从 Qwen3 数学到 CUDA Kernel：0.6B 推理运行时的完整实现与优化

0.1 FP32、BF16 和 INT8 各自负责什么#

本项目的主路径是 BF16 输入/输出，归约中需要更大动态范围的位置使用 FP32：

• RMSNorm 的平方和使用 FP32；
• attention score、最大值、指数和与输出累加使用 FP32；
• W8A16 的 INT8 权重先乘 FP32 scale，再累加到 FP32；
• 最终激活或 logits 再转换为 BF16。

INT8 不是“把所有 tensor 强制变成整数”。W8A16 表示权重 $W$ 是 8 bit，激活 $X$ 仍是 16 bit。量化只在通过质量门控的 LM head 启用。

1.1 参数量从 shape 逐项算出来#

单个 layer 的主要矩阵是：

因为 $H_qD=16\times128=2048$，Q 的宽度是 2048；因为 $H_{kv}D=8\times128=1024$，K/V 的宽度各是 1024。加上 norm 权重后，每层 15,730,944 个参数；28 层与 tied embedding 合计 596,049,920 个参数。

BF16 参数占用约

这个结果建立了第一个工程约束：6 GiB 显存可以容纳模型，但长 KV cache、CUDA Graph 静态 buffer、workspace 和框架开销仍需严格规划。

2.1 Input RMSNorm#

对 hidden row $x=(x_0,\ldots,x_{d-1})$，RMSNorm 定义为

$\gamma_i$ 是可学习缩放，$\epsilon=10^{-6}$ 防止分母为零。它没有减均值，所以不同于 LayerNorm。CUDA 实现需要先把 1024 个平方归约成一个 FP32 标量，再把标量广播给整行。

令

2.2 Packed QKV projection#

数学上有三次线性层：

运行时把权重按输出维拼接：

一次线性层得到 $[Q\mid K\mid V]$。这样不是减少 FLOP，而是减少 GEMM 调度和权重描述开销， 并让后续 split 只是 view。

2.3 Q/K Head RMSNorm#

Qwen3 对每个 Q/K head 的 128 维向量再次归一化。对 head 向量 $q_h$：

$r$ 是 head 内维度，$\odot$ 表示逐元素乘。K 同理，但 V 不做该归一化。

2.4 RoPE 将位置编码成二维旋转#

对位置 $p$ 和维度对 $i$，角频率为

相位为

把向量的两个半区配对，令 $a=q_i$、$b=q_{i+D/2}$：

K 做同样旋转。RoPE 不改变向量长度，但让 Q/K 点积包含相对位置信息。

2.5 GQA：16 个 Q head 共享 8 个 KV head#

每个 KV head 服务

个 query head。query head $h$ 使用

号 KV head。例如 Q head 0、1 都读 KV head 0；Q head 2、3 都读 KV head 1。 这个映射既减少 KV cache 容量，也给 CUDA kernel 提供同一 K/V 被两个 Q head 复用的机会。

2.6 Scaled dot-product attention#

对 batch $b$、query head $h$、query 位置 $t$ 和历史位置 $j$：

$s$ 是 score。causal mask 要求 $j\le t$。稳定 softmax 为

head 输出为

拼接 16 个 query head 后宽度为 2048，再经过 O projection：

第一条 residual：

2.7 SwiGLU MLP#

先做 post-attention RMSNorm：

gate 与 up projection 可以打包为一次 GEMM：

其中 $G,U\in\mathbb{R}^{B\times S\times3072}$。SiLU 定义为

SwiGLU 输出为

down projection 和第二条 residual 得到下一层输入：

到这里，模型数学已经闭环。下面才讨论为什么运行时不能按这些公式逐个调用通用算子。

3.1 Prefill#

prefill 一次处理 $S$ 个 prompt token。线性层中的

通常较大，因此是有充分行复用的 GEMM。同一权重 tile 被许多输入行复用，Tensor Core 更容易 发挥吞吐。

attention 需要处理整个 causal 三角区域，score 数量近似 $S^2/2$。若显式存储完整 score， 显存和带宽会按 $O(S^2)$ 增长，所以 production prefill 后端使用 fused scaled dot-product attention。

