diff --git a/README.md b/README.md index 6480177..5969fe5 100644 --- a/README.md +++ b/README.md @@ -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 全矩阵。 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 badba6d..a1f53a9 100644 --- a/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md @@ -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。 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 new file mode 100644 index 0000000..e47aac9 --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/SGLANG_DRAFT_PR.md @@ -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: + +- +- +- 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 new file mode 100644 index 0000000..10d3f4f --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/UPSTREAM_DUPLICATION_AUDIT.md @@ -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 服务级验收 | + +关键链接： + +- +- +- +- +- + +## 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 后通过，并非功能回归。 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 6540491..0947adc 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 @@ -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 暴露的兼容问题 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 new file mode 100644 index 0000000..3f733d8 --- /dev/null +++ b/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch @@ -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) diff --git a/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md index d99800c..e885e35 100644 --- a/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md +++ b/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md @@ -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`.