08 · Attention：把新 K/V 写入正确的缓存槽

Attention 层承担两件事

1. 计算当前 query 对上下文 key/value 的注意力输出。
2. 把当前新 token 对应的 K/V 写入持久 KV Cache，供以后 Decode 使用。

nano-vLLM 的 layers/attention.py 使用 FlashAttention，并定义一个很小的 Triton kernel：store_kvcache_kernel。

slot_mapping 是什么

控制面已经通过 Sequence 的 block table 决定每个 token 应该落在哪个物理槽位。ModelRunner 把这个结果整理成 slot_mapping：每个输入 token 对应一个整数 slot。

Triton kernel 对每个 token：

slot = slot_mapping[token_index]
cache[slot] = current_key_or_value

真实实现还需要处理 head 维度、stride 和无效 slot，但核心就是一次 scatter 写入。

Prefill Attention

Prefill 中多条请求长度不同。ModelRunner 把 token 拼接起来，并提供累计长度，让 flash_attn_varlen_func 知道每条序列的边界。新的 K/V 在计算前后写入缓存。

Decode Attention

Decode 只输入每条序列最新 token。Attention 通过 block_tables 和 context_lens 读取离散 KV block，调用带 KV Cache 的 FlashAttention 路径。

为什么这里适合 Triton

缓存写入操作很简单，但如果用 Python 循环逐 token 复制，会产生严重开销。Triton 允许用短小 kernel 并行处理 token 与 head 维度，代码仍比手写 CUDA C++ 更容易阅读。

阅读时重点检查形状

• key: [N, num_kv_heads, head_dim]
• value: 同形状
• slot_mapping: [N]
• 展平后的缓存每个 slot 容纳 num_kv_heads * head_dim 个元素

形状断言不是多余代码，它们保护了最危险的显存写入路径。

HTML INTERACTIVE LAB08 · Attention：把新 K/V 写入正确的缓存槽

单独打开 ↗

课后习题等待完成

`slot_mapping` 的主要作用是什么？

动手任务

假设 block_size=4、block_table=[3,8]，计算逻辑 token 位置 5 的物理 slot，并在互动实验中写入对应 slot。

下一节：Sampling