From dac1bb652d703ffa187c1d1ec90a17b8c1cc3cb1 Mon Sep 17 00:00:00 2001 From: Zhiyi Hong <2497491955@qq.com> Date: Fri, 14 Aug 2026 17:09:13 +0800 Subject: [PATCH] [Test] Reproduce Kimi SM120 SiTU contract gap --- README.md | 2 +- .../README.md | 161 +++++++++++++++++- .../reproduce_kimi_situ_sm120_mxfp4.py | 88 ++++++++++ 3 files changed, 243 insertions(+), 8 deletions(-) create mode 100644 experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/reproduce_kimi_situ_sm120_mxfp4.py diff --git a/README.md b/README.md index e118a63..3e4ae2b 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # sskj — 多平台大模型推理性能基准测试项目 -**更新（2026-08-14 15:47:54 CST）**：建立 Kimi-K3 / RTX PRO 6000D / SGLang FlashInfer MXFP4 MoE 独立开发分支并完成 Phase 0 上游与真机审计；确认 SM120 基础 MXFP8×MXFP4 CUTLASS MoE 可运行，但 Kimi SiTU 激活尚未接入 SM120 CUTLASS runner，详见 `experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md`。 +**更新（2026-08-14 17:08:00 CST）**：完成 Kimi-K3 / RTX PRO 6000D / SGLang FlashInfer MXFP4 MoE Phase 1 数据契约审计；已从真实 checkpoint 追到 SM120 CUTLASS 调用，确认 SiTU 激活、gate/up 权重变换和参数透传三处独立缺口，并在 6000D 上复现 C++ binding 的门控 shape 首错，详见 `experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md`。 > 历史更新见 `git log`。项目目的与工作流见下方。 > **项目目的**：当新显卡（GPU/NPU）到货时，用最短时间完成大模型在该平台上的推理性能评估与部署配置选型。 diff --git a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md index c7ce32c..3e5f561 100644 --- a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md @@ -85,14 +85,161 @@ csrc/nv_internal/tensorrt_llm/kernels/cutlass_kernels/include/common.h: InvalidType = 10 ``` -## 后续硬门槛 +## Phase 1 真实调用链审计 -Phase 1 必须先完成以下内容，才允许进入整模型实验： +状态：**完成**。本阶段只确认数据契约和最小失败边界，未修改 SGLang 或 FlashInfer 实现。 -1. 为 SM120 CUTLASS runner 建立 SiTU 数学参考与 MXFP8×MXFP4 正确性测试。 -2. 明确 Kimi checkpoint 的 gate/up 排列以及 `(4.0, 25.0)` 参数语义。 -3. 补齐 C++ activation enum、gated 判断、激活实现及 Python/C++ 枚举一致性检查。 -4. 在 SGLang 中只对满足 Kimi SiTU 条件的模型选择新路径；不影响 GPT-OSS、DeepSeek 和现有 fallback。 -5. 覆盖 token 数 1/2/8/32、不同 expert 命中、TP/EP 权重切分以及 CUDA Graph capture/replay。 +### Serving 调用链 + +```text +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 元数据： + +```text +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 的参考公式为： + +```text +gate = 4 * tanh(gate / 4) * sigmoid(gate) +up = 25 * tanh(up / 25) +out = gate * up +``` + +### TP/EP expert mapping + +SGLang 的关系为： + +```text +moe_tp_size = tp_size / moe_ep_size / moe_dp_size +num_local_experts = num_global_experts / moe_ep_size +``` + +当前四节点基线为 TP32 / EP32 / MoE-DP1，因此： + +```text +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 的复现文件： + +```text +reproduce_kimi_situ_sm120_mxfp4.py +``` + +它复用 FlashInfer 官方测试的数据构造，只把 activation 改为 SiTU 并传 Kimi 的 `4/25` 参数。当前 main 的预期首错是： + +```text +fc1_expert_weights inter size must be equal to fc2_expert_weights inter size +``` + +601 上的实际复现命令： + +```bash +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 +``` + +证据日志： + +```text +/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，至少覆盖： + +- token 数 1/2/8/32 以及 Kimi 真实 prefill/decode 代表 shape； +- 均匀、倾斜、空 expert 和 masked/global expert id； +- SiTU `4/25`、MXFP8 activation、MXFP4 packed weight 与 E8M0 scale； +- NaN/Inf、确定性、CUDA Graph capture/replay； +- TP1/EP1 micro case 与 TP32/EP32 的 local-expert 映射； +- candidate 与 reference 使用同一 routing、同一量化输入和同一 finalize 口径。 正确性未通过前，不做吞吐结论，不修改 checkpoint，不使用运行时 monkeypatch。 diff --git a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/reproduce_kimi_situ_sm120_mxfp4.py b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/reproduce_kimi_situ_sm120_mxfp4.py new file mode 100644 index 0000000..8025abb --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/reproduce_kimi_situ_sm120_mxfp4.py @@ -0,0 +1,88 @@ +#!/usr/bin/env python3 +"""Reproduce the missing SM120 CUTLASS SiTU contract with MXFP8 x MXFP4.""" + +from __future__ import annotations + +import argparse +import importlib +from pathlib import Path +import sys + +import torch + +from flashinfer import mxfp8_quantize +from flashinfer.fused_moe import cutlass_fused_moe +from flashinfer.tllm_enums import ActivationType + + +def load_official_helpers(flashinfer_source: Path): + sys.path.insert(0, str(flashinfer_source)) + return importlib.import_module( + "tests.moe.test_trtllm_cutlass_fused_moe" + ) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--flashinfer-source", + type=Path, + default=Path("/data/hzy/src/flashinfer-kimi-sm120"), + ) + args = parser.parse_args() + + helpers = load_official_helpers(args.flashinfer_source) + torch.manual_seed(42) + + tokens, hidden, experts, top_k, intermediate = 1, 128, 2, 2, 128 + dtype = torch.bfloat16 + x = torch.randn(tokens, hidden, dtype=dtype, device="cuda") + w31 = torch.randn( + experts, 2 * intermediate, hidden, dtype=dtype, device="cuda" + ) / 10 + w2 = torch.randn( + experts, hidden, intermediate, dtype=dtype, device="cuda" + ) / 10 + + x_q, x_sf = mxfp8_quantize(x, True, 32) + w31_q, w31_sf = helpers.quant_mxfp4_batches(w31, experts) + w2_q, w2_sf = helpers.quant_mxfp4_batches(w2, experts) + logits = torch.randn(tokens, experts, dtype=dtype, device="cuda") + weights, ids = helpers.compute_routing(logits, top_k) + + neutral_global_scale = torch.ones(experts, dtype=torch.float32, device="cuda") + situ_gate_beta = torch.full( + (experts,), 4.0, dtype=torch.float32, device="cuda" + ) + situ_linear_beta = torch.full( + (experts,), 25.0, dtype=torch.float32, device="cuda" + ) + output = torch.empty_like(x) + + cutlass_fused_moe( + input=x_q, + token_selected_experts=ids.to(torch.int32), + token_final_scales=weights, + fc1_expert_weights=w31_q.contiguous().view(torch.int64), + fc2_expert_weights=w2_q.contiguous().view(torch.int64), + output_dtype=dtype, + quant_scales=[ + w31_sf.view(torch.int32), + neutral_global_scale, + w2_sf.view(torch.int32), + neutral_global_scale, + ], + input_sf=x_sf, + swiglu_alpha=situ_gate_beta, + swiglu_beta=situ_linear_beta, + swiglu_limit=None, + activation_type=ActivationType.Situ, + use_mxfp8_act_scaling=True, + output=output, + ) + + raise SystemExit("unexpected success: SM120 CUTLASS SiTU gap is no longer present") + + +if __name__ == "__main__": + main()