[KernelSwift 算子创新大赛][T1, T2, T3, T4]tuolajilatuo

[xu-zhengzhong]

KernelSwift 算子创新大赛参赛信息

• 大赛名称：KernelSwift 算子创新大赛
• 参赛队伍：拖拉机拉拖
• 所选赛题：T1, T2, T3, T4
• 参赛成员：徐政中、赵世杰

PR 概述

本 PR 面向 LLM 推理场景，基于 NineToothed DSL 新增并优化以下四个算子：

1. block_scaled_fp8_mm
2. mxfp4_w4a16_grouped_mm
3. rms_norm_gated
4. fused_mla_rope_cache_write

本 PR 的改动均属于 A 部分，不包含 B 部分的编译器或后端修改，也未修改 NineToothed、Triton、CoreX 或 HIP 后端源码。

主要改动

Block-scaled FP8 MM

• 实现支持 BlockWise1x128 和 BlockWise128x128 scale 布局的 FP8 GEMM。
• 支持 FP16、BF16 和 FP32 输出，以及可选 bias 融合。
• 针对 CoreX、HIP 和具备原生 FP8 Tensor Core 的 CUDA 平台选择不同计算路径。
• 在 kernel tile 内完成 dot 和 scale 应用，不物化完整反量化矩阵。

MXFP4 W4A16 Grouped MM

• 支持 uniform 和 routed 两种 grouped GEMM 模式。
• 在 kernel 内融合 E2M1 权重解码、E8M0 scale 应用和矩阵乘法。
• 支持零 token 专家及不均匀 routed MoE workload。
• 不物化完整的 [G, K, N] 反量化权重。

Gated RMSNorm

• 融合 RMSNorm、门控和激活计算。
• 支持 norm_before_gate=True/False。
• 支持 SiLU 等激活方式、可选 weight 以及 grouped normalization。
• 使用 FP32 完成归约和归一化计算，并转换回输入 dtype。

Fused MLA RoPE Cache Write

• 融合 MLA RoPE 计算与 paged KV cache 写入。
• 支持 CUDA Graph padding、slot=-1 跳过语义和非连续输入。
• 直接写入压缩的 [kv_c | RoPE(k_pe)] cache entry，避免按 head 展开。
• 提供独立 PyTorch reference 用于正确性验证。

工程与评测

• 新增四个算子的定向测试和 benchmark。

• 新增统一构建、测试及评测脚本：

  ./scripts/build_and_evaluate.sh all

• 新增构建评测说明和技术报告。

• 在 ntops.kernels 与 ntops.torch 中导出新增公开 API。

正确性验证

技术报告记录的验证结果如下：

• 天数智芯 MR-V100：
  • 四个定向测试文件共 78 passed, 2 skipped。
  • 两个 skip 是由于对应 PyTorch 版本缺少原生 MXFP4 dtype；raw uint8 路径已实际执行并通过。
• 海光 BW/gfx936：
  • Block-scaled FP8、MXFP4、Gated RMSNorm 和 Fused MLA RoPE Cache Write 的代表场景正确性检查均已通过。

测试命令：

PYTHONPATH=src python -m pytest -q \
  tests/test_block_scaled_fp8_mm.py \
  tests/test_mxfp4_w4a16_grouped_mm.py \
  tests/test_rms_norm_gated.py \
  tests/test_fused_mla_rope_cache_write.py

性能结果摘要

在天数智芯 MR-V100 上：

• Gated RMSNorm 相对 PyTorch eager 获得约 5.27x–6.37x 加速。
• Fused MLA RoPE Cache Write 获得约 6.95x–7.74x 加速。
• MXFP4 uniform/routed 场景相对软件反量化 reference 获得约 6.47x–42.27x 加速。
• Block-scaled FP8 MM 在 decode 场景获得约 2.67x 加速。

在海光 BW/gfx936 上，MXFP4 uniform 场景已取得收益；其他部分路径仍受 codegen、wave 配置及 launch/lowering 开销影响，相关 负收益及适用边界已在技术报告中如实记录。

量化算子的性能基线为软件反量化 reference，不能解释为相对平台原生量化算子的性能收益。完整环境、测试方法和数据请参阅 docs/technical_report.md。

已知限制

• Block-scaled FP8 暂不支持 A 侧 128x128 scale、batch、grouped launch、N 尾块及 fast accumulation。
• MXFP4 当前仅支持 W4A16 BlockWise1x32，暂不支持 activation scale、bias、swizzle 和 grouped-K。
• Fused MLA RoPE Cache Write 暂不支持 FP8 cache layout。
• 非连续 cache 会使用临时连续副本。
• 部分海光场景仍有进一步优化空间。