在工作中如果我们写了一个 cuda kernel，需要去计算他的理论利用率和理论带宽，这样我们才能知道这个 kernel 还有多少的优化空间。

搞清楚两个定义，tflops 和 memorybandwidth 分别是什么？ memory bandwidth 是单位时间 HBM 到寄存器的访问量。

在 llm 中，最重要的一个 kernel 就是 attention，下面针对 causal LM，计算 self attention 在前向过程中的 tflops 和 bandwidth。

首先回顾一下 attention 的计算流程，假设 query，key 和 value 分别为 QKV，那么整体 attention 的计算过程如下：

1. 首先计算 attention score \(A= \frac{Q K^T}{\sqrt{d_k}}\)，其中 \(\sqrt{d_k}\) 一般是 head_size；
2. 接着通过 softmax 归一化将 attention scores 转换成概率 p = softmax(masked_tri(A))，注意对于 causal LM 的推理流程，我们取下三角;
3. 使用计算之后的 p 对 V 进行加权获得输出，O = p * V。

其中有两个矩阵操作...

一般来说，transformer 都是使用的 MHA(Multi-Head Attention)，这里可以对单个 head 进行分析，multi-head 只需要乘上 head_size 即可。

在 LLM 进行推理的过程中有两个不同的 stage，分别是 prefill 和 decode。在 prefill 阶段会一次推理处理所有的 token，而在 decode 阶段则是每次处理一个 token，下面分别分析这两个阶段的 tflops 和 bandwidth。

在开始分析之前，先约定一下超参数

如何计算 tflops

对于 single-head，假设 QKV 的 size 是 ‌[M, head_size]，通过 profiling 获得了 attention kernel 的计算耗时是 ms，计算 tflops 可以使用下面的公式

perf = lambda ms: (
    2  # two ops
    * 2  # q@k + p@v
    / 2  # causal
    * batch_size
    * num_heads
    * head_size
    * seq_len
    * seq_len
    * 1e-12
    / (ms * 1e-3)
)

下面来解释下为什么？

对于 \(QK^T\) 的计算，他的 flops 是 \(O(bs^2nd)\)，其中 b 是 b

如何计算 bandwidth

对于 pre-fill 阶段，计算 bandwidth 可以通过下面的公式

elem_size = q.element_size()
perf = lambda ms: (
    elem_size
    * batch_size
    * (num_q_heads + 2 * num_k_heads)
    * head_size
    * seq_len
    * 1e-9
    / (ms * 1e-3)
)

下面来解释一下为什么？

• elem_size 是每一个 data element 的大小，单位是 bytes，比如 fp16 就是 2，fp8 或者 int8 就是 1；

Reference

https://le.qun.ch/en/blog/2023/05/13/transformer-batching/