优化 MetaX 平台 HD64 GQA Paged Attention Decode 的 K/V 数据复用

[Qzong2020]

概述

本 PR 优化 MetaX 平台上 head_dim=64、GQA Ratio 为 8 的 Paged Attention Decode 路径。

在满足条件的场景下，将共享同一 K/V Head 的 Query Head 按每组 4 个进行分组，通过 CTA 内的线程协作加载 K/V tile，并在共享内存中复用，从而减少 Query Head 之间重复的 K/V 数据读取。

优化路径保留原有逐 token 的计算方式及 Split-KV 合并逻辑。不满足条件时继续使用原有实现，不修改公共 API 和默认 split 选择策略。

优化动机

原有实现中，不同 Query Head 由独立 CTA 处理。在 GQA 场景下，多个 Query Head 共享同一个 K/V Head，这会导致相同 K/V 数据被多个 CTA 重复加载。

本 PR 通过 Query Head 分组和共享内存复用，减少冗余 K/V 读取，改善目标 Decode workload 的执行性能。

主要修改

• 新增 MetaX HD64 GQA grouped Split-KV kernel。
• 每个 K/V Head 对应 8 个 Query Head；优化 kernel 以 4 个 Query Head 为一组，每个 CTA 使用 128 个线程处理一组 Query Head，并共享对应的 K/V tile。
• 对每个 sequence 和 KV split，同一个 K/V Head 对应两个 Query Head 分组；K/V tile 在各 CTA 内复用，不跨 CTA 共享。
• 使用 8-token K/V tile，通过 16 字节向量化加载写入共享内存。
• 保留原有 softmax 累积方式及 Split-KV combine kernel。
• 增加优化路径的布局、对齐和参数检查。
• 扩展已有 Paged Attention 测试，覆盖优化路径和部分 fallback 场景。

涉及文件：

• src/infiniop/ops/paged_attention/metax/paged_attention_hd64.maca
• src/infiniop/ops/paged_attention/metax/paged_attention_metax.maca
• test/infiniop/paged_attention.py

优化路径适用条件

优化路径仅在 Split-KV 路径内、满足以下条件时启用：

• Head Dimension 和 Value Dimension 均为 64。
• 数据类型为 FP16 或 BF16。
• Query Head 数量为 K/V Head 数量的 8 倍。
• Split 数量为 1、2、4 或 8。
• Page Block Size 为 16、32、64 或 256。
• 不使用 ALiBi。
• Q、输出、block table 和 K/V cache 的布局满足检查条件。
• K/V 地址及相关 stride 满足向量化加载的对齐要求。

不满足条件时，继续执行原有 kernel。

正确性验证

验证项	结果
原实现专项正确性验证	97/97 PASS
优化实现专项正确性验证	97/97 PASS
Paged Attention 框架测试，包含完整 API 调用路径	282/282 PASS
已测试的 fallback 场景：非 GQA Ratio 8、ALiBi、带 padding 的 KV 布局	PASS
优化版本专项 dispatch trace 检查	4/4 PASS
git diff --check	PASS

其中，282 项为 test/infiniop/paged_attention.py 中的测试用例，不代表全项目测试数量。

框架测试与专项验证共同覆盖：

• FP16、BF16 数据类型。
• I32、I64、U32 索引。
• Split 数量 1、2、4、8。
• 不同分页大小、空 shard、非整除 shard 和跨页边界。
• 与 CPU FP32 参考实现比较。
• 非 GQA Ratio 8、ALiBi 和带 padding 的 K/V 布局回退场景。

每个实现的 97 项专项验证包含：

• 24 项混合长度测试。
• 24 项带 padding 的 KV 布局测试。
• 24 项非 GQA Ratio 8 测试。
• 24 项 ALiBi 测试。
• 1 项 Graph 测试，执行 100 次重放并检查结果。

本 PR 未修改原有框架测试的误差阈值。独立专项验证脚本使用明确设置的容差，分别检查 CPU 参考结果、原实现与优化实现的输出及中间结果。

上述测试未穷举所有不满足优化条件的参数、地址对齐和布局组合。

模型集成验证

使用 TinyLlama BF16，在固定输入及后续 token 序列的条件下进行了模型集成验证：

• Context Length：1536。
• Batch：1、8。
• Split 数量：8。
• 关闭模型图编译。
• 使用一个 prompt 配置，每组比较 3 次前向调用保存的 logits。

Batch	Logits 张量形状	逐元素一致	最大绝对误差
1	[3, 1, 32000]	是	0
8	[3, 8, 32000]	是	0

保存的 logits 均为有限值。Profiler trace 同时确认：原实现运行触发 legacy kernel，优化实现运行触发 grouped kernel。

该结果限于上述模型和测试配置，不代表对任意模型、输入或长序列自由生成的一致性保证。

性能测试

测试环境：

项目	配置
GPU	MetaX C500
GPU 资源	25% 算力切片，16000 MiB 显存配额
MACA	3.5.3.20
驱动	3.8.30
PyTorch	2.8.0+metax3.5.3.9
性能表数据类型	BF16
Query / K/V Heads	32 / 4
Head Dimension	64
Page Block Size	256
Context Length	1536
Split 数量	8

以下为算子微基准结果，计时包含 Split-KV 主 kernel 和 combine kernel，不代表完整 API 调用开销或模型端到端加速比。

执行模式说明：

• Graph：在 MetaX PyTorch 环境中，通过 torch.cuda.CUDAGraph 接口捕获并重放 kernel 调用。
• Eager：不使用图捕获，逐次 launch kernel。

每个实现预热 50 次，随后交替测量 10 组，每组执行 100 次调用。通过 GPU event 记录整组耗时并计算平均单次延迟，表中报告 10 组平均单次延迟的中位数。计时运行未启用 profiler。

Batch	执行模式	原实现	优化后	加速比
1	Graph	177.41 μs	152.52 μs	1.16×
1	Eager	198.73 μs	160.34 μs	1.24×
8	Graph	426.47 μs	248.43 μs	1.72×
8	Eager	437.13 μs	257.43 μs	1.70×

以上绝对延迟来自 C500 25% 算力切片环境，主要用于比较同一环境下新旧 kernel 的相对性能，不代表完整 C500 的绝对性能。

上述结果在相同 Split 数量下比较原实现与优化实现，不包含调整 Split 数量带来的收益。FP16 已纳入正确性验证，本表仅列出最终记录中的 BF16 性能数据。

兼容性

• 不修改 Paged Attention 公共 API 和接口语义。
• 不修改默认 Split-KV 选择策略。
• 修改仅限 MetaX backend 实现，不改动其他 backend 的实现。
• 保留原有实现作为 fallback。
• 不引入 InfiniLM 源码修改或额外运行时依赖。

其他说明

框架回归测试包含在本 PR 中；97 项专项验证、性能微基准和模型集成验证使用独立实验脚本，结果来自留档记录，这些脚本未包含在本次提交中。

本 PR 仅包含两份 MetaX 算子实现文件和一份已有测试文件的修改，不包含实验日志、模型文件或构建产物。