3.2 Decode#

decode 每步只输入一个新 token：

线性层退化为 GEMV。权重几乎没有跨输入行复用，每步都要重新流过显存。attention 只计算一个 query，却必须读取长度 $T$ 的全部历史 K/V，所以成本随 context 线性增长。

3.3 为什么必须保存 KV Cache#

如果不缓存，第 $t$ 步会重新为前 $t$ 个 token 计算 K/V；生成 $N$ 步的重复工作近似 $O(N^2)$。缓存后，每层只追加当前 token 的 K/V，历史部分直接读取。

单层 BF16 KV cache 字节数为

第一个 2 表示 K 和 V，最后一个 2 表示 BF16 每元素 2 byte。代入 $B=1$、 $H_{kv}=8$、$D=128$：

28 层合计

因此 $T=8192$ 时约 896 MiB。decode 每步又要读取这些历史 K/V，这解释了为什么 prompt 越长， attention 占完整请求的比例越高。

4.1 Decode GEMV 的权重墙#

对 $M=1$ 的 BF16 线性层，每个权重参与一次乘加，产生约 2 FLOP，同时至少读取 2 byte：

这是低算术强度 workload。单纯改成 Tensor Core 指令不会消除权重读取下界。

28 层主要矩阵的 BF16 权重流量约 840 MiB/token；tied LM head 另有 296.75 MiB/token。 短 context 时 LM head 是显著的固定带宽成本，长 context 时 KV 读取逐渐追上并超过它。

4.2 Naive 实现的五类浪费#

1. RMSNorm、RoPE、cache append 分开，反复读写 Q/K。
2. 两个 Q head 共享同一 KV head，但 kernel 仍重复读取 K/V。
3. 短 context 分得太碎，长 context 又没有足够 CTA。
4. 每层、每 token 重复 Python、dispatcher 和 kernel launch。
5. 把 $M=1$ 的带宽瓶颈误判为“Tensor Core 用得不够”。

项目的优化顺序据此确定：先消除冗余数据移动与 launch，再做 shape dispatch，最后只在确有收益且 通过质量门控的位置量化。

5.1 PyTorch dispatcher 的完整调用链#

以 RMSNorm 为例：

1
def rms_norm(x, weight, eps):
2
    """调用已注册 CUDA 算子并保持模型层不依赖扩展细节。"""
3
    return torch.ops.qwen_cuda.rms_norm(x, weight, eps)

schema 描述接口；FakeTensor implementation 只推导输出 metadata；C++ wrapper 检查 device、 dtype、shape、contiguous 和同设备约束，再把当前 PyTorch stream 传给 kernel。kernel 不能偷偷 使用 default stream，否则 eager 可能偶尔正确，CUDA Graph capture 或异步流水会出现竞态。

5.2 为什么不在 hot path 分配显存#

KV pages、attention scratch、token history、context length 和 graph 输入输出都在初始化时分配。 decode 循环只覆盖静态 buffer，不反复 cudaMalloc。PyTorch caching allocator 充当显存池， CUDA Graph 又要求 replay 时地址不变，这两点在此处是一致的设计要求。

7.1 为什么一个 warp 正好处理一个 head#

$D=128$，一个 warp 有 32 lane。让 lane $l$ 负责

于是每 lane 连续处理 4 个值，全 warp 恰好覆盖 128 维。每 lane 先累加 4 个平方，warp shuffle 合并成 norm；随后所有 lane 使用同一个 inverse RMS。

RoPE 的 half-rotation 配对是 $(r,r+64)$。lane $l$ 同时持有 $l$、$l+64$，以及 $l+32$、$l+96$，不必通过 shared memory 交换。这个映射不是“warp 越多越快”，而是由 $D=128$ 与 RoPE 配对关系共同推导出来。

7.2 逻辑 token 到物理 cache 地址#

令当前位置为 $p$：

$\ell$ 是逻辑页，$o$ 是页内位置。block table 给出物理页：

融合 kernel 在得到 rotated K 后直接写物理页 $q$，避免中间 tensor 再走一次 cache append kernel。Q 只返回给当前 attention，K 同时返回给 prefill attention 并持久化到 cache。

