qqtang/sskj

Fork 0

Zhiyi Hong 08a35066d7 [Test] Complete Kimi-K3 Prefill MoE backend report

2026-08-19 16:54:45 +08:00

9.0 KiB

Raw Permalink Blame History

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 实现

为 Kimi-K3 在 SM120 上接入 FlashInfer CUTLASS MXFP4 MoE；
将 Kimi 连续 [gate; up] 权重与 scale 转换为 CUTLASS 所需布局；
独立透传 situ_beta=4.0 与 situ_linear_beta=25.0；
在 MXFP8 量化前处理非连续 hidden states；
排除 SM120 误入 tcgen05 attention-residual 路径；
保持 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

结果结论

FlashInfer 在 6 个正式 shape 上均优于 Marlin；
4K Chunk 的收益约为 +9.8% Input TPS / -9.0% TTFT；
8K Chunk 的收益约为 +11% Input TPS / -10% TTFT；
16K Chunk 的收益约为 +7.6% Input TPS / -7.1% TTFT；
同一 backend 下，8K Chunk 同时优于 4K 和 16K；4K 切分轮次过多，16K 单步过重；
C 从 8 增至 16 后吞吐基本不变、TTFT 约翻倍，C=8 已接近 Prefill 吞吐饱和；
当前推荐为 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 优化按独立任务推进：

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

9.0 KiB Raw Permalink Blame History Unescape Escape

Kimi-K3 SM120 FlashInfer MXFP4 MoE

目标

交付状态

交付内容

SGLang 实现

测试覆盖

上游边界

正确性结果

四节点真实 Serving

实验口径

实验结果

结果结论

Chunk=4K 补测结论

最终 PR 代表点

复现命令

证据路径

8K/16K 正式矩阵

4K 补测

最终 PR 代表点

单卡测试与 Profiler

后续方向

9.0 KiB

Raw Permalink Blame History