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

[Docs] Finalize Kimi SM120 SGLang draft PR

Browse Source
This commit is contained in:
Zhiyi Hong 2026-08-18 23:47:34 +08:00
parent e01df16667
commit a11c80b703
5 changed files with 139 additions and 102 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md
View File

@ -7,6 +7,13 @@
> Kimi-K3 SGLang 接入、精确 MXFP8×MXFP4+SiTU 测试和 6000D 实测证据。 > Kimi-K3 SGLang 接入、精确 MXFP8×MXFP4+SiTU 测试和 6000D 实测证据。
> 权威边界见 [`UPSTREAM_DUPLICATION_AUDIT.md`](UPSTREAM_DUPLICATION_AUDIT.md)。 > 权威边界见 [`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，并形成可回归、可性能比较、可提交上游的最小改动。 为 Kimi-K3 在 NVIDIA SM120（RTX PRO 6000D）上的 SGLang 部署补齐 FlashInfer MXFP8 activation × MXFP4 weight fused MoE，并形成可回归、可性能比较、可提交上游的最小改动。

189
experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/SGLANG_DRAFT_PR.md
View File

@ -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 ## Motivation
[Feature] Support Kimi-K3 FlashInfer MXFP4 MoE on SM120
```
## 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. - Kimi stores gate and up projections as contiguous halves, while the CUTLASS
> kernel consumes the opposite half order. This differs from GPT-OSS's
> Validated with FlashInfer built from PR #4460 on NVIDIA RTX 6000D. pair-interleaved layout.
> - Kimi-K3 uses SiTU with independent parameters
> This PR does not vendor or pin the unmerged FlashInfer kernel. `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 FlashInfer PR #4460 owns the public CUTLASS SiTU kernel and API. This PR only
contains only the Kimi-K3 SM120 integration. It must remain Draft until #4460 adds the SGLang-side Kimi-K3 integration and should remain Draft until #4460
lands in an official RC or release with a stable API. is available in an official RC or release.
Reviewable source diff against SGLang main `63d783bbe095`: ## Modifications
```text - Detect the public FlashInfer CUTLASS SiTU capability without making the new
patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch enum or function parameters mandatory at SGLang import time.
Draft HEAD: 775aa076b00 - 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
## Summary `[up; gate]` order while preserving the existing GPT-OSS interleaved path.
- Forward `situ_beta` and `situ_linear_beta` independently through the
- Dispatch Kimi-K3 MXFP8 activation by MXFP4 weight MoE to FlashInfer CUTLASS FlashInfer #4460 API.
on SM120 when `--moe-runner-backend flashinfer_mxfp4` is explicitly selected. - Materialize non-contiguous hidden states immediately before MXFP8
- 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
quantization. quantization.
- Prevent SM120 from entering the tcgen05 attention-residual path intended for - Exclude SM120 from the tcgen05 attention-residual architecture gate.
GB100/GB200-class architectures. - 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. ## Accuracy Tests
- 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.
## 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 ```bash
#4460 and the SGLang Draft source mounted into the container: pytest -q -s \
test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py \
```text test/registered/unit/layers/test_attn_residual_arch_gate.py
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
``` ```
The final review also added a regression proving that non-SiTU paths continue Result with the FlashInfer kernel enabled:
to work with a legacy FlashInfer enum. Capability detection, legacy
compatibility, and architecture gating were rerun without kernel JIT:
```text ```text
commit: 775aa076 8 passed, 17 warnings in 559.87s
pytest: 8 passed, 5 deselected, 17 warnings in 26.00s
log: /data/hzy/phase5_sglang_pr4460_draft_cpu_final.log
``` ```
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; After the final compatibility and CI-registration changes, the non-JIT subset
- Kimi gate/up weight and scale layout conversion; was rerun from the final commit:
- 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.
## Four-node serving evidence ```bash
pytest -q -s \
Environment: 601-604, 32 RTX 6000D GPUs, TP32/EP4/DP1, 16K input, one output test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py \
token, concurrency 8/16, chunked Prefill 8K/16K, three repetitions per point. 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 ```text
Run ID: kimi3-ep4-moe-full-20260818-151349 8 passed, 5 deselected, 17 warnings in 22.23s
Result: 24/24 points, every request 40/40 successful
Errors: no OOM, Traceback, NCCL error, or EngineDeadError
``` ```
With 8K chunks, FlashInfer improved Input TPS by 11.00%-11.08% and reduced ## Speed Tests and Profiling
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%.
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 | Chunk | C | Backend | Input TPS | TTFT P50 (s) | FlashInfer vs Marlin |
/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 | 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. ## Checklist
2. The Draft should then update its minimum supported FlashInfer version based
on the released package, rather than pinning a commit. - [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).
3. The tcgen05 SM120 guard is a separate commit in the same Draft so reviewers - [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).
can split it into a BugFix PR if preferred. - [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.)
4. Existing results from the historical candidate image remain performance - [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).
evidence; the #4460-based targeted run is the final correctness evidence. - [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: References:
- <https://github.com/flashinfer-ai/flashinfer/pull/4460> - https://github.com/flashinfer-ai/flashinfer/pull/4460
- <https://github.com/sgl-project/sglang/pull/24816> - https://github.com/sgl-project/sglang/pull/24816
- <https://github.com/sgl-project/sglang/pull/33997> - https://github.com/sgl-project/sglang/pull/33997

22
experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/UPSTREAM_DUPLICATION_AUDIT.md
View File

@ -103,10 +103,18 @@ SGLang Draft 的完整可审查 diff：
```text ```text
base: 63d783bbe0955237ec41f9ddabf7235ddf04673c base: 63d783bbe0955237ec41f9ddabf7235ddf04673c
head: 775aa076b00a47ced0f9784443899bdb6aa00159 head: 300c87a431ac40d3e7817246376b7fe20932db09
patch: patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch 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 ```text
@ -144,9 +152,15 @@ GPU: NVIDIA RTX 6000D, SM120
``` ```
最终审查发现普通 SwiGLU 路径也曾无条件访问 `ActivationType.Situ`，可能误伤 最终审查发现普通 SwiGLU 路径也曾无条件访问 `ActivationType.Situ`，可能误伤
旧 FlashInfer。提交 `775aa076` 改为动态加入 SiTU enum，并新增旧版兼容回归； 旧 FlashInfer。提交 `775aa076` 改为动态加入 SiTU enum，并新增旧版兼容回归。
能力探测、旧版兼容和架构 gate 为 `8 passed, 5 deselected in 26.00s`。GPU 最终提交 `300c87a` 为两个测试文件增加正式 CUDA CI 注册，并通过完整
kernel 路径继续由前述 8/8 测试覆盖。 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` 仍被视为启用、源码包缺少 前两次失败日志分别暴露 `FLASHINFER_DISABLE_JIT=0` 仍被视为启用、源码包缺少
生成的 `data/csrc`；第三次暴露只读源码不能生成 `_build_meta.py`。这些都是 生成的 `data/csrc`；第三次暴露只读源码不能生成 `_build_meta.py`。这些都是

7
experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/decision_log.md
View File

@ -29,9 +29,12 @@
`8 passed, 17 warnings in 559.87s`。正式日志： `8 passed, 17 warnings in 559.87s`。正式日志：
`/data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log`。 `/data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log`。
- 最终审查发现普通 SwiGLU 路径也曾无条件访问 `ActivationType.Situ`； - 最终审查发现普通 SwiGLU 路径也曾无条件访问 `ActivationType.Situ`；
`775aa076` 改为动态 enum 检测并补旧 FlashInfer 兼容回归。能力探测、 `775aa076` 改为动态 enum 检测并补旧 FlashInfer 兼容回归。最终提交
旧版兼容和 SM120 架构 gate 为 `8 passed, 5 deselected in 26.00s`；日志： `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`。 `/data/hzy/phase5_sglang_pr4460_draft_cpu_final.log`。
- 公开 Draft PR 正文遵循 SGLang 官方模板，不暴露内网原始路径；原始日志路径、
镜像 SHA、wheel SHA 和去重证据继续保留在本仓库审计文件中。
## 2026-08-18 Phase 5 ## 2026-08-18 Phase 5

16
experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch
View File

@ -320,7 +320,7 @@ index 12ed0cf2728d23d9ea6554d7ae713fd6814898e8..7946f685957b50914168b6813b3e1cb8
moe_tp_rank=layer.moe_tp_rank, moe_tp_rank=layer.moe_tp_rank,
moe_ep_size=layer.moe_ep_size, 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 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 --- a/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
+++ b/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 @@ -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): +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.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.moe_runner.base import MoeRunnerConfig
+ from sglang.srt.layers.moe.utils import MoeRunnerBackend + from sglang.srt.layers.moe.utils import MoeRunnerBackend
+ from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod + from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
@ -394,7 +394,7 @@ index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..26417ed028525ad170f9121457857cb3
config = MoeRunnerConfig( config = MoeRunnerConfig(
num_experts=num_experts, num_experts=num_experts,
num_local_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()) assert torch.equal(actual, expected[:, :hidden].contiguous())
@ -505,8 +505,7 @@ index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..26417ed028525ad170f9121457857cb3
+ dtype=torch.bfloat16, + dtype=torch.bfloat16,
+ device="cuda", + device="cuda",
+ generator=generator, + generator=generator,
+ ) + ).t()
+ .t()
+ * 0.1 + * 0.1
+ ) + )
+ assert not x.is_contiguous() + assert not x.is_contiguous()
@ -559,10 +558,10 @@ index 435c6d7fe682197282ed7f7cb8f336d7a70f4b0e..26417ed028525ad170f9121457857cb3
sys.exit(pytest.main([__file__, "-v"])) 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 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 new file mode 100644
index 0000000000000000000000000000000000000000..054d345218257583c20fda4a54a85a812e04d1fe index 0000000000000000000000000000000000000000..3c55b924312a7d2537f8221b5e985e7927e527d4
--- /dev/null --- /dev/null
+++ b/test/registered/unit/layers/test_attn_residual_arch_gate.py +++ 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.""" +"""Architecture gate tests for Kimi-K3's fused attention residual."""
+ +
+from unittest.mock import patch +from unittest.mock import patch
@ -570,6 +569,9 @@ index 0000000000000000000000000000000000000000..054d345218257583c20fda4a54a85a81
+import pytest +import pytest
+ +
+import sglang.srt.layers.attn_residual as attn_residual +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( +@pytest.mark.parametrize(