海光 DCU 上的长上下文 Attention：从 page784 到 page832 的 KV Cache 对齐与访存优化

摘要

在大模型推理中，Attention 的公式通常不是最难的部分，真正决定性能的往往是数据如何落在显存里，以及硬件一次能够以什么粒度把这些数据取出来。

本文以 Qwen3.5 的官方 prefix-off 长上下文路径为例，讨论一个看似很小、实际上贯穿 vLLM 配置、KV Cache、paged attention 和 DCU kernel dispatch 的问题：该路径的默认物理 KV Cache page 是 784 tokens，而海光 DCU gfx936 上的 vendor paged-attention 路径要求 page 长度满足 64-token 对齐。最终我们没有把 page size 当作 scheduler 参数硬编码，而是沿 vLLM 原有的混合 Attention/Mamba 页面计算逻辑，把它自动推导到最小兼容值 832。

这个改动本身只有一个数字级别的差异，但它改变了 kernel 能否被调用的前提。随后，q=4096 的 chunked-prefill 进入 vendor FlashAttention 路径，q=1 decode 继续使用 page-contiguous Triton 路径，避免了把两个完全不同的工作负载混成一个 kernel。

本文重点放在 DCU 的执行粒度、KV Cache 的物理布局、block table 的地址计算和端到端验证，不展开比赛提交、仓库地址或内部环境细节。

---

1. 为什么要从“页”开始看 Attention

对一个长度为 T 的上下文，Attention 需要访问历史 K/V。vLLM 不会为每个请求分配一整块连续的 K/V，而是把缓存切成固定大小的物理页，并通过 block table 把逻辑页映射到物理页。

可以把一个 token 的地址简化成：

logical_page = token_index // BLOCK_SIZE
page_offset  = token_index %  BLOCK_SIZE
physical_page = block_table[logical_page]

然后再结合 K/V cache 的 stride、KV head 和 head dimension 计算最终地址。

这种分页设计解决了显存碎片和动态请求管理问题，但也带来了一个硬件事实：BLOCK_SIZE 不只是调度器里的整数，它同时决定了 kernel 内部的页内循环、向量化访问和边界 mask。

如果硬件 kernel 希望以 64 tokens 为一个规整的访问单元，那么一个不能被 64 整除的物理页，通常会在每个 page 的尾部留下不规整的 tile 或额外的边界处理。对于短请求，这可能只是少量浪费；对于长上下文，页表解析和 K/V 读取会被重复很多次。

2. gfx936 的关键约束：wave64 与 64-token 访问粒度

海光 DCU 的 ROCm 兼容执行环境在 gfx936 路径上采用 wave64 执行粒度。这里不应把 wave64 简化成“线程数变成 64”这么一句话；它会影响：

• 一个 wave 内线程如何协同加载向量；
• K/V 的 token 维度如何分配给 lane；
• 连续地址能否形成规整的 memory transaction；
• kernel 对页尾部 mask 和跨页访问的处理方式。

在本实验使用的 vendor paged-attention 接口中，物理 page 长度需要满足 64-token 对齐。直接由源码和运行时日志确认的是这个接口约束；它与 wave64 的访问粒度相匹配，但本文不把 64 对齐唯一归因于 wave64。默认的 784 可以拆成：

784 = 12 × 64 + 16

而 832 可以拆成：

832 = 13 × 64

这意味着 page832 内部可以由 13 个完整的 64-token 访问组组成，不需要在每个物理页末尾留下一个 16-token 的非规整尾块。这里的收益不是来自减少 Attention 数学计算，而是来自让数据布局与硬件执行粒度一致。

需要强调的是，page832 不是“越大越快”的经验参数。1024 也满足 64 对齐，但它比 784 多出约 30.6% 的 token 容量；832 只增加约 6.1%，是满足该约束的最小兼容页，更适合显存紧张的 27B 混合模型。

3. Qwen3.5 的混合缓存为什么让这个问题更复杂

Qwen3.5 同时包含 full attention 和 GDN/Mamba 类状态路径。vLLM 需要让 Attention page 和 Mamba state page 在混合模型中保持兼容，否则同一个请求在不同层之间会遇到不同的缓存粒度和 padding 规则。

