sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md

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)

已确认：

1. FlashInfer 当前 main 的 SM120 CUTLASS fused MoE 能完成基础 MXFP8×MXFP4 运算；官方最小测试通过。
2. Kimi-K3 所需的 ActivationType.Situ 尚未接入该 SM120 CUTLASS C++ runner。
3. Python ActivationType.Situ 的值为 10，而 CUTLASS C++ ActivationType 的值 10 当前是 InvalidType。SiTU 调用因此被当作非门控激活，FC1 门控维度判断错误。
4. SGLang 当前 SM120 MXFP4 路径仍按 GPT-OSS 的 SwiGLU 权重布局和参数组织，未透传 Kimi SiTU 参数。
5. 旧目录 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_mxfp4 runner 位于普通 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 决策记录

1. 不在 SGLang 中用独立 Triton SiTU 对 CUTLASS 的错误输出做事后修补；CUTLASS 已有独立 activation kernel，正确位置在 FlashInfer。
2. 不复用 GPT-OSS 的奇偶行权重转换；Kimi loader 已明确给出连续 [gate; up] 半区。
3. 不新增 FlashInfer ABI 参数；现有两组 per-expert float 参数足以承载 SiTU 的 4/25，但字段命名需要在代码和测试中解释清楚。
4. 不把旧 flashinfer_cutedsl 实验当作本路径证据。
5. 不进入整模型性能测试，直到 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 决策记录

MXFP8 quantize/dequantize 后，注入的 NaN 被量化为有限值，而 Inf 保留。相同 Inf 行在 PyTorch GEMM 与 CUTLASS reduction 中产生不同的逐元素 NaN/Inf 混合，但两者影响同一 token row，未受影响行继续满足普通数值阈值。IEEE reduction 顺序不保证逐元素 NaN/Inf 分类一致，因此 harness 要求：

1. 相同 token rows 被标记为 non-finite；
2. 未受影响 rows 继续满足 rtol=0.2, atol=0.2；
3. candidate 不崩溃，且三次结果确定。

这不是放宽有限值正确性，有限 case 仍采用与 FlashInfer 官方 MXFP8×MXFP4 测试一致的阈值。

Phase 3 硬门槛

下一阶段只做最小实现：

1. FlashInfer C++ 增加与 Python 对齐的 Situ enum、gated classification 和 SiTU adaptor，复用已有 per-expert alpha/beta 参数与 grouped GEMM。
2. SGLang SM120 processor 按 gate_up_interleaved=False 处理 Kimi 连续半区，生成 CUTLASS [up; gate]，不得复用 GPT-OSS 奇偶行拆分。
3. SGLang runner 从 runner_config 透传 ActivationType.Situ 和 Kimi 4/25 参数。
4. 未满足支持条件时在服务初始化阶段明确选择 Marlin；不得在 kernel 失败后热切换，不修改 checkpoint，不使用 runtime monkeypatch。
5. 完成后必须以 --require-situ --suite all 通过 Phase 2 全矩阵，才能进入任何吞吐比较。