Support for Qualcomm AI Hub sliding window models

[baijumeswani]

Language models sourced from Qualcomm AI Hub have a sliding window structure. This pull-request adds support for sliding window models to the decoder-pipeline.

Working with the model sourced from the Qualcomm AI Hub, the model execution goes through two stages:

• Context (prompt) processing: The model runs through several iterations of the model processing window_size tokens at a time to compute the key-value cache and generate the first token. In each iteration, we slide the window (by window_size over the key-value cache buffer so the new key-value cache can be computed.
• Token generation: After the input prompt has been processed, the execution switches over to the token generation phase where new tokens are processed one at a time. In this phase, we slide the key-value cache buffer by 1 at each iteration.

For example, let's assume the following inputs:

• Prompt: ["What", "is", "the", "square", "root", "of", "16", "?"]
• Window size: 5
• Pad token: 0
• context length: 12

Given the prompt has 8 tokens (assuming a word is a token), to process the entire prompt, 2 iterations will be required.

• Iteration 1:

  • input_ids: [0, 0, "What", "is", "the"]
  • position_ids: [0, 0, 0, 1, 2]
  • attention_mask: [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1]

• Iteration 2:

  • input_ids: ["square", "root", "of", "16", "?"]
  • position_ids: [3, 4, 5, 6, 7]
  • attention_mask: [[0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]

The above two iterations complete the prompt processing phase. Subsequently, the token generation phase begins. The token generation phase processes 1 token at a time.

At each iteration, the key-value cache is updated by sliding the key-value cache buffer by window size tokens during the prompt processing phase and by 1 token during the token generation phase.

This pull-request introduces the above changes to the decoder-pipeline to support the sliding window models from the Qualcomm AI Hub

In addition to the above changes, the pull request also introduces the QNN_WITH_SHARED_MEMORY device type so memory can be allocated using the qnn allocator. Thanks to @edgchen1 for making these changes.

Open items:

• Continuous decoding support for these model types.
• Incorrect results when the generated tokens size becomes large enough.
• Performance.
• Tests

[baijumeswani]

Thank you for the review. :)