# Kimi-K3 SM120 FlashInfer MXFP4 MoE

## 目标

为 Kimi-K3 在 NVIDIA SM120（RTX 6000D）上的 SGLang 部署补齐
FlashInfer MXFP8 activation × MXFP4 weight fused MoE 支持，并形成：

- 可上游审查的最小 SGLang 改动；
- 精确覆盖 Kimi-K3 的正确性测试；
- RTX 6000D 单卡与四节点真实服务证据；
- Marlin 与 FlashInfer 的同口径性能结论。

本目录不修改已有 Kimi-K3 基线实验，也不提交与 FlashInfer 上游重复的 kernel。

## 交付状态

状态：**完成，SGLang Draft 代码与审计分支已推送。**

| 项目 | 结果 |
|---|---|
| GPU | NVIDIA RTX 6000D，SM120 |
| 模型 | Kimi-K3，MXFP4，SiTU `(4.0, 25.0)` |
| SGLang base | `593b1a9b8aeedb908ee21870c02a0fd1a3dc30a5` |
| SGLang Draft HEAD | `b95b5341ab9555ba5ea53f5828f6de203d749cce` |
| GitHub 分支 | `ACEEE-1222/sglang:hzy/kimi-k3-sm120-flashinfer-mxfp4` |
| FlashInfer 依赖 | PR #4460，implementation `b525c513fef7` |
| 默认行为 | 不改变默认 MoE backend，Marlin 仍可用 |

公开 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)

逐项验收：[`COMPLETION_AUDIT.md`](COMPLETION_AUDIT.md)

## 交付内容

### SGLang 实现

1. 为 Kimi-K3 在 SM120 上接入 FlashInfer CUTLASS MXFP4 MoE；
2. 将 Kimi 连续 `[gate; up]` 权重与 scale 转换为 CUTLASS 所需布局；
3. 独立透传 `situ_beta=4.0` 与 `situ_linear_beta=25.0`；
4. 在 MXFP8 量化前处理非连续 hidden states；
5. 排除 SM120 误入 tcgen05 attention-residual 路径；
6. 保持 GPT-OSS 路径、旧 FlashInfer import 和默认 Marlin 行为不变。

变更文件：

```text
python/sglang/srt/layers/attn_residual.py
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
python/sglang/srt/layers/quantization/mxfp4.py
test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py
test/registered/unit/layers/test_attn_residual_arch_gate.py
```

### 测试覆盖

- Kimi gate/up 权重和 scale 布局；
- SiTU `(4.0, 25.0)` 参数；
- MXFP8 activation × MXFP4 weight；
- 非连续输入；
- standard routing 与 EP expert IDs；
- eager、CUDA Graph capture/replay；
- SM120 attention-residual architecture gate；
- 缺少新 API 时的兼容性与明确报错。

## 上游边界

FlashInfer PR #4460 负责公共 CUTLASS SiTU kernel 和参数 API；本工作不提交
平行的 FlashInfer kernel PR。

SGLang Draft 只保留 Kimi-K3 的真实框架增量：权重/scale 布局、SiTU 参数、
非连续输入、SM120 guard 和测试。它不 vendor kernel、不 pin 未合并私人 commit、
不改变默认 backend。

完整去重审计见：
[`UPSTREAM_DUPLICATION_AUDIT.md`](UPSTREAM_DUPLICATION_AUDIT.md)

## 正确性结果

使用 FlashInfer #4460 构建的 wheel，在 RTX 6000D GPU6 上执行精确 Kimi 组合：

```text
8 passed, 17 warnings in 559.87s
```

覆盖 `SM120 + MXFP8 + MXFP4 + SiTU(4,25)`、Kimi 布局、非连续输入和
CUDA Graph。Kimi 真尺寸结果：

| Case | cosine | NRMSE |
|---|---:|---:|
| Decode M=1 | 0.9996223 | 2.748% |
| Prefill M=128 | 0.9996152 | 2.774% |

