qqtang/sskj
1
3
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4
History

README.md

Kimi-K3 SM120 FlashInfer MXFP4 MoE

目标

为 Kimi-K3 在 NVIDIA SM120（RTX 6000D）上的 SGLang 部署补齐 FlashInfer MXFP8 activation × MXFP4 weight fused MoE 支持，并形成：

  • 可上游审查的最小 SGLang 改动；
  • 精确覆盖 Kimi-K3 的正确性测试；
  • RTX 6000D 单卡与四节点真实服务证据；
  • Marlin 与 FlashInfer 的同口径性能结论。

本目录不修改已有 Kimi-K3 基线实验，也不提交与 FlashInfer 上游重复的 kernel。

交付状态

状态：完成，SGLang Draft 代码与审计分支已推送。

项目 结果
GPU NVIDIA RTX 6000D，SM120
模型 Kimi-K3，MXFP4，SiTU (4.0, 25.0)
SGLang base 593b1a9b8aeedb908ee21870c02a0fd1a3dc30a5
SGLang Draft HEAD b95b5341ab9555ba5ea53f5828f6de203d749cce
GitHub 分支 ACEEE-1222/sglang:hzy/kimi-k3-sm120-flashinfer-mxfp4
FlashInfer 依赖 PR #4460，implementation b525c513fef7
默认行为 不改变默认 MoE backend，Marlin 仍可用

公开 PR 正文：SGLANG_DRAFT_PR.md

可应用补丁： patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch

逐项验收：COMPLETION_AUDIT.md

交付内容

SGLang 实现

  1. 为 Kimi-K3 在 SM120 上接入 FlashInfer CUTLASS MXFP4 MoE；
  2. 将 Kimi 连续 [gate; up] 权重与 scale 转换为 CUTLASS 所需布局；
  3. 独立透传 situ_beta=4.0 与 situ_linear_beta=25.0；
  4. 在 MXFP8 量化前处理非连续 hidden states；
  5. 排除 SM120 误入 tcgen05 attention-residual 路径；
  6. 保持 GPT-OSS 路径、旧 FlashInfer import 和默认 Marlin 行为不变。

变更文件：

python/sglang/srt/layers/attn_residual.py
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
python/sglang/srt/layers/quantization/mxfp4.py
test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
test/registered/unit/layers/test_attn_residual_arch_gate.py

测试覆盖

  • Kimi gate/up 权重和 scale 布局；
  • SiTU (4.0, 25.0) 参数；
  • MXFP8 activation × MXFP4 weight；
  • 非连续输入；
  • standard routing 与 EP expert IDs；
  • eager、CUDA Graph capture/replay；
  • SM120 attention-residual architecture gate；
  • 缺少新 API 时的兼容性与明确报错。

上游边界

FlashInfer PR #4460 负责公共 CUTLASS SiTU kernel 和参数 API；本工作不提交 平行的 FlashInfer kernel PR。

SGLang Draft 只保留 Kimi-K3 的真实框架增量：权重/scale 布局、SiTU 参数、 非连续输入、SM120 guard 和测试。它不 vendor kernel、不 pin 未合并私人 commit、 不改变默认 backend。

完整去重审计见： UPSTREAM_DUPLICATION_AUDIT.md

正确性结果

使用 FlashInfer #4460 构建的 wheel，在 RTX 6000D GPU6 上执行精确 Kimi 组合：

8 passed, 17 warnings in 559.87s

覆盖 SM120 + MXFP8 + MXFP4 + SiTU(4,25)、Kimi 布局、非连续输入和 CUDA Graph。Kimi 真尺寸结果：

Case cosine NRMSE
Decode M=1 0.9996223 2.748%
Prefill M=128 0.9996152 2.774%

最终 Draft HEAD 的 compileall 和非 JIT 回归：

10 passed, 3 deselected, 17 warnings in 33.30s

五个变更文件已通过完整 pre-commit。同步 fork 后的 7 个提交与同步前版本在 git range-diff 中全部 patch-equivalent。

原始日志：

/data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log
/data/hzy/phase5_sglang_pr4460_draft_cpu_final.log

四节点真实 Serving

实验口径

维度 取值
节点 601-604，共 32 张 RTX 6000D
并行 TP32 / EP4 / DP1
输入/输出 16K / 1 token
并发 C=8、C=16
Backend Marlin、FlashInfer MXFP4
Chunk 4K、8K、16K
请求与重复 每点 40 请求，3 次重复
Cache 禁用 Radix Cache

正式 Run：

8K/16K: kimi3-ep4-moe-full-20260818-151349
4K:     kimi3-ep4-moe-chunk4k-20260819-150632

36/36 repeat 均完成 40/40 请求，输出长度均为 1，errors 为空；日志未发现 OOM、Traceback、NCCL error 或 EngineDeadError。下表为三次重复的中位数。

实验结果

