Fused Attention
=======

Attention (including MHA, GQA, MQA) is one of the key parts of transformers and also the performance critical in many scenarios. To implement various optimizations, a fused attention layer and corresponding utilities for the customized KV-cache it uses are introduced. As an example, fused attention can reduce the cost of MHA from 8521.276 ms (17.044 ms) to 248.586 ms (7.944 ms) of a 1975-token llama-7b first-token inference[[1]](#1).

Note that this doc assumes you have the basic knowledge of this cpp graph implementation including `ne_tensor::ne`, `ne_tensor::nb` etc. Model builder can enable fused attention with operators `ne_flash_attn*` defined in [`core/ne_layers.h`](core/ne_layers.h) while the implementation of fused attentions is mostly located in [`core/layers/mha_dense.cpp`](core/layers/mha_dense.cpp).

## KV-cache initialization
The memory for kv-cache is allocated in `/models/model_utils/model_utils.cpp`. As the fused attention implementation requires certain instruction extensions and potentially some other limitations, fused attention is enabled only if `bestla_reordered_attn_fp32_support()` aggress, denoting with `memory_type = NE_TYPE_BTLA`. Next, `get_batch_kv_elements_from_gpt_params()` will give the sizes (in terms of bytes if fused attention enabled, or in terms elements if fused attention disabled) of k-cache and v-cache respectively for each batch and each layer. The KV-cache is finally prepared with these 2 sizes by creating `ne_new_tensor` inside `model.kv_self.k/.v` or `model.layer[il].k_cache/.v_cache`.

## KV-cache Append
KV-cache is appended every time a new pair of K and V are generated by evaluating inner product for QKV (ne_mul_qkv). This operation append an additional K/V-tensor on the dimension of sequence length (i.e. resulting `n_past = n_past + N`, where `n_past` is number of previously cached tokens and `N` is the length of current tokens).

The tensor of K and V should be both permuted (with `ne_permute`) to `batch x N x head_num x head_size`. The permuted K/V tensors can then be passed to `ne_flash_attn_update_k()/_v()`, together with `n_past` as an offset in the dimension of sequence length, to concatenate to a "view" of the cached K/V of the current layer and current batch.

> Note: Currently the K and V tensor to append must be contiguous in the dimension of head size (i.e. `head_dim` in some implementations). The strides (`ne_tensor::nb`) of the other 3 dimensions are configurable.

## Fused Attention Computation
With the KV-cache of customized type and layout, the attention can be computed at its best performance. `ne_flash_attn` accepts a Q-tensor of `batch x head_num x N x head_size` and outputs a result tensor of `batch x N x head_num x head_size` (totally contiguous).

> Node: similar to the append operation, Q-tensor must be contiguous in the dimension of head and its strides (`ne_tensor::nb`) of the other 3 dimensions are configurable.

## Supported Models
Fused attention is designed to be able to easily support various models:

<table>
  <thead>
    <tr>
      <th>Models</th>
      <th align="center">Attention Type & Support</th>
      <th>Runtime ISA</th>
    </tr>
  </thead>
  <tbody>
    <tr>
      <td>

[LLaMA-7B](https://huggingface.co/decapoda-research/llama-7b-hf), [LLaMA-13B](https://huggingface.co/decapoda-research/llama-13b-hf), [LLaMA2-7B](https://huggingface.co/meta-llama/Llama-2-7b-chat-hf), [LLaMA2-13B](https://huggingface.co/meta-llama/Llama-2-13b-chat-hf)
</td>
      <td align="center">MHA ✅</td>
      <td rowspan=10>AMX_BF16</td>
    </tr>
    <tr>
      <td>

[LLaMA2-70B](https://huggingface.co/meta-llama/Llama-2-70b-chat-hf)</td>
      <td align="center">GQA ✅</td>
    </tr>
    <tr>
      <td>
[GPT-J-6B](https://huggingface.co/EleutherAI/gpt-j-6b)</td>
      <td align="center">MHA ✅</td>
    </tr>
    <tr>
      <td>
[GPT-NeoX-20B](https://huggingface.co/EleutherAI/gpt-neox-20b), [Dolly-v2-3B](https://huggingface.co/databricks/dolly-v2-3b)</td>
      <td align="center">MHA ✅</td>
    </tr>
    <tr>
      <td>
[Qwen-7B](https://huggingface.co/Qwen/Qwen-7B-Chat), [Qwen-14B](https://huggingface.co/Qwen/Qwen-14B-Chat)</td>
      <td align="center">MHA ✅</td>
    </tr>
    <tr>
      <td>
[MPT-7B](https://huggingface.co/mosaicml/mpt-7b), [MPT-30B](https://huggingface.co/mosaicml/mpt-30b)</td>
      <td align="center">MHA with [ALiBi](https://arxiv.org/abs/2108.12409) ✅</td>
    </tr>
    <tr>
      <td>
[Falcon-7B](https://huggingface.co/tiiuae/falcon-7b), [Falcon-40B](https://huggingface.co/tiiuae/falcon-40b)</td>
      <td align="center">MQA, GQA ✅</td>
    </tr>
    <tr>
      <td>
[BLOOM-7B](https://huggingface.co/bigscience/bloomz-7b1)</td>
      <td align="center">MHA with [ALiBi](https://arxiv.org/abs/2108.12409) 🚧</td>
    </tr>
    <tr>
      <td>
[OPT-125m](https://huggingface.co/facebook/opt-125m), [OPT-350m](https://huggingface.co/facebook/opt-350m), [OPT-1.3B](https://huggingface.co/facebook/opt-1.3b), [OPT-13B](https://huggingface.co/facebook/opt-13b)</td>
      <td align="center">MHA 🚧</td>
    </tr>
    <tr>
      <td>
[ChatGLM-6B](https://huggingface.co/THUDM/chatglm-6b), [ChatGLM2-6B](https://huggingface.co/THUDM/chatglm2-6b)</td>
      <td align="center">MHA with the autoregressive-blank-infilling pattern 🚧,<br>GQA ✅</td>
    </tr>
    <tr>
      <td>
[StarCoder-1B](https://huggingface.co/bigcode/starcoderbase-1b), [StarCoder-3B](https://huggingface.co/bigcode/starcoderbase-3b), [StarCoder-15.5B](https://huggingface.co/bigcode/starcoder)</td>
      <td align="center">MHA ✅</td>
    </tr>
  </tbody>
</table>

> ✅: Supported; 🚧: WIP

### Limitations
Currently the fused attention is only enabled when compiling the Neural Speed with GCC11+.

## Tips for parallelism
Thanks to the mathematical nature of attention, one can simply parallel the whole kv-cache operations and fused attention on commonly-parallelizable dimensions. Just pass each part to every KV-cache operations (and merge them together if needed).

## References
<details>
<summary><a id="1">[1]</a> The data was tested on a single socket of Intel(R) Xeon(R) Platinum 8480L on commit 01a809d. Details below.</summary>

|                  | 1st-token fused attn disabled | 1st-token fused attn enabled | 4th-token fused attn disabled | 4th-token fused attn enabled |
| :--------------- | ----------------------------: | ---------------------------: | ----------------------------: | ---------------------------: |
| total latency    |                     9748.26ms |                    1475.57ms |                       50.37ms |                      41.27ms |
| fused-attn lat   |                             / |         179.883ms + 68.703ms |                             / |            6.271ms + 1.673ms |
| est non-attn lat |                    1226.984ms |                   1226.984ms |                      33.326ms |                     33.326ms |
| MHA cost compare |                    8521.276ms |                    248.586ms |                      17.044ms |                      7.944ms |

(4th token is taking as an example of next-token performance)

Row logs:
```
# fused attn enabled
rm -rf bin && cmake .. -GNinja -DNS_BUILD_TESTS=ON -DNE_PROFILING=ON -DCMAKE_BUILD_TYPE=Release && ninja run_llama && env ENGINE_PROFILING=1 numactl -m 1 -C 56-111 bin/run_llama -m llama-7b-hf-pr447-q4j-sym-int8-fp32-g128.bin --seed 1234 -t 56 -b 2048 -c 2048 -n 4 --memory-auto -p "$(echo "$LUOYU_PROMPT" | cut -d' ' -f 1-1500)"
Welcome to use the llama on the ITREX!
...
=== GRAPH Profiling ===
perf_total_per_op_us[                     ADD] =  51.409 ms
perf_total_per_op_us[                     MUL] =  26.328 ms
perf_total_per_op_us[                RMS_NORM] =  42.445 ms
perf_total_per_op_us[                 MUL_MAT] = 127.810 ms
perf_total_per_op_us[                 RESHAPE] =   0.446 ms
perf_total_per_op_us[                    VIEW] =   0.997 ms
perf_total_per_op_us[                 PERMUTE] =   0.101 ms
perf_total_per_op_us[               TRANSPOSE] =   0.105 ms
perf_total_per_op_us[                GET_ROWS] =   8.342 ms
perf_total_per_op_us[                    ROPE] =  44.115 ms
perf_total_per_op_us[                 MUL_QKV] = 252.611 ms
perf_total_per_op_us[                FFN_SILU] = 668.217 ms
perf_total_per_op_us[              FLASH_ATTN] = 179.883 ms
perf_total_per_op_us[    FLASH_ATTN_KV_UPDATE] =  68.703 ms
perf_total_per_op_us[           INNER PRODUCT] =   0.000 ms
========================================
=== GRAPH Profiling ===
perf_total_per_op_us[                     ADD] =   0.420 ms
perf_total_per_op_us[                     MUL] =   0.447 ms
perf_total_per_op_us[                RMS_NORM] =   1.377 ms
perf_total_per_op_us[                 RESHAPE] =   0.432 ms
perf_total_per_op_us[                    VIEW] =   0.956 ms
perf_total_per_op_us[                 PERMUTE] =   0.126 ms
perf_total_per_op_us[               TRANSPOSE] =   0.105 ms
perf_total_per_op_us[                GET_ROWS] =   0.024 ms
perf_total_per_op_us[                    ROPE] =   1.992 ms
perf_total_per_op_us[                 MUL_QKV] =   6.311 ms
perf_total_per_op_us[                FFN_SILU] =  14.597 ms
perf_total_per_op_us[              FLASH_ATTN] =   6.425 ms
perf_total_per_op_us[    FLASH_ATTN_KV_UPDATE] =   1.717 ms
perf_total_per_op_us[           INNER PRODUCT] =   3.535 ms
========================================
=== GRAPH Profiling ===
perf_total_per_op_us[                     ADD] =   0.402 ms
perf_total_per_op_us[                     MUL] =   0.358 ms
perf_total_per_op_us[                RMS_NORM] =   1.281 ms
perf_total_per_op_us[                 RESHAPE] =   0.427 ms
perf_total_per_op_us[                    VIEW] =   1.058 ms
perf_total_per_op_us[                 PERMUTE] =   0.106 ms
perf_total_per_op_us[               TRANSPOSE] =   0.102 ms
perf_total_per_op_us[                GET_ROWS] =   0.024 ms
perf_total_per_op_us[                    ROPE] =   1.919 ms
perf_total_per_op_us[                 MUL_QKV] =   5.881 ms
perf_total_per_op_us[                FFN_SILU] =  14.522 ms
perf_total_per_op_us[              FLASH_ATTN] =   6.389 ms
perf_total_per_op_us[    FLASH_ATTN_KV_UPDATE] =   1.621 ms
perf_total_per_op_us[           INNER PRODUCT] =   3.339 ms
========================================
=== GRAPH Profiling ===
perf_total_per_op_us[                     ADD] =   0.327 ms
perf_total_per_op_us[                     MUL] =   0.361 ms
perf_total_per_op_us[                RMS_NORM] =   1.272 ms
perf_total_per_op_us[                 RESHAPE] =   0.422 ms
perf_total_per_op_us[                    VIEW] =   1.032 ms
perf_total_per_op_us[                 PERMUTE] =   0.110 ms
perf_total_per_op_us[               TRANSPOSE] =   0.101 ms
perf_total_per_op_us[                GET_ROWS] =   0.023 ms
perf_total_per_op_us[                    ROPE] =   1.967 ms
perf_total_per_op_us[                 MUL_QKV] =   6.034 ms
perf_total_per_op_us[                FFN_SILU] =  14.527 ms
perf_total_per_op_us[              FLASH_ATTN] =   6.271 ms
perf_total_per_op_us[    FLASH_ATTN_KV_UPDATE] =   1.673 ms
perf_total_per_op_us[           INNER PRODUCT] =   3.444 ms
========================================

model_print_timings:        load time =  2691.89 ms
model_print_timings:      sample time =     2.36 ms /     4 runs   (    0.59 ms per token)
model_print_timings: prompt eval time =  1475.57 ms /  1975 tokens (    0.75 ms per token)
model_print_timings:        eval time =   124.68 ms /     3 runs   (   41.56 ms per token)
model_print_timings:       total time =  2853.38 ms
========== eval time log of each prediction ==========
prediction   0, time: 1475.57ms
prediction   1, time: 42.19ms
prediction   2, time: 41.22ms
prediction   3, time: 41.27ms

# fused attn disabled
rm -rf bin && cmake .. -GNinja -DNS_BUILD_TESTS=ON -DNE_PROFILING=ON -DCMAKE_BUILD_TYPE=Release && ninja run_llama && env ENGINE_PROFILING=1 numactl -m 1 -C 56-111 bin/run_llama -m llama-7b-hf-pr447-q4j-sym-int8-fp32-g128.bin --seed 1234 -t 56 -b 2048 -c 2048 -n 4 --memory-f16 -p "$(echo "$LUOYU_PROMPT" | cut -d' ' -f 1-1500)"
Welcome to use the llama on the ITREX!
...
=== GRAPH Profiling ===
perf_total_per_op_us[                     ADD] =  55.300 ms
perf_total_per_op_us[                     MUL] =  40.209 ms
perf_total_per_op_us[                RMS_NORM] =  63.544 ms
perf_total_per_op_us[                 MUL_MAT] = 6698.093 ms
perf_total_per_op_us[                   SCALE] = 1325.542 ms
perf_total_per_op_us[                     CPY] = 273.083 ms
perf_total_per_op_us[                 RESHAPE] =   0.460 ms
perf_total_per_op_us[                    VIEW] =   0.734 ms
perf_total_per_op_us[                 PERMUTE] =   0.323 ms
perf_total_per_op_us[               TRANSPOSE] =   0.105 ms
perf_total_per_op_us[                GET_ROWS] =   8.467 ms
perf_total_per_op_us[           DIAG_MASK_INF] =  69.310 ms
perf_total_per_op_us[                SOFT_MAX] = 226.629 ms
perf_total_per_op_us[                    ROPE] =  44.610 ms
perf_total_per_op_us[                 MUL_QKV] = 264.430 ms
perf_total_per_op_us[                FFN_SILU] = 672.668 ms
perf_total_per_op_us[           INNER PRODUCT] =   0.000 ms
========================================
=== GRAPH Profiling ===
perf_total_per_op_us[                     ADD] =   0.445 ms
perf_total_per_op_us[                     MUL] =   0.405 ms
perf_total_per_op_us[                RMS_NORM] =   1.232 ms
perf_total_per_op_us[                 MUL_MAT] =  10.702 ms
perf_total_per_op_us[                   SCALE] =   0.952 ms
perf_total_per_op_us[                     CPY] =   3.040 ms
perf_total_per_op_us[                 RESHAPE] =   0.416 ms
perf_total_per_op_us[                    VIEW] =   0.792 ms
perf_total_per_op_us[                 PERMUTE] =   0.323 ms
perf_total_per_op_us[               TRANSPOSE] =   0.103 ms
perf_total_per_op_us[                GET_ROWS] =   0.023 ms
perf_total_per_op_us[           DIAG_MASK_INF] =   0.118 ms
perf_total_per_op_us[                SOFT_MAX] =   1.359 ms
perf_total_per_op_us[                    ROPE] =   1.888 ms
perf_total_per_op_us[                 MUL_QKV] =   6.133 ms
perf_total_per_op_us[                FFN_SILU] =  14.607 ms
perf_total_per_op_us[           INNER PRODUCT] =   3.504 ms
========================================
=== GRAPH Profiling ===
perf_total_per_op_us[                     ADD] =   0.324 ms
perf_total_per_op_us[                     MUL] =   0.402 ms
perf_total_per_op_us[                RMS_NORM] =   1.321 ms
perf_total_per_op_us[                 MUL_MAT] =  10.624 ms
perf_total_per_op_us[                   SCALE] =   0.954 ms
perf_total_per_op_us[                     CPY] =   3.104 ms
perf_total_per_op_us[                 RESHAPE] =   0.425 ms
perf_total_per_op_us[                    VIEW] =   0.748 ms
perf_total_per_op_us[                 PERMUTE] =   0.316 ms
perf_total_per_op_us[               TRANSPOSE] =   0.102 ms
perf_total_per_op_us[                GET_ROWS] =   0.021 ms
perf_total_per_op_us[           DIAG_MASK_INF] =   0.111 ms
perf_total_per_op_us[                SOFT_MAX] =   1.362 ms
perf_total_per_op_us[                    ROPE] =   1.874 ms
perf_total_per_op_us[                 MUL_QKV] =   6.001 ms
perf_total_per_op_us[                FFN_SILU] =  14.542 ms
perf_total_per_op_us[           INNER PRODUCT] =   3.314 ms
========================================
=== GRAPH Profiling ===
perf_total_per_op_us[                     ADD] =   0.354 ms
perf_total_per_op_us[                     MUL] =   0.391 ms
perf_total_per_op_us[                RMS_NORM] =   1.379 ms
perf_total_per_op_us[                 MUL_MAT] =  10.610 ms
perf_total_per_op_us[                   SCALE] =   0.964 ms
perf_total_per_op_us[                     CPY] =   3.115 ms
perf_total_per_op_us[                 RESHAPE] =   0.430 ms
perf_total_per_op_us[                    VIEW] =   0.866 ms
perf_total_per_op_us[                 PERMUTE] =   0.336 ms
perf_total_per_op_us[               TRANSPOSE] =   0.109 ms
perf_total_per_op_us[                GET_ROWS] =   0.022 ms
perf_total_per_op_us[           DIAG_MASK_INF] =   0.108 ms
perf_total_per_op_us[                SOFT_MAX] =   1.410 ms
perf_total_per_op_us[                    ROPE] =   1.959 ms
perf_total_per_op_us[                 MUL_QKV] =   5.826 ms
perf_total_per_op_us[                FFN_SILU] =  14.737 ms
perf_total_per_op_us[           INNER PRODUCT] =   3.378 ms
========================================

model_print_timings:        load time = 10987.95 ms
model_print_timings:      sample time =     2.38 ms /     4 runs   (    0.60 ms per token)
model_print_timings: prompt eval time =  9748.26 ms /  1975 tokens (    4.94 ms per token)
model_print_timings:        eval time =   150.93 ms /     3 runs   (   50.31 ms per token)
model_print_timings:       total time = 11175.74 ms
========== eval time log of each prediction ==========
prediction   0, time: 9748.26ms
prediction   1, time: 50.56ms
prediction   2, time: 50.00ms
prediction   3, time: 50.37ms
```
</details>
