INT8 KV-Cache Attention

原始题目：LeetGPU - INT8 KV-Cache Attention

题目描述

实现解码阶段的多头注意力，其中键和值缓存以 int8 格式存储并附有逐 token 缩放因子。相比 float32，这种布局将 KV-cache 带宽减半，被 TensorRT-LLM 和 vLLM 等生产级 LLM 服务系统广泛采用。

给定单个新 token 的查询张量 $Q$、int8 键缓存 $K_{int8}$、int8 值缓存 $V_{int8}$ 以及逐 token 缩放因子 $k\_scale$ 和 $v\_scale$，解量化缓存并计算缩放点积注意力输出。

解量化：$K_{float}[h, s, d] = K_{int8}[h, s, d] \times k\_scale[h, s]$（$V$ 同理）。注意力使用 $1/\sqrt{head\_dim}$ 作为缩放因子。

实现要求

• 实现 solve(Q, K_int8, V_int8, k_scale, v_scale, output, num_heads, seq_len, head_dim)。
• 不允许使用外部库。

示例

num_heads=1, seq_len=3, head_dim=4
Q = [[1,1,1,1]]
K_int8 = [[[100,0,0,0],[0,100,0,0],[0,0,100,0]]], k_scale = [0.1,0.1,0.1]
V_int8 = [[[10,20,30,40],...]], v_scale = [0.1,0.1,0.1]
K_float = [[1,0,0,0],[0,1,0,0],[0,0,1,0]], V_float = [[1,2,3,4],[5,6,7,8],[9,10,11,12]]
Scores = softmax([0.25,0.25,0.25]) = [1/3,1/3,1/3]
Output = [5,6,7,8]

约束条件

• $1 \le num\_heads \le 64$，$1 \le seq\_len \le 32{,}768$，$8 \le head\_dim \le 256$（8 的倍数）。
• $K_{int8}, V_{int8}$ 值在 $[-128, 127]$，缩放因子为正 float32。
• 性能测试在 $num\_heads=32, seq\_len=8{,}192, head\_dim=128$ 下进行。

解题思路

INT8 KV-cache 的核心价值在解码阶段——每个新 token 需要读取整个 KV-cache（$O(seq\_len)$），量化将带宽需求减半。解量化可以与注意力计算融合：在读入 int8 值的同时乘缩放因子，避免单独的解量化 kernel。逐 token 而非逐 channel 的缩放策略简化了实现，也是生产系统中的常见选择。欢迎在 GitHub Discussions 分享你的解法。