Kimi-K3 SM120 FlashInfer MXFP4 MoE
目标
为 Kimi-K3 在 NVIDIA SM120(RTX PRO 6000D)上的 SGLang 部署补齐 FlashInfer MXFP8 activation × MXFP4 weight fused MoE,并形成可回归、可性能比较、可提交上游的最小改动。
本目录不修改已有 Kimi-K3 基线实验。功能实现必须先通过算子正确性,再进入模型加载和性能测试。
Phase 0 审计结论
状态:完成。
审计环境:
| 对象 | 版本或提交 |
|---|---|
| GPU | NVIDIA RTX PRO 6000 Blackwell Server Edition,SM120 |
| SGLang | 827552bc1dda8507f2a1b11c7274340a83da49ce |
| FlashInfer | a9e03bf2f5fe,Python package 0.6.18 |
| PyTorch / CUDA | PyTorch 2.11,CUDA 13 |
| 构建镜像 | lmsysorg/sglang:kimi-k3-fiv617situ-warm |
| Kimi-K3 | hidden 7168,expert hidden 3584,896 experts,top-k 16,SiTU (alpha=4.0, beta=25.0) |
已确认:
- FlashInfer 当前 main 的 SM120 CUTLASS fused MoE 能完成基础 MXFP8×MXFP4 运算;官方最小测试通过。
- Kimi-K3 所需的
ActivationType.Situ尚未接入该 SM120 CUTLASS C++ runner。 - Python
ActivationType.Situ的值为 10,而 CUTLASS C++ActivationType的值 10 当前是InvalidType。SiTU 调用因此被当作非门控激活,FC1 门控维度判断错误。 - SGLang 当前 SM120 MXFP4 路径仍按 GPT-OSS 的 SwiGLU 权重布局和参数组织,未透传 Kimi SiTU 参数。
- 旧目录
kimi3_pro6000_sglang_tp32ep32_deepep_flashinfer_mxfp4实际使用flashinfer_cutedsl,并在模型加载前因量化配置识别失败;它不能作为 FlashInfer MXFP4 SM120 的成功或失败证据。
上游重叠审计:
- SGLang PR #30272 只实现 DeepSeek-V4 的 SM120 FlashInfer MXFP4 路径。
- FlashInfer issue #4223 跟踪 SM12x NVFP4 kernel,同一 issue 明确把 SiTU 标为当前目标模型“不需要”。
- 截至本次审计,没有找到覆盖 Kimi-K3 + SM120 + SiTU + MXFP4 fused MoE 的已合并实现或开放 PR。
因此,本任务有明确的独立增量,不是重复实现。
真机证据
FlashInfer 源码与 wheel:
/data/hzy/src/flashinfer-kimi-sm120
/data/hzy/artifacts/flashinfer-main-a9e03bf/flashinfer_python-0.6.18-py3-none-any.whl
/data/hzy/cache/flashinfer-sm120
官方 SM120 MXFP8×MXFP4 基线:
FLASHINFER_DISABLE_VERSION_CHECK=1 pytest -q \
tests/moe/test_trtllm_cutlass_fused_moe.py::test_moe_mxfp8_mxfp4\[False-None-None-None-otype1-128-2-2-128-1\] -s
实际结果:
1 passed, 3 warnings in 540.58s
SiTU 最小负例沿用官方 test_moe 数据构造,只将激活改为 ActivationType.Situ:
test_moe(1, 128, 2, 2, 128, ActivationType.Situ)
实际首个错误:
fc1_expert_weights inter size must be equal to fc2_expert_weights inter size
(256 vs. 128)
对应源码事实:
flashinfer/tllm_enums.py:
Situ = 10
InvalidType = 11
csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/common.h:
Identity = 9
InvalidType = 10
Phase 1 真实调用链审计
状态:完成。本阶段只确认数据契约和最小失败边界,未修改 SGLang 或 FlashInfer 实现。
Serving 调用链
KimiK3MoE.forward
-> TopK(select_experts): sigmoid + correction bias + top-k 16 + renormalize
-> KimiK3MoE._forward_routed
-> FusedMoE.run_moe_core
-> Mxfp4MoEMethod.apply
-> fused_experts_none_to_flashinfer_mxfp4
-> flashinfer.fused_moe.cutlass_fused_moe
-> flashinfer_cutlass_fused_moe_binding.cu
-> CUTLASS grouped GEMM1 -> gated activation -> GEMM2 -> finalize
关键源码:
| 层次 | 文件 | 已确认行为 |
|---|---|---|
| 模型 | sglang/srt/models/kimi_k3.py |
创建 FusedMoE 时传 activation="situ"、gemm1_alpha=4.0、gemm1_clamp_limit=25.0、gate_up_interleaved=False |
| 路由 | sglang/srt/layers/moe/topk.py |
K3 使用 sigmoid 分数;correction bias 只参与选 expert,最终权重从未加 bias 的 sigmoid 分数 gather,并做归一化 |
| 量化 | sglang/srt/layers/quantization/mxfp4.py |
checkpoint 的 packed MXFP4 权重与 E8M0 group-32 scale 被装入 w13/w2 |
| SGLang runner | sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py |
SM120 当前固定传 ActivationType.Swiglu,未读取 runner_config.activation |
| FlashInfer binding | csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu |
根据 C++ isGatedActivation 判断 FC1 是否为 GEMM2 intermediate 的两倍 |
| FlashInfer activation | csrc/fused_moe/cutlass_backend/cutlass_fused_moe_kernels.cuh |
gated activation 是独立 CUDA kernel;已有 per-expert 参数向量,可做最小 SiTU 扩展,无需重写 grouped GEMM |
真实 checkpoint 契约
模型配置来自 /data/hf_models/Kimi-K3/config.json:
| 项目 | 值 |
|---|---|
| hidden size | 7168 |
| routed latent width | 3584 |
| expert intermediate | 3072 |
| experts / active experts | 896 / 16 |
| layers | 93 |
| activation | SiTU |
| SiTU gate beta / linear beta | 4.0 / 25.0 |
| quantization | packed MXFP4 weight,group size 32,E8M0 scale |
第 3 层 expert 0 的真实 safetensors 元数据:
w1.weight_packed [3072, 1792] U8 w1.weight_scale [3072, 112] U8
w3.weight_packed [3072, 1792] U8 w3.weight_scale [3072, 112] U8
w2.weight_packed [3584, 1536] U8 w2.weight_scale [3584, 96] U8
1792=3584/2、1536=3072/2 是每 byte 两个 FP4 元素。SGLang loader 将独立 w1 与 w3 顺序装入 w13 的两个连续半区,因此加载后是 [gate; up],不是 GPT-OSS 的逐行 [gate0, up0, ...]。
SiTU 的参考公式为:
gate = 4 * tanh(gate / 4) * sigmoid(gate)
up = 25 * tanh(up / 25)
out = gate * up
TP/EP expert mapping
SGLang 的关系为:
moe_tp_size = tp_size / moe_ep_size / moe_dp_size
num_local_experts = num_global_experts / moe_ep_size
当前四节点基线为 TP32 / EP32 / MoE-DP1,因此:
moe_tp_size = 32 / 32 / 1 = 1
num_local_experts = 896 / 32 = 28
每个 rank 持有连续的 28 个完整专家,FC1/FC2 不再按 MoE-TP 切 intermediate。候选 kernel 的真实核心 shape 是 E_local=28, K=3584, N=3072, topK=16;输入 token 数随 prefill/decode batch 变化。
已确认的差异矩阵
| 契约 | Kimi-K3 正确值 | 当前 SM120 路径 | 影响 |
|---|---|---|---|
| Python/C++ enum 10 | Situ |
C++ 为 InvalidType |
SiTU 被当成非门控,首先触发 FC1 shape 校验失败 |
| 激活数学 | SiTU (4, 25) |
固定 Swiglu |
即使绕过 shape 校验也会产生错误输出 |
checkpoint w13 |
连续 [gate; up] |
SM120 processor 假定 GPT-OSS 交错行 | 0::2/1::2 会把 gate/up 行混合打乱 |
CUTLASS 期望 w31 |
连续 [up; gate] |
GPT-OSS 专用转换 | Kimi 需要按半区交换,不能按奇偶行拆分 |
| activation 参数 | gate beta=4,linear beta=25 | 固定 SwiGLU alpha=1.702、beta=1、limit=7 | 参数语义完全不同 |
| routing | 已计算的标准 top-k ids/weights | CUTLASS runner 支持标准 routing | 这一层无需重写 |
| EP mapping | global ids + EP32/rank -> local 28 experts | CUTLASS ABI 已有 TP/EP 参数 | 需要 correctness 覆盖,但不是当前首错 |
因此这不是单一“参数名不匹配”。至少要同时修复 FlashInfer SiTU activation 和 SGLang Kimi SM120 权重/参数适配,缺一项都不能得到正确结果。
CUDA Graph 与 fallback
flashinfer_mxfp4runner 位于普通 model forward 内,eager 与 CUDA Graph 使用同一条 MoE 调用,不存在另一个自动回退到 Marlin 的 graph 专用分支。- Decode 只在 batch 能被已捕获 bucket 覆盖且满足 graph 条件时 replay;超过 capture max、宽度不匹配或其他 eligibility 失败时,
ModelRunner._forward_raw转入 eager runner。 - Graph replay 会把实际 batch pad 到最近的 captured batch。
tune_max_num_tokens=next_power_of_2(x.shape[0])因而看到的是静态 padded token 数。 - 后续必须分别验证 eager、capture 和 replay;“eager 正确”不能代表 CUDA Graph replay 正确。
--moe-runner-backend flashinfer_mxfp4当前没有运行时按 shape 回退 Marlin 的机制。生产 fallback 应在服务初始化选择阶段明确决定,不能等 kernel 运行失败后热切换。
最小复现
目标 MXFP8 activation × MXFP4 weight 的复现文件:
reproduce_kimi_situ_sm120_mxfp4.py
它复用 FlashInfer 官方测试的数据构造,只把 activation 改为 SiTU 并传 Kimi 的 4/25 参数。当前 main 的预期首错是:
fc1_expert_weights inter size must be equal to fc2_expert_weights inter size
601 上的实际复现命令:
cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4
CUDA_VISIBLE_DEVICES=2 \
PYTHONPATH=/tmp/hzy-fi-main-phase1 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
/data/hzy/envs/sparse-opd-sm120/bin/python \
reproduce_kimi_situ_sm120_mxfp4.py \
> phase1_reproduce_sm120_situ.log 2>&1
证据日志:
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/phase1_reproduce_sm120_situ.log
2026-08-14 实测进入 flashinfer.cutlass_fused_moe 后,在
flashinfer_cutlass_fused_moe_binding.cu:369 以 256 vs. 128 命中上述检查。
这证明失败发生在 FlashInfer SM120 CUTLASS ABI/activation 边界,不依赖完整模型、SGLang scheduler、DeepEP 或多节点网络。
Phase 1 决策记录
- 不在 SGLang 中用独立 Triton SiTU 对 CUTLASS 的错误输出做事后修补;CUTLASS 已有独立 activation kernel,正确位置在 FlashInfer。
- 不复用 GPT-OSS 的奇偶行权重转换;Kimi loader 已明确给出连续
[gate; up]半区。 - 不新增 FlashInfer ABI 参数;现有两组 per-expert float 参数足以承载 SiTU 的
4/25,但字段命名需要在代码和测试中解释清楚。 - 不把旧
flashinfer_cutedsl实验当作本路径证据。 - 不进入整模型性能测试,直到 Phase 2 的 BF16/PyTorch 或 Marlin 参考正确性通过。
Phase 2 correctness harness
状态:完成。本阶段建立验证框架并校准 reference,但 SiTU candidate 按预期尚未通过;未产生性能晋级结论。
唯一入口:
validate_kimi_sm120_mxfp4.py
对称性约束
- reference 与 candidate 共用同一
Case、随机种子、routing IDs/weights、量化后的 MXFP8 input、packed MXFP4 weights 和 E8M0 scales;reference 不读取量化前的 BF16 权重。 - FC1 在生成时显式采用 CUTLASS oracle 布局
[up; gate]。这使 harness 同时约束 SGLang 后续必须把 Kimi checkpoint 的[gate; up]正确交换。 - SwiGLU 作为当前 kernel 已支持的控制组;SiTU
(4, 25)使用完全相同的数据与调用,只改变 activation contract。 - correctness 固定
use_fused_finalize=False,避免非结合 atomic reduction 干扰确定性;benchmark 也由同一个 CUDA Event 计时函数驱动,但本阶段设置 iterations=0,不在共享 GPU 负载下生成性能数字。
覆盖矩阵:
| 类别 | 覆盖 |
|---|---|
| token shape | M=1/2/8/32 |
| routing | uniform、skewed、空 expert、-1 masked slot |
| EP | EP2 global/local ID 过滤;Kimi EP32/rank7/E_local=28 |
| 数值 | finite、NaN/Inf row-level propagation、三次 bitwise determinism |
| execution | eager、CUDA Graph capture、两次 replay |
| Kimi 真尺寸 | H=3584, N=3072, E_local=28, topK=16;decode M=1、prefill M=128 |
执行命令
cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4
# 已支持的 SwiGLU 控制组;本阶段最终证据统一使用 GPU 6
CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/tmp/hzy-fi-main-phase1 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
/data/hzy/envs/sparse-opd-sm120/bin/python validate_kimi_sm120_mxfp4.py \
--suite smoke --activation swiglu \
--output results/phase2_control_swiglu_smoke.json
# Kimi SiTU 目标真尺寸
CUDA_VISIBLE_DEVICES=6 PYTHONPATH=/tmp/hzy-fi-main-phase1 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
/data/hzy/envs/sparse-opd-sm120/bin/python validate_kimi_sm120_mxfp4.py \
--suite real --activation situ \
--output results/phase2_target_situ_real_gpu6.json
实现完成后的 CI-style gate 必须增加 --require-situ;若仍命中已知缺口,入口返回非零。
结果与证据
| 组别 | 结果 | 关键证据 |
|---|---|---|
| SwiGLU smoke control | 6/6 PASS | cosine 0.999625–0.999812,NRMSE 1.50%–2.75% |
| SwiGLU Kimi real shape | 2/2 PASS | decode/prefill cosine 0.999597/0.999621,NRMSE 2.85%/2.75% |
| SiTU smoke target | 6/6 EXPECTED_MISSING_SITU | 全部命中 binding 256 vs. 128 首错 |
| SiTU Kimi real shape | 2/2 EXPECTED_MISSING_SITU | GPU6 上 decode/prefill 均命中同一首错 |
| determinism | PASS | non-fused finalize 三次 bitwise equal |
| CUDA Graph | PASS | control replay 与 eager max_abs=0,两次 replay bitwise equal |
原始证据:
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/phase2_control_swiglu_smoke_gpu6.log
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/phase2_target_situ_smoke_gpu6.log
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/phase2_control_swiglu_real_gpu6.log
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/phase2_target_situ_real_gpu6.log
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/results/phase2_*.json
NaN/Inf 决策记录
2026-08-17 的统一 gate 发现旧 oracle 不成立:原始 BF16 NaN/Inf 经
mxfp8_quantize 后,E4M3 payload 全部有限,scale 最大值为 0xFE。
NVIDIA E8M0 定义中 0x00..0xFE 表示有限的二次幂,只有 0xFF 表示
NaN;转换采用 saturating finite 行为。旧 reference 把 0xFE scale
重新展开到 BF16 时发生溢出,再要求 Tensor Core candidate 传播相同
non-finite row,不是有效的 MXFP8 oracle。
修正后的 nonfinite case 同时要求:
- 量化后的 E4M3 payload 全部有限;
- E8M0 scale 不得出现
0xFF; - candidate 输出全部有限且三次 bitwise deterministic;
- BF16 dequant reference 的 overflow rows 继续记录为 diagnostic,不参与数值判定。
这项修正没有改变任何 finite case 的 reference、容差或性能计时。有限值仍
使用同一量化输入、权重、scale 与 routing,并满足 rtol=0.2, atol=0.2。
官方格式依据:
https://docs.nvidia.com/cuda/cuda-math-api/cuda_math_api/struct____nv__fp8__e8m0.html。
Phase 3 最小实现
状态:完成;有限值正确性与定义明确的 MXFP8 nonfinite 边界通过,可以进入 Phase 4 profiling。尚未形成整模型性能结论。
源码提交
| 仓库 | 分支 | 提交 |
|---|---|---|
| FlashInfer | hzy/kimi-k3-sm120-situ-mxfp4 |
7d1a4ee [Feature] Add CUTLASS SiTU fused MoE support on SM120 |
| SGLang | hzy/kimi-k3-sm120-flashinfer-mxfp4 |
fb929bb [Feature] Support Kimi-K3 FlashInfer MXFP4 MoE on SM120 |
FlashInfer 的最小增量:
- C++
ActivationType增加与 Python ABI 对齐的Situ; - 把 SiTU 标为 gated activation;
- 增加
4*tanh(gate/4)*sigmoid(gate) * 25*tanh(up/25)adaptor,复用已有 grouped GEMM 和 per-expert 参数; - Python 暴露显式 capability marker,避免仅凭早已存在的 enum 误判 kernel 可用。
SGLang 的最小增量:
- 同时支持 GPT-OSS 交错
[gate_i, up_i]与 Kimi 连续[gate; up]checkpoint 布局,统一生成 FlashInfer[up; gate]; - 从
MoeRunnerConfig透传Situ、alpha4、beta25,不再硬编码 SwiGLU; - bias-free Kimi 路径不向 kernel 传零占位 bias;
- FlashInfer 缺少 SiTU capability marker 时在 runner 初始化阶段明确报错,并提示以
--moe-runner-backend marlin重启;不做 kernel 失败后的热切换。
构建产物
/data/hzy/artifacts/flashinfer-kimi-sm120-situ-phase3-v3/flashinfer_python-0.6.18-py3-none-any.whl
SHA256 9298d34b6f9cec9b8bf535d1688fc6e8c73f7db65f64e0c496580e881c3d997b
正确性与集成结果
统一 Phase 2/3 gate 命令:
cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4
CUDA_VISIBLE_DEVICES=6 \
PYTHONPATH=/tmp/hzy-fi-situ-phase3-v3 \
FLASHINFER_DISABLE_VERSION_CHECK=1 \
/data/hzy/envs/sparse-opd-sm120/bin/python validate_kimi_sm120_mxfp4.py \
--suite all --activation situ --require-situ \
--warmup 2 --benchmark-iterations 3 \
--output results/phase3_target_situ_all_gpu6_v3.json
结果:8/8 PASS。
| 覆盖 | 结果 |
|---|---|
| M=1/2/8/32,uniform/skewed/empty/masked | PASS |
| EP2 global/local IDs | PASS |
| MXFP8 NaN/Inf saturation boundary | PASS,candidate finite、3-run bitwise |
| CUDA Graph capture/replay | PASS,replay 与 eager max_abs=0 |
Kimi decode M=1, H=3584, N=3072, E=28, topK=16 |
cosine 0.9996223,NRMSE 2.748% |
Kimi prefill M=128 同真尺寸 |
cosine 0.9996152,NRMSE 2.774% |
SGLang 集成测试:6 passed, 18 warnings in 563.41s。warnings 为 asyncio
配置、PyTorch deprecation 和只读 pytest cache,不含断言失败。
原始证据:
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/results/phase3_target_situ_all_gpu6_v3.log
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/results/phase3_target_situ_all_gpu6_v3.json
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/results/phase3_sglang_integration_gpu6_final.log
Phase 3 决策
- 当前缺口同时位于 FlashInfer SiTU activation 与 SGLang Kimi layout/参数集成;两个补丁都必要。
- 没有修改 checkpoint,没有 runtime monkeypatch,没有通过放宽 finite tolerance 获得通过。
- microbenchmark latency 仅用于检查 candidate 可执行,不作为生产性能结论。
- 下一阶段必须用 Nsight Systems/Compute 和三轮交错 A/B 区分 wrapper、同步、权重搬运、dispatch 与 kernel 本体;在此之前不宣称优于 Marlin。