[Docs] Prepare Kimi SM120 SGLang draft PR
This commit is contained in:
Zhiyi Hong
parent
ec7b604a50
commit
e01df16667
@ -1,5 +1,9 @@
|
||||
# sskj — 多平台大模型推理性能基准测试项目
|
||||
|
||||
**更新（2026-08-18 23:00:42 CST）**：完成依赖 FlashInfer #4460 的 SGLang Draft PR 收敛。确认不提交任何 FlashInfer PR，也不 vendor 或 pin 未合并 kernel；在 601 GPU6 上从 #4460 源码构建 FlashInfer 0.6.18 wheel，SGLang 定向测试 8/8 通过。Draft 标题、依赖、兼容策略、测试和四机 EP4 数据见 `experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/SGLANG_DRAFT_PR.md`。
|
||||
|
||||
**更新（2026-08-18 22:21:40 CST）**：完成 Kimi-K3 SM120 FlashInfer MXFP4 去重审计。确认 CUTLASS SiTU kernel 已由未合并的 FlashInfer PR #4460 实现，不再提交平行 kernel；SGLang 贡献收敛为 Kimi gate/up 与 scale 布局、SiTU 4.0/25.0 参数映射、非连续输入和 SM120 attention-residual guard，并保留 601-604 EP4 的全部服务级结果。详见 `experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/UPSTREAM_DUPLICATION_AUDIT.md`。
|
||||
|
||||
**更新（2026-08-18 18:38:17 CST）**：完成 Kimi-K3 四节点 TP32×EP4 的真实 Prefill MoE backend 验收；24/24 重复均 40/40 成功，FlashInfer MXFP4 在 8K Chunk、C=8/16 下相对 Marlin 提升约 11% Input TPS、降低约 10% TTFT，并确认 8K 优于 16K。实验默认 EP 改为 4，补齐非连续 Prefill 输入与 FlashInfer 官方 SiTU 回归，详见 `experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md`。
|
||||
|
||||
**更新（2026-08-18 13:55:30 CST）**：Kimi-K3 四节点 MoE backend 验收入口支持通过环境变量覆盖 EP size 及 backend/chunk/concurrency 子矩阵，用于先执行 TP32×EP4、16K Chunk、16K→1、C=16 的最大压力容量验证，同时保留默认 TP32×EP32 全矩阵。
|
||||
|
||||
@ -1,5 +1,12 @@
|
||||
# Kimi-K3 SM120 FlashInfer MXFP4 MoE
|
||||
|
||||
> **2026-08-18 去重审计更正**：下文 Phase 0-4 保留当时的调查与实验
|
||||
> 时间线，但“没有开放 PR、需要独立实现 FlashInfer SiTU kernel”的判断已失效。
|
||||
> 当前 kernel 依赖方向是尚未合并的 FlashInfer PR #4460；我们不提交
|
||||
> FlashInfer PR，SGLang Draft 只保留
|
||||
> Kimi-K3 SGLang 接入、精确 MXFP8×MXFP4+SiTU 测试和 6000D 实测证据。
|
||||
> 权威边界见 [`UPSTREAM_DUPLICATION_AUDIT.md`](UPSTREAM_DUPLICATION_AUDIT.md)。
|
||||
|
||||
## 目标
|
||||
|
||||
为 Kimi-K3 在 NVIDIA SM120（RTX PRO 6000D）上的 SGLang 部署补齐 FlashInfer MXFP8 activation × MXFP4 weight fused MoE，并形成可回归、可性能比较、可提交上游的最小改动。
|
||||
@ -29,13 +36,16 @@
|
||||
4. SGLang 当前 SM120 MXFP4 路径仍按 GPT-OSS 的 SwiGLU 权重布局和参数组织，未透传 Kimi SiTU 参数。
|
||||
5. 旧目录 `kimi3_pro6000_sglang_tp32ep32_deepep_flashinfer_mxfp4` 实际使用 `flashinfer_cutedsl`，并在模型加载前因量化配置识别失败；它不能作为 FlashInfer MXFP4 SM120 的成功或失败证据。
|
||||
|
||||
上游重叠审计：
|
||||
当时的上游重叠审计遗漏了后续出现的 FlashInfer #4460。最新复核结论是：
|
||||
|
||||
- SGLang PR #30272 只实现 DeepSeek-V4 的 SM120 FlashInfer MXFP4 路径。
|
||||
- FlashInfer issue #4223 跟踪 SM12x NVFP4 kernel，同一 issue 明确把 SiTU 标为当前目标模型“不需要”。
|
||||
- 截至本次审计，没有找到覆盖 Kimi-K3 + SM120 + SiTU + MXFP4 fused MoE 的已合并实现或开放 PR。
|
||||
- FlashInfer #4460 已实现 CUTLASS SiTU kernel 和独立 `situ_beta` /
|
||||
`situ_linear_beta` ABI，当前开放、未合并。
|
||||
- SGLang #32668 已合并 GPT-OSS SM120 接入，但不覆盖 Kimi 的连续 gate/up
|
||||
半区布局和 SiTU 参数。
|
||||
- 605 的可运行镜像也是移植 #4460，不是 0.6.17 原版直接支持。
|
||||
|
||||
因此，本任务有明确的独立增量，不是重复实现。
|
||||
因此独立增量仅位于 Kimi SGLang 接入、精确测试和 6000D 服务证据；不再
|
||||
包含平行 FlashInfer kernel。完整文件级审计见 `UPSTREAM_DUPLICATION_AUDIT.md`。
|
||||
|
||||
## 真机证据
|
||||
|
||||
@ -674,3 +684,14 @@ MXFP8×MXFP4 参数组，覆盖 FP16/BF16 与 autotune on/off：
|
||||
/data/hzy/phase5_sglang_noncontiguous_gpu6_v2.log
|
||||
/data/hzy/phase5_flashinfer_official_situ_gpu6.log
|
||||
```
|
||||
|
||||
去重后又在 FlashInfer #4460 的正式参数接口上完成最终组合验证：从 #4460
|
||||
源码构建 `flashinfer-python 0.6.18` wheel，加载当前 SGLang Draft，结果为：
|
||||
|
||||
```text
|
||||
8 passed, 17 warnings in 559.87s
|
||||
/data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log
|
||||
```
|
||||
|
||||
该 FlashInfer 测试提交仅保留为本地验证，不提交 FlashInfer PR。SGLang Draft
|
||||
不 vendor kernel、不 pin 未合并 commit、也不改变默认 MoE backend。
|
||||
|
||||
@ -0,0 +1,120 @@
|
||||
# SGLang Draft PR
|
||||
|
||||
## Title
|
||||
|
||||
```text
|
||||
[Feature] Support Kimi-K3 FlashInfer MXFP4 MoE on SM120
|
||||
```
|
||||
|
||||
## Dependency status
|
||||
|
||||
> 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.
|
||||
|
||||
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.
|
||||
|
||||
Reviewable source diff against SGLang main `63d783bbe095`:
|
||||
|
||||
```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
|
||||
quantization.
|
||||
- Prevent SM120 from entering the tcgen05 attention-residual path intended for
|
||||
GB100/GB200-class architectures.
|
||||
|
||||
## Compatibility and 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.
|
||||
|
||||
## Tests
|
||||
|
||||
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
|
||||
```
|
||||
|
||||
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:
|
||||
|
||||
```text
|
||||
commit: 775aa076
|
||||
pytest: 8 passed, 5 deselected, 17 warnings in 26.00s
|
||||
log: /data/hzy/phase5_sglang_pr4460_draft_cpu_final.log
|
||||
```
|
||||
|
||||
Coverage includes:
|
||||
|
||||
- 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.
|
||||
|
||||
## 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.
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
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%.
|
||||
|
||||
Raw evidence:
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
## Known limitations and merge order
|
||||
|
||||
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.
|
||||
|
||||
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>
|
||||
@ -0,0 +1,153 @@
|
||||
# Kimi-K3 SM120 FlashInfer MXFP4 去重审计
|
||||
|
||||
审计时间：2026-08-18 22:21:40 CST。
|
||||
|
||||
## 结论
|
||||
|
||||
FlashInfer CUTLASS SiTU kernel 不再作为我们的独立实现提交。真正的上游
|
||||
kernel 工作是尚未合并的 FlashInfer PR #4460。我们不提交 FlashInfer PR，
|
||||
只在本地将精确 Kimi 测试作为依赖验证：SM120、MXFP8 activation、MXFP4
|
||||
weight、SiTU `4.0/25.0` 和 Kimi shape。
|
||||
|
||||
SGLang 仍有独立增量：Kimi-K3 的 gate/up 权重与 scale 布局、SiTU 参数映射、
|
||||
非连续 hidden states、SM120 backend 分发，以及阻止 SM120 误入 tcgen05
|
||||
attention-residual 路径。SGLang 改动必须显式依赖 FlashInfer #4460，不能通过
|
||||
runtime monkey patch 或私有 capability marker 假装 0.6.17 已完整支持。
|
||||
|
||||
## 上游已有内容
|
||||
|
||||
| 项目 | 状态 | 已覆盖 | 未覆盖或边界 |
|
||||
|---|---|---|---|
|
||||
| FlashInfer #4180 | 已合并 | TRTLLM-gen MoE 的 SiTU | 不是 CUTLASS SM120 MXFP8×MXFP4 路径 |
|
||||
| FlashInfer #4209 | 开放、未合并 | 早期参数化 SiTU 方案 | 与 #4460 竞争；未作为当前依赖方向 |
|
||||
| FlashInfer #4460 | 开放、未合并；维护者已 approve，CI 22/22 | CUTLASS `ActivationType.Situ`、`situ_beta`、`situ_linear_beta`，覆盖 SM89/90/100/103/120 | 现有 `test_moe_mxfp8_mxfp4` 未明确覆盖 Kimi SiTU 4/25 |
|
||||
| SGLang #32668 | 已合并 | GPT-OSS SM120 FlashInfer MXFP4 | GPT-OSS 是 pair-interleaved gate/up；不是 Kimi 连续半区布局，也没有 Kimi SiTU 接入 |
|
||||
| SGLang #32650 | 已关闭、未合并 | Kimi SM90 参数化 SiTU 草案 | 采用 #4209 风格接口，不是当前 #4460 ABI；没有 SM120 服务级验收 |
|
||||
|
||||
关键链接：
|
||||
|
||||
- <https://github.com/flashinfer-ai/flashinfer/pull/4460>
|
||||
- <https://github.com/flashinfer-ai/flashinfer/pull/4209>
|
||||
- <https://github.com/flashinfer-ai/flashinfer/pull/4180>
|
||||
- <https://github.com/sgl-project/sglang/pull/32668>
|
||||
- <https://github.com/sgl-project/sglang/pull/32650>
|
||||
|
||||
## 605 已有方案
|
||||
|
||||
605 使用的不是“官方 FlashInfer 0.6.17 开箱即用”，而是旧 SGLang 镜像加
|
||||
运行时移植：
|
||||
|
||||
```text
|
||||
profile: /data/yy/sskj/deploy/profiles/pro6000/kimi3_pro6000_sglang_tp32ep32_flashinfer.env
|
||||
image: lmsysorg/sglang:kimi-k3-fiv617situ-warm
|
||||
image: sha256:66f5e4ba1da183621c3ca01ebb2b699e04373a2680085291f5840b01bdfe0e2c
|
||||
SGLang: 0.5.16
|
||||
flashinfer-python: 0.6.17
|
||||
flashinfer-cubin: 0.6.15.post1
|
||||
flashinfer-jit-cache: 0.6.15.post1+cu130
|
||||
```
|
||||
|
||||
它设置 `FLASHINFER_DISABLE_VERSION_CHECK=1`，并在启动前执行：
|
||||
|
||||
```text
|
||||
/tmp/patch_k3_sm120.py
|
||||
/tmp/patch_k3_flashinfer_sm120.py
|
||||
```
|
||||
|
||||
补丁 SHA256：
|
||||
|
||||
```text
|
||||
patch_k3_sm120.py:
|
||||
30975b43bd9abda98eefb5d9f5a6a39c319a8d1cce9c3b298a73a42ce3db5166
|
||||
patch_k3_flashinfer_sm120.py:
|
||||
202fe03ff800dab06b45fde24ef8b870d5e1021fe53ea145e5ce8df1eee50ecb
|
||||
```
|
||||
|
||||
`patch_k3_flashinfer_sm120.py` 明确写明移植 #4460，并额外处理 Kimi 布局、
|
||||
非连续输入和 SM120 autotune workaround。`patch_k3_sm120.py` 修复
|
||||
attention residual 架构 gate。它们证明功能方向可行，但 runtime 文本替换、
|
||||
混合版本包和跳过 autotune 不适合作为上游最终实现。
|
||||
|
||||
605 模型配置实测为：
|
||||
|
||||
```text
|
||||
activation_situ_beta = 4.0
|
||||
activation_situ_linear_beta = 25.0
|
||||
```
|
||||
|
||||
## 我们保留的新增内容
|
||||
|
||||
| 层次 | 新增内容 | 证据或提交 |
|
||||
|---|---|---|
|
||||
| SGLang | 按 `gate_up_interleaved` 区分 GPT-OSS pair-interleaved 与 Kimi `[gate; up]` 连续半区，并转换为 CUTLASS `[up; gate]` | `cab839c` |
|
||||
| SGLang | Kimi SiTU 映射到 #4460 的 `ActivationType.Situ`、`situ_beta=4.0`、`situ_linear_beta=25.0` | `2adc099` |
|
||||
| SGLang | `mxfp8_quantize` 前 materialize 非连续 Prefill hidden states | `35a9612` |
|
||||
| SGLang | 禁止 SM120 进入 GB100/GB200/GB300 tcgen05 attention-residual 路径 | `5e61458` |
|
||||
| 本地依赖验证 | 在 #4460 之上运行精确 MXFP8×MXFP4 + Kimi SiTU 4/25 + SM120 case | 本地提交 `542310c`；不提交 FlashInfer PR |
|
||||
| 6000D 验收 | 601-604、TP32/EP4、16K→1、C8/C16、Chunk 8K/16K、三次重复 | `kimi3-ep4-moe-full-20260818-151349` |
|
||||
|
||||
## 明确排除的重复内容
|
||||
|
||||
- FlashInfer 历史提交 `7d1a4ee` 是平行 SiTU kernel，实现与 #4460 高度
|
||||
重合，不进入 PR。
|
||||
- 历史测试提交 `e0bc653` 使用旧 SwiGLU 字段承载 4/25，已由基于 #4460
|
||||
独立字段的 `542310c` 取代。
|
||||
- 605 的 `skip_ops` autotune workaround 不默认进入 SGLang PR。我们的候选
|
||||
镜像完成 autotune 和 24/24 服务矩阵；只有在可复现崩溃时才另报问题。
|
||||
- 不把 FlashInfer 0.6.17 的版本号当成功能证明。完整组合依赖未合并的
|
||||
#4460 或等价后续合并版本。
|
||||
|
||||
## 固定实验身份与原始证据
|
||||
|
||||
SGLang Draft 的完整可审查 diff：
|
||||
|
||||
```text
|
||||
base: 63d783bbe0955237ec41f9ddabf7235ddf04673c
|
||||
head: 775aa076b00a47ced0f9784443899bdb6aa00159
|
||||
patch: patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch
|
||||
```
|
||||
|
||||
候选镜像：
|
||||
|
||||
```text
|
||||
local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5
|
||||
sha256:0b8e8ddceccecc0daf10b59ddd4177036affc81facc8bf87ea9e2456ab1e96b9
|
||||
SGLang 0.5.16
|
||||
flashinfer-python 0.6.18（历史候选源码 wheel）
|
||||
flashinfer-cubin 0.6.15.post1
|
||||
flashinfer-jit-cache 0.6.15.post1+cu130
|
||||
```
|
||||
|
||||
原始服务级结果：
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
原始 GPU6 correctness 日志：
|
||||
|
||||
```text
|
||||
/data/hzy/phase5_sglang_noncontiguous_gpu6_v2.log
|
||||
/data/hzy/phase5_flashinfer_official_situ_gpu6.log
|
||||
/data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log
|
||||
/data/hzy/phase5_sglang_pr4460_draft_cpu_final.log
|
||||
```
|
||||
|
||||
最终组合验证已在 601 的 GPU6 完成。测试从本地 #4460 工作树构建
|
||||
`flashinfer-python 0.6.18` wheel，再加载当前 SGLang Draft 源码：
|
||||
|
||||
```text
|
||||
wheel sha256: deafdac2dc44b331441daba64e5b5dd28f591fe44f087af0aac3521c6ddbfa1f
|
||||
result: 8 passed, 17 warnings in 559.87s
|
||||
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_DISABLE_JIT=0` 仍被视为启用、源码包缺少
|
||||
生成的 `data/csrc`；第三次暴露只读源码不能生成 `_build_meta.py`。这些都是
|
||||
测试装配问题，最终 v4 采用可写临时源码树构建 wheel 后通过，并非功能回归。
|
||||
@ -2,6 +2,37 @@
|
||||
|
||||
本文件记录可恢复的关键决策、失败尝试和证据路径。成功结论见 `README.md`。
|
||||
|
||||
## 2026-08-18 上游去重审计
|
||||
|
||||
- FlashInfer #4460 已覆盖 CUTLASS SiTU kernel、`ActivationType.Situ` 和独立
|
||||
`situ_beta/situ_linear_beta`；历史独立 kernel 提交 `7d1a4ee` 不提交。
|
||||
- FlashInfer 侧只在本地保留基于 #4460 的精确 Kimi
|
||||
MXFP8×MXFP4+SiTU 4/25 依赖验证；
|
||||
旧测试 `e0bc653` 由新测试提交 `542310c` 取代。
|
||||
- SGLang 接入改为公开 API 能力检查，不再依赖自造
|
||||
`CUTLASS_FUSED_MOE_SUPPORTS_SITU` marker，也不再借用 SwiGLU 字段传 4/25。
|
||||
- 605 的 runtime patches 作为历史可行性证据，不直接移植；其中 autotune
|
||||
skip 只有在重新复现崩溃后才考虑单独修复。
|
||||
- EP4 24/24 服务实验不丢弃、不重跑 EP32。它证明候选组合可用且相对
|
||||
Marlin 有收益，但最终上游 correctness 必须在 #4460 ABI 上重跑。
|
||||
- 详细矩阵、版本、SHA256 和原始路径见 `UPSTREAM_DUPLICATION_AUDIT.md`。
|
||||
|
||||
### 最终 PR 边界与 #4460 验证
|
||||
|
||||
- 不提交 FlashInfer PR；本地 `542310c` 仅用于验证 #4460 对 Kimi 精确组合的
|
||||
correctness，不作为上游贡献。
|
||||
- SGLang 只开 Draft PR，明确依赖 FlashInfer #4460；不 vendor kernel、不在依赖
|
||||
文件 pin 未合并 commit，也不改变默认 MoE backend。
|
||||
- 缺少 CUTLASS SiTU API 时，普通 SGLang import 和默认 Marlin 路径不受影响；
|
||||
只有显式选择 Kimi SM120 `flashinfer_mxfp4` 时给出升级或回退错误。
|
||||
- 601 GPU6 使用 #4460 构建的 FlashInfer 0.6.18 wheel 跑当前 SGLang Draft，
|
||||
`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`；日志：
|
||||
`/data/hzy/phase5_sglang_pr4460_draft_cpu_final.log`。
|
||||
|
||||
## 2026-08-18 Phase 5
|
||||
|
||||
### 真实 serving 暴露的兼容问题
|
||||
|
||||
@ -0,0 +1,608 @@
|
||||
diff --git a/python/sglang/srt/layers/attn_residual.py b/python/sglang/srt/layers/attn_residual.py
|
||||
index 2e3cc94e4206208264fa49331c4201482f9c8ee6..72f0f7b3c6f97a1aefc8453ebd7871403972f0a4 100644
|
||||
--- a/python/sglang/srt/layers/attn_residual.py
|
||||
+++ b/python/sglang/srt/layers/attn_residual.py
|
||||
@@ -8,7 +8,7 @@
|
||||
# fast — warp-specialized TMA kernel: cp.async.bulk producer +
|
||||
# online-softmax consumers over a double-buffered chunk ring, out
|
||||
# norm fused, per-nvb tuned launch config, one persistent CTA per
|
||||
-# SM. Taken on SM100+ with H=7168.
|
||||
+# SM. Taken on GB100/GB200/GB300 (SM100/SM103/SM110) with H=7168.
|
||||
# hip — single Triton kernel, everything in one launch; taken on ROCm
|
||||
# within its register budget.
|
||||
# fused — Triton 2-kernel pipeline with full H-parallelism; the fallback
|
||||
@@ -34,14 +34,17 @@ _HIP_SHAPE_GATE = None
|
||||
|
||||
|
||||
def _use_fast(hidden_size: int) -> bool:
|
||||
- """The TMA kernel needs SM100+ (tcgen05, cp.async.bulk) and its H=7168
|
||||
- template instantiation; everything else takes the triton pipeline."""
|
||||
+ """Select the tcgen05 TMA kernel only on supported data-center Blackwell.
|
||||
+
|
||||
+ RTX 6000D is SM120 but does not support this GB100/GB200/GB300-specific
|
||||
+ kernel. It must use the Triton pipeline.
|
||||
+ """
|
||||
global _FAST_SUPPORTED
|
||||
if is_npu():
|
||||
return False
|
||||
if _FAST_SUPPORTED is None:
|
||||
major, _ = torch.cuda.get_device_capability()
|
||||
- _FAST_SUPPORTED = major >= 10
|
||||
+ _FAST_SUPPORTED = major in (10, 11)
|
||||
return _FAST_SUPPORTED and hidden_size == 7168
|
||||
|
||||
|
||||
diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
|
||||
index 6a592fb5ebbda136fa13ac423f6bd36455af1a87..d824a8a3dd4f3938682f7d14f1ac5ff04f48d669 100644
|
||||
--- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
|
||||
+++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
|
||||
@@ -8,6 +8,7 @@ Quantization methods prepare a small quant_info payload and route through
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
+import inspect
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
|
||||
@@ -89,6 +90,10 @@ class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo):
|
||||
swiglu_beta: Optional[torch.Tensor] = None
|
||||
swiglu_limit: Optional[torch.Tensor] = None
|
||||
|
||||
+ # Optional per-expert SiTU scales, fp32 [E].
|
||||
+ situ_beta: Optional[torch.Tensor] = None
|
||||
+ situ_linear_beta: Optional[torch.Tensor] = None
|
||||
+
|
||||
# TP/EP topology (forwarded to the FlashInfer kernel)
|
||||
moe_tp_size: int = 1
|
||||
moe_tp_rank: int = 0
|
||||
@@ -111,22 +116,44 @@ def _flashinfer_cutlass_fused_moe():
|
||||
return cutlass_fused_moe, ActivationType
|
||||
|
||||
|
||||
+def flashinfer_cutlass_supports_situ() -> bool:
|
||||
+ """Return whether the installed FlashInfer exposes the CUTLASS SiTU API."""
|
||||
+ try:
|
||||
+ cutlass_fused_moe, activation_type = _flashinfer_cutlass_fused_moe()
|
||||
+ parameters = inspect.signature(cutlass_fused_moe).parameters
|
||||
+ except (ImportError, RuntimeError, TypeError, ValueError):
|
||||
+ return False
|
||||
+ return (
|
||||
+ hasattr(activation_type, "Situ")
|
||||
+ and "situ_beta" in parameters
|
||||
+ and "situ_linear_beta" in parameters
|
||||
+ )
|
||||
+
|
||||
+
|
||||
def _activation_type(runner_config: MoeRunnerConfig):
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_trtllm import get_activation_type
|
||||
|
||||
_, ActivationType = _flashinfer_cutlass_fused_moe()
|
||||
- activation = ActivationType(
|
||||
- get_activation_type(
|
||||
- runner_config.activation,
|
||||
- is_gated=runner_config.is_gated,
|
||||
+ situ_type = getattr(ActivationType, "Situ", None)
|
||||
+ if runner_config.activation == "situ" and runner_config.is_gated:
|
||||
+ if situ_type is None:
|
||||
+ raise RuntimeError("FlashInfer CUTLASS SiTU support is not available.")
|
||||
+ activation = situ_type
|
||||
+ else:
|
||||
+ activation = ActivationType(
|
||||
+ get_activation_type(
|
||||
+ runner_config.activation,
|
||||
+ is_gated=runner_config.is_gated,
|
||||
+ )
|
||||
)
|
||||
- )
|
||||
supported = {
|
||||
ActivationType.Swiglu,
|
||||
ActivationType.Geglu,
|
||||
ActivationType.Relu2,
|
||||
ActivationType.Identity,
|
||||
}
|
||||
+ if situ_type is not None:
|
||||
+ supported.add(situ_type)
|
||||
assert activation in supported, (
|
||||
f"Activation {runner_config.activation!r} "
|
||||
f"(is_gated={runner_config.is_gated}) maps to {activation.name}, "
|
||||
@@ -308,7 +335,7 @@ def fused_experts_none_to_flashinfer_mxfp4(
|
||||
quant_info, FlashInferCutlassMxfp4MoeQuantInfo
|
||||
), f"Unexpected quant_info type for flashinfer_mxfp4: {type(quant_info)}"
|
||||
|
||||
- flashinfer_cutlass_fused_moe, ActivationType = _flashinfer_cutlass_fused_moe()
|
||||
+ flashinfer_cutlass_fused_moe, _ = _flashinfer_cutlass_fused_moe()
|
||||
|
||||
x = dispatch_output.hidden_states
|
||||
topk_output = dispatch_output.topk_output
|
||||
@@ -342,6 +369,7 @@ def fused_experts_none_to_flashinfer_mxfp4(
|
||||
if weight_global_scale is not None:
|
||||
from flashinfer import mxfp8_quantize
|
||||
|
||||
+ x = x.contiguous()
|
||||
x, input_sf = mxfp8_quantize(
|
||||
x,
|
||||
is_sf_swizzled_layout=True,
|
||||
@@ -380,13 +408,15 @@ def fused_experts_none_to_flashinfer_mxfp4(
|
||||
swiglu_alpha=quant_info.swiglu_alpha,
|
||||
swiglu_beta=quant_info.swiglu_beta,
|
||||
swiglu_limit=quant_info.swiglu_limit,
|
||||
+ situ_beta=quant_info.situ_beta,
|
||||
+ situ_linear_beta=quant_info.situ_linear_beta,
|
||||
tp_size=quant_info.moe_tp_size,
|
||||
tp_rank=quant_info.moe_tp_rank,
|
||||
ep_size=quant_info.moe_ep_size,
|
||||
ep_rank=quant_info.moe_ep_rank,
|
||||
use_w4_group_scaling=not use_mxfp8_act_scaling,
|
||||
use_mxfp8_act_scaling=use_mxfp8_act_scaling,
|
||||
- activation_type=ActivationType.Swiglu,
|
||||
+ activation_type=_activation_type(runner_config),
|
||||
tune_max_num_tokens=next_power_of_2(x.shape[0]),
|
||||
output=out,
|
||||
use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(),
|
||||
diff --git a/python/sglang/srt/layers/quantization/mxfp4.py b/python/sglang/srt/layers/quantization/mxfp4.py
|
||||
index 12ed0cf2728d23d9ea6554d7ae713fd6814898e8..7946f685957b50914168b6813b3e1cb85b0cfd09 100644
|
||||
--- a/python/sglang/srt/layers/quantization/mxfp4.py
|
||||
+++ b/python/sglang/srt/layers/quantization/mxfp4.py
|
||||
@@ -1153,15 +1153,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def _process_weights_for_sm120_cutlass(self, layer):
|
||||
- """Prepare GPT-OSS MXFP4 experts for FlashInfer CUTLASS on SM120.
|
||||
-
|
||||
- GPT-OSS stores gate/up rows pair-wise as
|
||||
- ``[gate_0, up_0, gate_1, up_1, ...]``. FlashInfer's fused MoE consumes
|
||||
- two contiguous halves in ``[up; gate]`` order. Build that layout after
|
||||
- checkpoint loading so padding cannot move the split, pad both GEMMs to
|
||||
- CUTLASS's 128-element alignment, and swizzle the native E8M0 scales for
|
||||
- the SM120 MXFP8-by-MXFP4 kernels. Packed FP4 weight bytes themselves do
|
||||
- not need an SM120 permutation.
|
||||
+ """Prepare MXFP4 experts for FlashInfer CUTLASS on SM120.
|
||||
+
|
||||
+ FlashInfer consumes two contiguous halves in ``[up; gate]`` order.
|
||||
+ GPT-OSS checkpoints store pair-interleaved ``[gate_i, up_i]`` rows,
|
||||
+ while Kimi-K3 stores contiguous ``[gate; up]`` halves. Normalize either
|
||||
+ layout after loading, pad both GEMMs to CUTLASS's 128-element alignment,
|
||||
+ and swizzle the native E8M0 scales. Packed FP4 bytes need no additional
|
||||
+ SM120 permutation.
|
||||
"""
|
||||
from flashinfer import block_scale_interleave
|
||||
|
||||
@@ -1173,9 +1172,15 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
E = layer.num_local_experts
|
||||
device = layer.w13_weight.device
|
||||
|
||||
+ gate_up_interleaved = self.moe_runner_config.gate_up_interleaved
|
||||
+
|
||||
+ def _split_gate_up(unpadded):
|
||||
+ if gate_up_interleaved:
|
||||
+ return unpadded[:, 0::2, :], unpadded[:, 1::2, :]
|
||||
+ return unpadded[:, :N_un, :], unpadded[:, N_un : 2 * N_un, :]
|
||||
+
|
||||
def _stack_up_gate_w13(unpadded, last_pad, last_un):
|
||||
- gate_rows = unpadded[:, 0::2, :]
|
||||
- up_rows = unpadded[:, 1::2, :]
|
||||
+ gate_rows, up_rows = _split_gate_up(unpadded)
|
||||
out = torch.zeros(
|
||||
E, 2 * N_pad, last_pad, dtype=unpadded.dtype, device=device
|
||||
)
|
||||
@@ -1192,8 +1197,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
|
||||
bias_dtype = layer.w13_weight_bias.dtype
|
||||
w13_bias_padded = torch.zeros(E, 2 * N_pad, dtype=bias_dtype, device=device)
|
||||
- w13_bias_padded[:, :N_un] = layer.w13_weight_bias.data[:, 1::2]
|
||||
- w13_bias_padded[:, N_pad : N_pad + N_un] = layer.w13_weight_bias.data[:, 0::2]
|
||||
+ if gate_up_interleaved:
|
||||
+ gate_bias = layer.w13_weight_bias.data[:, 0::2]
|
||||
+ up_bias = layer.w13_weight_bias.data[:, 1::2]
|
||||
+ else:
|
||||
+ gate_bias = layer.w13_weight_bias.data[:, :N_un]
|
||||
+ up_bias = layer.w13_weight_bias.data[:, N_un : 2 * N_un]
|
||||
+ w13_bias_padded[:, :N_un] = up_bias
|
||||
+ w13_bias_padded[:, N_pad : N_pad + N_un] = gate_bias
|
||||
|
||||
def _pad_w2_3d(unpadded, last_pad, last_un):
|
||||
out = torch.zeros(E, K_pad, last_pad, dtype=unpadded.dtype, device=device)
|
||||
@@ -1225,17 +1236,59 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
layer.w13_weight_bias = Parameter(w13_bias_padded, requires_grad=False)
|
||||
layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False)
|
||||
|
||||
- layer.swiglu_alpha = Parameter(
|
||||
- torch.full((E,), 1.702, dtype=torch.float32, device=device),
|
||||
- requires_grad=False,
|
||||
+ activation = self.moe_runner_config.activation
|
||||
+ alpha = self.moe_runner_config.gemm1_alpha
|
||||
+ beta = self.moe_runner_config.gemm1_beta
|
||||
+ limit = self.moe_runner_config.gemm1_clamp_limit
|
||||
+ if activation == "situ":
|
||||
+ situ_beta = 4.0 if alpha is None else alpha
|
||||
+ situ_linear_beta = 25.0 if limit is None else limit
|
||||
+ alpha = beta = limit = None
|
||||
+ else:
|
||||
+ situ_beta = situ_linear_beta = None
|
||||
+ alpha = 1.702 if alpha is None else alpha
|
||||
+ beta = 1.0 if beta is None else beta
|
||||
+ limit = 7.0 if limit is None else limit
|
||||
+
|
||||
+ layer.swiglu_alpha = (
|
||||
+ None
|
||||
+ if alpha is None
|
||||
+ else Parameter(
|
||||
+ torch.full((E,), alpha, dtype=torch.float32, device=device),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
)
|
||||
- layer.swiglu_beta = Parameter(
|
||||
- torch.ones(E, dtype=torch.float32, device=device),
|
||||
- requires_grad=False,
|
||||
+ layer.swiglu_beta = (
|
||||
+ None
|
||||
+ if beta is None
|
||||
+ else Parameter(
|
||||
+ torch.full((E,), beta, dtype=torch.float32, device=device),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
)
|
||||
- layer.swiglu_limit = Parameter(
|
||||
- torch.full((E,), 7.0, dtype=torch.float32, device=device),
|
||||
- requires_grad=False,
|
||||
+ layer.swiglu_limit = (
|
||||
+ None
|
||||
+ if limit is None
|
||||
+ else Parameter(
|
||||
+ torch.full((E,), limit, dtype=torch.float32, device=device),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
+ )
|
||||
+ layer.situ_beta = (
|
||||
+ None
|
||||
+ if situ_beta is None
|
||||
+ else Parameter(
|
||||
+ torch.full((E,), situ_beta, dtype=torch.float32, device=device),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
+ )
|
||||
+ layer.situ_linear_beta = (
|
||||
+ None
|
||||
+ if situ_linear_beta is None
|
||||
+ else Parameter(
|
||||
+ torch.full((E,), situ_linear_beta, dtype=torch.float32, device=device),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
)
|
||||
# The MXFP4 ABI uses a neutral global weight scale for each GEMM.
|
||||
layer.mxfp4_weight_global_scale = Parameter(
|
||||
@@ -1278,6 +1331,22 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
"cutlass_sm90",
|
||||
"cutlass_sm120",
|
||||
):
|
||||
+ if (
|
||||
+ self._fi_kernel == "cutlass_sm120"
|
||||
+ and moe_runner_config.activation == "situ"
|
||||
+ ):
|
||||
+ from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
|
||||
+ flashinfer_cutlass_supports_situ,
|
||||
+ )
|
||||
+
|
||||
+ if not flashinfer_cutlass_supports_situ():
|
||||
+ raise RuntimeError(
|
||||
+ "Kimi-K3 FlashInfer MXFP4 MoE on SM120 requires a "
|
||||
+ "FlashInfer build with CUTLASS SiTU support "
|
||||
+ "(ActivationType.Situ plus situ_beta/situ_linear_beta). "
|
||||
+ "Upgrade FlashInfer or restart with "
|
||||
+ "--moe-runner-backend marlin."
|
||||
+ )
|
||||
# Register the fused func at runner construction so the FusedOpPool
|
||||
# lookup at `MoeRunner.__init__` finds it.
|
||||
import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass # noqa: F401
|
||||
@@ -1338,7 +1407,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
)
|
||||
|
||||
def _apply_sm120_cutlass(self, layer, dispatch_output):
|
||||
- """SM120 GPT-OSS MXFP8 x MXFP4 MoE via FlashInfer CUTLASS."""
|
||||
+ """SM120 MXFP8 x MXFP4 MoE via FlashInfer CUTLASS."""
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
|
||||
FlashInferCutlassMxfp4MoeQuantInfo,
|
||||
)
|
||||
@@ -1349,11 +1418,13 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
||||
w13_weight_scale=layer.w13_weight_scale,
|
||||
w2_weight_scale=layer.w2_weight_scale,
|
||||
mxfp4_weight_global_scale=layer.mxfp4_weight_global_scale,
|
||||
- w13_bias=layer.w13_weight_bias,
|
||||
- w2_bias=layer.w2_weight_bias,
|
||||
+ w13_bias=layer.w13_weight_bias if self.with_bias else None,
|
||||
+ w2_bias=layer.w2_weight_bias if self.with_bias else None,
|
||||
swiglu_alpha=layer.swiglu_alpha,
|
||||
swiglu_beta=layer.swiglu_beta,
|
||||
swiglu_limit=layer.swiglu_limit,
|
||||
+ situ_beta=layer.situ_beta,
|
||||
+ situ_linear_beta=layer.situ_linear_beta,
|
||||
moe_tp_size=layer.moe_tp_size,
|
||||
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
|
||||
--- 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
|
||||
import importlib
|
||||
import sys
|
||||
from contextlib import nullcontext
|
||||
+from enum import IntEnum
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
@@ -105,6 +106,54 @@ def test_dsv4_sm120_load_contract(monkeypatch):
|
||||
assert captured["fp4_scale_dtype"] == torch.float8_e8m0fnu
|
||||
|
||||
|
||||
+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
|
||||
+ 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
|
||||
+
|
||||
+ monkeypatch.setattr(
|
||||
+ mxfp4_module,
|
||||
+ "get_moe_runner_backend",
|
||||
+ lambda: MoeRunnerBackend.FLASHINFER_MXFP4,
|
||||
+ )
|
||||
+ monkeypatch.setattr(
|
||||
+ runner_module, "flashinfer_cutlass_supports_situ", lambda: False
|
||||
+ )
|
||||
+
|
||||
+ method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod)
|
||||
+ method._fi_kernel = "cutlass_sm120"
|
||||
+ config = MoeRunnerConfig(activation="situ", is_gated=True)
|
||||
+ with pytest.raises(RuntimeError, match="--moe-runner-backend marlin"):
|
||||
+ method.create_moe_runner(SimpleNamespace(), config)
|
||||
+
|
||||
+
|
||||
+def test_non_situ_activation_supports_legacy_flashinfer(monkeypatch):
|
||||
+ import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as runner_module
|
||||
+ import sglang.srt.layers.moe.moe_runner.flashinfer_trtllm as trtllm_module
|
||||
+
|
||||
+ class LegacyActivationType(IntEnum):
|
||||
+ Swiglu = 0
|
||||
+ Geglu = 1
|
||||
+ Relu2 = 2
|
||||
+ Identity = 3
|
||||
+
|
||||
+ monkeypatch.setattr(
|
||||
+ runner_module,
|
||||
+ "_flashinfer_cutlass_fused_moe",
|
||||
+ lambda: (object(), LegacyActivationType),
|
||||
+ )
|
||||
+ monkeypatch.setattr(
|
||||
+ trtllm_module,
|
||||
+ "get_activation_type",
|
||||
+ lambda *args, **kwargs: LegacyActivationType.Swiglu.value,
|
||||
+ )
|
||||
+
|
||||
+ config = SimpleNamespace(activation="silu", is_gated=True)
|
||||
+ assert runner_module._activation_type(config) == LegacyActivationType.Swiglu
|
||||
+
|
||||
+
|
||||
def test_dsv4_sm120_matches_direct_flashinfer(monkeypatch):
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required")
|
||||
@@ -317,6 +366,7 @@ def test_gpt_oss_sm120_padding_layout_and_kernel(monkeypatch):
|
||||
method.intermediate_size_per_partition = intermediate
|
||||
method._padded_hidden = padded_hidden
|
||||
method._padded_intermediate = padded_intermediate
|
||||
+ method.with_bias = True
|
||||
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):
|
||||
assert torch.equal(actual, expected[:, :hidden].contiguous())
|
||||
|
||||
|
||||
+def test_kimi_k3_sm120_situ_layout_and_noncontiguous_input(monkeypatch):
|
||||
+ if not torch.cuda.is_available():
|
||||
+ pytest.skip("CUDA required")
|
||||
+ if torch.cuda.get_device_capability() != (12, 0):
|
||||
+ pytest.skip("SM120 required")
|
||||
+ pytest.importorskip("flashinfer.fused_moe")
|
||||
+
|
||||
+ from flashinfer import block_scale_interleave, mxfp8_quantize
|
||||
+ from flashinfer.fused_moe import cutlass_fused_moe
|
||||
+ from flashinfer.fused_moe.core import ActivationType
|
||||
+
|
||||
+ import sglang.srt.layers.moe.moe_runner.flashinfer_cutlass as runner_module
|
||||
+ from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
+ from sglang.srt.layers.moe.moe_runner.runner import MoeRunner
|
||||
+ from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
|
||||
+ from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
+ from sglang.srt.layers.moe.utils import MoeRunnerBackend
|
||||
+ from sglang.srt.layers.quantization.mxfp4 import Mxfp4MoEMethod
|
||||
+
|
||||
+ if not runner_module.flashinfer_cutlass_supports_situ():
|
||||
+ pytest.skip("FlashInfer CUTLASS SiTU support required")
|
||||
+
|
||||
+ monkeypatch.setattr(
|
||||
+ runner_module, "use_symmetric_memory", lambda *args, **kwargs: nullcontext()
|
||||
+ )
|
||||
+ monkeypatch.setattr(runner_module, "is_allocation_symmetric", lambda: False)
|
||||
+ monkeypatch.setattr(runner_module, "get_tp_group", lambda: None)
|
||||
+
|
||||
+ num_experts, hidden, intermediate = 4, 256, 256
|
||||
+ w13, w2, w13_scale, w2_scale = _random_weights(num_experts, hidden, intermediate)
|
||||
+ gate, up = w13.chunk(2, dim=1)
|
||||
+ gate_scale, up_scale = w13_scale.chunk(2, dim=1)
|
||||
+ generator = torch.Generator(device="cuda").manual_seed(3)
|
||||
+ layer = SimpleNamespace(
|
||||
+ w13_weight=torch.nn.Parameter(w13.view(torch.uint8), requires_grad=False),
|
||||
+ w2_weight=torch.nn.Parameter(w2.view(torch.uint8), requires_grad=False),
|
||||
+ w13_weight_scale=torch.nn.Parameter(
|
||||
+ w13_scale.view(torch.uint8), requires_grad=False
|
||||
+ ),
|
||||
+ w2_weight_scale=torch.nn.Parameter(
|
||||
+ w2_scale.view(torch.uint8), requires_grad=False
|
||||
+ ),
|
||||
+ # CUDA weight creation keeps zero placeholders even for bias-free models.
|
||||
+ w13_weight_bias=torch.nn.Parameter(
|
||||
+ torch.zeros(
|
||||
+ num_experts,
|
||||
+ 2 * intermediate,
|
||||
+ dtype=torch.bfloat16,
|
||||
+ device="cuda",
|
||||
+ ),
|
||||
+ requires_grad=False,
|
||||
+ ),
|
||||
+ w2_weight_bias=torch.nn.Parameter(
|
||||
+ torch.zeros(num_experts, hidden, dtype=torch.bfloat16, device="cuda"),
|
||||
+ requires_grad=False,
|
||||
+ ),
|
||||
+ num_local_experts=num_experts,
|
||||
+ moe_tp_size=1,
|
||||
+ moe_tp_rank=0,
|
||||
+ moe_ep_size=1,
|
||||
+ moe_ep_rank=0,
|
||||
+ )
|
||||
+
|
||||
+ method = Mxfp4MoEMethod.__new__(Mxfp4MoEMethod)
|
||||
+ method._fi_kernel = "cutlass_sm120"
|
||||
+ method.num_experts = num_experts
|
||||
+ method.hidden_size = hidden
|
||||
+ method.intermediate_size_per_partition = intermediate
|
||||
+ method._padded_hidden = hidden
|
||||
+ method._padded_intermediate = intermediate
|
||||
+ method.with_bias = False
|
||||
+ config = MoeRunnerConfig(
|
||||
+ num_experts=num_experts,
|
||||
+ num_local_experts=num_experts,
|
||||
+ hidden_size=hidden,
|
||||
+ intermediate_size_per_partition=intermediate,
|
||||
+ top_k=2,
|
||||
+ activation="situ",
|
||||
+ is_gated=True,
|
||||
+ gemm1_alpha=4.0,
|
||||
+ gemm1_clamp_limit=25.0,
|
||||
+ gate_up_interleaved=False,
|
||||
+ )
|
||||
+ method.moe_runner_config = config
|
||||
+ method.runner = MoeRunner(MoeRunnerBackend.FLASHINFER_MXFP4, config)
|
||||
+ method._process_weights_for_sm120_cutlass(layer)
|
||||
+
|
||||
+ expected_w13 = torch.cat((up, gate), dim=1).view(torch.uint8)
|
||||
+ expected_w13_scale = block_scale_interleave(
|
||||
+ torch.cat((up_scale.view(torch.uint8), gate_scale.view(torch.uint8)), dim=1)
|
||||
+ ).reshape_as(layer.w13_weight_scale)
|
||||
+ assert torch.equal(layer.w13_weight, expected_w13)
|
||||
+ assert torch.equal(layer.w13_weight_scale, expected_w13_scale)
|
||||
+ assert layer.swiglu_alpha is None
|
||||
+ assert layer.swiglu_beta is None
|
||||
+ assert layer.swiglu_limit is None
|
||||
+ assert torch.all(layer.situ_beta == 4.0)
|
||||
+ assert torch.all(layer.situ_linear_beta == 25.0)
|
||||
+ assert runner_module._activation_type(config) == ActivationType.Situ
|
||||
+
|
||||
+ x = (
|
||||
+ torch.randn(
|
||||
+ hidden,
|
||||
+ 8,
|
||||
+ dtype=torch.bfloat16,
|
||||
+ device="cuda",
|
||||
+ generator=generator,
|
||||
+ )
|
||||
+ .t()
|
||||
+ * 0.1
|
||||
+ )
|
||||
+ assert not x.is_contiguous()
|
||||
+ logits = torch.randn(
|
||||
+ 8,
|
||||
+ num_experts,
|
||||
+ dtype=torch.float32,
|
||||
+ device="cuda",
|
||||
+ generator=generator,
|
||||
+ )
|
||||
+ topk_weights, topk_ids = torch.topk(torch.softmax(logits, dim=-1), 2, dim=-1)
|
||||
+ topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
|
||||
+ dispatch_output = StandardDispatchOutput(
|
||||
+ x,
|
||||
+ None,
|
||||
+ StandardTopKOutput(topk_weights, topk_ids.to(torch.int32), logits),
|
||||
+ )
|
||||
+ actual = method._apply_sm120_cutlass(layer, dispatch_output).hidden_states
|
||||
+
|
||||
+ x_quant, x_scale = mxfp8_quantize(
|
||||
+ x.contiguous(), is_sf_swizzled_layout=True, alignment=32
|
||||
+ )
|
||||
+ expected = torch.empty_like(x, memory_format=torch.contiguous_format)
|
||||
+ cutlass_fused_moe(
|
||||
+ input=x_quant,
|
||||
+ token_selected_experts=topk_ids.to(torch.int32),
|
||||
+ token_final_scales=topk_weights,
|
||||
+ fc1_expert_weights=layer.w13_weight.view(torch.int64),
|
||||
+ fc2_expert_weights=layer.w2_weight.view(torch.int64),
|
||||
+ output_dtype=torch.bfloat16,
|
||||
+ quant_scales=[
|
||||
+ layer.w13_weight_scale.view(torch.int32),
|
||||
+ layer.mxfp4_weight_global_scale,
|
||||
+ layer.w2_weight_scale.view(torch.int32),
|
||||
+ layer.mxfp4_weight_global_scale,
|
||||
+ ],
|
||||
+ input_sf=x_scale,
|
||||
+ situ_beta=layer.situ_beta,
|
||||
+ situ_linear_beta=layer.situ_linear_beta,
|
||||
+ use_w4_group_scaling=False,
|
||||
+ use_mxfp8_act_scaling=True,
|
||||
+ activation_type=ActivationType.Situ,
|
||||
+ tune_max_num_tokens=8,
|
||||
+ output=expected,
|
||||
+ )
|
||||
+ assert torch.equal(actual, expected)
|
||||
+
|
||||
+
|
||||
if __name__ == "__main__":
|
||||
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
|
||||
--- /dev/null
|
||||
+++ b/test/registered/unit/layers/test_attn_residual_arch_gate.py
|
||||
@@ -0,0 +1,43 @@
|
||||
+"""Architecture gate tests for Kimi-K3's fused attention residual."""
|
||||
+
|
||||
+from unittest.mock import patch
|
||||
+
|
||||
+import pytest
|
||||
+
|
||||
+import sglang.srt.layers.attn_residual as attn_residual
|
||||
+
|
||||
+
|
||||
+@pytest.mark.parametrize(
|
||||
+ ("capability", "expected"),
|
||||
+ [
|
||||
+ ((10, 0), True),
|
||||
+ ((10, 3), True),
|
||||
+ ((11, 0), True),
|
||||
+ ((12, 0), False),
|
||||
+ ((9, 0), False),
|
||||
+ ],
|
||||
+)
|
||||
+def test_fast_attn_residual_arch_gate(capability, expected):
|
||||
+ with (
|
||||
+ patch.object(attn_residual, "_FAST_SUPPORTED", None),
|
||||
+ patch.object(attn_residual, "is_npu", return_value=False),
|
||||
+ patch.object(
|
||||
+ attn_residual.torch.cuda,
|
||||
+ "get_device_capability",
|
||||
+ return_value=capability,
|
||||
+ ),
|
||||
+ ):
|
||||
+ assert attn_residual._use_fast(7168) is expected
|
||||
+
|
||||
+
|
||||
+def test_fast_attn_residual_requires_kimi_hidden_size():
|
||||
+ with (
|
||||
+ patch.object(attn_residual, "_FAST_SUPPORTED", None),
|
||||
+ patch.object(attn_residual, "is_npu", return_value=False),
|
||||
+ patch.object(
|
||||
+ attn_residual.torch.cuda,
|
||||
+ "get_device_capability",
|
||||
+ return_value=(10, 0),
|
||||
+ ),
|
||||
+ ):
|
||||
+ assert not attn_residual._use_fast(4096)
|
||||
@ -112,3 +112,24 @@ The candidate image preserves the Kimi base image's dependency set. It installs
|
||||
the patched FlashInfer wheel/AOT libraries and applies a minimal diff against
|
||||
the image's exact two SGLang source files; it does not replace the complete
|
||||
SGLang Python tree or require a newer `sglang-kernel` package.
|
||||
|
||||
## Audited software identity and upstream boundary
|
||||
|
||||
The accepted run used the following immutable candidate identity:
|
||||
|
||||
```text
|
||||
image: local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5
|
||||
image ID: sha256:0b8e8ddceccecc0daf10b59ddd4177036affc81facc8bf87ea9e2456ab1e96b9
|
||||
SGLang: 0.5.16
|
||||
flashinfer-python: 0.6.18 (candidate source wheel)
|
||||
flashinfer-cubin: 0.6.15.post1
|
||||
flashinfer-jit-cache: 0.6.15.post1+cu130
|
||||
Kimi SiTU: beta=4.0, linear_beta=25.0
|
||||
```
|
||||
|
||||
These results remain valid for that candidate. They do not mean stock
|
||||
FlashInfer 0.6.17 supports the complete Kimi SM120 combination. The CUTLASS
|
||||
SiTU kernel is being upstreamed in FlashInfer PR #4460. The final SGLang change
|
||||
depends on that API and retains only Kimi layout/parameter integration,
|
||||
noncontiguous input handling, the SM120 attention-residual guard, and tests.
|
||||
See `../kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/UPSTREAM_DUPLICATION_AUDIT.md`.
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user