9.1 为什么分区结果仍能精确合并#

长 context 被切成 partition。第 $p$ 个 partition 输出 $(m_p,\ell_p,o_p)$。全局最大值

把每个局部状态重标定到 $M$：

最终结果 $O/L$。这不是近似，也不是省略 softmax；它只是利用指数平移不改变归一化概率的性质， 把一次长归约变成可并行的两阶段归约。

对应 CUDA 主循环保留 FP32 状态：

1
/**
2
 * 将一个新 score/value 合并进 online-softmax 状态。
3
 * max_value、denom 和 out_accum 始终以同一个指数基准表示。
4
 */
5
__device__ void update_online_state(
6
    float score,
7
    const float* value,
8
    int head_dim,
9
    float& max_value,
10
    float& denom,
11
    float* out_accum) {
12
    const float next_max = fmaxf(max_value, score);
13
    const float old_scale = expf(max_value - next_max);
14
    const float new_scale = expf(score - next_max);
15

16
    denom = denom * old_scale + new_scale;
17
    for (int r = 0; r < head_dim; ++r) {
18
        out_accum[r] = out_accum[r] * old_scale + value[r] * new_scale;
19
    }
20
    max_value = next_max;
21
}

真实 kernel 会让 lane 分担 $D=128$ 的分量并用 warp primitive 归约 score；这里的标量版本用于 展示数学与代码状态的一一对应。

10.1 短路径：grouped single CTA#

grid 映射为 $(b,\text{kv head})$，一个 CTA 同时处理共享该 KV head 的两个 Q head。K/V 从 paged cache 读取一次，两个 query 分别维护 online-softmax 状态。

适用条件是 context 较短，单 CTA 可以在可接受时间内扫描完。收益是：

1. 避免 partition scratch；
2. 避免第二次 reduction launch；
3. GQA 的 K/V 复用发生在同一 CTA；
4. 固定 launch 数更少。

代价是 context 增长后单 CTA 工作量过大，SM 间并行不足。

10.2 长路径：64-token partition 加 reduction#

第一阶段 grid 映射为 $(b,\text{kv head},\text{partition})$。每个 warp 处理 64 个历史 token， 输出 $(m_p,\ell_p,o_p)$ 到预分配 FP32 scratch。第二阶段按前节公式合并 partition。

收益是把 context 维切到多个 CTA 并行；代价是 scratch 写入、额外 launch 和 reduction。 因此 dispatch threshold 必须 benchmark，而不能凭感觉固定。

10.3 Native 与 production 后端的诚实边界#

仓库的 native attention 是完整、可运行、可测试的本项目 CUDA 实现，适合学习、调试和继续优化。 production attention 明确复用 vLLM 0.23.0 已注册的 paged_attention_v1。端到端“与 vLLM 持平”的结论使用 production 后端，不能把它写成 native kernel 单独达到同样水平。

这是工程上的正确取舍：底层学习实现负责暴露机制，production 路径负责给出可靠基线。

12.1 K 维双缓冲#

把 $K$ 分成多个 $K_{tile}$。stage 0 计算时，stage 1 预取下一块：

双缓冲隐藏 global-to-shared 延迟，但 shared memory 翻倍，寄存器地址状态增加，可能降低 occupancy。只有计算足够长时，隐藏延迟的收益才覆盖资源代价。

12.2 对齐、向量化和 bank conflict#

packed load 要求地址满足对应对齐，例如 16-byte load 的起始地址应 16-byte 对齐；尾部不满足 tile 的 shape 必须 mask 或回退。shared memory 若让一个 warp 的多个 lane 命中同一 bank 的 不同地址，会串行化；layout 常通过 padding 或 swizzle 改变 bank 映射。

occupancy 也不是越高越好。较大 tile 可提高复用，却消耗更多寄存器/shared memory，减少驻留 CTA。正确选择来自 shape-specific benchmark，而不是固定最大 block。

12.3 精度路径的取舍#

12.4 为什么自写 WMMA 被删除#

