Skip to content

Add competition operators for KernelSwift operator innovation competition - #104

Open
logicore-code wants to merge 2 commits into
InfiniTensor:masterfrom
logicore-code:kernelswift-competition-chongli
Open

Add competition operators for KernelSwift operator innovation competition#104
logicore-code wants to merge 2 commits into
InfiniTensor:masterfrom
logicore-code:kernelswift-competition-chongli

Conversation

@logicore-code

Copy link
Copy Markdown

KernelSwift 算子创新大赛参赛作品

KernelSwift 算子创新大赛

参赛队伍: 崇理
队长: 蒋泽宇 (jiangzeyu-2026)
队员: 蒋光荣
所选赛题: T1, T2, T3, T4


修改内容与动机

本PR提交了四个高性能LLM推理算子的完整实现:

T1: MXFP4 W4A16分组专家矩阵乘

  • 即时反量化(On-the-fly Dequantization)减少显存带宽
  • 分组专家并行,单kernel处理多个专家
  • 支持E2M1格式的快速查找表解码
  • 对齐参考: PyTorch scaled_grouped_mm, vLLM MXFP4编码

T2: Block-scaled FP8矩阵乘

  • 动态在线缩放防止FP32累加器溢出
  • 双缓冲流水线隐藏延迟
  • 128×128分块匹配Tensor Core计算单元
  • 对齐参考: PyTorch scaled_mm

T3: Gated RMSNorm融合算子

  • 单遍完成归一化和门控计算
  • 支持Sigmoid/SiLU/GELU/Tanh多种门控激活
  • 消除中间结果显存读写,提升KDA等架构效率
  • 对齐参考: vLLM RMSNormGated

T4: MLA RoPE与压缩KV Cache写入融合

  • 三阶段融合:投影+RoPE+Cache写入
  • Paged KV Cache高效写入,支持动态序列长度
  • 实现DeepSeek-V2/V3的MLA低秩压缩注意力
  • 对齐参考: vLLM concat_and_cache_mla

方案与实现

所有算子基于九齿(NineToothed) DSL实现,遵循arrange-and-apply编程范式。

关键优化策略:

  1. 算子融合减少kernel launch和显存带宽
  2. 即时量化/反量化避免额外数据搬运
  3. 硬件感知的分块和内存布局优化
  4. 数值稳定的在线缩放算法

工程适用边界

  • 支持海光DCU (50%) 和 天数智芯Iluvatar (50%)
  • 接口设计遵循最小变更原则,保持向后兼容
  • constexpr参数传达编译期常量,支持编译优化
  • 测试覆盖正确性验证、数值稳定性、边界条件

测试

pytest tests/test_mxfp4_grouped_gemm.py tests/test_block_scaled_fp8_gemm.py tests/test_gated_rmsnorm.py tests/test_mla_rope_kv_cache.py -v

pytest output: All 21 tests passed

技术报告

详见 docs/COMPETITION_REPORT_CHONGLI.md


参赛队伍: 崇理
成员: 蒋泽宇, 蒋光荣
日期: 2026年8月6日

jiangzeyu-2026 added 2 commits August 6, 2026 23:45
…tion

Implement four high-performance LLM inference operators:

- MXFP4 W4A16 grouped expert matrix multiplication (T1)
- Block-scaled FP8 matrix multiplication (T2)
- Gated RMSNorm fusion operator (T3)
- MLA RoPE and compressed KV Cache fusion (T4)

Team: Chongli
Members: Zeyu Jiang, Guangrong Jiang

All operators support Hygon and Iluvatar platforms.
pytest output: tests pass locally
Detailed technical report covering all four competition problems:
- Optimization strategies and implementation details
- Performance analysis and expected speedups
- Ablation studies and design decisions

Team: Chongli (Zeyu Jiang, Guangrong Jiang)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant