[RFC]: NixlConnector: 将标准 KV 页面转换为解码端的 ROCm AITER 随机排序页面

理由。 TL;DR

• 在 ROCm 上，AITER 的随机 KV 页在高批量下提供了最快的解码速度。预填并没有从这些页中受益，而不同的后端则是预填的最快速度。- P/D 分解可以使每个方向使用最佳格式，但 NixlConnector 要求两个方向都使用相同的格式。- 我建议让 ROCm 解码器接受标准页，并在传输后将其转换为随机页。- 我现在需要的决策是：- 连接器是否应该比较页面格式标签并调用后端拥有的转换器；- 预填和解码是否可以在其页面格式进行协商后使用不同的后端；- 解码器的步骤路径上的转换是否可接受。- 决定是否合并转换的测量结果将在本次讨论之后提供（请参阅 go/no-go 计划）。- 这取决于 #59110（bug：NixlConnector 无法看到 KV 页面格式）。该问题的两个修复是：- 在 NixlAgentMetadata 中添加区域级页面格式标签；- 在 NIXL 描述中将 K 和 V 位置分开。背景。 使用 VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1 (#29887)，ROCM_AITER_FA 将每个块的 K 写为 [H, D/x, N, x]，其 V 为 [H, N/x, D, x]，其中 x = 16 // element_size。AITER 的 asm 和 gluon 解码内核读取此格式。其他后端写入标准页面，其中每个头的标记都保持 K 和 V 并排（[H, N, 2*D]）。这些是 CUDA FlashAttention、未随机的 ROCM_AITER_FA 和 ROCM_AITER_UNIFIED_ATTN。这两个内核家族都不读取对方的页面。解码：高批量下随机页面胜出。 我在一台 MI300X 上测试了完整的解码步骤（Qwen3-8B，BF16，TP 1，16 标记页面）。随机臂使用了非 gluon asm/ll4mi 路径，且 KV 缓存布局为 LHBNC。每个单元是随机加速与 ROCM_AITER_FA 相比，这是 ROCm 默认选择的 KV 连接器。括号中的数字表示在该形状下与最快的未随机后端相比的加速率（TRITON_ATTN、ROCM_AITER_UNIFIED_ATTN 或 ROCM_AITER_FA）：

上下文 \ 批量	1	8	32	64	128
1k	+2.6% (+0.2%)	+4.3% (+0.8%)	+5.7% (+5.7%)	+6.0% (+6.0%)	+9.2% (+3.6%)
4k	+3.3% (−0.5%)	+2.5% (+2.5%)	+5.3% (+3.2%)	+5.5% (+2.3%)	+7.7% (−1.5%)
16k	+2.4% (+1.2%)	+1.4% (−0.4%)	+10.2% (+5.1%)	–	–
32k	+3.7% (+1.3%)	+1.2% (−2.0%)	–	–	–

一个短横线标记了超出硬件 95% KV 容量限制的形状。每个配置运行一次，因此将 2% 以下的差距视为平局。- 与 ROCM_AITER_FA 相比，随着批量的增加，AITER 的 asm 解码内核接管。对于 MI300X 上的 Qwen3-8B，这发生在批量超过 19 时，即批量 × 查询头超过 CU 个数的两倍。- 与最快的未随机后端相比，随机页面在 7 个形状中的 6 个形状中比最快的未随机后端快 2.3–6.0%。在批量 128 时，差距缩小到 +3.6% 在 1k 和 −1.5% 在 4k。- 在所有 15 个形状中，与最快的未随机后端相比的中位数加速率为