最终 Draft HEAD 的 compileall 和非 JIT 回归：

```text
10 passed, 3 deselected, 17 warnings in 33.30s
```

五个变更文件已通过完整 pre-commit。同步 fork 后的 7 个提交与同步前版本在
`git range-diff` 中全部 patch-equivalent。

原始日志：

```text
/data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log
/data/hzy/phase5_sglang_pr4460_draft_cpu_final.log
```

## 四节点真实 Serving

### 实验口径

| 维度 | 取值 |
|---|---|
| 节点 | 601-604，共 32 张 RTX 6000D |
| 并行 | TP32 / EP4 / DP1 |
| 输入/输出 | 16K / 1 token |
| 并发 | C=8、C=16 |
| Backend | Marlin、FlashInfer MXFP4 |
| Chunk | 4K、8K、16K |
| 请求与重复 | 每点 40 请求，3 次重复 |
| Cache | 禁用 Radix Cache |

正式 Run：

```text
8K/16K: kimi3-ep4-moe-full-20260818-151349
4K:     kimi3-ep4-moe-chunk4k-20260819-150632
```

36/36 repeat 均完成 40/40 请求，输出长度均为 1，`errors` 为空；日志未发现
OOM、Traceback、NCCL error 或 EngineDeadError。下表为三次重复的中位数。

### 实验结果

| Chunk | C | Backend | Input TPS | TTFT P50 | TTFT P95 | FlashInfer vs Marlin |
|---:|---:|---|---:|---:|---:|---:|
| 4K | 8 | Marlin | 2694.53 | 48.86 s | 49.20 s | baseline |
| 4K | 8 | FlashInfer | **2957.23** | **44.48 s** | **44.86 s** | +9.75% TPS / -8.96% TTFT |
| 4K | 16 | Marlin | 2692.60 | 96.25 s | 97.96 s | baseline |
| 4K | 16 | FlashInfer | **2960.41** | **87.61 s** | **89.04 s** | +9.95% TPS / -8.98% TTFT |
| 8K | 8 | Marlin | 2935.02 | 43.51 s | 46.22 s | baseline |
| 8K | 8 | FlashInfer | **3257.96** | **39.19 s** | **41.62 s** | +11.00% TPS / -9.92% TTFT |
| 8K | 16 | Marlin | 2934.90 | 87.09 s | 89.86 s | baseline |
| 8K | 16 | FlashInfer | **3260.14** | **78.39 s** | **80.88 s** | +11.08% TPS / -9.99% TTFT |
| 16K | 8 | Marlin | 2812.91 | 45.47 s | 45.86 s | baseline |
| 16K | 8 | FlashInfer | **3027.50** | **42.27 s** | **42.63 s** | +7.63% TPS / -7.04% TTFT |
| 16K | 16 | Marlin | 2812.33 | 90.97 s | 91.36 s | baseline |
| 16K | 16 | FlashInfer | **3027.34** | **84.52 s** | **84.85 s** | +7.65% TPS / -7.09% TTFT |

### 结果结论

1. FlashInfer 在 6 个正式 shape 上均优于 Marlin；
2. 4K Chunk 的收益约为 `+9.8% Input TPS / -9.0% TTFT`；
3. 8K Chunk 的收益约为 `+11% Input TPS / -10% TTFT`；
4. 16K Chunk 的收益约为 `+7.6% Input TPS / -7.1% TTFT`；
5. 同一 backend 下，8K Chunk 同时优于 4K 和 16K；4K 切分轮次过多，16K 单步过重；
6. C 从 8 增至 16 后吞吐基本不变、TTFT 约翻倍，C=8 已接近 Prefill 吞吐饱和；
7. 当前推荐为 `EP4 + flashinfer_mxfp4 + chunked_prefill_size=8192`。