Chunk C Backend Input TPS TTFT P50 TTFT P95 FlashInfer vs Marlin
4K 8 Marlin 2694.53 48.86 s 49.20 s baseline
4K 8 FlashInfer 2957.23 44.48 s 44.86 s +9.75% TPS / -8.96% TTFT
4K 16 Marlin 2692.60 96.25 s 97.96 s baseline
4K 16 FlashInfer 2960.41 87.61 s 89.04 s +9.95% TPS / -8.98% TTFT
8K 8 Marlin 2935.02 43.51 s 46.22 s baseline
8K 8 FlashInfer 3257.96 39.19 s 41.62 s +11.00% TPS / -9.92% TTFT
8K 16 Marlin 2934.90 87.09 s 89.86 s baseline
8K 16 FlashInfer 3260.14 78.39 s 80.88 s +11.08% TPS / -9.99% TTFT
16K 8 Marlin 2812.91 45.47 s 45.86 s baseline
16K 8 FlashInfer 3027.50 42.27 s 42.63 s +7.63% TPS / -7.04% TTFT
16K 16 Marlin 2812.33 90.97 s 91.36 s baseline
16K 16 FlashInfer 3027.34 84.52 s 84.85 s +7.65% TPS / -7.09% TTFT

结果结论

  1. FlashInfer 在 6 个正式 shape 上均优于 Marlin；
  2. 4K Chunk 的收益约为 +9.8% Input TPS / -9.0% TTFT；
  3. 8K Chunk 的收益约为 +11% Input TPS / -10% TTFT；
  4. 16K Chunk 的收益约为 +7.6% Input TPS / -7.1% TTFT；
  5. 同一 backend 下，8K Chunk 同时优于 4K 和 16K；4K 切分轮次过多，16K 单步过重；
  6. C 从 8 增至 16 后吞吐基本不变、TTFT 约翻倍，C=8 已接近 Prefill 吞吐饱和；
  7. 当前推荐为 EP4 + flashinfer_mxfp4 + chunked_prefill_size=8192。

该优化有效但不是数量级变化。最佳 TTFT P50 仍约 39 秒，后续需要继续优化 TP/EP 通信与 Prefill 并行结构。

Chunk=4K 补测结论

4K 补测复用了 8K/16K 矩阵的镜像、模型、服务参数和入口。12/12 repeat 全部有效，共 480/480 请求成功。

FlashInfer 在 C=8/16 下仍分别提升 9.75%/9.95% Input TPS，但 4K 的绝对 Input TPS 比 8K 低约 9%，TTFT 也更高。该结果补齐了 Chunk 曲线并确认 8K 不是网格边界造成的偶然最优点。

最终 PR 代表点

最终又用等价 SGLang patch 300c87a 和 FlashInfer #4460 implementation b525c513 重跑 8K/C8 代表点：

Run ID: kimi3-sm120-pr-representative-20260819-115530
TP32 / EP4 / DP1
ISL=16384 / OSL=1 / C=8 / requests=40 / chunk=8192
Marlin 3 repeats + FlashInfer MXFP4 3 repeats
Backend Input TPS TTFT P50 TTFT P95
Marlin 2930.03 43.587 s 46.287 s
FlashInfer MXFP4 3251.25 39.267 s 41.694 s
相对变化 +10.96% -9.91% -9.92%

六轮均为 40/40 成功，所有 output_lens=1，无 OOM、Traceback、NCCL error 或 EngineDeadError。

复现命令

四节点矩阵只在 601 执行：

cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill

read -rsp 'sudo password: ' SUDO_PASSWORD; echo
export SUDO_PASSWORD
export RUN_ID=kimi3-moe-prefill-$(date +%Y%m%d-%H%M%S)

bash run_moe_backend_prefill.sh run

可通过以下环境变量缩小矩阵：

EP_SIZE
BACKENDS_CSV
CHUNKED_PREFILL_SIZES_CSV
CONCURRENCIES_CSV
REPEATS

4K 补测使用：

EP_SIZE=4 \
BACKENDS_CSV=marlin,flashinfer_mxfp4 \
CHUNKED_PREFILL_SIZES_CSV=4096 \
CONCURRENCIES_CSV=8,16 \
REPEATS=3 \
bash run_moe_backend_prefill.sh run

证据路径

8K/16K 正式矩阵

/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/results/kimi3-ep4-moe-full-20260818-151349/
/data/hzy/kimi3-ep4-moe-full-20260818-151349.log

results/phase5_ep4_moe_backend_prefill_summary.json
results/phase5_ep4_moe_backend_prefill_results.csv

4K 补测

/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/results/kimi3-ep4-moe-chunk4k-20260819-150632/
/data/hzy/kimi3-ep4-moe-chunk4k-20260819-150632.log

results/phase5_ep4_moe_backend_chunk4k_summary.json
results/phase5_ep4_moe_backend_chunk4k_results.csv

最终 PR 代表点

/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/results/kimi3-sm120-pr-representative-20260819-115530/
/data/hzy/kimi3-sm120-pr-representative-20260819-115530.log

results/pr_representative_16k_c8_chunk8k_summary.json
results/pr_representative_16k_c8_chunk8k_results.csv

单卡测试与 Profiler

results/phase3_target_situ_all_gpu6_v3.json
results/phase4_mscan_ep32_5rounds_gpu6.json
results/phase4_nsys_ab_gpu6.nsys-rep
results/phase4_ncu_2026_2_1_ab_gpu6.ncu-rep

后续方向

MoE backend 选型完成后，不再扩大同类 backend 网格。后续 Prefill 优化按独立 任务推进：

  1. TP Reduce Scatter / Attention TP input scattered；
  2. MoE A2A；
  3. Prefill Pipeline Parallelism；
  4. 对最优新配置重新做 Nsight 时间线验证。