本项目曾实现 BF16 WMMA 与 fused SwiGLU。实测 0.11264 ms，而 cuBLAS 加单独 SwiGLU 是 0.07885 ms，只达到 0.70x。原因包括 tile 边界、epilogue、occupancy 和 cuBLAS 成熟调度优势。

“使用 Tensor Core”只是实现事实，不是性能证明。该路径没有达到门槛，因此从正式代码删除，只在 实验记录中保留结论。

13.1 CUDA lane 映射#

一个 warp 负责一个输出词表行，一个 block 8 warp。每个 128-value group 中，32 lane 各加载 4 个 packed INT8 和 4 个 BF16：

lane 局部完成 4 次乘加，warp shuffle 合并，lane 0 乘该 group scale 并累加 FP32 total。

1
/**
2
 * 计算一个 128-value group 的局部 W8A16 点积。
3
 * packed_weight 每个 int32 存 4 个 int8；packed_input 存 4 个 bf16。
4
 */
5
float partial = 0.0F;
6
#pragma unroll
7
for (int item = 0; item < 4; ++item) {
8
    partial = fmaf(
9
        static_cast<float>(weight_values[item]),
10
        __bfloat162float(input_values[item]),
11
        partial);
12
}
13
partial = warp_sum(partial);
14
if (lane == 0) {
15
    total = fmaf(partial, group_scales[row * groups + group], total);
16
}

这个 kernel 使用普通 CUDA core 做反量化与 FP32 FMA，不宣称是 INT8 Tensor Core GEMM。 $M=1$ 时它的核心收益来自权重字节减少，而不是更高理论 TOPS。

13.2 为什么只量化 LM head#

全层 W8A16 实验的 teacher-forced top-1 agreement 只有 0.9625，未通过质量门槛；LM-head-only 模式达到 0.99375，因此正式路径只保留后者。量化范围是由质量证据决定的，不是由“8 bit 更快” 决定的。

14.1 环境与方法#

benchmark 前检查 nvidia-smi，确认没有其他计算进程占用 GPU。计时覆盖同一离线请求的 prefill 与 128-token decode，不包含模型加载、tokenize 和 detokenize。

完整请求输出吞吐定义为

它不是多请求服务的 aggregate throughput，也不是只算 decode loop 的 token/s。

14.2 端到端结果#

必须定义三个不同的比值：

这里 W8A16/BF16 才是启用量化本身带来的实测增量：短 prompt 10.4%，8K prompt 2.6%。 W8A16/vLLM 同时包含本项目其他路径与外部 baseline 的差异，不能把 1.317x 写成 W8 kernel 单独加速 31.7%。

14.3 为什么 context 越长，W8 收益越小#

LM head 是固定成本；attention KV 读取随 $T$ 增长。短 prompt 时压缩 296.75 MiB LM head 占比高，端到端收益明显。长 prompt 时每步接近 896 MiB KV 读取，LM head 占比下降，所以 W8A16 microbenchmark 的约 1.95x 不会等比例变成完整请求加速。

这也是为什么必须同时报告 microbenchmark 和端到端数据。

14.4 算子与质量结果#

W8A16 microbenchmark：

LM-head-only 质量门控使用 5 个 prompt、每个 32 个 teacher-forced token：

这些指标说明该测试集上 token 决策高度一致，但不等价于“所有任务绝对无损”。改变模型、量化范围、 group size 或采样方式后必须重做质量验证。

14.5 Native attention 的优化链#

native 后端 decode token/s 从初版到 dual graph：

短 prompt 因 launch 占比高而受益；8K 出现轻微回退，说明 graph 并不自动让所有 shape 更快。 生产路径仍使用 vLLM 注册 attention，native 数据不能替代生产端到端表。

14.6 为什么旧表和新表不能直接相除#

旧 baseline 的最大 prompt 是 1024，运行时按 max_model_len=1168 配置；新表测试到 8192，使用 max_model_len=8336。静态 cache、graph buffer、attention dispatch 和 工作集都变化了。因此跨两轮结果的差异不能全部归因于某一个新 kernel。

正确做法是同一 commit、同一配置、同一计时脚本做 A/B；本轮 W8A16/BF16 比值就是这样的 同配置对照。