课程 04 · 计算内核

KV 到底写到了哪一格

这一课把 slot_mapping 这个 int32 数组追到底：谁构造它、谁消费它、为什么 prefill 和 decode 的 FlashAttention 调用长得完全不一样。

约 60 分钟 源码：layers/attention.py · model_runner.py · layers/sampler.py 前置：02、03

把维度变化拆到每一步

1. QKᵀ中的转置只交换最后两根轴，沿d求和。
2. 将逻辑矩阵与真实backend布局分开读，再追踪slot_mapping。

打开元素级图解：矩阵、轴、索引与可点击反例 →

1一次前向需要的五样东西

模型本身不关心调度。它只问：这次要算哪些 token、它们在序列里的位置是什么、KV 写到哪、历史 KV 去哪读。prepare_prefill()[L129–170] 把调度器的状态翻译成这五样：

张量	形状 / 类型	作用
input_ids	[总 token 数] int64	把所有序列的 token 首尾相接压平
positions	[总 token 数] int64	每个 token 的绝对位置，喂给 RoPE
cu_seqlens_q / k	[bs+1] int32	前缀和，标记每条序列在压平数组里的区间
slot_mapping	[总 token 数] int32	本课主角：每个 token 的 KV 写到哪个物理槽
block_tables	[bs, 最大块数] int32	每条序列的块表，padding 用 -1

cu_seqlens 是变长 batch 的通用表示法：序列 i 的 token 落在 [cu_seqlens[i], cu_seqlens[i+1])。这样一次前向就能处理长度各异的序列，不需要 padding。

2slot_mapping 是怎么算出来的

model_runner.py · L151–161start_block = start // self.block_size
end_block = (end + self.block_size - 1) // self.block_size
for i in range(start_block, end_block):
    slot_start = seq.block_table[i] * self.block_size
    if i == start_block:
        slot_start += start % self.block_size      # 第一块可能从中间开始
    if i != end_block - 1:
        slot_end = seq.block_table[i] * self.block_size + self.block_size
    else:
        slot_end = seq.block_table[i] * self.block_size + end - i * self.block_size
    slot_mapping.extend(range(slot_start, slot_end))

KV cache 在显存里是一个扁平的大张量：

self.kv_cache = torch.empty(2, num_layers, num_blocks, block_size, num_kv_heads, head_dim)

所以「第 b 个块的第 j 格」对应的槽号就是 b * block_size + j。slot_mapping 就是把每个待写入的 token 映射到这个一维槽号——它是一条序列的块表在时间维度上的展开。

3写 KV：一个 30 行的 Triton kernel

layers/attention.py · L10–30@triton.jit
def store_kvcache_kernel(key_ptr, key_stride, value_ptr, value_stride,
                         k_cache_ptr, v_cache_ptr, slot_mapping_ptr, D: tl.constexpr):
    idx = tl.program_id(0)                 # 第 idx 个 token
    slot = tl.load(slot_mapping_ptr + idx)
    if slot == -1: return                  # padding 槽，跳过
    key_offsets = idx * key_stride + tl.arange(0, D)
    key = tl.load(key_ptr + key_offsets)
    value = tl.load(value_ptr + value_offsets)
    cache_offsets = slot * D + tl.arange(0, D)
    tl.store(k_cache_ptr + cache_offsets, key)
    tl.store(v_cache_ptr + cache_offsets, value)

一个 program 处理一个 token，每个 program 内部并行搬运 D = num_heads * head_dim 个数。为什么值得写 Triton？因为这是纯搬运：PyTorch 的索引赋值 k_cache[slots] = k 会触发一次 gather/scatter 的算子调度，而这里合并成一次 kernel，且没有中间张量。

4两次 FlashAttention，两种形状

layers/attention.py · L59–75def forward(self, q, k, v):
    context = get_context()
    if k_cache.numel() and v_cache.numel():
        store_kvcache(k, v, k_cache, v_cache, context.slot_mapping)   # 先写
    if context.is_prefill:
        if context.block_tables is not None:      # 前缀缓存命中
            k, v = k_cache, v_cache               # 读整块 paged cache
        o = flash_attn_varlen_func(q, k, v,
                max_seqlen_q=context.max_seqlen_q, cu_seqlens_q=context.cu_seqlens_q,
                max_seqlen_k=context.max_seqlen_k, cu_seqlens_k=context.cu_seqlens_k,
                softmax_scale=self.scale, causal=True, block_table=context.block_tables)
    else:
        o = flash_attn_with_kvcache(q.unsqueeze(1), k_cache, v_cache,
                cache_seqlens=context.context_lens, block_table=context.block_tables,
                softmax_scale=self.scale, causal=True)

三个关键差异：

