[Docs] Finalize Kimi SM120 SGLang draft PR
This commit is contained in:
Zhiyi Hong
parent
e01df16667
commit
a11c80b703
@ -7,6 +7,13 @@
|
||||
> Kimi-K3 SGLang 接入、精确 MXFP8×MXFP4+SiTU 测试和 6000D 实测证据。
|
||||
> 权威边界见 [`UPSTREAM_DUPLICATION_AUDIT.md`](UPSTREAM_DUPLICATION_AUDIT.md)。
|
||||
|
||||
当前 SGLang Draft HEAD 为 `300c87a431ac`，已推送到
|
||||
`ACEEE-1222/sglang:hzy/kimi-k3-sm120-flashinfer-mxfp4`。公开 PR 正文见
|
||||
[`SGLANG_DRAFT_PR.md`](SGLANG_DRAFT_PR.md)，可审查补丁见
|
||||
[`patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch`](patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch)。
|
||||
五个变更文件已通过完整 pre-commit；最终非 JIT 回归为
|
||||
`8 passed, 5 deselected, 17 warnings in 22.23s`。
|
||||
|
||||
## 目标
|
||||
|
||||
为 Kimi-K3 在 NVIDIA SM120（RTX PRO 6000D）上的 SGLang 部署补齐 FlashInfer MXFP8 activation × MXFP4 weight fused MoE，并形成可回归、可性能比较、可提交上游的最小改动。
|
||||
|
||||
@ -1,120 +1,131 @@
|
||||
# SGLang Draft PR
|
||||
# [Feature] Support Kimi-K3 FlashInfer MXFP4 MoE on SM120
|
||||
|
||||
## Title
|
||||
> Draft: depends on flashinfer-ai/flashinfer#4460. The dependency is not
|
||||
> vendored or pinned by this PR, and the default MoE backend is unchanged.
|
||||
|
||||
```text
|
||||
[Feature] Support Kimi-K3 FlashInfer MXFP4 MoE on SM120
|
||||
```
|
||||
## Motivation
|
||||
|
||||
## Dependency status
|
||||
SGLang already supports the FlashInfer CUTLASS MXFP4 MoE path for GPT-OSS on
|
||||
SM120, but Kimi-K3 needs additional model-specific integration:
|
||||
|
||||
> Draft: depends on flashinfer-ai/flashinfer#4460.
|
||||
>
|
||||
> Validated with FlashInfer built from PR #4460 on NVIDIA RTX 6000D.
|
||||
>
|
||||
> This PR does not vendor or pin the unmerged FlashInfer kernel.
|
||||
- Kimi stores gate and up projections as contiguous halves, while the CUTLASS
|
||||
kernel consumes the opposite half order. This differs from GPT-OSS's
|
||||
pair-interleaved layout.
|
||||
- Kimi-K3 uses SiTU with independent parameters
|
||||
`situ_beta=4.0` and `situ_linear_beta=25.0`.
|
||||
- Chunked Prefill can pass non-contiguous hidden states to the MXFP8
|
||||
quantizer.
|
||||
- RTX 6000D (SM120) must not enter the tcgen05 fused attention-residual path
|
||||
intended for GB100/GB200/GB300-class architectures.
|
||||
|
||||
FlashInfer #4460 owns the public CUTLASS SiTU implementation. This SGLang PR
|
||||
contains only the Kimi-K3 SM120 integration. It must remain Draft until #4460
|
||||
lands in an official RC or release with a stable API.
|
||||
FlashInfer PR #4460 owns the public CUTLASS SiTU kernel and API. This PR only
|
||||
adds the SGLang-side Kimi-K3 integration and should remain Draft until #4460
|
||||
is available in an official RC or release.
|
||||
|
||||
Reviewable source diff against SGLang main `63d783bbe095`:
|
||||
## Modifications
|
||||
|
||||
```text
|
||||
patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch
|
||||
Draft HEAD: 775aa076b00
|
||||
```
|
||||
|
||||
## Summary
|
||||
|
||||
- Dispatch Kimi-K3 MXFP8 activation by MXFP4 weight MoE to FlashInfer CUTLASS
|
||||
on SM120 when `--moe-runner-backend flashinfer_mxfp4` is explicitly selected.
|
||||
- Convert Kimi's contiguous `[gate; up]` weight and scale halves to the
|
||||
`[up; gate]` layout consumed by the CUTLASS kernel, without changing the
|
||||
existing GPT-OSS pair-interleaved path.
|
||||
- Forward Kimi's SiTU parameters as independent `situ_beta=4.0` and
|
||||
`situ_linear_beta=25.0` values through the #4460 API.
|
||||
- Materialize non-contiguous Prefill hidden states before FlashInfer MXFP8
|
||||
- Detect the public FlashInfer CUTLASS SiTU capability without making the new
|
||||
enum or function parameters mandatory at SGLang import time.
|
||||
- Dispatch Kimi-K3 MXFP8-activation-by-MXFP4-weight MoE to FlashInfer CUTLASS
|
||||
on SM120 only when `--moe-runner-backend flashinfer_mxfp4` is selected.
|
||||
- Convert Kimi's contiguous `[gate; up]` weight and scale halves to CUTLASS's
|
||||
`[up; gate]` order while preserving the existing GPT-OSS interleaved path.
|
||||
- Forward `situ_beta` and `situ_linear_beta` independently through the
|
||||
FlashInfer #4460 API.
|
||||
- Materialize non-contiguous hidden states immediately before MXFP8
|
||||
quantization.
|
||||
- Prevent SM120 from entering the tcgen05 attention-residual path intended for
|
||||
GB100/GB200-class architectures.
|
||||
- Exclude SM120 from the tcgen05 attention-residual architecture gate.
|
||||
- Add registered unit tests for capability fallback, legacy FlashInfer
|
||||
compatibility, Kimi/GPT-OSS layouts, SiTU parameters, non-contiguous input,
|
||||
numerical parity, and the architecture gate.
|
||||
|
||||
## Compatibility and fallback
|
||||
The PR does not copy FlashInfer C++ code, pin an unmerged commit, alter the
|
||||
default MoE backend, or remove the Marlin fallback.
|
||||
|
||||
- The default MoE backend is unchanged; Marlin remains available.
|
||||
- Importing SGLang does not require the new FlashInfer API.
|
||||
- Capability detection checks `ActivationType.Situ` and both public function
|
||||
parameters only when the Kimi SM120 FlashInfer path is selected.
|
||||
- An older FlashInfer then raises a targeted error suggesting an upgrade or
|
||||
`--moe-runner-backend marlin`.
|
||||
- No FlashInfer source is copied into SGLang and no dependency file pins an
|
||||
unmerged commit.
|
||||
## Accuracy Tests
|
||||
|
||||
## Tests
|
||||
Target hardware: NVIDIA RTX 6000D, compute capability 12.0. FlashInfer was
|
||||
built from a snapshot of PR #4460 fetched on 2026-08-18.
|
||||
|
||||
Targeted validation on one RTX 6000D GPU (SM120), with FlashInfer built from
|
||||
#4460 and the SGLang Draft source mounted into the container:
|
||||
|
||||
```text
|
||||
flashinfer-python: 0.6.18 (built from #4460)
|
||||
wheel sha256: deafdac2dc44b331441daba64e5b5dd28f591fe44f087af0aac3521c6ddbfa1f
|
||||
pytest: 8 passed, 17 warnings in 559.87s
|
||||
log: /data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log
|
||||
```bash
|
||||
pytest -q -s \
|
||||
test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py \
|
||||
test/registered/unit/layers/test_attn_residual_arch_gate.py
|
||||
```
|
||||
|
||||
The final review also added a regression proving that non-SiTU paths continue
|
||||
to work with a legacy FlashInfer enum. Capability detection, legacy
|
||||
compatibility, and architecture gating were rerun without kernel JIT:
|
||||
Result with the FlashInfer kernel enabled:
|
||||
|
||||
```text
|
||||
commit: 775aa076
|
||||
pytest: 8 passed, 5 deselected, 17 warnings in 26.00s
|
||||
log: /data/hzy/phase5_sglang_pr4460_draft_cpu_final.log
|
||||
8 passed, 17 warnings in 559.87s
|
||||
```
|
||||
|
||||
Coverage includes:
|
||||
The GPU test covers the exact Kimi combination of SM120, MXFP8 activations,
|
||||
MXFP4 weights, SiTU `(4.0, 25.0)`, Kimi gate/up layout, and non-contiguous
|
||||
input. It compares the SGLang adapter output directly with the public
|
||||
FlashInfer CUTLASS API.
|
||||
|
||||
- missing-capability fallback/error behavior;
|
||||
- Kimi gate/up weight and scale layout conversion;
|
||||
- independent SiTU 4.0/25.0 parameter forwarding;
|
||||
- non-contiguous hidden-state input;
|
||||
- direct SGLang-to-FlashInfer numerical comparison on SM120;
|
||||
- SM120 attention-residual architecture gating.
|
||||
After the final compatibility and CI-registration changes, the non-JIT subset
|
||||
was rerun from the final commit:
|
||||
|
||||
## Four-node serving evidence
|
||||
|
||||
Environment: 601-604, 32 RTX 6000D GPUs, TP32/EP4/DP1, 16K input, one output
|
||||
token, concurrency 8/16, chunked Prefill 8K/16K, three repetitions per point.
|
||||
```bash
|
||||
pytest -q -s \
|
||||
test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py \
|
||||
test/registered/unit/layers/test_attn_residual_arch_gate.py \
|
||||
-k "not matches_direct_flashinfer and not padding_layout_and_kernel and not situ_layout_and_noncontiguous_input"
|
||||
```
|
||||
|
||||
```text
|
||||
Run ID: kimi3-ep4-moe-full-20260818-151349
|
||||
Result: 24/24 points, every request 40/40 successful
|
||||
Errors: no OOM, Traceback, NCCL error, or EngineDeadError
|
||||
8 passed, 5 deselected, 17 warnings in 22.23s
|
||||
```
|
||||
|
||||
With 8K chunks, FlashInfer improved Input TPS by 11.00%-11.08% and reduced
|
||||
TTFT P50 by 9.92%-9.99% versus Marlin. With 16K chunks, Input TPS improved
|
||||
7.63%-7.65% and TTFT P50 fell 7.04%-7.09%.
|
||||
## Speed Tests and Profiling
|
||||
|
||||
Raw evidence:
|
||||
Four-node serving validation used 32 RTX 6000D GPUs with TP32/EP4/DP1,
|
||||
16K input tokens, one output token, concurrency 8/16, and chunked Prefill
|
||||
sizes 8K/16K. Each point was repeated three times; all 24 runs completed with
|
||||
40/40 successful requests and no OOM, Traceback, NCCL error, or
|
||||
`EngineDeadError`.
|
||||
|
||||
```text
|
||||
/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
|
||||
```
|
||||
| Chunk | C | Backend | Input TPS | TTFT P50 (s) | FlashInfer vs Marlin |
|
||||
|---:|---:|---|---:|---:|---:|
|
||||
| 8K | 8 | Marlin | 2935.02 | 43.51 | baseline |
|
||||
| 8K | 8 | FlashInfer | 3257.96 | 39.19 | +11.00% TPS, -9.92% TTFT |
|
||||
| 8K | 16 | Marlin | 2934.90 | 87.09 | baseline |
|
||||
| 8K | 16 | FlashInfer | 3260.14 | 78.39 | +11.08% TPS, -9.99% TTFT |
|
||||
| 16K | 8 | Marlin | 2812.91 | 45.47 | baseline |
|
||||
| 16K | 8 | FlashInfer | 3027.50 | 42.27 | +7.63% TPS, -7.04% TTFT |
|
||||
| 16K | 16 | Marlin | 2812.33 | 90.97 | baseline |
|
||||
| 16K | 16 | FlashInfer | 3027.34 | 84.52 | +7.65% TPS, -7.09% TTFT |
|
||||
|
||||
## Known limitations and merge order
|
||||
The service benchmark used the previously validated candidate integration
|
||||
with the same Kimi layout, SiTU, non-contiguous-input, and SM120 guard fixes.
|
||||
The final public #4460 API is covered by the accuracy test above; no claim is
|
||||
made that an unreleased #4460 wheel was used for this four-node benchmark.
|
||||
|
||||
1. FlashInfer #4460 must merge and publish an RC or release first.
|
||||
2. The Draft should then update its minimum supported FlashInfer version based
|
||||
on the released package, rather than pinning a commit.
|
||||
3. The tcgen05 SM120 guard is a separate commit in the same Draft so reviewers
|
||||
can split it into a BugFix PR if preferred.
|
||||
4. Existing results from the historical candidate image remain performance
|
||||
evidence; the #4460-based targeted run is the final correctness evidence.
|
||||
## Checklist
|
||||
|
||||
- [x] Format your code according to the [Format code with pre-commit](https://docs.sglang.io/developer_guide/contribution_guide.html#format-code-with-pre-commit).
|
||||
- [x] Add unit tests according to the [Run and add unit tests](https://docs.sglang.io/developer_guide/contribution_guide.html#run-and-add-unit-tests).
|
||||
- [x] Update documentation according to [Write documentations](https://docs.sglang.io/developer_guide/contribution_guide.html#write-documentations). (N/A for this Draft: it does not expose a stable released capability or change defaults; user-facing docs should be added after #4460 is released.)
|
||||
- [x] Provide accuracy and speed benchmark results according to [Test the accuracy](https://docs.sglang.io/developer_guide/contribution_guide.html#test-the-accuracy) and [Benchmark the speed](https://docs.sglang.io/developer_guide/contribution_guide.html#benchmark-the-speed).
|
||||
- [x] Follow the SGLang code style [guidance](https://docs.sglang.io/developer_guide/contribution_guide.html#code-style-guidance).
|
||||
|
||||
Full pre-commit validation was run on all five changed files and every hook
|
||||
passed. The two test files are registered with CUDA CI in `base-b` using the
|
||||
`1-gpu-small` runner.
|
||||
|
||||
## Dependency and Merge Order
|
||||
|
||||
1. FlashInfer #4460 merges and publishes an RC or release with a stable API.
|
||||
2. This Draft updates its version gate to that released package if required.
|
||||
3. SGLang CI runs with the released FlashInfer dependency.
|
||||
4. The Draft becomes Ready for Review.
|
||||
|
||||
The SM120 attention-residual guard is kept as a separate commit so reviewers
|
||||
can split it into a `[BugFix]` PR if preferred.
|
||||
|
||||
References:
|
||||
|
||||
- <https://github.com/flashinfer-ai/flashinfer/pull/4460>
|
||||
- <https://github.com/sgl-project/sglang/pull/24816>
|
||||
- <https://github.com/sgl-project/sglang/pull/33997>
|
||||
- https://github.com/flashinfer-ai/flashinfer/pull/4460
|
||||
- https://github.com/sgl-project/sglang/pull/24816
|
||||
- https://github.com/sgl-project/sglang/pull/33997
|
||||
|
||||
@ -103,10 +103,18 @@ SGLang Draft 的完整可审查 diff：
|
||||
|
||||
```text
|
||||
base: 63d783bbe0955237ec41f9ddabf7235ddf04673c
|
||||
head: 775aa076b00a47ced0f9784443899bdb6aa00159
|
||||
head: 300c87a431ac40d3e7817246376b7fe20932db09
|
||||
patch: patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch
|
||||
```
|
||||
|
||||
GitHub 交付分支：
|
||||
|
||||
```text
|
||||
fork: ACEEE-1222/sglang
|
||||
branch: hzy/kimi-k3-sm120-flashinfer-mxfp4
|
||||
compare: https://github.com/sgl-project/sglang/compare/main...ACEEE-1222:sglang:hzy/kimi-k3-sm120-flashinfer-mxfp4?expand=1
|
||||
```
|
||||
|
||||
候选镜像：
|
||||
|
||||
```text
|
||||
@ -144,9 +152,15 @@ GPU: NVIDIA RTX 6000D, SM120
|
||||
```
|
||||
|
||||
最终审查发现普通 SwiGLU 路径也曾无条件访问 `ActivationType.Situ`，可能误伤
|
||||
旧 FlashInfer。提交 `775aa076` 改为动态加入 SiTU enum，并新增旧版兼容回归；
|
||||
能力探测、旧版兼容和架构 gate 为 `8 passed, 5 deselected in 26.00s`。GPU
|
||||
kernel 路径继续由前述 8/8 测试覆盖。
|
||||
旧 FlashInfer。提交 `775aa076` 改为动态加入 SiTU enum，并新增旧版兼容回归。
|
||||
最终提交 `300c87a` 为两个测试文件增加正式 CUDA CI 注册，并通过完整
|
||||
pre-commit。最终 HEAD 的能力探测、旧版兼容和架构 gate 为
|
||||
`8 passed, 5 deselected, 17 warnings in 22.23s`；GPU kernel 路径继续由前述
|
||||
8/8 测试覆盖。
|
||||
|
||||
完整 pre-commit 命令覆盖全部五个变更文件；isort、ruff、black、codespell、
|
||||
registered-tests 检查及其他 hook 全部通过。下载 GitHub hook 时只对该命令使用
|
||||
HTTPS-to-SSH URL rewrite，没有修改全局 Git 配置。
|
||||
|
||||
前两次失败日志分别暴露 `FLASHINFER_DISABLE_JIT=0` 仍被视为启用、源码包缺少
|
||||
生成的 `data/csrc`；第三次暴露只读源码不能生成 `_build_meta.py`。这些都是
|
||||
|
||||
@ -29,9 +29,12 @@
|
||||
`8 passed, 17 warnings in 559.87s`。正式日志：
|
||||
`/data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log`。
|
||||
- 最终审查发现普通 SwiGLU 路径也曾无条件访问 `ActivationType.Situ`；
|
||||
`775aa076` 改为动态 enum 检测并补旧 FlashInfer 兼容回归。能力探测、
|
||||
旧版兼容和 SM120 架构 gate 为 `8 passed, 5 deselected in 26.00s`；日志：
|
||||
`775aa076` 改为动态 enum 检测并补旧 FlashInfer 兼容回归。最终提交
|
||||
`300c87a` 增加 CUDA CI 注册并通过全部 pre-commit hook。能力探测、旧版兼容
|
||||
和 SM120 架构 gate 为 `8 passed, 5 deselected, 17 warnings in 22.23s`；日志：
|
||||
`/data/hzy/phase5_sglang_pr4460_draft_cpu_final.log`。
|
||||
- 公开 Draft PR 正文遵循 SGLang 官方模板，不暴露内网原始路径；原始日志路径、
|
||||
镜像 SHA、wheel SHA 和去重证据继续保留在本仓库审计文件中。
|
||||
|
||||
## 2026-08-18 Phase 5
|
||||
|
||||
|
||||
@ -320,7 +320,7 @@ index 12ed0cf2728d23d9ea6554d7ae713fd6814898e8..7946f685957b50914168b6813b3e1cb8
|
||||
moe_tp_rank=layer.moe_tp_rank,
|
||||
moe_ep_size=layer.moe_ep_size,
|
||||
diff --git a/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py b/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
|
||||
index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..26417ed028525ad170f9121457857cb3580bb0b9 100644
|
||||
index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..0eed38e03f980fc9919f84a95265e0b6b8e55df9 100644
|
||||
--- a/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
|
||||
+++ b/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
|
||||
@@ -6,6 +6,7 @@ import builtins
|
||||
@ -336,8 +336,8 @@ index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..26417ed028525ad170f9121457857cb3
|
||||
|
||||
|
||||
+def test_kimi_k3_sm120_situ_requires_flashinfer_cutlass_situ_api(monkeypatch):
|
||||
+ import sglang.srt.layers.quantization.mxfp4 as mxfp4_module
|
||||
+ import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as runner_module
|
||||
+ import sglang.srt.layers.quantization.mxfp4 as mxfp4_module
|
||||
+ from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
+ from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
||||
+ from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
|
||||
@ -394,7 +394,7 @@ index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..26417ed028525ad170f9121457857cb3
|
||||
config = MoeRunnerConfig(
|
||||
num_experts=num_experts,
|
||||
num_local_experts=num_experts,
|
||||
@@ -433,5 +483,162 @@ def test_gpt_oss_sm120_padding_layout_and_kernel(monkeypatch):
|
||||
@@ -433,5 +483,161 @@ def test_gpt_oss_sm120_padding_layout_and_kernel(monkeypatch):
|
||||
assert torch.equal(actual, expected[:, :hidden].contiguous())
|
||||
|
||||
|
||||
@ -505,8 +505,7 @@ index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..26417ed028525ad170f9121457857cb3
|
||||
+ dtype=torch.bfloat16,
|
||||
+ device="cuda",
|
||||
+ generator=generator,
|
||||
+ )
|
||||
+ .t()
|
||||
+ ).t()
|
||||
+ * 0.1
|
||||
+ )
|
||||
+ assert not x.is_contiguous()
|
||||
@ -559,10 +558,10 @@ index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..26417ed028525ad170f9121457857cb3
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
diff --git a/test/registered/unit/layers/test_attn_residual_arch_gate.py b/test/registered/unit/layers/test_attn_residual_arch_gate.py
|
||||
new file mode 100644
|
||||
index 0000000000000000000000000000000000000000..054d345218257583c20fda4a54a85a812e04d1fe
|
||||
index 0000000000000000000000000000000000000000..3c55b924312a7d2537f8221b5e985e7927e527d4
|
||||
--- /dev/null
|
||||
+++ b/test/registered/unit/layers/test_attn_residual_arch_gate.py
|
||||
@@ -0,0 +1,43 @@
|
||||
@@ -0,0 +1,46 @@
|
||||
+"""Architecture gate tests for Kimi-K3's fused attention residual."""
|
||||
+
|
||||
+from unittest.mock import patch
|
||||
@ -570,6 +569,9 @@ index 0000000000000000000000000000000000000000..054d345218257583c20fda4a54a85a81
|
||||
+import pytest
|
||||
+
|
||||
+import sglang.srt.layers.attn_residual as attn_residual
|
||||
+from sglang.test.ci.ci_register import register_cuda_ci
|
||||
+
|
||||
+register_cuda_ci(est_time=10, stage="base-b", runner_config="1-gpu-small")
|
||||
+
|
||||
+
|
||||
+@pytest.mark.parametrize(
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user