该优化有效但不是数量级变化。最佳 TTFT P50 仍约 39 秒，后续需要继续优化
TP/EP 通信与 Prefill 并行结构。

### Chunk=4K 补测结论

4K 补测复用了 8K/16K 矩阵的镜像、模型、服务参数和入口。12/12 repeat
全部有效，共 480/480 请求成功。

FlashInfer 在 C=8/16 下仍分别提升 `9.75%/9.95%` Input TPS，但 4K 的绝对
Input TPS 比 8K 低约 9%，TTFT 也更高。该结果补齐了 Chunk 曲线并确认 8K
不是网格边界造成的偶然最优点。

## 最终 PR 代表点

最终又用等价 SGLang patch `300c87a` 和 FlashInfer #4460 implementation
`b525c513` 重跑 8K/C8 代表点：

```text
Run ID: kimi3-sm120-pr-representative-20260819-115530
TP32 / EP4 / DP1
ISL=16384 / OSL=1 / C=8 / requests=40 / chunk=8192
Marlin 3 repeats + FlashInfer MXFP4 3 repeats
```

| Backend | Input TPS | TTFT P50 | TTFT P95 |
|---|---:|---:|---:|
| Marlin | 2930.03 | 43.587 s | 46.287 s |
| FlashInfer MXFP4 | 3251.25 | 39.267 s | 41.694 s |
| 相对变化 | +10.96% | -9.91% | -9.92% |

六轮均为 40/40 成功，所有 `output_lens=1`，无 OOM、Traceback、NCCL error
或 EngineDeadError。

## 复现命令

四节点矩阵只在 601 执行：

```bash
cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill

read -rsp 'sudo password: ' SUDO_PASSWORD; echo
export SUDO_PASSWORD
export RUN_ID=kimi3-moe-prefill-$(date +%Y%m%d-%H%M%S)

bash run_moe_backend_prefill.sh run
```

可通过以下环境变量缩小矩阵：

```text
EP_SIZE
BACKENDS_CSV
CHUNKED_PREFILL_SIZES_CSV
CONCURRENCIES_CSV
REPEATS
```

4K 补测使用：

```bash
EP_SIZE=4 \
BACKENDS_CSV=marlin,flashinfer_mxfp4 \
CHUNKED_PREFILL_SIZES_CSV=4096 \
CONCURRENCIES_CSV=8,16 \
REPEATS=3 \
bash run_moe_backend_prefill.sh run
```

## 证据路径

### 8K/16K 正式矩阵

```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

results/phase5_ep4_moe_backend_prefill_summary.json
results/phase5_ep4_moe_backend_prefill_results.csv
```

### 4K 补测

```text
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/results/kimi3-ep4-moe-chunk4k-20260819-150632/
/data/hzy/kimi3-ep4-moe-chunk4k-20260819-150632.log

results/phase5_ep4_moe_backend_chunk4k_summary.json
results/phase5_ep4_moe_backend_chunk4k_results.csv
```

### 最终 PR 代表点

```text
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/results/kimi3-sm120-pr-representative-20260819-115530/
/data/hzy/kimi3-sm120-pr-representative-20260819-115530.log

results/pr_representative_16k_c8_chunk8k_summary.json
results/pr_representative_16k_c8_chunk8k_results.csv
```

### 单卡测试与 Profiler

```text
results/phase3_target_situ_all_gpu6_v3.json
results/phase4_mscan_ep32_5rounds_gpu6.json
results/phase4_nsys_ab_gpu6.nsys-rep
results/phase4_ncu_2026_2_1_ab_gpu6.ncu-rep
```

## 后续方向

MoE backend 选型完成后，不再扩大同类 backend 网格。后续 Prefill 优化按独立
任务推进：

1. TP Reduce Scatter / Attention TP input scattered；
2. MoE A2A；
3. Prefill Pipeline Parallelism；
4. 对最优新配置重新做 Nsight 时间线验证。