在 prefix caching 关闭的路径上，页面选择可以抽象为（省略了现有配置中的最小值比较和后续 padding 分支）：

attn_page_size_1_token = FullAttentionSpec(...).page_size_bytes
required_page = ceil(mamba_page_bytes /
                     (alignment_tokens * attn_page_size_1_token))
                 * alignment_tokens

对于这个 Qwen3.5/DCU 形状，backend 适配把 alignment_tokens 从通用的 16 提升到 64；沿原有的 verify_and_update_config 逻辑计算后，默认页从 784 向上取整到 832。

这和直接在 scheduler 中写一句 block_size = 832 是两件事：

1. 计算仍然由现有的模型配置和缓存规格数据流完成；
2. 页面大小仍然作为 KV Cache 的物理布局结果传递给后续组件；
3. scheduler 的 token budget、chunk 数量和请求调度语义不变；
4. 不满足 Qwen3.5、ROCm、head size、KV head 数量和 prefix 状态条件时，保持原有逻辑。

这类改动的核心是“让 backend 的物理约束进入既有布局计算”，而不是用一个新参数绕过 vLLM 的缓存模型。

4. Qwen3.5 的真实页形状

本文实验中的 full-attention 形状具有很强的确定性：

query heads: 24
KV heads:     4
GQA ratio:    6
head size:    256
dtype:        BF16

vLLM 的完整 K/V cache tensor 还包含一个 K/V 维度，形状可以写成：

[num_physical_pages, 2, BLOCK_SIZE, 4, 256]

在 unbind(1) 之后，单独的 K 或 V cache 形状才是：

[num_physical_pages, BLOCK_SIZE, 4, 256]

K 和 V 分开存储时，单个 page 的数据量近似为：

BLOCK_SIZE × 4 × 256 × sizeof(BF16)

K、V 两者合计则再乘以 2。换成 page832 后，单页会多放一些 token，但页内每个 64-token group 都是完整的。这个增加的容量成本必须和 kernel 访存效率一起评估，不能只看单页耗时。

5. 从 page table 到 coalesced load

通用 Triton paged-attention kernel 的核心循环可以抽象成下面的形式：

for tile in key_tiles:
    token = tile * TILE_SIZE + lane_offset
    logical_page = token // BLOCK_SIZE
    page_offset = token % BLOCK_SIZE
    physical_page = block_table[logical_page]

    K = load(key_cache[physical_page, page_offset, kv_head, dim])
    V = load(value_cache[physical_page, page_offset, kv_head, dim])
    score = dot(Q, K) * scale
    update_online_softmax(score, V)

在 page784 上，token % BLOCK_SIZE 的结果在页尾会留下一个无法与 64-token 访问组对齐的区间。通用 kernel 可以通过 mask 保证正确性，但 mask 并不会让不规则地址变成规则地址；这里描述的是访存布局的机制推断，不是额外的硬件 trace 计数。

在 page832 上，页内 token 区间可以自然地分成完整的 64-token 组。vendor kernel 仍然需要处理最后一个请求的真实长度，但不需要在每个物理页内部面对一个固定的 16-token 尾块。

这里的优化点不是“去掉 causal mask”。因果约束仍然完整保留；改变的是 page 内循环与物理地址布局，使向量加载更适合 DCU 的执行粒度。

6. 为什么 vendor kernel 只接管严格的单序列 prefill 形状

长上下文服务不是一种单一 workload。q=1 decode、q=4096 chunked-prefill 和最后不足 4096 的尾块，在并行度、访存方向和图捕获方式上都不同。

因此最终 Tail-Q 版本的 vendor dispatch 使用严格条件，而不是全局替换：

1 < num_actual_tokens <= 4096
max_seqlen_q == num_actual_tokens
single sequence
decoder attention
BF16 Q/K/V/output
Qwen3.5: 24 Q heads, 4 KV heads, head size 256
page size = 832
GQA ratio = 6
causal = true
no sliding window
no ALiBi / sinks / softcap
no multimodal prefix range

满足条件时，调用已有的 vendor varlen_fwd_unified 接口：

