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 a1f53a9..dc2af7b 100644 --- a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md @@ -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，并形成可回归、可性能比较、可提交上游的最小改动。 diff --git a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/SGLANG_DRAFT_PR.md b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/SGLANG_DRAFT_PR.md index e47aac9..fac3171 100644 --- a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/SGLANG_DRAFT_PR.md +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/SGLANG_DRAFT_PR.md @@ -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 diff --git a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/UPSTREAM_DUPLICATION_AUDIT.md b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/UPSTREAM_DUPLICATION_AUDIT.md index 10d3f4f..986522f 100644 --- a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/UPSTREAM_DUPLICATION_AUDIT.md +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/UPSTREAM_DUPLICATION_AUDIT.md @@ -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`。这些都是 diff --git a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/decision_log.md b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/decision_log.md index 0947adc..7b6c323 100644 --- a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/decision_log.md +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/decision_log.md @@ -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 diff --git a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch index 3f733d8..c62463b 100644 --- a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch @@ -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(