1. prefill 用 varlen：一次算 N 个 query，用 cu_seqlens 描述变长 batch，Q 和 K 是同一次前向里算出来的。
2. decode 用 with_kvcache：Q 只有 1 个 token（所以要 unsqueeze(1) 变成 [bs, 1, heads, dim]），K/V 全部从 cache 读。
3. 前缀缓存命中时 prefill 也读 cache：k, v = k_cache, v_cache 把「这次算的 K/V」换成「整块 cache」，再靠 block_table 让 kernel 只读有效部分——那些已经算过的前缀就不必重算了。

5为什么 prefill 只算最后一个 token 的 logits

layers/embed_head.py · L56–66def forward(self, x):
    context = get_context()
    if context.is_prefill:
        last_indices = context.cu_seqlens_q[1:] - 1     # 每条序列的最后一个位置
        x = x[last_indices].contiguous()
    logits = F.linear(x, self.weight)
    ...

因果注意力下，位置 i 的隐藏状态只依赖 [0, i]。生成第一个新 token 只需要整段 prompt 的最后一个位置的输出。如果对 4096 个位置都算 vocab_size 维的 logits，等于把 15 万 × 4096 的矩阵乘白做一遍——在 0.6B 模型上这只是浪费，在 100B 模型上就是灾难。

这也是为什么 ParallelLMHead 在 prefill 时先 gather 再取：每个 rank 只持有词表的一部分，最后拼起来得到完整 logits。

6采样：一次「指数竞赛」

layers/sampler.py · L7–12@torch.compile
def forward(self, logits, temperatures):
    logits = logits.float().div_(temperatures.unsqueeze(dim=1))
    probs = torch.softmax(logits, dim=-1)
    sample_tokens = probs.div_(torch.empty_like(probs).exponential_(1).clamp_min_(1e-10)).argmax(dim=-1)

看起来像技巧，其实是一个精确的数学等价：

若 E_i ~ Exp(1) 独立同分布，则 argmax_i (p_i / E_i) 的分布恰好是 Multinomial(p)。

这叫 exponential race（也叫 Gumbel-max 的变体）：把每个候选的「等待时间」定义为 E_i / p_i，谁先到就是谁。p_i 越大，等待时间越短，被选中的概率越大——而且严格等于 p_i。

为什么不用 torch.multinomial？两个原因：multinomial 在 CUDA 上不保证与 torch.compile 兼容的确定性实现；而这种纯张量运算可以整体编译进一张图，省掉一次 kernel launch 和一个随机数生成器调用。

7练习 A · 手算 slot_mapping

block_size = 4。序列块表 [7, 9]，已缓存 0 个 token，本轮要算 6 个 token（整段 prompt）。写出 slot_mapping。

展开：答案

start = 0, end = 6。第 0 块（物理块 7）覆盖 token 0–3 → 槽 7*4+0 … 7*4+3 = 28,29,30,31；第 1 块（物理块 9）覆盖 token 4–5 → 槽 9*4+0, 9*4+1 = 36,37。

结果：[28,29,30,31,36,37]。注意物理上不连续——这正是分页的意义。

如果这轮是 chunked prefill 的第二块（已缓存 4 个 token，本轮算 2 个），则 start=4, end=6，start_block=1，start % block_size = 0，结果同样是 [36,37]。

8练习 B · 闭卷测验

9练习 C · 代码寻宝

1. prepare_decode 的 slot_mapping 为什么只有一个元素？

decode 每序列只写 1 个 token 的 KV，所以每条序列对应一个槽：block_table[-1] * block_size + last_block_num_tokens - 1（model_runner L181）。last_block_num_tokens 是最后一块里已有的 token 数，减 1 得到这个 token 的 0-based 槽内偏移。

2. 如果 block_size 和 FlashAttention 的 block 粒度不一致，会发生什么？

FlashAttention 的 block_table 参数会按它自己的粒度去索引 KV cache。如果两者对不上，读到的是错位的 K/V——不会报错，只会默默输出胡话。这就是 config.py 里 assert self.kvcache_block_size % 256 == 0 存在的原因。

3. 为什么采样前要 logits.float()？

logits 可能是 bf16，直接除以 temperature 会损失精度，softmax 在低精度下更容易溢出或产生 0 概率。转成 fp32 再算，代价很小（只在最后一步做）。

10面试复述卡

「PagedAttention 的 kernel 怎么知道 KV 在哪？」

展开：30 秒版本

「推理时 KV 存在一个大张量里，按块切分、物理上不连续。注意力 kernel 拿到的是 block_table——每条序列的逻辑块到物理块映射，加上这次的 slot_mapping——每个待写 token 到物理槽的一维索引。写入时 Triton kernel 按 slot 逐 token 搬运；读取时 FlashAttention 用 block_table 做间接寻址，所以它能在分页的显存上直接算注意力，不需要先拼成连续内存。」