q, k, v
cu_seqlens_q
seqused_k
block_table
max_seqlen_q / max_seqlen_k
causal = true

不满足条件就回到原有 unified Triton attention。这个设计有两个工程意义：

• vendor kernel 只面对它真正擅长的固定形状；
• 新路径即使在某个边界上不适合，也不会污染通用 attention。

最初的候选只覆盖 q=4096；最终 Tail-Q 版本覆盖 1 < q <= 4096 且 max_seqlen_q == q，明确排除 q=1 decode。这样最后一个不满 4096 的 prefill chunk 也能进入同一类完整 causal attention，而不会改变 scheduler 的切块逻辑。

7. decode 为什么要单独维护 page-contiguous 路径

prefill 的 q 很大，适合让 vendor FlashAttention 处理一整个 query chunk；decode 的 q 通常为 1，主要问题变成从大量历史 page 中读取 K/V，并以低 launch 开销完成在线 softmax。

page-contiguous decode kernel 的循环更接近：

for page_idx in pages_for_sequence:
    physical_page = block_table[page_idx]
    for local_tile in tiles_inside_page:
        page_offset = local_tile * TILE_SIZE + offsets
        K = load(K_cache[physical_page, page_offset, ...])
        V = load(V_cache[physical_page, page_offset, ...])
        update_online_softmax(Q, K, V)

这个路径仍然解析 block table，没有跳过任何逻辑 page，也没有改变因果范围。它只是针对 q=1、GQA6、head size 256 的访问几何减少通用调度开销，并同时支持 page784 和 page832。

在 page832 上，独立 CUDAGraph microbench 的 decode 加速约为：

KV 长度	通用 page832	page-contiguous	加速比
8192	0.115989 ms	0.080010 ms	1.45x
16384	0.154312 ms	0.122411 ms	1.26x
32768	0.405256 ms	0.278832 ms	1.45x

这组结果是 bitwise exact 的；但独立 kernel 的加速比不能直接当成端到端服务加速比，因为 decode 还会被 Linear、采样、图 replay 和 Python/调度边界包围。

8. 实测结果：kernel 加速如何传到服务层

在同一容器、同一模型、同一原生扩展和同一数据集下，baseline 使用官方 prefix-off page784 路径，candidate 使用 page832、vendor prefill 和 page-contiguous decode。三档每侧均完成 10/10、failed=0。运行时日志还确认了四个关键事实：

baseline: prefix caching = False, block size = 784
candidate: prefix caching = False, block size = 832
vendor q=4096 prefill marker: hit
page832 decode marker: hit

本地三档 A/B 的相对变化为：

上下文档位	吞吐变化
4K-8K	+0.62%
8K-16K	+2.67%
16K-32K	+13.78%

q=4096 vendor prefill standalone 相对 page784 Triton 路径约为 2.75x-2.92x；完整服务的提升明显小于这个数字。长档收益更明显，是因为该测试 workload 中 q=4096 chunk 和长序列 paged-KV 访问占比更高；这是由 workload/profile 推出的解释，不应当理解为单独 trace 已经证明了某一种访存事务的精确占比。

这也说明为什么不能只拿一个 kernel 表格作为结论：在 page832、prefix-off、all-Tail-Q 的官方 8K profile 中，q=4096 prefill 的 Attention 总桶从 681.071 ms 降到 238.718 ms；这不是单个 vendor kernel 的 standalone 耗时，端到端服务还要支付 GEMM、GDN、decode 和调度成本。

这里的数字口径需要分开：2.75x-2.92x 是 standalone kernel 对比，+0.62%/+2.67%/+13.78% 是同容器端到端 A/B，681.071/238.718 ms 是 profile 的 Attention 总桶。它们不能互相替代，也不能把本地 A/B 倍率直接当作另一台机器的官方绝对吞吐。

9. 数值正确性：快不等于逐位相同

vendor BF16 attention 与 Triton attention 计算的是同一个完整的因果 Attention，但并行归约顺序可能不同。因此：

• vendor standalone 输出有限且重复运行稳定；
• page832 Triton 与 page784 Triton 可以逐位一致；
• vendor 与 Triton 的相对 L2 误差约为 0.0058-0.0064；
• 生成文本不保证逐请求 bitwise 相同。

