Skip to content

Qualcomm AI Engine Direct - documentation for KV cache update #8134

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 5, 2025
Merged

Qualcomm AI Engine Direct - documentation for KV cache update #8134

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 57 additions & 3 deletions examples/qualcomm/oss_scripts/llama/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,10 @@ This file provides you the instructions to run LLAMA model with different parame
1. LLAMA2 Stories 110M
2. LLAMA3.2 1B
3. LLAMA3.2 3B (WIP)

We offer the following modes to execute the model:

Prefill Mode: This is also known as batch prefill mode, where the model takes in a list of tokens as input and generates the next token along with the key-value (KV) cache for all tokens. This mode is efficient for generating the initial sequence of tokens (usually the user's prompt).
Prefill Mode: This is also known as batch prefill mode, where the model takes in a list of tokens as input and generates the next token along with the key-value (KV) cache for all tokens. This mode is efficient for encoding the user's prompt.

KV Cache Mode: In KV Cache mode, the model takes in a single previous token and generates the next predicted token along with its KV cache. It is efficient for generating subsequent tokens after the initial prompt.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Let's rename it to generate mode, as this term isn't used very common and a bit confusing


Expand Down Expand Up @@ -41,7 +42,7 @@ python -m extension.llm.tokenizer.tokenizer -t tokenizer.model -o tokenizer.bin
echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps": 1e-05, "vocab_size": 32000}' > params.json
```

#### LLAMA3.2
#### LLAMA3.2
Follow the [instructions](https://www.llama.com/) to download models.
At the end of this step, users should have the following files ready: `consolidated.00.pth`, `params.json`, and `tokenizer.model`.

Expand All @@ -58,6 +59,53 @@ Default example using hybrid mode.
python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --ptq 16a4w --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --llama_model llama3_2 --model_mode hybrid --prefill_seq_len 32 --kv_seq_len 128 --prompt "what is 1+1"
```

### KV Cache update mechanism
We have two distinct mechanisms for updating the key-value (KV) cache, which can be selected at runtime. Shift Pointer and Smart Mask.

#### Shift Pointer mechanism

<figure>
<img src="./assets/ShiftPointer.png" alt="Shift Pointer mechanism"> <figcaption>
The figure illustrates the process of updating the key and value caches during each inference step. In key cache update process, we initially allocate memory for each layer with <code>num_head</code> size of <code>(head_dim + 1) * (seq_len - 1)</code>. After a single inference, the new key cache is copied from the key output pointer <code>k_out</code> and appended to the key cache. Subsequently, the buffer start pointer of the key cache <code>k_in</code> moves to the next token, making the previous position of the buffer start pointer unused. This process is repeated for each subsequent inference step.
For the value cache update process, we first allocate a contiguous memory of size <code>(num_head + 1) * head_dim * (seq_len - 1)</code> for each layer, with the last head reserved for I/O shifting, After the first inference, the cache is updated by simply shifting the pointers of all heads to the next token position, making only the previous <code>head_dim * 1</code> section of the buffer start pointer <code>v_in</code> of the first head unused. This process is repeated for each subsequent inference step.</figcaption>
</figure>

#### Smart Mask mechanism:
<figure>
<img src="./assets/SmartMask.png" alt="Smart Mask mechanism">
<figcaption>The Smart Mask mechanism streamlines the process of updating tokens in the cache. Unlike the Shift Pointer mechanism, which requires moving the buffer start pointer <code>k_in</code>/<code>v_in</code> of the cache, the Smart Mask mechanism updates only the new token at the specified position. This approach eliminates the need to adjust the buffer start pointer. This mechanism is beneficial for shared buffers but requires CPU memory copying. </figcaption>
</figure>

#### Analysis KV Cache Update Mechanism for each Layer each inference
<table>
<tr>
<th>Mechanism</th>
<th colspan="2" style="text-align:center;">Time Complexity</th>
<th colspan="2" style="text-align:center;">Space Complexity</th>
</tr>
<tr>
<th></th>
<th style="text-align:center;">K</th>
<th style="text-align:center;">V</th>
<th style="text-align:center;">K</th>
<th style="text-align:center;">V</th>
</tr>
<tr>
<td style="text-align:center;">Shift Pointer</td>
<td style="text-align:center;">num_head * head_dim</td>
<td style="text-align:center;">1</td>
<td style="text-align:center;">num_head * (head_dim + 1) * seq_len</td>
<td style="text-align:center;">(num_head + 1) * head_dim * (seq_len - 1)</td>
</tr>
<tr>
<td style="text-align:center;">Smart Mask</td>
<td style="text-align:center;">num_head * head_dim</td>
<td style="text-align:center;">num_head * head_dim</td>
<td style="text-align:center;">num_head * seq_len * head_dim</td>
<td style="text-align:center;">num_head * seq_len * head_dim</td>
</tr>
</table>

### Additional Configs when running the script
If you would like to compile the model only, we have provided the flag `--compile_only`. Taking LLAMA3.2 as an example:
```bash
Expand All @@ -67,4 +115,10 @@ python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -m ${SOC_MO
On the other hand, if you already have a pre-compiled .pte model, you can perform inference by providing the flag `--pre_gen_pte` and specifying the folder that contains the .pte model. Taking LLAMA3.2 as an example:
```bash
python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --ptq 16a4w --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --llama_model llama3_2 --model_mode hybrid --prefill_seq_len 32 --kv_seq_len 128 --prompt "what is 1+1" --pre_gen_pte ${FOLDER_TO_PRE_GEN_PTE}
```
```

You can select the KV Cache update mechanism at runtime by setting the `KV_UPDATER` variable to either "shift_pointer" or "smart_mask". By default, it is set to "smart_mask".
`KV_UPDATER` = "shift_pointer"
```bash
python examples/qualcomm/oss_scripts/llama/llama.py -b build-android -s ${SERIAL_NUM} -m ${SOC_MODEL} --ptq 16a4w --checkpoint consolidated.00.pth --params params.json --tokenizer_model tokenizer.model --llama_model llama3_2 --model_mode hybrid --prefill_seq_len 32 --kv_seq_len 128 --prompt "what is 1+1" --kv_updator ${KV_UPDATER}
```
Binary file added examples/qualcomm/oss_scripts/llama/assets/ShiftPointer.png
View file
Open in desktop
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file added examples/qualcomm/oss_scripts/llama/assets/SmartMask.png
View file
Open in desktop
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Loading