这不是可以忽略的细节。BF16 reduction 的微小差异可能在长序列生成中被放大，所以必须同时做：

1. kernel 输出 finite 检查；
2. 重复运行一致性检查；
3. 相对误差与参考 kernel 对比；
4. 四类任务级 accuracy smoke；
5. 逐请求记录生成文本和 token 数。

在这条路线中，没有跳过 Q/K/V，没有缩短 causal 范围，没有做稀疏注意力，也没有复用跨请求 KV；变化是物理页面布局和等价 kernel 选择。

10. 一个重要的失败案例：page832 曾经“看起来无效”

这条路线最有价值的发现并不是第一次跑出 2.9x，而是后来发现旧结论的测试路径并不完整。

旧实验只提交了 attention dispatch 改动，却没有同时提交 page-layout 对齐逻辑。官方 loader 关闭 prefix caching，实际运行时仍是 page784；由于 784 不满足 vendor kernel 的 64-token 页要求，vendor 路径根本没有命中。

另一方面，旧 fast-start 脚本打开了 prefix caching，Mamba align 又把页面抬到了 1024。本地看到的结果因此是 page1024，不是官方路径的 page832。

于是出现了一个典型矛盾：

standalone page832 vendor：很快
旧官方提交：几乎没收益

正确结论不是“standalone 不可信”，也不是“官方一定不适合”，而是要先检查：

源码是否真的包含布局改动？
运行时 block size 是多少？
vendor marker 是否出现？
prefix caching 是否改变了页面？

只有这些条件对齐，性能数字才属于同一条路径。

11. 复现与验收清单

如果要在另一台 DCU 或另一个 vLLM 版本复现，建议按以下顺序：

11.1 先确认硬件和软件约束

确认 gfx936/对应 DCU 目标架构
确认 wave64 相关 kernel 约束
确认 vendor FlashAttention 的 page 对齐要求
确认 BF16、GQA6、head size 256 的真实模型形状

11.2 再确认页面计算

打印 prefix caching 状态
打印最终 cache block size
打印 attention page bytes 与 Mamba page bytes
确认 page832 来自对齐计算，而不是脚本里的隐藏参数

11.3 最后确认 kernel 命中

打印 vendor prefill marker
打印 page-contiguous decode marker
检查不满足条件时是否回退到 unified Triton
检查 block table 和 K/V stride 与 page size 一致

11.4 端到端验收

同一容器、同一模型、同一数据集
充分 warmup，避免首次编译污染
三档吞吐，failed = 0
记录 TTFT、TPOT、输出 token 数和文本
accuracy smoke

12. 结论：真正的优化发生在“硬件约束”和“框架语义”的交界处

page832 不是一个孤立的调参结果，它同时涉及：

• gfx936 的 wave64 执行粒度；
• vendor attention 对 64-token 页的要求；
• vLLM 混合 Attention/Mamba 的页面兼容计算；
• paged KV 的 block table 和 stride 地址公式；
• q=4096 prefill 与 q=1 decode 的不同访问几何；
• BF16 归约顺序带来的数值验证问题。

只改 kernel 而不改页面布局，kernel 可能根本不会命中；只改页面大小而不改 decode 访问，长上下文也可能收益有限；只看 microbench 而不看输出和 official loader，结论还可能被错误路径污染。

对 DCU 推理优化而言，最值得复用的方法是：

先把硬件的访问粒度翻译成数据布局约束，再把布局约束接入框架已有的语义路径，最后用运行时 marker 和端到端结果证明 kernel 确实命中。

这比“换一个 BLOCK_SIZE 看看”更接近真正的 kernel 工程，也更能解释为什么一个 48-token 的页面差异，会在每层 attention 和每个长请求的多个物理页访问中反复出现并影响整体服务性能。

---

本文讨论的是公开可解释的硬件、布局和 kernel 方法；实现描述不包含内部仓库地址、账号信息、容器凭据或私有基础设施细节。性能数字来自固定形状的同容器实验，实际结果会随 DCU 型号、驱动、DTK、编译缓存和运行时负载变化。