Compare commits
38 Commits
main
...
hzy-kimi-k
| Author | SHA1 | Date | |
|---|---|---|---|
|
a58b931cc8 | ||
|
|
0433fcc3ee | ||
|
|
d790df39b2 | ||
|
|
96ffea6d37 | ||
|
|
d7381abe84 | ||
|
|
60f77cd4ef | ||
|
|
6fac5ad567 | ||
|
|
2a3b12fa78 | ||
|
|
08a35066d7 | ||
|
|
74ec19dd48 | ||
|
|
0fdcab9927 | ||
|
|
a5248ed80e | ||
|
|
fc336a3c7b | ||
|
|
63f2327a90 | ||
|
|
a90c898683 | ||
|
|
ab9a5422f6 | ||
|
|
7f67dfe6b3 | ||
|
|
e8ff3ce1e8 | ||
|
|
d28db48e4b | ||
|
|
39f692caae | ||
|
|
a9206ff105 | ||
|
|
a11c80b703 | ||
|
|
e01df16667 | ||
|
|
ec7b604a50 | ||
|
|
27b8be09cb | ||
|
|
be9d6bfe3a | ||
|
|
d0863501ca | ||
|
|
13944079fa | ||
|
|
e3974e2352 | ||
|
|
b50de8fe99 | ||
|
|
c14f8aa43a | ||
|
|
daeffd147b | ||
|
|
5454fb984e | ||
|
|
c8f30ab7dc | ||
|
|
6493798ad5 | ||
|
|
a1c18d736b | ||
|
|
dac1bb652d | ||
|
|
0684d269df |
28
README.md
28
README.md
@ -1,5 +1,31 @@
|
||||
# sskj — 多平台大模型推理性能基准测试项目
|
||||
|
||||
**更新（2026-08-20 10:13:01 CST）**：完成 Kimi-K3 Prefill TP Reduce Scatter 可行性审计并停止该方向。K3 的 MLA 输出门控仍依赖完整 7168 维 hidden，且 69/93 层为 KDA；保留或恢复 gate hidden 后，原 hidden All-Reduce 无法消除并新增 2112 维 latent All-Gather，估算通信量反增约 14.7%。研究原型仅保留为否决证据，不进入四节点实验或上游 PR；后续转向 MoE A2A 与 Pipeline Parallelism。详见 `experiments/pro6000/kimi3_pro6000_sglang_tp_reduce_scatter_prefill/README.md`。
|
||||
|
||||
**更新（2026-08-19 16:51:43 CST）**：完成 Kimi-K3 / RTX 6000D / SGLang Prefill MoE backend 交付文档收敛与 Chunk=4K 补测。四节点 TP32×EP4、16K→1、C=8/16 下 12/12 重复、480/480 请求成功；FlashInfer 相对 Marlin 提升约 9.8% Input TPS、降低约 9.0% TTFT。结合 4K/8K/16K 全部结果，最终推荐保持 `flashinfer_mxfp4 + chunked_prefill_size=8192`。
|
||||
|
||||
**更新（2026-08-19 14:23:52 CST）**：601 的 GitLab SSH key 已生效，内部 Kimi-K3 SM120 审计分支已推送至 `qqtang/sskj:hzy-kimi-k3-sm120-flashinfer-mxfp4`。远端已有普通分支 `hzy`，因此不能创建 `hzy/...` ref，最终采用无层级冲突的分支名；GitHub SGLang feature 分支保持不变。
|
||||
|
||||
**更新（2026-08-19 13:58:50 CST）**：用户同步 GitHub fork 后，将 Kimi-K3 SM120 SGLang 的 7 个提交无冲突重放到最新 `main` `593b1a9`，`range-diff` 逐提交确认补丁等价；最终 Draft HEAD `b95b534` 已推送至 `ACEEE-1222/sglang:hzy/kimi-k3-sm120-flashinfer-mxfp4`。在 601 统一验证镜像中重跑最终 HEAD 的语法与定向回归，结果为 `10 passed, 3 deselected`；代表性四机性能数据保留，不重复消耗 GPU。
|
||||
|
||||
**更新（2026-08-19 13:07:48 CST）**：完成 Kimi-K3 SM120 SGLang Draft PR 的最终代表点验收。基于 SGLang `300c87a` 与 FlashInfer #4460 `b525c51`，在 601-604 四节点 TP32×EP4 上完成 16K→1、C=8、Chunk=8K 的 Marlin/FlashInfer 各 3 次 A/B；240/240 请求成功。FlashInfer 的 Input TPS 中位数提升 10.96%，TTFT P50/P95 分别下降 9.91%/9.92%。公开 PR Checklist 已全部勾选，轻量汇总和原始证据索引已归档。
|
||||
|
||||
**更新（2026-08-19 10:45:18 CST）**：新增 Kimi-K3 SM120 SGLang Draft PR 的单点代表性验收入口。使用精确 SGLang `300c87a`、FlashInfer #4460 实现 `b525c51` 和统一镜像，在 601-604 四节点 TP32×EP4 上仅复测 16K→1、C=8、Chunk=8K 的 Marlin/FlashInfer A/B，各 3 次重复；完整命令、镜像身份、原始日志和汇总统一落入单个 Run 目录。
|
||||
|
||||
**更新（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 全矩阵。
|
||||
|
||||
**更新（2026-08-18 13:03:33 CST）**：Phase 5 候选镜像改为针对 Kimi 基础镜像原生 SGLang 源码应用最小兼容补丁，仅补齐 SM120 FlashInfer MXFP4 backend 分发、权重布局与 SiTU 激活映射，避免整体替换新版 Python 包造成 `sglang-kernel` 版本错配。
|
||||
|
||||
**更新（2026-08-18 12:24:19 CST）**：新增 Kimi-K3 / 601-604 / SGLang TP32×EP32 的真实 Prefill MoE backend 验收实验；固定比较 Marlin 与 FlashInfer MXFP4 在 16K→1、C=8/16、Chunk=8K/16K 下的 TTFT、E2E 与 Input TPS，并提供四节点一致候选镜像构建、原始证据和自动汇总，详见 `experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/README.md`。
|
||||
|
||||
**更新（2026-08-17 11:56:47 CST）**：完成 Kimi-K3 / RTX PRO 6000D / SGLang FlashInfer MXFP4 MoE Phase 2 correctness harness；SwiGLU 控制矩阵 8/8 通过，覆盖 edge routing、EP global ID、NaN/Inf、确定性、CUDA Graph 和 Kimi 真尺寸，SiTU 目标矩阵 8/8 稳定复现同一 C++ 缺口，详见 `experiments/pro6000/kimi3_pro6000_sglang_sm120_flashinfer_mxfp4/README.md`。
|
||||
|
||||
> 历史更新见 `git log`。项目目的与工作流见下方。
|
||||
> **项目目的**：当新显卡（GPU/NPU）到货时，用最短时间完成大模型在该平台上的推理性能评估与部署配置选型。
|
||||
> 当前模型：DeepSeek-V4-Flash（FP8 / INT8）；后续接入 GLM5.2，**完全复用**本项目的实验与报告流程。
|
||||
@ -45,6 +71,8 @@
|
||||
| `experiments/pro6000/dsv4_pro6000_vllm_tp_dp_matrix/` | RTX 6000D + vLLM |
|
||||
| `experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/` | RTX 6000D + SGLang |
|
||||
| `experiments/pro6000/kimi3_pro6000_sglang_tp32ep32/` | RTX 6000D + SGLang，Kimi-K3（TP32×EP32，部署手册见 docs/KIMI_K3_DEPLOY.md） |
|
||||
| `experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/` | RTX 6000D + SGLang，Kimi-K3 真实 Prefill 的 Marlin/FlashInfer MXFP4 与 Chunk 8K/16K 对照 |
|
||||
| `experiments/pro6000/kimi3_pro6000_sglang_tp_reduce_scatter_prefill/` | Kimi-K3 Prefill TP Reduce Scatter 可行性审计；因 MLA gate 与 KDA 依赖完整 hidden，结论为不推荐继续 |
|
||||
| `experiments/pro6000/kimi3_pro6000_pd_rdma/` | RTX 6000D + SGLang，Kimi-K3 PD 分离（MoonCake RDMA，8 节点，见 deploy_pd.sh + docs/KIMI_K3_DEPLOY.md 附录 B） |
|
||||
|
||||
每个目录内：`run_bench.sh` 跑固定并发矩阵；`run_adaptive_concurrency.sh` 从 C=1 指数倍增搜饱和点；`run_adaptive_concurrency_add16.sh` 从 C=16 线性 +16 步进、带 TTFT SLO 停止与回退（当前主力用法，见 `experiments/ADAPTIVE_CONCURRENCY_USAGE.md`）。
|
||||
|
||||
@ -0,0 +1,454 @@
|
||||
# Kimi K3 推理优化 - 优化 Prefill - Deep PP
|
||||
|
||||
## 1. 工作背景
|
||||
|
||||
Kimi-K3 在 4 台 RTX 6000D、共 32 张 GPU 上部署时，最直接的方式是
|
||||
`PP1 / TP32 / EP4`：完整模型只有一个 Pipeline Stage，Tensor Parallel
|
||||
覆盖全部 32 张卡，MoE Expert Parallel 使用 4 张卡。
|
||||
|
||||
这套配置可以正确运行，但 16K 长输入的 Prefill 延迟很高。既有 Nsight
|
||||
Systems Profile 显示，在 `16K -> 1、C=8、chunk=8K` 场景中：
|
||||
|
||||
- 稳定 GPU 的 Prefill 窗口中，NCCL 暴露时间占比中位数为 **57.27%**；
|
||||
- 单个 Profile 窗口约有 **557 次 AllReduce**；
|
||||
- AllReduce p50 约为 **7.10 ms**；
|
||||
- 通信与计算重叠率接近 **0%**；
|
||||
- 其中 MoE 相关 AllReduce 约占已分类 AllReduce 时间的 **58.73%**。
|
||||
|
||||
这说明 PP1/TP32 并不只是“32 张卡一起算得更快”。Kimi-K3 每层中频繁出现
|
||||
的 TP Collective 会跨越 4 台服务器，在没有 NVLink 的 RTX 6000D 上通过
|
||||
RoCE 执行；等待通信的时间已经超过实际计算时间。
|
||||
|
||||
## 2. 为什么尝试 Deep PP
|
||||
|
||||
本实验所说的 Deep PP，是把一个模型划分为更多 Pipeline Stage，同时相应
|
||||
缩小每个 Stage 内的 TP Group：
|
||||
|
||||
```text
|
||||
总 GPU 数 = PP × TP = 32
|
||||
```
|
||||
|
||||
在固定 32 张 GPU 的前提下，本实验考察：
|
||||
|
||||
| 配置 | 每个 Stage 使用的 GPU | Stage 与节点的关系 |
|
||||
|---|---:|---|
|
||||
| PP1 / TP32 | 32 | 一个 Stage 跨 4 个节点 |
|
||||
| PP2 / TP16 | 16 | 一个 Stage 跨 2 个节点 |
|
||||
| PP4 / TP8 | 8 | 一个 Stage 完整落在 1 个节点 |
|
||||
| PP8 / TP4 | 4 | 每个节点放置 2 个 Stage |
|
||||
| PP16 / TP2 | 2 | 每个节点放置 4 个 Stage |
|
||||
|
||||
PP 加深后会新增 Pipeline Activation 的 Stage-to-Stage 传输，也会产生流水线
|
||||
填充和排空造成的 bubble，因此 PP 并不是越大越快。它可能带来收益的原因是：
|
||||
|
||||
1. 每层高频 TP AllReduce 的通信域从 32 卡逐步缩小到 16、8、4 卡；
|
||||
2. PP4 开始，TP Group 完整落在单机内，不再让每次 TP AllReduce 都跨节点；
|
||||
3. PP8 中 TP4/EP4 都位于单机，跨节点通信主要集中在少数 Pipeline 边界；
|
||||
4. 只要 C=8/16 提供了足够多的请求，Scheduler 就可能用不同请求交错填充
|
||||
Pipeline，抵消一部分 bubble。
|
||||
|
||||
因此，这一阶段要回答的不是“PP 能不能启动”，而是：在 Kimi-K3、6000D
|
||||
和固定 Prefill 压力下，缩小 TP 通信域的收益能否超过 Pipeline 开销，以及
|
||||
拐点出现在 PP2、PP4、PP8 还是 PP16。
|
||||
|
||||
## 3. 实现方式
|
||||
|
||||
### 3.1 唯一实验入口
|
||||
|
||||
实验只保留一个入口脚本：
|
||||
|
||||
```text
|
||||
run_pp_baseline_search.sh
|
||||
```
|
||||
|
||||
脚本负责：
|
||||
|
||||
1. 检查 601-604 的模型、镜像与 `/dev/infiniband`；
|
||||
2. 依次启动 PP2/TP16、PP4/TP8、PP8/TP4、PP16/TP2 服务；
|
||||
3. 通过 `/get_server_info` 校验实际 PP、TP、EP 和 MoE backend；
|
||||
4. 每个配置执行 C=8/16，各重复 3 次；
|
||||
5. 校验每轮是否完成 40/40 请求以及是否存在请求错误；
|
||||
6. 保存完整 Docker 命令、服务日志、benchmark 日志、原始 JSONL 和 GPU 快照；
|
||||
7. 生成逐轮 `results.csv` 和中位数 `summary.csv`；
|
||||
8. 无论成功或失败，最终清理 4 台服务器上的实验容器。
|
||||
|
||||
### 3.2 SGLang PP 路径
|
||||
|
||||
本轮使用镜像：
|
||||
|
||||
```text
|
||||
local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5
|
||||
```
|
||||
|
||||
服务直接通过 SGLang 官方参数启动：
|
||||
|
||||
```bash
|
||||
python3 -m sglang.launch_server \
|
||||
--model-path /data/hf_models/Kimi-K3 \
|
||||
--tp-size "${TP_SIZE}" \
|
||||
--pp-size "${PP_SIZE}" \
|
||||
--ep-size "${EP_SIZE}" \
|
||||
--nnodes 4 \
|
||||
--node-rank "${NODE_RANK}" \
|
||||
--dist-init-addr 174.1.60.1:20000 \
|
||||
--moe-runner-backend flashinfer_mxfp4 \
|
||||
--chunked-prefill-size 8192
|
||||
```
|
||||
|
||||
脚本默认取 `EP_SIZE=min(4, TP_SIZE)`。因此 PP2/4/8 使用 EP4，PP16 因 TP2
|
||||
必须使用 EP2，满足 SGLang 的 `TP % EP == 0` 约束。
|
||||
|
||||
实验没有复制同事旧代码，也没有注入 `patch_pp_group.py`。准备阶段曾验证过
|
||||
一版运行时 PP Group patch，但它会让 Python Object 通信和 Tensor/P2P 通信
|
||||
复用错误的 Gloo Group，出现消息大小不匹配。回到当前 SGLang 上游实现后，
|
||||
PP2、PP4 和 PP8 均能正常启动，因此该 patch 被完全移除，没有进入正式数据。
|
||||
|
||||
### 3.3 网络配置
|
||||
|
||||
4 个节点都向容器映射 `/dev/infiniband`，并使用：
|
||||
|
||||
```bash
|
||||
NCCL_SOCKET_IFNAME=bond0
|
||||
GLOO_SOCKET_IFNAME=bond0
|
||||
NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3
|
||||
NCCL_IB_GID_INDEX=3
|
||||
```
|
||||
|
||||
`bond0` 承担 Bootstrap/控制连接，`mlx5_0~3` 是 NCCL 的 RoCE HCA。正式运行
|
||||
中没有出现 NCCL、Gloo、RDMA、connection refused 或 communicator failure。
|
||||
|
||||
## 4. 实验设计
|
||||
|
||||
为了让结果只反映 PP/TP 变化，其他主要变量保持固定：
|
||||
|
||||
| 项目 | 固定值 |
|
||||
|---|---|
|
||||
| 模型 | Kimi-K3 |
|
||||
| GPU | 601-604，32 × RTX 6000D |
|
||||
| MoE runner | FlashInfer MXFP4 |
|
||||
| A2A backend | none |
|
||||
| EP | PP1/2/4/8 为 4；PP16 为 2 |
|
||||
| 输入/输出 | 16,384 / 1 token |
|
||||
| Chunked Prefill | 8,192 token |
|
||||
| 并发 | C=8、C=16 |
|
||||
| 每轮请求数 | 40 |
|
||||
| 重复次数 | 3 |
|
||||
| 数据集 | random-ids，输入长度比例固定为 1.0 |
|
||||
| Prefix Cache | `--disable-radix-cache` |
|
||||
| 显存比例 | 0.88 |
|
||||
|
||||
输出长度固定为 1，是为了让 TTFT 和 Input TPS 主要反映 Prefill，而不是被长
|
||||
Decode 混入。每个点重复 3 次，最终使用中位数，避免首轮 JIT、缓存状态或系统
|
||||
瞬时抖动主导结果。
|
||||
|
||||
PP1 没有重复跑。它来自前一阶段已经完成的 FlashInfer MXFP4、EP4、chunk8K
|
||||
实验。经逐项核对，PP1 与本轮使用相同镜像、模型、MoE backend、EP、输入输出、
|
||||
并发、请求数、显存比例和网络参数，区别只有 `PP1/TP32`，因此可以直接并入。
|
||||
|
||||
## 5. 一个关键实现差异：PP Micro Batch
|
||||
|
||||
同事旧 PP4 配置显式设置过：
|
||||
|
||||
```bash
|
||||
--pp-max-micro-batch-size 1
|
||||
```
|
||||
|
||||
当前 SGLang 在该参数未设置时，会在 Scheduler 初始化后自动计算：
|
||||
|
||||
```python
|
||||
pp_max_micro_batch_size = max(max_running_requests // pp_size, 1)
|
||||
```
|
||||
|
||||
接纳新请求时又会执行：
|
||||
|
||||
```python
|
||||
allocatable_requests = pp_max_micro_batch_size - running_batch_size
|
||||
```
|
||||
|
||||
因此这里的 `1` 不只是一个普通的 Kernel batch 参数，它会把 PP Scheduler
|
||||
允许同时运行的请求数压到 1。对于 C=8/16，这会让请求大量排队，Pipeline
|
||||
难以通过请求交错保持各 Stage 忙碌，bubble、P2P 和调度固定成本都会被放大。
|
||||
|
||||
本实验不传这个参数，使用 SGLang 当前的自动值。旧实验还混合了 Marlin、
|
||||
chunk16K、EP1/EP8、PP16、旧镜像和运行时 PP patch，因此旧数据不能作为
|
||||
“PP 天生更慢”的单变量证据。
|
||||
|
||||
## 6. 实验过程与正确性
|
||||
|
||||
正式 Run ID：
|
||||
|
||||
```text
|
||||
kimi3-pp-upstream-20260821-120208
|
||||
kimi3-pp16-20260821-161721
|
||||
```
|
||||
|
||||
服务启动时间分别为：
|
||||
|
||||
| 配置 | 健康检查通过时间 |
|
||||
|---|---:|
|
||||
| PP2 / TP16 / EP4 | 417 s |
|
||||
| PP4 / TP8 / EP4 | 405 s |
|
||||
| PP8 / TP4 / EP4 | 338 s |
|
||||
| PP16 / TP2 / EP2 | 538 s |
|
||||
|
||||
PP2、PP4、PP8 共形成 18 个正式样本；补充的 PP16 形成 6 个正式样本。
|
||||
每个样本均完成 40/40 请求，原始结果中的 errors 为空。服务和 benchmark
|
||||
日志没有发现：
|
||||
|
||||
- CUDA OOM / OutOfMemory；
|
||||
- Traceback / Exception；
|
||||
- NCCL / Gloo / RDMA failure；
|
||||
- Engine failure；
|
||||
- 请求连接失败或提前终止。
|
||||
|
||||
PP8 在 C=8 和 C=16 下三次 Input TPS 的极差都小于 0.7%，说明结果不是某
|
||||
一次运行的偶然峰值。
|
||||
|
||||
PP16 首次启动在 decode CUDA Graph 捕获时触发 custom all-reduce IPC buffer
|
||||
注册错误；关闭 PP16 的 decode graph 后，16-stage warmup 又超过默认 300 秒
|
||||
watchdog。正式 PP16 使用 `--disable-decode-cuda-graph --watchdog-timeout
|
||||
3000000`，538 秒后健康。该差异不改变本轮 Prefill 路径：输出仅 1 token，
|
||||
没有迭代 Decode；但 PP16 的 Decode 性能不能由本实验推断。PP16 三次 Input
|
||||
TPS 极差不超过 0.7%，结果同样稳定。
|
||||
|
||||
## 7. 性能结果
|
||||
|
||||
下表均为 3 次正式运行的中位数：
|
||||
|
||||
| PP / TP / EP | C | Input TPS | TTFT p50 | TTFT p95 | Input TPS 相对 PP1 | TTFT p50 相对 PP1 |
|
||||
|---|---:|---:|---:|---:|---:|---:|
|
||||
| 1 / 32 / 4 | 8 | 3,257.96 | 39.19 s | 41.62 s | 基线 | 基线 |
|
||||
| 2 / 16 / 4 | 8 | 5,015.95 | 25.22 s | 26.69 s | +53.96% | -35.65% |
|
||||
| 4 / 8 / 4 | 8 | 7,155.20 | 17.28 s | 18.62 s | +119.62% | -55.90% |
|
||||
| 8 / 4 / 4 | 8 | **8,809.50** | **13.47 s** | **16.45 s** | **+170.40%** | **-65.64%** |
|
||||
| 16 / 2 / 2 | 8 | 6,701.04 | 17.94 s | 21.39 s | +105.68% | -54.22% |
|
||||
| 1 / 32 / 4 | 16 | 3,260.14 | 78.39 s | 80.88 s | 基线 | 基线 |
|
||||
| 2 / 16 / 4 | 16 | 5,018.38 | 50.70 s | 52.14 s | +53.93% | -35.32% |
|
||||
| 4 / 8 / 4 | 16 | 7,152.58 | 34.84 s | 35.95 s | +119.39% | -55.56% |
|
||||
| 8 / 4 / 4 | 16 | **8,824.12** | **27.31 s** | **29.77 s** | **+170.67%** | **-65.16%** |
|
||||
| 16 / 2 / 2 | 16 | 6,700.96 | 35.88 s | 39.49 s | +105.54% | -54.23% |
|
||||
|
||||
从 PP1 到 PP8，Input TPS 沿着 PP2、PP4、PP8 单调上升，TTFT 也同步下降；
|
||||
PP16 则出现反转。相对 PP8，PP16 在 C=8/16 的 Input TPS 分别下降 23.93%
|
||||
和 24.06%，TTFT p50 分别增加 33.18% 和 31.37%。这说明继续把 TP4 缩到
|
||||
TP2 的收益已无法抵消 16-stage bubble、更多 P2P 边界和调度开销。
|
||||
|
||||
PP16 同时把 EP4 改为 EP2，因此它不是严格的纯 PP 单变量对照。但它准确代表
|
||||
32 卡上可实际部署的 `PP16/TP2` 组合，足以判定该组合不应替代 PP8 基线。
|
||||
|
||||
C=16 的 TTFT 大致是 C=8 的两倍，而同一 PP 下 Input TPS 基本不变。这表明
|
||||
这批 Prefill 工作已接近该配置的稳定吞吐区：增加并发主要增加排队等待，而不是
|
||||
继续提高输入吞吐。PP 改善的是服务处理这批 token 的速率，因此两个并发点都
|
||||
出现了接近一致的 TPS 增益。
|
||||
|
||||
## 8. 显存结果
|
||||
|
||||
下表取每个配置 benchmark 完成后，4 个节点 32 张卡中的最高显存：
|
||||
|
||||
| PP / TP / EP | 峰值显存 | 相对 85,651 MiB 的剩余空间 |
|
||||
|---|---:|---:|
|
||||
| 1 / 32 / 4 | 81,355 MiB | 4,296 MiB |
|
||||
| 2 / 16 / 4 | 80,335 MiB | 5,316 MiB |
|
||||
| 4 / 8 / 4 | 82,417 MiB | 3,234 MiB |
|
||||
| 8 / 4 / 4 | 79,407 MiB | 6,244 MiB |
|
||||
| 16 / 2 / 2 | 79,989 MiB | 5,662 MiB |
|
||||
|
||||
PP 的 Stage 不会获得完全相同的层、输出头和状态，因此各 PP 配置的显存不再
|
||||
像 PP1 那样均匀。PP4 的最重 Stage 已使用约 82.4 GiB；PP8 的最重 Stage
|
||||
约 79.4 GiB，在本轮反而保留了最多的最坏 Rank 余量。
|
||||
|
||||
这一点对下一阶段 A2A 很重要，因为 DeepEP 或 FlashInfer A2A 都可能申请额外
|
||||
dispatch/combine buffer。只看 TPS 而忽略最重 Rank，可能出现基线可跑、打开
|
||||
A2A 后启动或 Prefill OOM 的情况。
|
||||
|
||||
## 9. 通信结果如何解释
|
||||
|
||||
### 9.1 累计收益变大，但每次增加 PP 的边际收益在下降
|
||||
|
||||
第 7 节相对 PP1 的累计数字会让人感觉“PP 越大，收益增长越快”。逐级比较后，
|
||||
实际规律相反。以 C=8 为例：
|
||||
|
||||
| 变化 | Input TPS 边际变化 | TTFT p50 边际变化 |
|
||||
|---|---:|---:|
|
||||
| PP1 -> PP2 | +53.96% | -35.65% |
|
||||
| PP2 -> PP4 | +42.65% | -31.48% |
|
||||
| PP4 -> PP8 | +23.12% | -22.05% |
|
||||
| PP8 -> PP16 | -23.93% | +33.18% |
|
||||
|
||||
所以 PP1 到 PP8 的累计收益持续扩大，但每次把 PP 翻倍带来的新增收益已经逐步
|
||||
减小；PP16 则越过最优点。C=16 也呈现相同趋势。
|
||||
|
||||
### 9.2 节点内 PCIe 与节点间 RoCE 的实测
|
||||
|
||||
在 601--604 上使用 NCCL 2.28.9 + CUDA 13 的 `all_reduce_perf`，分别测量与
|
||||
PP Stage 对应的 TP4、TP8、TP16 和 TP32 通信域。所有测试均为 0 个错误值：
|
||||
|
||||
| 通信域 | 实际拓扑 | 128 MiB 延迟 | 128 MiB busbw | 128/256 MiB 平均 busbw |
|
||||
|---|---|---:|---:|---:|
|
||||
| TP4 | 单节点 GPU0--3，同一 PCIe/NUMA 岛 | 4.93 ms | 40.85 GB/s | 41.00 GB/s |
|
||||
| TP8 | 单节点 8 卡，跨两个 NUMA 岛 | 6.00 ms | 39.13 GB/s | 38.89 GB/s |
|
||||
| TP16 | 两节点 | 6.24 ms | 40.31 GB/s | 40.60 GB/s |
|
||||
| TP32 | 四节点 | 6.59 ms | 39.47 GB/s | 40.08 GB/s |
|
||||
|
||||
大消息带宽都在约 39--41 GB/s，说明 PP 收益不是因为节点间 400G RoCE 的
|
||||
持续带宽远低于节点内 PCIe。真正变化更明显的是 collective 的参与 Rank 数、
|
||||
固定延迟和同步范围：128 MiB AllReduce 从 TP32 的 6.59 ms 降到 TP4 的
|
||||
4.93 ms。
|
||||
|
||||
两个拓扑边界尤其重要：
|
||||
|
||||
- PP2 -> PP4 把每个 Stage 的 TP 域从两节点 TP16 缩到单节点 TP8，逐层
|
||||
AllReduce 不再跨节点；
|
||||
- PP4 -> PP8 把 TP8 缩到 TP4，每个 Stage 可落在同一个 PCIe/NUMA 岛，避免
|
||||
单节点内跨 CPU Root Complex，同时每卡得到更大的 GEMM 分片，减轻 TP32
|
||||
过度切分造成的小矩阵低效率。
|
||||
|
||||
### 9.3 PP1 与 PP8 的 Nsight 对照
|
||||
|
||||
为避免把 Profiler 开销混入第 7 节性能结果，另起同场景短窗口 Nsight：
|
||||
`16K -> 1, C=8, chunk=8K`。以下是稳定 Rank 的中位数：
|
||||
|
||||
| 每个活跃 Rank/Stage 的指标 | PP1 / TP32 | PP8 / TP4 | 变化 |
|
||||
|---|---:|---:|---:|
|
||||
| 普通 AllReduce 调用数 | 557 | 66 | -88.2% |
|
||||
| 普通 AllReduce 累计时间 | 4,160.6 ms | 392.3 ms | -90.6% |
|
||||
| 单次 AllReduce p50 | 7.10 ms | 5.83 ms | -17.9% |
|
||||
| 计算与 NCCL 时间线重叠 | 约 0 | 约 0 | 未形成有效重叠 |
|
||||
|
||||
这里的 557 -> 66 是**单个 Rank/Stage**的 Profile 窗口，不是一个请求在整个
|
||||
模型上的 collective 总数。PP8 把模型层分给 8 个 Stage，一个请求仍要依次经过
|
||||
全部 Stage；它减少的是每个活跃 Rank 在一个 Stage 内负责的层数和同步工作，
|
||||
并让不同请求的 Stage 可以流水并行，而不是凭空删除模型全部通信。
|
||||
|
||||
PP8 Trace 还出现约 3.43 s 的非 AllReduce NCCL 区间，主要对应 Pipeline
|
||||
Send/Recv 及 Stage 等待。该值包含异步 Stage 等待，不能直接当作 3.43 s 的
|
||||
网络传输时间，也不能把 PP8 的 NCCL wall 百分比与 PP1 做等口径比较。可靠的
|
||||
对照是上表中的调用数、普通 AllReduce 累计时间和单次延迟。
|
||||
|
||||
### 9.4 为什么 PP8 最优，PP16 反而退化
|
||||
|
||||
PP1 -> PP8 的收益由四项共同产生：
|
||||
|
||||
1. TP 通信域由 32 Rank 缩到 4 Rank，单次 collective 延迟降低；
|
||||
2. PP4/PP8 的 TP collective 保持在节点内，PP8 进一步保持在同一 PCIe/NUMA
|
||||
岛；
|
||||
3. 每个 Rank 只执行所在 Stage 的层，C=8/16 的不同请求能在多个 Stage 并行，
|
||||
将通信和计算流水化；
|
||||
4. TP32 过度切分的 GEMM 在 TP4 下变成更大的本地矩阵，Kernel 效率更高。
|
||||
|
||||
收益没有继续延伸到 PP16：C=8 时请求数甚至少于 16 个 Stage，流水线无法填满；
|
||||
C=16 也只有一波请求。更多 Stage 边界、P2P Send/Recv、调度和 bubble 超过了
|
||||
TP4 -> TP2 的剩余收益，而且 PP16 还因整除约束从 EP4 改成 EP2。因此当前证据
|
||||
支持的是“PP8 在这个 Prefill 压力点达到通信、计算粒度与流水线开销的平衡”，
|
||||
而不是“PP 越大越快”。
|
||||
|
||||
## 10. 最终判断与后续使用
|
||||
|
||||
在本次固定场景中，`PP8 / TP4 / EP4` 同时得到：
|
||||
|
||||
- 最高的 Input TPS；
|
||||
- 最低的 TTFT p50/p95；
|
||||
- 三次重复中稳定的结果；
|
||||
- 比 PP1 和 PP4 更充足的最坏 Rank 显存余量；
|
||||
- 单节点内的 TP4/EP4 通信域。
|
||||
|
||||
因此，下一阶段 MoE A2A 的首个兼容性和显存 Smoke 应固定使用：
|
||||
|
||||
```text
|
||||
PP8 / TP4 / EP4
|
||||
16K -> 1
|
||||
C=8/16
|
||||
chunked_prefill_size=8192
|
||||
FlashInfer MXFP4
|
||||
```
|
||||
|
||||
`PP4 / TP8 / EP4` 保留为回退对照：如果某个 A2A backend 对 PP8/TP4 有当前
|
||||
代码未覆盖的约束，可以快速判断问题来自 backend 兼容性，还是 A2A 本身。
|
||||
|
||||
`PP16 / TP2 / EP2` 不进入下一阶段默认矩阵。它相对 PP1 仍有明显收益，但比
|
||||
PP8 慢约 24%，并引入 decode graph 兼容性、超长启动 warmup 和 EP2 变量。
|
||||
|
||||
这项结果不能外推到低并发 C=1、长 Decode、在线混合流量或更多节点。当前
|
||||
TTFT 即使在 PP8 下仍为 13.47 s（C8）和 27.31 s（C16），距离 5 s 目标仍有
|
||||
明显差距。Deep PP 解决了一个主要通信瓶颈，但不是 Prefill 优化的终点。
|
||||
|
||||
## 11. 复现命令与证据
|
||||
|
||||
在 601 执行：
|
||||
|
||||
```bash
|
||||
cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_pp_baseline_search
|
||||
RUN_ID=kimi3-pp-$(date +%Y%m%d-%H%M%S) \
|
||||
SUDO_PASSWORD_FILE=/path/to/password-file \
|
||||
bash run_pp_baseline_search.sh run
|
||||
```
|
||||
|
||||
PP2/4/8 正式结果：
|
||||
|
||||
```text
|
||||
/data/hzy/sskj/experiments/pro6000/
|
||||
kimi3_pro6000_sglang_pp_baseline_search/
|
||||
results/kimi3-pp-upstream-20260821-120208/
|
||||
```
|
||||
|
||||
其中：
|
||||
|
||||
```text
|
||||
summary.csv 6 个 PP×C 中位数
|
||||
results.csv 18 个逐轮结果
|
||||
metadata/ 实际服务配置与 Run manifest
|
||||
service/ 四节点 Docker 命令和完整服务日志
|
||||
bench/ benchmark 文本输出
|
||||
raw/ benchmark JSONL
|
||||
gpu/ 四节点启动前、健康、完成和清理后的显存快照
|
||||
```
|
||||
|
||||
PP1 原始结果：
|
||||
|
||||
```text
|
||||
/data/hzy/sskj/experiments/pro6000/
|
||||
kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/
|
||||
results/kimi3-ep4-moe-full-20260818-151349/
|
||||
```
|
||||
|
||||
PP16 正式结果：
|
||||
|
||||
```text
|
||||
/data/hzy/sskj/experiments/pro6000/
|
||||
kimi3_pro6000_sglang_pp_baseline_search/
|
||||
results/kimi3-pp16-20260821-161721/
|
||||
```
|
||||
|
||||
PP16 启动兼容性证据（未产生 benchmark 数据）：
|
||||
|
||||
```text
|
||||
results/kimi3-pp16-20260821-155734/ # custom AR decode graph IPC 错误
|
||||
results/kimi3-pp16-20260821-160555/ # 默认 300s watchdog 超时
|
||||
```
|
||||
|
||||
PP1 Nsight 通信归因：
|
||||
|
||||
```text
|
||||
/data/hzy/sskj/experiments/pro6000/
|
||||
kimi3_pro6000_sglang_prefill_communication_profile/
|
||||
results/kimi3-prefill-comm-20260820-143749/nsys_analysis.json
|
||||
```
|
||||
|
||||
PP8 Nsight 通信归因与四组 NCCL 微基准：
|
||||
|
||||
```text
|
||||
/data/hzy/sskj/experiments/pro6000/
|
||||
kimi3_pro6000_sglang_prefill_communication_profile/
|
||||
results/kimi3-prefill-comm-pp8-20260821-172717/
|
||||
|
||||
nsys_analysis.json # PP8 各 Rank CUDA/NCCL 汇总
|
||||
nsys_device_metrics.csv # PP8 各设备时间线指标
|
||||
nsys/node0/prefill_node0.nsys-rep # 601 原始 Nsight Trace
|
||||
nsys/node1/prefill_node1.nsys-rep # 602 原始 Nsight Trace
|
||||
nsys/node2/prefill_node2.nsys-rep # 603 原始 Nsight Trace
|
||||
nsys/node3/prefill_node3.nsys-rep # 604 原始 Nsight Trace
|
||||
nccl_tp4_intra_node.log # 单节点同 PCIe/NUMA 岛
|
||||
nccl_tp8_intra_node.log # 单节点跨 NUMA
|
||||
nccl_tp16_2node.log # 两节点
|
||||
nccl_tp32_4node.log # 四节点
|
||||
```
|
||||
@ -0,0 +1,379 @@
|
||||
#!/usr/bin/env bash
|
||||
# Search Kimi-K3 prefill pipeline-parallel baselines on 601-604.
|
||||
set -Eeuo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
# shellcheck source=/dev/null
|
||||
source "${ROOT_DIR}/scripts/common/lib.sh"
|
||||
|
||||
ACTION="${1:-run}"
|
||||
EXPERIMENT="kimi3_pro6000_sglang_pp_baseline_search"
|
||||
RUN_ID="${RUN_ID:-kimi3-pp-baseline-$(date '+%Y%m%d-%H%M%S')}"
|
||||
RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
|
||||
|
||||
MODEL_PATH="${MODEL_PATH:-/data/hf_models/Kimi-K3}"
|
||||
SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-kimi-k3}"
|
||||
DOCKER_IMAGE="${DOCKER_IMAGE:-local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5}"
|
||||
PORT="${PORT:-30000}"
|
||||
HEAD_HOST="${HEAD_HOST:-174.1.60.1}"
|
||||
NODE_SSH_USER="${NODE_SSH_USER:-user}"
|
||||
NODE_HOSTS=(174.1.60.1 174.1.60.2 174.1.60.3 174.1.60.4)
|
||||
SSH_OPTS=(-o BatchMode=yes -o StrictHostKeyChecking=no -o ConnectTimeout=10)
|
||||
DIST_PORT="${DIST_PORT:-20000}"
|
||||
|
||||
MAX_EP_SIZE=4
|
||||
read -r -a PP_SIZES <<< "${PP_SIZES_OVERRIDE:-2 4 8 16}"
|
||||
CONCURRENCIES=(8 16)
|
||||
INPUT_LEN=16384
|
||||
OUTPUT_LEN=1
|
||||
CHUNKED_PREFILL_SIZE=8192
|
||||
NUM_PROMPTS=40
|
||||
REPEATS="${REPEATS:-3}"
|
||||
WARMUP_REQUESTS=2
|
||||
HEALTH_WAIT_S="${HEALTH_WAIT_S:-2400}"
|
||||
CONTAINER_PREFIX="${EXPERIMENT}"
|
||||
|
||||
mkdir -p "${RESULT_ROOT}"/{service,raw,bench,gpu,metadata}
|
||||
log_init "${RESULT_ROOT}/orchestrator.log"
|
||||
|
||||
usage() {
|
||||
printf 'Usage: SUDO_PASSWORD=... %s {dry-run|run|summarize|stop}\n' "$0"
|
||||
}
|
||||
|
||||
require_password() {
|
||||
if [[ -z "${SUDO_PASSWORD:-}" && -n "${SUDO_PASSWORD_FILE:-}" ]]; then
|
||||
[[ -r "$SUDO_PASSWORD_FILE" ]] || {
|
||||
echo "ERROR: cannot read SUDO_PASSWORD_FILE=$SUDO_PASSWORD_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
IFS= read -r SUDO_PASSWORD <"$SUDO_PASSWORD_FILE"
|
||||
fi
|
||||
[[ -n "${SUDO_PASSWORD:-}" ]] || {
|
||||
echo "ERROR: set SUDO_PASSWORD or SUDO_PASSWORD_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
}
|
||||
|
||||
is_head() {
|
||||
[[ "$1" == "$HEAD_HOST" ]]
|
||||
}
|
||||
|
||||
sudo_host() {
|
||||
local host="$1"
|
||||
shift
|
||||
require_password
|
||||
if is_head "$host"; then
|
||||
printf '%s\n' "$SUDO_PASSWORD" | sudo -S -p '' -- "$@"
|
||||
return
|
||||
fi
|
||||
local remote_cmd
|
||||
printf -v remote_cmd '%q ' "$@"
|
||||
printf '%s\n' "$SUDO_PASSWORD" | ssh "${SSH_OPTS[@]}" \
|
||||
"${NODE_SSH_USER}@${host}" "sudo -S -p '' -- ${remote_cmd}"
|
||||
}
|
||||
|
||||
container_name() {
|
||||
printf '%s_node%s' "$CONTAINER_PREFIX" "$1"
|
||||
}
|
||||
|
||||
tp_for_pp() {
|
||||
local pp="$1"
|
||||
(( 32 % pp == 0 )) || { echo "ERROR: PP=$pp does not divide 32" >&2; return 1; }
|
||||
echo $((32 / pp))
|
||||
}
|
||||
|
||||
ep_for_tp() {
|
||||
local tp="$1"
|
||||
local ep="$MAX_EP_SIZE"
|
||||
if (( tp < ep )); then
|
||||
ep="$tp"
|
||||
fi
|
||||
(( tp % ep == 0 )) || {
|
||||
echo "ERROR: TP=$tp must be divisible by EP=$ep" >&2
|
||||
return 1
|
||||
}
|
||||
echo "$ep"
|
||||
}
|
||||
|
||||
preflight() {
|
||||
require_password
|
||||
[[ -d "$MODEL_PATH" ]] || { echo "ERROR: missing model $MODEL_PATH" >&2; exit 2; }
|
||||
local host
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
sudo_host "$host" docker image inspect "$DOCKER_IMAGE" >/dev/null
|
||||
sudo_host "$host" test -e /dev/infiniband/uverbs0
|
||||
done
|
||||
{
|
||||
echo "run_id=$RUN_ID"
|
||||
echo "image=$DOCKER_IMAGE"
|
||||
echo "model=$MODEL_PATH"
|
||||
echo "pp_sizes=${PP_SIZES[*]}"
|
||||
local pp tp ep
|
||||
for pp in "${PP_SIZES[@]}"; do
|
||||
tp="$(tp_for_pp "$pp")"
|
||||
ep="$(ep_for_tp "$tp")"
|
||||
echo "layout=pp${pp}_tp${tp}_ep${ep}"
|
||||
done
|
||||
echo "shape=${INPUT_LEN}->${OUTPUT_LEN}"
|
||||
echo "chunked_prefill_size=$CHUNKED_PREFILL_SIZE"
|
||||
echo "concurrencies=${CONCURRENCIES[*]}"
|
||||
echo "num_prompts=$NUM_PROMPTS"
|
||||
echo "repeats=$REPEATS"
|
||||
} >"${RESULT_ROOT}/metadata/manifest.env"
|
||||
}
|
||||
|
||||
stop_service() {
|
||||
local rank host name
|
||||
for rank in 0 1 2 3; do
|
||||
host="${NODE_HOSTS[$rank]}"
|
||||
name="$(container_name "$rank")"
|
||||
sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true
|
||||
done
|
||||
}
|
||||
|
||||
collect_service_logs() {
|
||||
local label="$1" rank host name
|
||||
for rank in 0 1 2 3; do
|
||||
host="${NODE_HOSTS[$rank]}"
|
||||
name="$(container_name "$rank")"
|
||||
sudo_host "$host" docker logs "$name" \
|
||||
>"${RESULT_ROOT}/service/${label}_node${rank}.log" 2>&1 || true
|
||||
done
|
||||
}
|
||||
|
||||
collect_gpu_snapshot() {
|
||||
local label="$1" rank host
|
||||
for rank in 0 1 2 3; do
|
||||
host="${NODE_HOSTS[$rank]}"
|
||||
sudo_host "$host" nvidia-smi \
|
||||
--query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu,power.draw \
|
||||
--format=csv,noheader,nounits \
|
||||
>"${RESULT_ROOT}/gpu/${label}_node${rank}.csv" 2>&1 || true
|
||||
done
|
||||
}
|
||||
|
||||
start_node() {
|
||||
local rank="$1" pp="$2" tp="$3" ep="$4"
|
||||
local host="${NODE_HOSTS[$rank]}" name bootstrap pp16_args=""
|
||||
name="$(container_name "$rank")"
|
||||
# PP16/TP2 currently fails while registering custom-all-reduce IPC buffers
|
||||
# for decode graph capture. This prefill-only OSL=1 experiment does not use
|
||||
# iterative decode. Its 16-stage startup also exceeds the default 300-second
|
||||
# watchdog while lazily creating adjacent-stage P2P communicators.
|
||||
if (( pp == 16 )); then
|
||||
pp16_args="--disable-decode-cuda-graph --watchdog-timeout 3000000"
|
||||
fi
|
||||
bootstrap="export SGLANG_HOST_IP=174.1.60.$((rank + 1)) && exec python3 -m sglang.launch_server --model-path ${MODEL_PATH} --served-model-name ${SERVED_MODEL_NAME} --tp-size ${tp} --pp-size ${pp} --ep-size ${ep} --nnodes 4 --node-rank ${rank} --dist-init-addr ${HEAD_HOST}:${DIST_PORT} --trust-remote-code --moe-runner-backend flashinfer_mxfp4 --chunked-prefill-size ${CHUNKED_PREFILL_SIZE} --mem-fraction-static 0.88 --cuda-graph-max-bs-decode 16 ${pp16_args} --mamba-radix-cache-strategy extra_buffer_lazy --disable-radix-cache --dist-timeout 3600 --mamba-full-memory-ratio 0.36 --host 0.0.0.0 --port ${PORT}"
|
||||
local -a cmd=(
|
||||
docker run -d --name "$name"
|
||||
--gpus all --network host --ipc=host --ulimit memlock=-1
|
||||
--device /dev/infiniband --shm-size 32g --entrypoint bash
|
||||
-v "${MODEL_PATH}:${MODEL_PATH}:ro"
|
||||
-e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
|
||||
-e NCCL_SOCKET_IFNAME=bond0 -e GLOO_SOCKET_IFNAME=bond0
|
||||
-e NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3
|
||||
-e NCCL_IB_GID_INDEX=3 -e NCCL_IB_TIMEOUT=22 -e NCCL_IB_RETRY_CNT=7
|
||||
-e NCCL_CUMEM_ENABLE=1 -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
-e SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK=0 -e SGLANG_MOE_FUSED_GATE_RADIX=1
|
||||
-e FLASHINFER_DISABLE_JIT=1 -e FLASHINFER_DISABLE_VERSION_CHECK=1
|
||||
"$DOCKER_IMAGE" -lc "$bootstrap"
|
||||
)
|
||||
printf '%q ' "${cmd[@]}" >"${RESULT_ROOT}/service/pp${pp}_node${rank}.cmd.txt"
|
||||
printf '\n' >>"${RESULT_ROOT}/service/pp${pp}_node${rank}.cmd.txt"
|
||||
sudo_host "$host" "${cmd[@]}" >/dev/null
|
||||
}
|
||||
|
||||
wait_health() {
|
||||
local pp="$1" i
|
||||
for ((i = 1; i <= HEALTH_WAIT_S; i++)); do
|
||||
if curl --fail --silent --max-time 5 "http://${HEAD_HOST}:${PORT}/health" >/dev/null 2>&1; then
|
||||
log "service healthy PP=${pp} wait_s=${i}"
|
||||
return 0
|
||||
fi
|
||||
if (( i % 30 == 0 )); then
|
||||
log "waiting for PP=${pp} service elapsed_s=${i}"
|
||||
collect_service_logs "pp${pp}_starting"
|
||||
if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|RuntimeError|AssertionError' \
|
||||
"${RESULT_ROOT}/service/pp${pp}_starting_node"*.log; then
|
||||
log "ERROR: PP=${pp} startup log contains a fatal error"
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
verify_server() {
|
||||
local pp="$1" tp="$2" ep="$3" info="${RESULT_ROOT}/metadata/pp${pp}_server_info.json"
|
||||
curl --fail --silent "http://${HEAD_HOST}:${PORT}/get_server_info" >"$info"
|
||||
python3 - "$info" "$pp" "$tp" "$ep" <<'PY'
|
||||
import json, sys
|
||||
|
||||
path, pp, tp, ep = sys.argv[1], *map(int, sys.argv[2:])
|
||||
info = json.load(open(path, encoding="utf-8"))
|
||||
expected = {"pp_size": pp, "tp_size": tp, "ep_size": ep,
|
||||
"moe_runner_backend": "flashinfer_mxfp4"}
|
||||
bad = [f"{k}={info.get(k)!r}, expected={v!r}" for k, v in expected.items()
|
||||
if info.get(k) != v]
|
||||
if bad:
|
||||
raise SystemExit("server configuration mismatch: " + "; ".join(bad))
|
||||
print(expected)
|
||||
PY
|
||||
}
|
||||
|
||||
start_service() {
|
||||
local pp="$1" tp ep
|
||||
tp="$(tp_for_pp "$pp")"
|
||||
ep="$(ep_for_tp "$tp")"
|
||||
stop_service
|
||||
log "starting PP=${pp} TP=${tp} EP=${ep} on 601-604"
|
||||
start_node 1 "$pp" "$tp" "$ep"
|
||||
start_node 2 "$pp" "$tp" "$ep"
|
||||
start_node 3 "$pp" "$tp" "$ep"
|
||||
sleep 5
|
||||
start_node 0 "$pp" "$tp" "$ep"
|
||||
if ! wait_health "$pp"; then
|
||||
collect_service_logs "pp${pp}_startup_failed"
|
||||
stop_service
|
||||
return 1
|
||||
fi
|
||||
verify_server "$pp" "$tp" "$ep"
|
||||
collect_service_logs "pp${pp}_healthy"
|
||||
collect_gpu_snapshot "pp${pp}_healthy"
|
||||
}
|
||||
|
||||
run_one_bench() {
|
||||
local pp="$1" concurrency="$2" repeat="$3"
|
||||
local stem="pp${pp}_c${concurrency}_r${repeat}"
|
||||
local output_file="${RESULT_ROOT}/raw/${stem}.jsonl"
|
||||
local bench_log="${RESULT_ROOT}/bench/${stem}.log"
|
||||
rm -f "$output_file"
|
||||
log "bench PP=${pp} C=${concurrency} repeat=${repeat}/${REPEATS}"
|
||||
sudo_host "$HEAD_HOST" docker run --rm --network host \
|
||||
-v "${MODEL_PATH}:${MODEL_PATH}:ro" \
|
||||
-v "${RESULT_ROOT}:${RESULT_ROOT}" \
|
||||
-e PYTHONUNBUFFERED=1 --entrypoint python3 "$DOCKER_IMAGE" \
|
||||
-m sglang.benchmark.serving \
|
||||
--backend sglang --host "$HEAD_HOST" --port "$PORT" \
|
||||
--tokenizer "$MODEL_PATH" --dataset-name random-ids \
|
||||
--random-input-len "$INPUT_LEN" --random-output-len "$OUTPUT_LEN" \
|
||||
--random-range-ratio 1.0 --num-prompts "$NUM_PROMPTS" \
|
||||
--max-concurrency "$concurrency" --request-rate 10000 \
|
||||
--warmup-requests "$WARMUP_REQUESTS" --output-file "$output_file" \
|
||||
--output-details --disable-tqdm >"$bench_log" 2>&1
|
||||
python3 - "$output_file" "$NUM_PROMPTS" <<'PY'
|
||||
import json, sys
|
||||
|
||||
path, expected = sys.argv[1], int(sys.argv[2])
|
||||
rows = [json.loads(line) for line in open(path, encoding="utf-8") if line.strip()]
|
||||
assert len(rows) == 1, (path, len(rows))
|
||||
assert rows[0].get("completed") == expected, rows[0].get("completed")
|
||||
assert not any(rows[0].get("errors", [])), "benchmark contains request errors"
|
||||
PY
|
||||
}
|
||||
|
||||
summarize() {
|
||||
python3 - "$RESULT_ROOT" <<'PY'
|
||||
import csv, json, re, statistics, sys
|
||||
from pathlib import Path
|
||||
|
||||
root = Path(sys.argv[1])
|
||||
pattern = re.compile(r"pp(\d+)_c(\d+)_r(\d+)\.jsonl$")
|
||||
metrics = [
|
||||
"request_throughput", "input_throughput", "output_throughput",
|
||||
"total_throughput", "mean_ttft_ms", "median_ttft_ms", "p95_ttft_ms",
|
||||
"p99_ttft_ms", "mean_e2e_latency_ms", "median_e2e_latency_ms",
|
||||
"p95_e2e_latency_ms", "p99_e2e_latency_ms",
|
||||
]
|
||||
rows = []
|
||||
for path in sorted((root / "raw").glob("*.jsonl")):
|
||||
match = pattern.match(path.name)
|
||||
if not match:
|
||||
continue
|
||||
data = next(json.loads(line) for line in path.open(encoding="utf-8") if line.strip())
|
||||
pp, concurrency, repeat = map(int, match.groups())
|
||||
tp = 32 // pp
|
||||
row = {"pp_size": pp, "tp_size": tp, "ep_size": min(4, tp),
|
||||
"concurrency": concurrency, "repeat": repeat,
|
||||
"completed": data.get("completed")}
|
||||
row.update({name: data.get(name) for name in metrics})
|
||||
rows.append(row)
|
||||
|
||||
fields = list(rows[0]) if rows else []
|
||||
with (root / "results.csv").open("w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=fields)
|
||||
if fields:
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
groups = {}
|
||||
for row in rows:
|
||||
groups.setdefault((row["pp_size"], row["concurrency"]), []).append(row)
|
||||
summary = []
|
||||
for (pp, concurrency), group in sorted(groups.items()):
|
||||
tp = 32 // pp
|
||||
item = {"pp_size": pp, "tp_size": tp, "ep_size": min(4, tp),
|
||||
"concurrency": concurrency, "repeats": len(group)}
|
||||
for name in metrics:
|
||||
values = [float(row[name]) for row in group if row.get(name) is not None]
|
||||
item[f"median_{name}"] = statistics.median(values) if values else None
|
||||
summary.append(item)
|
||||
|
||||
summary_fields = list(summary[0]) if summary else []
|
||||
with (root / "summary.csv").open("w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=summary_fields)
|
||||
if summary_fields:
|
||||
writer.writeheader()
|
||||
writer.writerows(summary)
|
||||
print(f"wrote {len(rows)} result rows and {len(summary)} summary rows")
|
||||
PY
|
||||
}
|
||||
|
||||
run_matrix() {
|
||||
preflight
|
||||
collect_gpu_snapshot before_all
|
||||
local pp concurrency repeat
|
||||
for pp in "${PP_SIZES[@]}"; do
|
||||
if ! start_service "$pp"; then
|
||||
log "ERROR: PP=${pp} failed to start; preserving evidence and continuing"
|
||||
continue
|
||||
fi
|
||||
for concurrency in "${CONCURRENCIES[@]}"; do
|
||||
for ((repeat = 1; repeat <= REPEATS; repeat++)); do
|
||||
run_one_bench "$pp" "$concurrency" "$repeat"
|
||||
done
|
||||
done
|
||||
collect_service_logs "pp${pp}_completed"
|
||||
collect_gpu_snapshot "pp${pp}_completed"
|
||||
stop_service
|
||||
sleep 5
|
||||
done
|
||||
collect_gpu_snapshot after_all
|
||||
summarize
|
||||
}
|
||||
|
||||
dry_run() {
|
||||
local pp tp ep
|
||||
for pp in "${PP_SIZES[@]}"; do
|
||||
tp="$(tp_for_pp "$pp")"
|
||||
ep="$(ep_for_tp "$tp")"
|
||||
echo "PP=${pp} TP=${tp} EP=${ep} ${INPUT_LEN}->${OUTPUT_LEN} C=${CONCURRENCIES[*]} chunk=${CHUNKED_PREFILL_SIZE}"
|
||||
done
|
||||
bash -n "$0"
|
||||
}
|
||||
|
||||
case "$ACTION" in
|
||||
dry-run) dry_run ;;
|
||||
run)
|
||||
trap stop_service EXIT INT TERM
|
||||
run_matrix
|
||||
;;
|
||||
summarize) summarize ;;
|
||||
stop)
|
||||
require_password
|
||||
stop_service
|
||||
;;
|
||||
*) usage; exit 2 ;;
|
||||
esac
|
||||
@ -0,0 +1,241 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Summarize CUDA/NCCL time and overlap from stage-scoped Nsight SQLite files."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import math
|
||||
import sqlite3
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def merge_intervals(intervals: list[tuple[int, int]]) -> list[tuple[int, int]]:
|
||||
merged: list[list[int]] = []
|
||||
for start, end in sorted(intervals):
|
||||
if end <= start:
|
||||
continue
|
||||
if not merged or start > merged[-1][1]:
|
||||
merged.append([start, end])
|
||||
else:
|
||||
merged[-1][1] = max(merged[-1][1], end)
|
||||
return [(start, end) for start, end in merged]
|
||||
|
||||
|
||||
def interval_length(intervals: list[tuple[int, int]]) -> int:
|
||||
return sum(end - start for start, end in intervals)
|
||||
|
||||
|
||||
def intersection_length(
|
||||
left: list[tuple[int, int]], right: list[tuple[int, int]]
|
||||
) -> int:
|
||||
i = j = total = 0
|
||||
while i < len(left) and j < len(right):
|
||||
start = max(left[i][0], right[j][0])
|
||||
end = min(left[i][1], right[j][1])
|
||||
if end > start:
|
||||
total += end - start
|
||||
if left[i][1] <= right[j][1]:
|
||||
i += 1
|
||||
else:
|
||||
j += 1
|
||||
return total
|
||||
|
||||
|
||||
def percentile(values: list[float], pct: float) -> float:
|
||||
if not values:
|
||||
return math.nan
|
||||
ordered = sorted(values)
|
||||
pos = (len(ordered) - 1) * pct
|
||||
low = math.floor(pos)
|
||||
high = math.ceil(pos)
|
||||
if low == high:
|
||||
return ordered[low]
|
||||
return ordered[low] * (high - pos) + ordered[high] * (pos - low)
|
||||
|
||||
|
||||
def analyze_sqlite(path: Path) -> tuple[list[dict], dict[str, dict]]:
|
||||
connection = sqlite3.connect(path)
|
||||
rows = connection.execute(
|
||||
"""
|
||||
SELECT k.deviceId, k.start, k.end, s.value
|
||||
FROM CUPTI_ACTIVITY_KIND_KERNEL AS k
|
||||
JOIN StringIds AS s ON k.shortName = s.id
|
||||
ORDER BY k.deviceId, k.start
|
||||
"""
|
||||
).fetchall()
|
||||
connection.close()
|
||||
|
||||
by_device: dict[int, list[tuple[int, int, str]]] = defaultdict(list)
|
||||
top_names: dict[str, dict[str, float]] = defaultdict(
|
||||
lambda: {"count": 0, "duration_ns": 0}
|
||||
)
|
||||
for device, start, end, name in rows:
|
||||
by_device[device].append((start, end, name))
|
||||
top_names[name]["count"] += 1
|
||||
top_names[name]["duration_ns"] += end - start
|
||||
|
||||
output = []
|
||||
node = path.parent.name
|
||||
for device, kernels in sorted(by_device.items()):
|
||||
all_intervals = merge_intervals([(start, end) for start, end, _ in kernels])
|
||||
nccl_rows = [row for row in kernels if row[2].lower().startswith("nccl")]
|
||||
compute_rows = [row for row in kernels if not row[2].lower().startswith("nccl")]
|
||||
nccl_intervals = merge_intervals([(start, end) for start, end, _ in nccl_rows])
|
||||
compute_intervals = merge_intervals(
|
||||
[(start, end) for start, end, _ in compute_rows]
|
||||
)
|
||||
|
||||
wall_ns = max(end for _, end, _ in kernels) - min(start for start, _, _ in kernels)
|
||||
kernel_union_ns = interval_length(all_intervals)
|
||||
nccl_union_ns = interval_length(nccl_intervals)
|
||||
compute_union_ns = interval_length(compute_intervals)
|
||||
overlap_ns = intersection_length(nccl_intervals, compute_intervals)
|
||||
allreduce_ms = [
|
||||
(end - start) / 1e6
|
||||
for start, end, name in nccl_rows
|
||||
if "allreduce" in name.lower()
|
||||
]
|
||||
normal_allreduce_ms = [value for value in allreduce_ms if value < 50.0]
|
||||
# The 8K Prefill trace has a clean gap between the two known payloads:
|
||||
# [M, 7168] ~= 112 MiB and M * (3584 + 7168) ~= 168 MiB.
|
||||
hidden_allreduce_ms = [value for value in normal_allreduce_ms if value < 7.5]
|
||||
moe_allreduce_ms = [value for value in normal_allreduce_ms if value >= 7.5]
|
||||
normal_allreduce_total_ms = sum(normal_allreduce_ms)
|
||||
|
||||
output.append(
|
||||
{
|
||||
"node": node,
|
||||
"local_gpu": device,
|
||||
"wall_ms": wall_ns / 1e6,
|
||||
"gpu_busy_union_ms": kernel_union_ns / 1e6,
|
||||
"gpu_busy_pct": 100.0 * kernel_union_ns / wall_ns,
|
||||
"nccl_union_ms": nccl_union_ns / 1e6,
|
||||
"nccl_wall_pct": 100.0 * nccl_union_ns / wall_ns,
|
||||
"compute_union_ms": compute_union_ns / 1e6,
|
||||
"compute_wall_pct": 100.0 * compute_union_ns / wall_ns,
|
||||
"nccl_compute_overlap_ms": overlap_ns / 1e6,
|
||||
"nccl_overlap_pct": (
|
||||
100.0 * overlap_ns / nccl_union_ns if nccl_union_ns else 0.0
|
||||
),
|
||||
"nccl_exposed_ms": (nccl_union_ns - overlap_ns) / 1e6,
|
||||
"nccl_exposed_wall_pct": 100.0
|
||||
* (nccl_union_ns - overlap_ns)
|
||||
/ wall_ns,
|
||||
"allreduce_count": len(allreduce_ms),
|
||||
"allreduce_normal_count": len(normal_allreduce_ms),
|
||||
"allreduce_mean_ms": (
|
||||
sum(normal_allreduce_ms) / len(normal_allreduce_ms)
|
||||
if normal_allreduce_ms
|
||||
else math.nan
|
||||
),
|
||||
"allreduce_p50_ms": percentile(normal_allreduce_ms, 0.50),
|
||||
"allreduce_p95_ms": percentile(normal_allreduce_ms, 0.95),
|
||||
"allreduce_max_ms": max(allreduce_ms, default=math.nan),
|
||||
"hidden_allreduce_count": len(hidden_allreduce_ms),
|
||||
"hidden_allreduce_sum_ms": sum(hidden_allreduce_ms),
|
||||
"hidden_allreduce_p50_ms": percentile(hidden_allreduce_ms, 0.50),
|
||||
"moe_allreduce_count": len(moe_allreduce_ms),
|
||||
"moe_allreduce_sum_ms": sum(moe_allreduce_ms),
|
||||
"moe_allreduce_p50_ms": percentile(moe_allreduce_ms, 0.50),
|
||||
"moe_allreduce_time_share_pct": (
|
||||
100.0 * sum(moe_allreduce_ms) / normal_allreduce_total_ms
|
||||
if normal_allreduce_total_ms
|
||||
else 0.0
|
||||
),
|
||||
}
|
||||
)
|
||||
return output, top_names
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("result_root", type=Path)
|
||||
args = parser.parse_args()
|
||||
|
||||
sqlite_paths = sorted((args.result_root / "nsys").glob("node*/*.sqlite"))
|
||||
if len(sqlite_paths) != 4:
|
||||
raise SystemExit(f"expected four SQLite files, found {len(sqlite_paths)}")
|
||||
|
||||
devices: list[dict] = []
|
||||
global_names: dict[str, dict[str, float]] = defaultdict(
|
||||
lambda: {"count": 0, "duration_ns": 0}
|
||||
)
|
||||
for path in sqlite_paths:
|
||||
device_rows, names = analyze_sqlite(path)
|
||||
devices.extend(device_rows)
|
||||
for name, values in names.items():
|
||||
global_names[name]["count"] += values["count"]
|
||||
global_names[name]["duration_ns"] += values["duration_ns"]
|
||||
|
||||
stable_devices = [
|
||||
row for row in devices if row["allreduce_max_ms"] < 50.0
|
||||
]
|
||||
metric_names = [
|
||||
"wall_ms",
|
||||
"gpu_busy_pct",
|
||||
"nccl_union_ms",
|
||||
"nccl_wall_pct",
|
||||
"compute_union_ms",
|
||||
"compute_wall_pct",
|
||||
"nccl_compute_overlap_ms",
|
||||
"nccl_overlap_pct",
|
||||
"nccl_exposed_ms",
|
||||
"nccl_exposed_wall_pct",
|
||||
"allreduce_count",
|
||||
"allreduce_mean_ms",
|
||||
"allreduce_p50_ms",
|
||||
"allreduce_p95_ms",
|
||||
"hidden_allreduce_count",
|
||||
"hidden_allreduce_sum_ms",
|
||||
"hidden_allreduce_p50_ms",
|
||||
"moe_allreduce_count",
|
||||
"moe_allreduce_sum_ms",
|
||||
"moe_allreduce_p50_ms",
|
||||
"moe_allreduce_time_share_pct",
|
||||
]
|
||||
stable_medians = {
|
||||
name: percentile([float(row[name]) for row in stable_devices], 0.50)
|
||||
for name in metric_names
|
||||
}
|
||||
|
||||
top_kernels = [
|
||||
{
|
||||
"name": name,
|
||||
"count": int(values["count"]),
|
||||
"duration_ms": values["duration_ns"] / 1e6,
|
||||
}
|
||||
for name, values in sorted(
|
||||
global_names.items(),
|
||||
key=lambda item: item[1]["duration_ns"],
|
||||
reverse=True,
|
||||
)[:40]
|
||||
]
|
||||
payload = {
|
||||
"result_root": str(args.result_root),
|
||||
"sqlite_files": [str(path) for path in sqlite_paths],
|
||||
"devices": devices,
|
||||
"stable_device_rule": "allreduce_max_ms < 50; excludes profiler trigger lanes",
|
||||
"stable_device_count": len(stable_devices),
|
||||
"stable_medians": stable_medians,
|
||||
"top_kernels": top_kernels,
|
||||
}
|
||||
|
||||
output_json = args.result_root / "nsys_analysis.json"
|
||||
output_csv = args.result_root / "nsys_device_metrics.csv"
|
||||
output_json.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
with output_csv.open("w", newline="", encoding="utf-8") as handle:
|
||||
writer = csv.DictWriter(handle, fieldnames=list(devices[0]))
|
||||
writer.writeheader()
|
||||
writer.writerows(devices)
|
||||
|
||||
print(json.dumps(payload["stable_medians"], indent=2))
|
||||
print(f"stable devices: {len(stable_devices)}/{len(devices)}")
|
||||
print(f"wrote {output_json}")
|
||||
print(f"wrote {output_csv}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -0,0 +1,265 @@
|
||||
#!/usr/bin/env bash
|
||||
# Capture a short, stage-scoped Nsight Systems trace for Kimi-K3 prefill.
|
||||
set -Eeuo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
RUN_ID="${RUN_ID:-kimi3-prefill-comm-$(date '+%Y%m%d-%H%M%S')}"
|
||||
RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/results}"
|
||||
RESULT_ROOT="${RESULT_BASE}/${RUN_ID}"
|
||||
|
||||
MODEL_PATH="${MODEL_PATH:-/data/hf_models/Kimi-K3}"
|
||||
IMAGE="${IMAGE:-local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5}"
|
||||
HEAD_HOST="${HEAD_HOST:-174.1.60.1}"
|
||||
NODE_USER="${NODE_USER:-user}"
|
||||
NODE_HOSTS=(174.1.60.1 174.1.60.2 174.1.60.3 174.1.60.4)
|
||||
PORT="${PORT:-30000}"
|
||||
DIST_PORT="${DIST_PORT:-20000}"
|
||||
PROFILE_STEPS="${PROFILE_STEPS:-3}"
|
||||
CONCURRENCY="${CONCURRENCY:-8}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-40}"
|
||||
PP_SIZE="${PP_SIZE:-1}"
|
||||
TP_SIZE="${TP_SIZE:-32}"
|
||||
EP_SIZE="${EP_SIZE:-4}"
|
||||
CONTAINER_PREFIX="kimi3_prefill_comm_profile"
|
||||
SSH_OPTS=(-o BatchMode=yes -o StrictHostKeyChecking=no -o ConnectTimeout=10)
|
||||
|
||||
mkdir -p "${RESULT_ROOT}"/{bench,nsys,service,gpu}
|
||||
|
||||
log() {
|
||||
printf '[%(%F %T)T] %s\n' -1 "$*" | tee -a "${RESULT_ROOT}/orchestrator.log"
|
||||
}
|
||||
|
||||
require_password() {
|
||||
if [[ -z "${SUDO_PASSWORD:-}" && -n "${SUDO_PASSWORD_FILE:-}" ]]; then
|
||||
IFS= read -r SUDO_PASSWORD <"${SUDO_PASSWORD_FILE}"
|
||||
fi
|
||||
[[ -n "${SUDO_PASSWORD:-}" ]] || {
|
||||
echo "ERROR: set SUDO_PASSWORD or SUDO_PASSWORD_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
}
|
||||
|
||||
is_head() {
|
||||
[[ "$1" == "$HEAD_HOST" ]]
|
||||
}
|
||||
|
||||
sudo_host() {
|
||||
local host="$1"
|
||||
shift
|
||||
require_password
|
||||
if is_head "$host"; then
|
||||
printf '%s\n' "$SUDO_PASSWORD" | sudo -S -p '' -- "$@"
|
||||
else
|
||||
local command
|
||||
printf -v command '%q ' "$@"
|
||||
printf '%s\n' "$SUDO_PASSWORD" | ssh "${SSH_OPTS[@]}" "${NODE_USER}@${host}" \
|
||||
"sudo -S -p '' -- ${command}"
|
||||
fi
|
||||
}
|
||||
|
||||
run_host() {
|
||||
local host="$1"
|
||||
shift
|
||||
if is_head "$host"; then
|
||||
"$@"
|
||||
else
|
||||
ssh "${SSH_OPTS[@]}" "${NODE_USER}@${host}" "$@"
|
||||
fi
|
||||
}
|
||||
|
||||
container_name() {
|
||||
printf '%s_node%s' "$CONTAINER_PREFIX" "$1"
|
||||
}
|
||||
|
||||
prepare_nodes() {
|
||||
local host rank=0
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
run_host "$host" mkdir -p "${RESULT_ROOT}/nsys/node${rank}"
|
||||
rank=$((rank + 1))
|
||||
done
|
||||
}
|
||||
|
||||
stop_service() {
|
||||
local host rank=0
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
sudo_host "$host" docker rm -f "$(container_name "$rank")" >/dev/null 2>&1 || true
|
||||
rank=$((rank + 1))
|
||||
done
|
||||
}
|
||||
|
||||
start_node() {
|
||||
local rank="$1" host name output launch
|
||||
host="${NODE_HOSTS[$rank]}"
|
||||
name="$(container_name "$rank")"
|
||||
output="${RESULT_ROOT}/nsys/node${rank}/prefill_node${rank}"
|
||||
launch="export SGLANG_HOST_IP=174.1.60.$((rank + 1)); exec nsys profile --trace=cuda,nvtx,nccl --nccl-trace=api,group,gpu,coll,kernel-launch --sample=none --cpuctxsw=none --capture-range=cudaProfilerApi --capture-range-end=stop --cuda-graph-trace=node --force-overwrite=true --stats=true --output ${output} python3 -m sglang.launch_server --model-path ${MODEL_PATH} --served-model-name kimi-k3 --tp-size ${TP_SIZE} --pp-size ${PP_SIZE} --ep-size ${EP_SIZE} --nnodes 4 --node-rank ${rank} --dist-init-addr ${HEAD_HOST}:${DIST_PORT} --trust-remote-code --moe-runner-backend flashinfer_mxfp4 --chunked-prefill-size 8192 --mem-fraction-static 0.88 --cuda-graph-max-bs-decode 16 --mamba-radix-cache-strategy extra_buffer_lazy --disable-radix-cache --dist-timeout 3600 --mamba-full-memory-ratio 0.36 --host 0.0.0.0 --port ${PORT}"
|
||||
|
||||
local -a command=(
|
||||
docker run -d --name "$name"
|
||||
--gpus all --network host --ipc=host --ulimit memlock=-1
|
||||
--device /dev/infiniband --shm-size 32g --entrypoint bash
|
||||
-v "${MODEL_PATH}:${MODEL_PATH}:ro"
|
||||
-v "${RESULT_ROOT}:${RESULT_ROOT}"
|
||||
-e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
|
||||
-e NCCL_SOCKET_IFNAME=bond0 -e GLOO_SOCKET_IFNAME=bond0
|
||||
-e NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3
|
||||
-e NCCL_IB_GID_INDEX=3 -e NCCL_IB_TIMEOUT=22 -e NCCL_IB_RETRY_CNT=7
|
||||
-e NCCL_CUMEM_ENABLE=1 -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
-e SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK=0 -e SGLANG_MOE_FUSED_GATE_RADIX=1
|
||||
-e FLASHINFER_DISABLE_JIT=1 -e FLASHINFER_DISABLE_VERSION_CHECK=1
|
||||
"$IMAGE" -lc "$launch"
|
||||
)
|
||||
printf '%q ' "${command[@]}" >"${RESULT_ROOT}/service/node${rank}.cmd.txt"
|
||||
printf '\n' >>"${RESULT_ROOT}/service/node${rank}.cmd.txt"
|
||||
sudo_host "$host" "${command[@]}" >/dev/null
|
||||
}
|
||||
|
||||
wait_health() {
|
||||
local elapsed
|
||||
for ((elapsed = 1; elapsed <= 2400; elapsed++)); do
|
||||
if curl --fail --silent --max-time 5 "http://${HEAD_HOST}:${PORT}/health" >/dev/null 2>&1; then
|
||||
log "service healthy after ${elapsed}s"
|
||||
return 0
|
||||
fi
|
||||
if (( elapsed % 30 == 0 )); then
|
||||
log "waiting for service: ${elapsed}s"
|
||||
collect_logs starting
|
||||
if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|RuntimeError' \
|
||||
"${RESULT_ROOT}/service/starting_node"*.log; then
|
||||
log "ERROR: startup log contains a fatal error"
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
collect_logs() {
|
||||
local label="$1" host rank=0
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
sudo_host "$host" docker logs "$(container_name "$rank")" \
|
||||
>"${RESULT_ROOT}/service/${label}_node${rank}.log" 2>&1 || true
|
||||
rank=$((rank + 1))
|
||||
done
|
||||
}
|
||||
|
||||
collect_gpu() {
|
||||
local label="$1" host rank=0
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
sudo_host "$host" nvidia-smi \
|
||||
--query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu,power.draw \
|
||||
--format=csv,noheader,nounits \
|
||||
>"${RESULT_ROOT}/gpu/${label}_node${rank}.csv" 2>&1 || true
|
||||
rank=$((rank + 1))
|
||||
done
|
||||
}
|
||||
|
||||
run_bench() {
|
||||
local label="$1" prompts="$2" concurrency="$3"
|
||||
local output="${RESULT_ROOT}/bench/${label}.jsonl"
|
||||
sudo_host "$HEAD_HOST" docker run --rm --network host \
|
||||
-v "${MODEL_PATH}:${MODEL_PATH}:ro" -v "${RESULT_ROOT}:${RESULT_ROOT}" \
|
||||
-e PYTHONUNBUFFERED=1 --entrypoint python3 "$IMAGE" \
|
||||
-m sglang.benchmark.serving --backend sglang --host "$HEAD_HOST" --port "$PORT" \
|
||||
--tokenizer "$MODEL_PATH" --dataset-name random-ids \
|
||||
--random-input-len 16384 --random-output-len 1 --random-range-ratio 1.0 \
|
||||
--num-prompts "$prompts" --max-concurrency "$concurrency" --request-rate 10000 \
|
||||
--warmup-requests 0 --output-file "$output" --output-details --disable-tqdm \
|
||||
>"${RESULT_ROOT}/bench/${label}.log" 2>&1
|
||||
}
|
||||
|
||||
trigger_profile() {
|
||||
local payload response
|
||||
payload="{\"activities\":[\"CUDA_PROFILER\"],\"num_steps\":${PROFILE_STEPS},\"profile_by_stage\":true,\"profile_id\":\"${RUN_ID}\",\"profile_prefix\":\"prefill\"}"
|
||||
printf '%s\n' "$payload" >"${RESULT_ROOT}/service/start_profile_request.json"
|
||||
response="$(curl --fail --silent --show-error -X POST \
|
||||
"http://${HEAD_HOST}:${PORT}/start_profile" \
|
||||
-H 'Content-Type: application/json' -d "$payload")"
|
||||
printf '%s\n' "$response" | tee "${RESULT_ROOT}/service/start_profile_response.json"
|
||||
}
|
||||
|
||||
copy_worker_artifacts() {
|
||||
local host rank
|
||||
for rank in 1 2 3; do
|
||||
host="${NODE_HOSTS[$rank]}"
|
||||
printf '%s\n' "$SUDO_PASSWORD" | ssh "${SSH_OPTS[@]}" "${NODE_USER}@${host}" \
|
||||
"sudo -S -p '' -- tar -C '${RESULT_ROOT}' -cf - 'nsys/node${rank}'" | \
|
||||
tar -C "${RESULT_ROOT}" -xf -
|
||||
done
|
||||
}
|
||||
|
||||
write_manifest() {
|
||||
{
|
||||
printf 'run_id=%s\n' "$RUN_ID"
|
||||
printf 'image=%s\n' "$IMAGE"
|
||||
printf 'model=%s\n' "$MODEL_PATH"
|
||||
printf 'pp=%s\ntp=%s\nep=%s\ndp=1\n' "$PP_SIZE" "$TP_SIZE" "$EP_SIZE"
|
||||
printf 'moe_runner_backend=flashinfer_mxfp4\nmoe_a2a_backend=none\n'
|
||||
printf 'input_len=16384\noutput_len=1\nconcurrency=%s\n' "$CONCURRENCY"
|
||||
printf 'chunked_prefill_size=8192\nprofile_steps=%s\n' "$PROFILE_STEPS"
|
||||
} >"${RESULT_ROOT}/manifest.env"
|
||||
}
|
||||
|
||||
run_all() {
|
||||
require_password
|
||||
(( PP_SIZE * TP_SIZE == 32 )) || {
|
||||
echo "ERROR: PP_SIZE * TP_SIZE must equal 32" >&2
|
||||
return 2
|
||||
}
|
||||
(( TP_SIZE % EP_SIZE == 0 )) || {
|
||||
echo "ERROR: TP_SIZE must be divisible by EP_SIZE" >&2
|
||||
return 2
|
||||
}
|
||||
prepare_nodes
|
||||
write_manifest
|
||||
stop_service
|
||||
collect_gpu before
|
||||
log "starting PP${PP_SIZE}/TP${TP_SIZE}/EP${EP_SIZE}/FlashInfer service under Nsight"
|
||||
start_node 1
|
||||
start_node 2
|
||||
start_node 3
|
||||
sleep 5
|
||||
start_node 0
|
||||
wait_health
|
||||
collect_logs healthy
|
||||
collect_gpu healthy
|
||||
|
||||
log "warming kernels without profiling"
|
||||
run_bench warmup 2 2
|
||||
log "arming ${PROFILE_STEPS} Prefill steps"
|
||||
trigger_profile
|
||||
run_bench profile_c8 "$NUM_PROMPTS" "$CONCURRENCY"
|
||||
collect_gpu profiled
|
||||
collect_logs profiled
|
||||
|
||||
log "stopping services to finalize reports"
|
||||
stop_service
|
||||
sleep 10
|
||||
copy_worker_artifacts
|
||||
collect_gpu after
|
||||
find "${RESULT_ROOT}/nsys" -type f -printf '%p %s bytes\n' | sort \
|
||||
>"${RESULT_ROOT}/nsys/files.txt"
|
||||
log "capture complete: ${RESULT_ROOT}"
|
||||
}
|
||||
|
||||
cleanup() {
|
||||
collect_logs cleanup 2>/dev/null || true
|
||||
stop_service 2>/dev/null || true
|
||||
}
|
||||
|
||||
case "${1:-run}" in
|
||||
run)
|
||||
trap cleanup EXIT INT TERM
|
||||
run_all
|
||||
trap - EXIT INT TERM
|
||||
;;
|
||||
stop)
|
||||
require_password
|
||||
stop_service
|
||||
;;
|
||||
*)
|
||||
echo "Usage: SUDO_PASSWORD=... bash $0 [run|stop]" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
@ -0,0 +1,63 @@
|
||||
# Kimi-K3 SM120 FlashInfer MXFP4 完成审计
|
||||
|
||||
审计时间：2026-08-19。本文只记录 goal 的逐项验收状态；技术结论见
|
||||
`UPSTREAM_DUPLICATION_AUDIT.md`，公开 PR 正文见 `SGLANG_DRAFT_PR.md`。
|
||||
|
||||
## 验收矩阵
|
||||
|
||||
| 要求 | 状态 | 权威证据 |
|
||||
|---|---|---|
|
||||
| 固定 605 历史环境身份 | 已证明 | profile 为 `/data/yy/sskj/deploy/profiles/pro6000/kimi3_pro6000_sglang_tp32ep32_flashinfer.env`；镜像 `lmsysorg/sglang:kimi-k3-fiv617situ-warm`；两个 runtime patch SHA256 见去重审计 |
|
||||
| 区分官方 0.6.17 与实际能力 | 已证明 | 605 profile 设置 `FLASHINFER_DISABLE_VERSION_CHECK=1` 并启动前执行两个补丁；报告未再声称 0.6.17 原版直接支持完整组合 |
|
||||
| 逐文件比较 605 补丁、FlashInfer #4460、SGLang 上游 | 已证明 | `UPSTREAM_DUPLICATION_AUDIT.md` 的“上游已有/605 已有/我们新增/明确排除”四部分；605 补丁目标已从原文件重新核验 |
|
||||
| 不提交重复 FlashInfer kernel | 已满足 | FlashInfer 历史实现与测试仅作为本地研究记录；SGLang Draft 不含 FlashInfer C++、依赖 pin 或 private capability marker |
|
||||
| SGLang 仅保留真实增量 | 已满足 | SGLang base `593b1a9b8aee` 到 head `b95b5341ab95` 仅修改 5 个文件；补丁位于 `patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch` |
|
||||
| 可应用 patch 与 Draft HEAD 一致 | 已证明 | 在全新 detached base `593b1a9b8aee` 应用仓库 patch 并 `git add -A` 后，暂存树与 `b95b5341ab95` 的 diff 为空 |
|
||||
| Kimi gate/up 权重和 scale 布局 | 已证明 | 注册测试覆盖连续 `[gate; up]` 到 CUTLASS `[up; gate]`，并保留 GPT-OSS pair-interleaved 分支 |
|
||||
| SiTU 4.0/25.0 参数透传 | 已证明 | #4460 ABI 的 SM120 GPU correctness 测试覆盖独立 `situ_beta` 与 `situ_linear_beta` |
|
||||
| 非连续 hidden states | 已证明 | GPU correctness 测试输入为非连续 view，runner 在 MXFP8 quantize 边界 materialize contiguous tensor |
|
||||
| attn_residual SM120 guard | 已证明 | 注册架构测试覆盖 SM100/103/110 为 true、SM120/SM90 为 false |
|
||||
| autotune 稳定性边界 | 已证明且不纳入 PR | 605 的 skip workaround 未进入 Draft；正式 24 点候选矩阵未复现 autotune 崩溃，因此没有用历史 workaround 扩大 PR |
|
||||
| 最终定向 correctness | 已通过 | RTX 6000D GPU6、#4460 snapshot：`8 passed, 17 warnings in 559.87s`；wheel SHA256 `deafdac2dc44b331441daba64e5b5dd28f591fe44f087af0aac3521c6ddbfa1f` |
|
||||
| fork 同步与补丁等价性 | 已通过 | 旧 base/head `63d783b..300c87a` 重放到 `593b1a9..b95b534`，7/7 commits 在 `git range-diff` 中均为 `=` |
|
||||
| 最终兼容回归 | 已通过 | Draft HEAD `b95b534`：compileall 通过，`10 passed, 3 deselected, 17 warnings in 33.30s` |
|
||||
| 格式和 CI 注册 | 已通过 | 五个变更文件完整 pre-commit 全部通过；两个测试文件注册到 CUDA CI `base-b/1-gpu-small` |
|
||||
| 保留 EP4 服务级结果 | 已证明 | 601 原始目录有 24 个正式 JSONL 点；每组 3 repeats、每点 40/40 成功；错误扫描为空；未重跑 EP32 |
|
||||
| 当前 Draft 代表点 A/B | 已通过 | Run `kimi3-sm120-pr-representative-20260819-115530`；Marlin/FlashInfer 各 3 次，240/240 请求成功；FlashInfer Input TPS +10.96%，TTFT P50/P95 -9.91%/-9.92% |
|
||||
| 公开 PR 标题、正文、依赖和已知限制 | Draft 正文已完成 | `SGLANG_DRAFT_PR.md` 遵循官方模板，Checklist 全部勾选；标题为 `[Feature] Support Kimi-K3 FlashInfer MXFP4 MoE on SM120`；Speed Tests 只报告一个目标硬件代表点，不作普遍加速声明 |
|
||||
| GitHub fork 交付 | 已完成 | `ACEEE-1222/sglang:hzy/kimi-k3-sm120-flashinfer-mxfp4`，Draft HEAD `b95b534`，远端 ref 已复核 |
|
||||
| 内部仓库独立分支 | 已完成 | 601 本地分支 `hzy/feature-kimi-k3-sm120-flashinfer-mxfp4` 已推送到 `qqtang/sskj:hzy-kimi-k3-sm120-flashinfer-mxfp4`；远端已有普通 `hzy` ref，因此不能使用 `hzy/...` 子分支名 |
|
||||
|
||||
## 原始证据位置
|
||||
|
||||
```text
|
||||
# EP4 服务级原始结果（601）
|
||||
/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
|
||||
|
||||
# #4460 ABI correctness（601）
|
||||
/data/hzy/phase5_sglang_pr4460_draft_gpu6_v4.log
|
||||
/data/hzy/phase5_sglang_pr4460_draft_cpu_final.log
|
||||
|
||||
# 当前 SGLang Draft worktree（601）
|
||||
/data/hzy/src/sglang-kimi-sm120-draft-wt
|
||||
|
||||
# 当前 Draft 代表点（601）
|
||||
/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
|
||||
```
|
||||
|
||||
## 交付确认
|
||||
|
||||
```text
|
||||
GitHub code branch:
|
||||
ACEEE-1222/sglang:hzy/kimi-k3-sm120-flashinfer-mxfp4
|
||||
|
||||
Internal audit branch:
|
||||
qqtang/sskj:hzy-kimi-k3-sm120-flashinfer-mxfp4
|
||||
|
||||
Internal compare:
|
||||
https://git.meta-stone.net/qqtang/sskj/compare/main...hzy-kimi-k3-sm120-flashinfer-mxfp4
|
||||
```
|
||||
|
||||
代码、代表点、原始证据索引、公开 PR 正文、可应用 patch 和两处远端分支均已完成。
|
||||
@ -0,0 +1,29 @@
|
||||
ARG BASE_IMAGE=lmsysorg/sglang:kimi-k3-fiv617situ-warm
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG SGLANG_COMMIT
|
||||
ARG FLASHINFER_COMMIT
|
||||
|
||||
ENV FLASHINFER_DISABLE_VERSION_CHECK=1 \
|
||||
PYTHONPATH=/sgl-workspace/sglang/python \
|
||||
PYTHONUNBUFFERED=1
|
||||
|
||||
COPY flashinfer_python-0.6.18-py3-none-any.whl /tmp/
|
||||
RUN python3 -m pip install --no-deps --force-reinstall \
|
||||
/tmp/flashinfer_python-0.6.18-py3-none-any.whl && \
|
||||
rm -f /tmp/*.whl
|
||||
|
||||
# Replace the image's older Python package with the exact Draft tree. Keep the
|
||||
# base image's ABI-matched torch and sglang-kernel binaries: the Draft changes
|
||||
# only the Python Kimi/FlashInfer integration and does not require a new kernel.
|
||||
RUN rm -rf /sgl-workspace/sglang/python/sglang
|
||||
COPY sglang/ /sgl-workspace/sglang/python/sglang/
|
||||
COPY test_mxfp4_sm120_cutlass.py /opt/pr-tests/test_mxfp4_sm120_cutlass.py
|
||||
COPY source_identity.txt /opt/pr-build/source_identity.txt
|
||||
|
||||
RUN python3 -m compileall -q /sgl-workspace/sglang/python/sglang && \
|
||||
python3 -c "import inspect; from flashinfer.fused_moe import cutlass_fused_moe; from flashinfer.fused_moe.core import ActivationType; assert hasattr(ActivationType, 'Situ'); assert 'situ_beta' in inspect.signature(cutlass_fused_moe).parameters; import sglang.srt.layers.quantization.mxfp4"
|
||||
|
||||
LABEL ai.meta-stone.purpose="Kimi-K3 SM120 SGLang Draft representative validation" \
|
||||
ai.meta-stone.sglang.commit="${SGLANG_COMMIT}" \
|
||||
ai.meta-stone.flashinfer.commit="${FLASHINFER_COMMIT}"
|
||||
@ -0,0 +1,281 @@
|
||||
# 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 时间线验证。
|
||||
@ -0,0 +1,145 @@
|
||||
# [Feature] Support Kimi-K3 FlashInfer MXFP4 MoE on SM120
|
||||
|
||||
> Draft: depends on flashinfer-ai/flashinfer#4460. The dependency is not
|
||||
> vendored or pinned by this PR, and the default MoE backend is unchanged.
|
||||
|
||||
## Motivation
|
||||
|
||||
SGLang already supports the FlashInfer CUTLASS MXFP4 MoE path for GPT-OSS on
|
||||
SM120, but Kimi-K3 needs additional model-specific integration:
|
||||
|
||||
- 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 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.
|
||||
|
||||
## Modifications
|
||||
|
||||
- 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.
|
||||
- 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.
|
||||
|
||||
The PR does not copy FlashInfer C++ code, pin an unmerged commit, alter the
|
||||
default MoE backend, or remove the Marlin fallback.
|
||||
|
||||
## Accuracy Tests
|
||||
|
||||
Target hardware: NVIDIA RTX 6000D, compute capability 12.0. FlashInfer was
|
||||
built from a snapshot of PR #4460 fetched on 2026-08-18.
|
||||
|
||||
```bash
|
||||
pytest -q -s \
|
||||
test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py \
|
||||
test/registered/unit/layers/test_attn_residual_arch_gate.py
|
||||
```
|
||||
|
||||
Result with the FlashInfer kernel enabled:
|
||||
|
||||
```text
|
||||
8 passed, 17 warnings in 559.87s
|
||||
```
|
||||
|
||||
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.
|
||||
|
||||
After syncing the fork and rebasing the seven-patch series onto SGLang main
|
||||
`593b1a9b8aeedb908ee21870c02a0fd1a3dc30a5`, the non-JIT subset was rerun
|
||||
from the final Draft head `b95b5341ab9555ba5ea53f5828f6de203d749cce`:
|
||||
|
||||
```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
|
||||
10 passed, 3 deselected, 17 warnings in 33.30s
|
||||
```
|
||||
|
||||
## Speed Tests and Profiling
|
||||
|
||||
This is a compatibility PR, so the following result is a target-hardware
|
||||
validation point rather than a general speed claim. Marlin and
|
||||
`flashinfer_mxfp4` used the same SGLang commit, image, model, service flags,
|
||||
requests, and three-repeat protocol.
|
||||
|
||||
```text
|
||||
Hardware: 4 nodes / 32 NVIDIA RTX 6000D GPUs (SM120)
|
||||
Model: Kimi-K3
|
||||
Parallelism: TP32 / EP4 / DP1
|
||||
Workload: ISL=16384, OSL=1, concurrency=8, 40 requests
|
||||
Chunked Prefill: 8192 tokens
|
||||
SGLang patch series used for the run: 300c87a431ac40d3e7817246376b7fe20932db09
|
||||
Final rebased Draft head: b95b5341ab9555ba5ea53f5828f6de203d749cce
|
||||
Final SGLang main base: 593b1a9b8aeedb908ee21870c02a0fd1a3dc30a5
|
||||
FlashInfer: PR #4460 implementation b525c513fef7c1d162efd3ce8636fcc8760de66a
|
||||
Repeats: 3 per backend
|
||||
```
|
||||
|
||||
The service benchmark was collected before the fork sync. `git range-diff`
|
||||
reports all seven commits as patch-equivalent after the rebase, and the
|
||||
directed tests above were rerun from the final Draft head. The performance
|
||||
result is therefore retained without claiming that the benchmark container
|
||||
contained the new commit object.
|
||||
|
||||
| Backend | Median Input TPS | Median TTFT P50 | Median TTFT P95 |
|
||||
|---|---:|---:|---:|
|
||||
| Marlin | 2930.03 | 43.587 s | 46.287 s |
|
||||
| FlashInfer MXFP4 | 3251.25 | 39.267 s | 41.694 s |
|
||||
| Change | +10.96% | -9.91% | -9.92% |
|
||||
|
||||
All six repeats completed 40/40 requests, all outputs contained exactly one
|
||||
token, and no request errors, OOM, traceback, NCCL error, or engine-death
|
||||
event was observed. The result only establishes this exact configuration;
|
||||
the default backend remains unchanged.
|
||||
|
||||
## 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 documentation 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
|
||||
@ -0,0 +1,178 @@
|
||||
# Kimi-K3 SM120 FlashInfer MXFP4 去重审计
|
||||
|
||||
审计时间：2026-08-19 13:07:48 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` |
|
||||
| 当前 Draft 代表点 | 精确 SGLang `300c87a` + #4460 `b525c513`，TP32/EP4、16K→1、C8、Chunk 8K，各三次 | `kimi3-sm120-pr-representative-20260819-115530` |
|
||||
|
||||
## 明确排除的重复内容
|
||||
|
||||
- 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: 593b1a9b8aeedb908ee21870c02a0fd1a3dc30a5
|
||||
head: b95b5341ab9555ba5ea53f5828f6de203d749cce
|
||||
patch: patches/sglang-kimi-k3-sm120-flashinfer-mxfp4-draft.patch
|
||||
```
|
||||
|
||||
该 7 提交系列由旧范围 `63d783b..300c87a` 重放到用户新同步的 main；
|
||||
`git range-diff` 逐项均为 `=`。在新 base 的独立 worktree 应用上述 patch、
|
||||
暂存全部文件后，与 `b95b534` 比较的 diff 为空。
|
||||
|
||||
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
|
||||
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，并新增旧版兼容回归。
|
||||
同步前提交 `300c87a` 为两个测试文件增加正式 CUDA CI 注册，并通过完整
|
||||
pre-commit。同步后的最终 HEAD `b95b534` 在 601 通过 compileall；能力探测、
|
||||
旧版兼容和架构 gate 为 `10 passed, 3 deselected, 17 warnings in 33.30s`；
|
||||
GPU kernel 路径继续由前述 8/8 测试覆盖。
|
||||
|
||||
完整 pre-commit 命令覆盖全部五个变更文件；isort、ruff、black、codespell、
|
||||
registered-tests 检查及其他 hook 全部通过。下载 GitHub hook 时只对该命令使用
|
||||
HTTPS-to-SSH URL rewrite，没有修改全局 Git 配置。
|
||||
|
||||
最终代表点使用同步前等价 patch `300c87a` 和 #4460 implementation commit 的统一组合，
|
||||
六轮 240/240 请求成功。FlashInfer 在所有 32 rank 完成 autotune，因此 605 的
|
||||
skip-autotune workaround 仍不属于 SGLang PR 增量。该固定点的 FlashInfer
|
||||
Input TPS 为 Marlin 的 `+10.96%`，TTFT P50/P95 为 `-9.91%/-9.92%`；这只作为
|
||||
SM120 集成的服务级证据，不扩展为所有 shape 的性能结论。
|
||||
|
||||
前两次失败日志分别暴露 `FLASHINFER_DISABLE_JIT=0` 仍被视为启用、源码包缺少
|
||||
生成的 `data/csrc`；第三次暴露只读源码不能生成 `_build_meta.py`。这些都是
|
||||
测试装配问题，最终 v4 采用可写临时源码树构建 wheel 后通过，并非功能回归。
|
||||
@ -0,0 +1,197 @@
|
||||
# Kimi-K3 SM120 FlashInfer MXFP4 Decision Log
|
||||
|
||||
本文件记录可恢复的关键决策、失败尝试和证据路径。成功结论见 `README.md`。
|
||||
|
||||
## 2026-08-19 fork 同步与最终交付
|
||||
|
||||
- 用户重新同步 `ACEEE-1222/sglang:main` 后，最终 base 为 `593b1a9b8aee`。
|
||||
- 原 7 提交系列从 `63d783b..300c87a` 无冲突重放为
|
||||
`593b1a9..b95b534`；`git range-diff` 的 7 项全部为 `=`，没有语义改动。
|
||||
- 最终 HEAD 在 601 统一验证镜像中通过 compileall 和定向非 JIT 回归：
|
||||
`10 passed, 3 deselected, 17 warnings in 33.30s`。
|
||||
- GitHub feature 分支通过精确 `force-with-lease` 从 `300c87a` 更新到
|
||||
`b95b534`，随后用 `ls-remote` 复核远端 ref。
|
||||
- 四机代表点不重跑。其容器记录的是同步前 `300c87a`，但对应 patch 与最终
|
||||
HEAD 逐提交等价；PR 正文同时列出实验 patch SHA 和最终 HEAD，不混淆对象。
|
||||
|
||||
## 2026-08-19 最终 PR 代表点
|
||||
|
||||
- 只补跑一个代表点，不重跑 EP32：TP32/EP4、16K→1、C=8、Chunk=8K，
|
||||
Marlin/FlashInfer 各三次，Run ID
|
||||
`kimi3-sm120-pr-representative-20260819-115530`。
|
||||
- 当前组合固定为 SGLang `300c87a431ac`、FlashInfer #4460 实现
|
||||
`b525c513fef7`；精确 FlashInfer wheel SHA256 为
|
||||
`5c67841b0bc5aa72d04910ee1cb727ad5fc9c13b5b80c14e26ac40de1583ce88`。
|
||||
- 六轮 240/240 请求成功，输出长度全部为 1；错误扫描为空。
|
||||
- 三次中位数：Marlin Input TPS `2930.03`、TTFT P50/P95
|
||||
`43.587/46.287 s`；FlashInfer 为 `3251.25`、`39.267/41.694 s`。
|
||||
- 该点证明当前 Draft 组合在目标硬件可用，并在这一固定 shape 上相对 Marlin
|
||||
为 Input TPS `+10.96%`、TTFT P50/P95 `-9.91%/-9.92%`；不外推到其他
|
||||
shape，也不改变默认 backend。
|
||||
- 第一次离线汇总误报 `expected 3 repeats`，根因是 benchmark JSON 的
|
||||
`backend=sglang` 覆盖了文件名解析出的 MoE backend；修复字段合并顺序后，
|
||||
原始六轮数据直接生成最终汇总，没有重跑或删除结果。
|
||||
- 公开 PR 的 Checklist 全部勾选；Speed Tests 仅报告该代表点，仍把 PR 定位为
|
||||
compatibility support。
|
||||
|
||||
## 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 兼容回归。最终提交
|
||||
`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 和去重证据继续保留在本仓库审计文件中。
|
||||
- Checklist 在最终代表点和内部档案完成后全部勾选。公开 PR 仅报告一个固定
|
||||
shape 的目标硬件结果，不做普遍提速声明；用户文档仍等待 #4460 发布稳定 API。
|
||||
|
||||
## 2026-08-18 Phase 5
|
||||
|
||||
### 真实 serving 暴露的兼容问题
|
||||
|
||||
- Kimi 基线镜像早于通用 `MoeRunnerConfig.gemm1_beta` 字段，首次
|
||||
FlashInfer EP4 启动在权重处理阶段抛出 `AttributeError`。镜像兼容补丁
|
||||
改为 `getattr(..., "gemm1_beta", None)`，SiTU 的 beta 仍从 Kimi 已有的
|
||||
`gemm1_clamp_limit=25` 读取；没有改变新版本源码的通用字段语义。
|
||||
- 第二次启动完成权重和 KV cache 初始化后，autotune 收到非连续 Prefill
|
||||
hidden states，`mxfp8_quantize` 抛出 `RuntimeError: input must be contiguous`。
|
||||
修复位于 FlashInfer runner 边界，在量化前执行 `x.contiguous()`；正式
|
||||
24 点矩阵验证了该路径。
|
||||
- benchmark 容器最初尝试访问 Hugging Face。入口增加
|
||||
`--tokenizer /data/hf_models/Kimi-K3`，正式矩阵完全使用本地 tokenizer。
|
||||
|
||||
### EP4 正式矩阵
|
||||
|
||||
- Run ID：`kimi3-ep4-moe-full-20260818-151349`。
|
||||
- 口径：TP32/EP4/DP1，16K 输入、1 token 输出、40 requests，
|
||||
Marlin/FlashInfer MXFP4，chunk 8K/16K，C8/C16，每点三次。
|
||||
- 结果：24/24 raw JSONL 均为 40/40 成功且 errors 为空；服务与 bench 日志
|
||||
未发现 OOM、Traceback、NCCL error 或 EngineDeadError。
|
||||
- 8K chunk：FlashInfer Input TPS 提升 11.00%-11.08%，TTFT P50 降低
|
||||
9.92%-9.99%。
|
||||
- 16K chunk：FlashInfer Input TPS 提升 7.63%-7.65%，TTFT P50 降低
|
||||
7.04%-7.09%。
|
||||
- 决策：该 16K/1 Prefill 工作负载采用 EP4 + FlashInfer MXFP4 + 8K chunk。
|
||||
|
||||
### EP32 证据处理
|
||||
|
||||
- 不补跑 EP32。仅复用已有 `kimi3-moe-prefill-20260818-130900` 中
|
||||
Marlin/8K/C8 两次结果。
|
||||
- EP32 中位数为 Input TPS 2531.43、TTFT P50 50.52 s；EP4 三次中位数为
|
||||
2935.02、43.51 s，即 EP4 吞吐 +15.94%、TTFT P50 -13.88%。
|
||||
- EP32 只有两次重复，因此只作为方向性同口径对照，不替代 EP4 正式矩阵。
|
||||
|
||||
### 原始证据与清理
|
||||
|
||||
```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
|
||||
```
|
||||
|
||||
仓库内保留小型汇总：
|
||||
|
||||
```text
|
||||
results/phase5_ep4_moe_backend_prefill_summary.json
|
||||
results/phase5_ep4_moe_backend_prefill_results.csv
|
||||
```
|
||||
|
||||
运行结束后 601-604 无实验容器，32 张 GPU 均为 0 MiB / 0% utilization。
|
||||
|
||||
### PR 回归补充
|
||||
|
||||
- SGLang runner 增加非连续输入回归。第一次测试把
|
||||
`expected=torch.empty_like(noncontiguous_x)` 也创建成非连续 stride，底层
|
||||
kernel 按连续输出 ABI 写入后被测试错误解读，导致数值错位。reference
|
||||
输出改为 `memory_format=torch.contiguous_format` 后通过。
|
||||
- SGLang 最终结果：`1 passed, 5 deselected`，提交 `8a2d501`。
|
||||
- FlashInfer 官方 `test_trtllm_cutlass_fused_moe.py` 增加 SiTU
|
||||
`(alpha=4, beta=25)` 的 MXFP8×MXFP4 case，覆盖 FP16/BF16 和 autotune
|
||||
on/off；最终 `4 passed, 153 deselected`，提交 `e0bc653`。
|
||||
- 证据日志：
|
||||
|
||||
```text
|
||||
/data/hzy/phase5_sglang_noncontiguous_gpu6.log
|
||||
/data/hzy/phase5_sglang_noncontiguous_gpu6_v2.log
|
||||
/data/hzy/phase5_flashinfer_official_situ_gpu6.log
|
||||
```
|
||||
|
||||
## 2026-08-18 Phase 4
|
||||
|
||||
### 对称 A/B
|
||||
|
||||
- 决策：baseline/candidate 都走 SGLang production runner，共享输入、权重、
|
||||
scale、routing、SiTU 参数和计时方式。
|
||||
- 结果：M1/M128 FlashInfer 略慢；扩展五轮扫描中 M512/2048/8192 分别比
|
||||
Marlin 快约 1.46x/2.08x/2.13x。
|
||||
- 限制：M 是单卡本地 runner token 数，不是 HTTP ISL 或全局 batch。
|
||||
- 证据：`results/phase4_ab_5rounds_gpu6.json`、
|
||||
`results/phase4_mscan_ep32_5rounds_gpu6.json`。
|
||||
|
||||
### GPU 绑定事故
|
||||
|
||||
- 现象：最初 Nsight 命令同时使用 `--gpus device=6` 和 `--privileged`，UUID
|
||||
guard 检测到实际设备为物理 GPU0 并终止。
|
||||
- 根因：privileged 重新暴露全部 NVIDIA device node，破坏容器 GPU 隔离。
|
||||
- 处理：移除 privileged；固定 `--gpus device=6`、
|
||||
`NVIDIA_VISIBLE_DEVICES=6` 和预期 UUID。未继续占用 GPU0。
|
||||
|
||||
### JIT 与 AOT
|
||||
|
||||
- 现象：纯 Python FlashInfer wheel 在新进程重新 JIT。NCU profiler 注入
|
||||
JIT 子进程后产生大量挂起 `nvcc`。
|
||||
- 处理：将已验证的 `fused_moe_120`、`mxfp8_quantization_sm100`、
|
||||
`fp4_quantization_120f` `.so` 写入 wheel 的 AOT 目录；设置
|
||||
`FLASHINFER_DISABLE_JIT=1`；持久化 SGLang JIT cache。
|
||||
- 验证：`results/phase4_aot_nojit_smoke_gpu6.json`。
|
||||
|
||||
### Profiler 工具
|
||||
|
||||
- NCU 2025.3.1：RTX 6000D/SM120 不受支持，报
|
||||
`Profiling is not supported on device 0`。
|
||||
- NCU 2026.2.1：从 NVIDIA 官方 CUDA repository 下载并校验 SHA256，解压到
|
||||
`/data/hzy/tools/nsight-compute-2026.2.1/`，正式捕获成功。
|
||||
- Nsight Systems：正式捕获成功，原始文件为
|
||||
`results/phase4_nsys_ab_gpu6.nsys-rep`。
|
||||
|
||||
### 大输出统计
|
||||
|
||||
- 现象：M8192 输出约 2936 万元素，`torch.quantile` 报 input too large。
|
||||
- 处理：P95/P99 改用精确 `torch.kthvalue`。未抽样，未改正确性阈值。
|
||||
|
||||
### Warmup 敏感性
|
||||
|
||||
- M512 pilot 使用 warmup=1，得到与正式结果相反的排序。
|
||||
- 正式五轮使用 warmup=5、每轮 30 次，五轮速度比为
|
||||
`1.428x-1.474x`，采用正式结果。
|
||||
- pilot 日志保留，不作为性能结论。
|
||||
|
||||
### 阶段出口
|
||||
|
||||
- 不继续 synthetic EP4 扫描。
|
||||
- 下一步需要 601-604 共 32 卡运行真实 TP32/EP32 Kimi-K3 serving。
|
||||
- 未经资源授权不清理现有作业、不启动 Phase 5。
|
||||
@ -0,0 +1,610 @@
|
||||
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..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
|
||||
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.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
|
||||
+
|
||||
+ 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,161 @@ 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..3c55b924312a7d2537f8221b5e985e7927e527d4
|
||||
--- /dev/null
|
||||
+++ b/test/registered/unit/layers/test_attn_residual_arch_gate.py
|
||||
@@ -0,0 +1,46 @@
|
||||
+"""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
|
||||
+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(
|
||||
+ ("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)
|
||||
@ -0,0 +1,583 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Profile Kimi-K3 MXFP4 MoE on SM120: SGLang Marlin vs FlashInfer.
|
||||
|
||||
Both backends consume the same BF16 input, packed MXFP4 checkpoint values,
|
||||
E8M0 scales, global top-k ids/weights, SiTU parameters, and EP32/rank7 mapping.
|
||||
Backend-specific post-load layouts are derived once before correctness checks,
|
||||
warmup, timing, and profiling.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import asdict
|
||||
from datetime import datetime
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
import statistics
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from flashinfer import (
|
||||
SfLayout,
|
||||
block_scale_interleave,
|
||||
mxfp4_dequantize,
|
||||
mxfp4_quantize,
|
||||
mxfp8_quantize,
|
||||
)
|
||||
from sglang.srt.distributed import init_distributed_environment
|
||||
from sglang.srt.distributed.parallel_state import (
|
||||
destroy_model_parallel,
|
||||
initialize_model_parallel,
|
||||
)
|
||||
from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
|
||||
FlashInferCutlassMxfp4MoeQuantInfo,
|
||||
fused_experts_none_to_flashinfer_mxfp4,
|
||||
)
|
||||
from sglang.srt.layers.moe.moe_runner.marlin import (
|
||||
MarlinMoeQuantInfo,
|
||||
fused_experts_none_to_marlin,
|
||||
)
|
||||
from sglang.srt.layers.moe.token_dispatcher.standard import StandardDispatchOutput
|
||||
from sglang.srt.layers.moe.topk import StandardTopKOutput
|
||||
from sglang.srt.layers.quantization.marlin_utils_fp4 import (
|
||||
prepare_moe_mxfp4_layer_for_marlin,
|
||||
)
|
||||
|
||||
from validate_kimi_sm120_mxfp4 import (
|
||||
Case,
|
||||
QuantizedCase,
|
||||
dequantize_mxfp8_input,
|
||||
make_routing,
|
||||
run_reference,
|
||||
)
|
||||
|
||||
|
||||
SHAPES = {
|
||||
"decode": Case(
|
||||
"kimi_decode_m1_ep32",
|
||||
1,
|
||||
3584,
|
||||
3072,
|
||||
28,
|
||||
16,
|
||||
"ep_mixed",
|
||||
ep_size=32,
|
||||
ep_rank=7,
|
||||
),
|
||||
"prefill": Case(
|
||||
"kimi_prefill_m128_ep32",
|
||||
128,
|
||||
3584,
|
||||
3072,
|
||||
28,
|
||||
16,
|
||||
"ep_mixed",
|
||||
ep_size=32,
|
||||
ep_rank=7,
|
||||
),
|
||||
"prefill_m512": Case(
|
||||
"kimi_prefill_m512_ep32",
|
||||
512,
|
||||
3584,
|
||||
3072,
|
||||
28,
|
||||
16,
|
||||
"ep_mixed",
|
||||
ep_size=32,
|
||||
ep_rank=7,
|
||||
),
|
||||
"prefill_m2048": Case(
|
||||
"kimi_prefill_m2048_ep32",
|
||||
2048,
|
||||
3584,
|
||||
3072,
|
||||
28,
|
||||
16,
|
||||
"ep_mixed",
|
||||
ep_size=32,
|
||||
ep_rank=7,
|
||||
),
|
||||
"prefill_m8192": Case(
|
||||
"kimi_prefill_m8192_ep32",
|
||||
8192,
|
||||
3584,
|
||||
3072,
|
||||
28,
|
||||
16,
|
||||
"ep_mixed",
|
||||
ep_size=32,
|
||||
ep_rank=7,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _quantize_experts_linear(
|
||||
weight: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
packed = []
|
||||
scales = []
|
||||
dequantized = []
|
||||
for expert_weight in weight:
|
||||
expert_packed, expert_scale = mxfp4_quantize(
|
||||
expert_weight,
|
||||
sfLayout=SfLayout.layout_linear,
|
||||
)
|
||||
packed.append(expert_packed)
|
||||
scales.append(expert_scale)
|
||||
dequantized.append(
|
||||
mxfp4_dequantize(
|
||||
expert_packed.cpu(),
|
||||
expert_scale.cpu(),
|
||||
sfLayout=SfLayout.layout_linear,
|
||||
)
|
||||
)
|
||||
return (
|
||||
torch.stack(packed),
|
||||
torch.stack(scales),
|
||||
torch.stack(dequantized).to(device=weight.device, dtype=weight.dtype),
|
||||
)
|
||||
|
||||
|
||||
def _prepare_linear_checkpoint_case(
|
||||
case: Case,
|
||||
seed: int,
|
||||
device: torch.device,
|
||||
) -> QuantizedCase:
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
dtype = torch.bfloat16
|
||||
x = torch.randn(
|
||||
case.tokens,
|
||||
case.hidden,
|
||||
generator=generator,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
gate = torch.randn(
|
||||
case.local_experts,
|
||||
case.intermediate,
|
||||
case.hidden,
|
||||
generator=generator,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
) * (1.0 / math.sqrt(case.hidden))
|
||||
up = torch.randn(
|
||||
case.local_experts,
|
||||
case.intermediate,
|
||||
case.hidden,
|
||||
generator=generator,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
) * (1.0 / math.sqrt(case.hidden))
|
||||
down = torch.randn(
|
||||
case.local_experts,
|
||||
case.hidden,
|
||||
case.intermediate,
|
||||
generator=generator,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
) * (1.0 / math.sqrt(case.intermediate))
|
||||
input_q, input_sf = mxfp8_quantize(
|
||||
x,
|
||||
is_sf_swizzled_layout=True,
|
||||
alignment=32,
|
||||
)
|
||||
input_dq = dequantize_mxfp8_input(input_q, input_sf, dtype)
|
||||
fc1_q, fc1_sf, fc1_dq = _quantize_experts_linear(
|
||||
torch.cat((up, gate), dim=1).contiguous()
|
||||
)
|
||||
fc2_q, fc2_sf, fc2_dq = _quantize_experts_linear(down)
|
||||
topk_ids, topk_weights = make_routing(case, device)
|
||||
return QuantizedCase(
|
||||
case=case,
|
||||
input_q=input_q,
|
||||
input_sf=input_sf,
|
||||
input_dq=input_dq,
|
||||
fc1_q=fc1_q,
|
||||
fc1_sf=fc1_sf,
|
||||
fc1_dq=fc1_dq,
|
||||
fc2_q=fc2_q,
|
||||
fc2_sf=fc2_sf,
|
||||
fc2_dq=fc2_dq,
|
||||
topk_ids=topk_ids,
|
||||
topk_weights=topk_weights,
|
||||
expert_scale=torch.ones(
|
||||
case.local_experts, dtype=torch.float32, device=device
|
||||
),
|
||||
situ_gate_beta=torch.full(
|
||||
(case.local_experts,), 4.0, dtype=torch.float32, device=device
|
||||
),
|
||||
situ_linear_beta=torch.full(
|
||||
(case.local_experts,), 25.0, dtype=torch.float32, device=device
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _init_sglang_distributed() -> None:
|
||||
"""Create the rank-one groups that production MoE wrappers expect."""
|
||||
torch.cuda.set_device(0)
|
||||
init_distributed_environment(
|
||||
world_size=1,
|
||||
rank=0,
|
||||
local_rank=0,
|
||||
distributed_init_method="tcp://127.0.0.1:29541",
|
||||
backend="nccl",
|
||||
)
|
||||
initialize_model_parallel(
|
||||
tensor_model_parallel_size=1,
|
||||
expert_model_parallel_size=1,
|
||||
)
|
||||
|
||||
|
||||
def _destroy_sglang_distributed() -> None:
|
||||
destroy_model_parallel()
|
||||
if torch.distributed.is_initialized():
|
||||
torch.distributed.destroy_process_group()
|
||||
|
||||
|
||||
def _quantile_exact(values: torch.Tensor, quantile: float) -> float:
|
||||
"""Return an exact quantile without torch.quantile's 2^24 element limit."""
|
||||
rank = max(1, math.ceil(quantile * values.numel()))
|
||||
return float(torch.kthvalue(values, rank).values.item())
|
||||
|
||||
|
||||
def _metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict:
|
||||
ref = reference.float().reshape(-1)
|
||||
got = candidate.float().reshape(-1)
|
||||
diff = got - ref
|
||||
abs_diff = diff.abs()
|
||||
denominator = torch.sqrt(torch.mean(ref.square())).clamp_min(1e-12)
|
||||
return {
|
||||
"mean_abs": float(abs_diff.mean().item()),
|
||||
"p95_abs": _quantile_exact(abs_diff, 0.95),
|
||||
"p99_abs": _quantile_exact(abs_diff, 0.99),
|
||||
"max_abs": float(abs_diff.max().item()),
|
||||
"nrmse": float((torch.sqrt(torch.mean(diff.square())) / denominator).item()),
|
||||
"cosine": float(F.cosine_similarity(ref, got, dim=0).item()),
|
||||
"reference_rms": float(denominator.item()),
|
||||
"reference_max_abs": float(ref.abs().max().item()),
|
||||
"finite": bool(torch.isfinite(candidate).all()),
|
||||
}
|
||||
|
||||
|
||||
def _assert_close(name: str, reference: torch.Tensor, candidate: torch.Tensor) -> dict:
|
||||
metrics = _metrics(reference, candidate)
|
||||
if not metrics["finite"]:
|
||||
raise AssertionError(f"{name} produced non-finite output: {metrics}")
|
||||
torch.testing.assert_close(candidate, reference, rtol=0.2, atol=0.2)
|
||||
return metrics
|
||||
|
||||
|
||||
def _checkpoint_gate_up_from_flashinfer_up_gate(tensor: torch.Tensor) -> torch.Tensor:
|
||||
up, gate = tensor.chunk(2, dim=1)
|
||||
return torch.cat((gate, up), dim=1).contiguous()
|
||||
|
||||
|
||||
def _build_expert_map(case: Case, device: torch.device) -> torch.Tensor:
|
||||
expert_map = torch.full(
|
||||
(case.total_experts,), -1, dtype=torch.int32, device=device
|
||||
)
|
||||
start = case.local_expert_start
|
||||
expert_map[start : start + case.local_experts] = torch.arange(
|
||||
case.local_experts, dtype=torch.int32, device=device
|
||||
)
|
||||
return expert_map
|
||||
|
||||
|
||||
class BackendPair:
|
||||
def __init__(self, case: Case, seed: int):
|
||||
self.case = case
|
||||
self.data = _prepare_linear_checkpoint_case(
|
||||
case,
|
||||
seed,
|
||||
torch.device("cuda"),
|
||||
)
|
||||
self.config = MoeRunnerConfig(
|
||||
num_experts=case.total_experts,
|
||||
num_local_experts=case.local_experts,
|
||||
hidden_size=case.hidden,
|
||||
intermediate_size_per_partition=case.intermediate,
|
||||
top_k=case.top_k,
|
||||
activation="situ",
|
||||
is_gated=True,
|
||||
gemm1_alpha=4.0,
|
||||
gemm1_clamp_limit=25.0,
|
||||
gate_up_interleaved=False,
|
||||
)
|
||||
logits = torch.zeros(
|
||||
case.tokens,
|
||||
case.total_experts,
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
self.expert_map = _build_expert_map(case, torch.device("cuda"))
|
||||
self.flashinfer_dispatch = StandardDispatchOutput(
|
||||
self.data.input_dq.contiguous(),
|
||||
None,
|
||||
StandardTopKOutput(
|
||||
self.data.topk_weights,
|
||||
self.data.topk_ids,
|
||||
logits,
|
||||
),
|
||||
)
|
||||
marlin_topk_ids = self.expert_map[self.data.topk_ids]
|
||||
self.marlin_dispatch = StandardDispatchOutput(
|
||||
self.data.input_dq.contiguous(),
|
||||
None,
|
||||
StandardTopKOutput(
|
||||
self.data.topk_weights,
|
||||
marlin_topk_ids,
|
||||
logits,
|
||||
),
|
||||
)
|
||||
self.reference = run_reference(self.data, "situ")
|
||||
self.fi_quant = self._build_flashinfer_quant_info()
|
||||
self.marlin_quant = self._build_marlin_quant_info()
|
||||
|
||||
def _build_flashinfer_quant_info(self) -> FlashInferCutlassMxfp4MoeQuantInfo:
|
||||
case = self.case
|
||||
w13_scale = block_scale_interleave(self.data.fc1_sf).reshape_as(
|
||||
self.data.fc1_sf
|
||||
)
|
||||
w2_scale = block_scale_interleave(self.data.fc2_sf).reshape_as(
|
||||
self.data.fc2_sf
|
||||
)
|
||||
return FlashInferCutlassMxfp4MoeQuantInfo(
|
||||
w13_weight=self.data.fc1_q,
|
||||
w2_weight=self.data.fc2_q,
|
||||
w13_weight_scale=w13_scale.contiguous(),
|
||||
w2_weight_scale=w2_scale.contiguous(),
|
||||
mxfp4_weight_global_scale=self.data.expert_scale,
|
||||
w13_bias=None,
|
||||
w2_bias=None,
|
||||
swiglu_alpha=self.data.situ_gate_beta,
|
||||
swiglu_beta=self.data.situ_linear_beta,
|
||||
swiglu_limit=None,
|
||||
moe_tp_size=1,
|
||||
moe_tp_rank=0,
|
||||
moe_ep_size=case.ep_size,
|
||||
moe_ep_rank=case.ep_rank,
|
||||
padded_hidden=case.hidden,
|
||||
)
|
||||
|
||||
def _build_marlin_quant_info(self) -> MarlinMoeQuantInfo:
|
||||
case = self.case
|
||||
layer = SimpleNamespace(
|
||||
w13_weight=torch.nn.Parameter(
|
||||
_checkpoint_gate_up_from_flashinfer_up_gate(self.data.fc1_q.clone()),
|
||||
requires_grad=False,
|
||||
),
|
||||
w2_weight=torch.nn.Parameter(
|
||||
self.data.fc2_q.clone(), requires_grad=False
|
||||
),
|
||||
w13_weight_scale=torch.nn.Parameter(
|
||||
_checkpoint_gate_up_from_flashinfer_up_gate(
|
||||
self.data.fc1_sf.clone()
|
||||
),
|
||||
requires_grad=False,
|
||||
),
|
||||
w2_weight_scale=torch.nn.Parameter(
|
||||
self.data.fc2_sf.clone(), requires_grad=False
|
||||
),
|
||||
orig_dtype=torch.bfloat16,
|
||||
)
|
||||
prepare_moe_mxfp4_layer_for_marlin(layer)
|
||||
return MarlinMoeQuantInfo(
|
||||
w13_qweight=layer.w13_weight,
|
||||
w2_qweight=layer.w2_weight,
|
||||
w13_scales=layer.w13_weight_scale,
|
||||
w2_scales=layer.w2_weight_scale,
|
||||
w13_g_idx_sort_indices=None,
|
||||
w2_g_idx_sort_indices=None,
|
||||
weight_bits=4,
|
||||
is_k_full=True,
|
||||
expert_map=self.expert_map,
|
||||
global_num_experts=case.total_experts,
|
||||
)
|
||||
|
||||
def flashinfer(self) -> torch.Tensor:
|
||||
return fused_experts_none_to_flashinfer_mxfp4(
|
||||
self.flashinfer_dispatch, self.fi_quant, self.config
|
||||
).hidden_states
|
||||
|
||||
def marlin(self) -> torch.Tensor:
|
||||
return fused_experts_none_to_marlin(
|
||||
self.marlin_dispatch, self.marlin_quant, self.config
|
||||
).hidden_states[:, : self.case.hidden]
|
||||
|
||||
def validate(self) -> dict:
|
||||
fi = self.flashinfer()
|
||||
marlin = self.marlin()
|
||||
torch.cuda.synchronize()
|
||||
metrics = {
|
||||
"flashinfer_vs_reference": _metrics(self.reference, fi),
|
||||
"marlin_vs_reference": _metrics(self.reference, marlin),
|
||||
"flashinfer_vs_marlin": _metrics(marlin, fi),
|
||||
}
|
||||
print("CORRECTNESS " + json.dumps(metrics, sort_keys=True), flush=True)
|
||||
_assert_close("flashinfer", self.reference, fi)
|
||||
_assert_close("marlin", self.reference, marlin)
|
||||
_assert_close("cross_backend", marlin, fi)
|
||||
return metrics
|
||||
|
||||
|
||||
def _measure(fn: Callable[[], torch.Tensor], iterations: int) -> dict:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
torch.cuda.synchronize()
|
||||
wall_start = time.perf_counter()
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
fn()
|
||||
end.record()
|
||||
torch.cuda.synchronize()
|
||||
wall_ms = (time.perf_counter() - wall_start) * 1000.0 / iterations
|
||||
return {
|
||||
"gpu_ms": start.elapsed_time(end) / iterations,
|
||||
"wall_ms": wall_ms,
|
||||
}
|
||||
|
||||
|
||||
def _warmup(pair: BackendPair, iterations: int) -> None:
|
||||
for _ in range(iterations):
|
||||
pair.marlin()
|
||||
pair.flashinfer()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def run_benchmark(args: argparse.Namespace) -> dict:
|
||||
report = {
|
||||
"created_at": datetime.now().astimezone().isoformat(),
|
||||
"device": torch.cuda.get_device_name(),
|
||||
"compute_capability": list(torch.cuda.get_device_capability()),
|
||||
"cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
|
||||
"rounds": args.rounds,
|
||||
"iterations": args.iterations,
|
||||
"shapes": {},
|
||||
}
|
||||
for shape_name in args.shapes:
|
||||
pair = BackendPair(SHAPES[shape_name], args.seed)
|
||||
correctness = pair.validate()
|
||||
_warmup(pair, args.warmup)
|
||||
rounds = []
|
||||
for round_index in range(args.rounds):
|
||||
order = (
|
||||
("marlin", "flashinfer")
|
||||
if round_index % 2 == 0
|
||||
else ("flashinfer", "marlin")
|
||||
)
|
||||
result = {"round": round_index + 1, "order": list(order)}
|
||||
for backend in order:
|
||||
result[backend] = _measure(getattr(pair, backend), args.iterations)
|
||||
rounds.append(result)
|
||||
summary = {}
|
||||
for backend in ("marlin", "flashinfer"):
|
||||
summary[backend] = {
|
||||
metric: statistics.median(
|
||||
result[backend][metric] for result in rounds
|
||||
)
|
||||
for metric in ("gpu_ms", "wall_ms")
|
||||
}
|
||||
summary["speedup_marlin_over_flashinfer"] = {
|
||||
metric: summary["marlin"][metric] / summary["flashinfer"][metric]
|
||||
for metric in ("gpu_ms", "wall_ms")
|
||||
}
|
||||
report["shapes"][shape_name] = {
|
||||
"case": asdict(pair.case),
|
||||
"correctness": correctness,
|
||||
"rounds": rounds,
|
||||
"summary": summary,
|
||||
}
|
||||
del pair
|
||||
torch.cuda.empty_cache()
|
||||
return report
|
||||
|
||||
|
||||
def run_profile(args: argparse.Namespace) -> None:
|
||||
shape_names = tuple(SHAPES) if args.shape == "all" else (args.shape,)
|
||||
pairs = [
|
||||
(shape_name, BackendPair(SHAPES[shape_name], args.seed))
|
||||
for shape_name in shape_names
|
||||
]
|
||||
for shape_name, pair in pairs:
|
||||
print(
|
||||
"PROFILE_CORRECTNESS "
|
||||
+ json.dumps({shape_name: pair.validate()}, sort_keys=True),
|
||||
flush=True,
|
||||
)
|
||||
_warmup(pair, args.warmup)
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.cudart().cudaProfilerStart()
|
||||
backends = (
|
||||
(args.backend,)
|
||||
if args.backend != "both"
|
||||
else ("marlin", "flashinfer")
|
||||
)
|
||||
for shape_name, pair in pairs:
|
||||
for backend in backends:
|
||||
for index in range(args.profile_iterations):
|
||||
label = f"{backend}_{shape_name}_iteration_{index + 1}"
|
||||
context = (
|
||||
torch.cuda.nvtx.range(label)
|
||||
if torch.cuda.is_available()
|
||||
else nullcontext()
|
||||
)
|
||||
with context:
|
||||
getattr(pair, backend)()
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.cudart().cudaProfilerStop()
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--mode", choices=("benchmark", "profile"), default="benchmark")
|
||||
parser.add_argument("--shapes", nargs="+", choices=tuple(SHAPES), default=list(SHAPES))
|
||||
parser.add_argument("--shape", choices=(*SHAPES, "all"), default="decode")
|
||||
parser.add_argument("--backend", choices=("marlin", "flashinfer", "both"), default="both")
|
||||
parser.add_argument("--rounds", type=int, default=3)
|
||||
parser.add_argument("--iterations", type=int, default=30)
|
||||
parser.add_argument("--warmup", type=int, default=5)
|
||||
parser.add_argument("--profile-iterations", type=int, default=3)
|
||||
parser.add_argument("--seed", type=int, default=20260817)
|
||||
parser.add_argument("--output", type=Path)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA is required")
|
||||
expected_uuid = os.environ.get("EXPECTED_GPU_UUID")
|
||||
actual_uuid = str(torch.cuda.get_device_properties(0).uuid)
|
||||
if expected_uuid and actual_uuid.lower() != expected_uuid.lower():
|
||||
raise SystemExit(
|
||||
f"GPU binding mismatch: expected {expected_uuid}, got {actual_uuid}"
|
||||
)
|
||||
if torch.cuda.get_device_capability() != (12, 0):
|
||||
raise SystemExit(f"SM120 is required, got {torch.cuda.get_device_capability()}")
|
||||
_init_sglang_distributed()
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
if args.mode == "profile":
|
||||
run_profile(args)
|
||||
return
|
||||
report = run_benchmark(args)
|
||||
finally:
|
||||
_destroy_sglang_distributed()
|
||||
rendered = json.dumps(report, indent=2, sort_keys=True)
|
||||
print(rendered)
|
||||
if args.output is not None:
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(rendered + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -0,0 +1,88 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Reproduce the missing SM120 CUTLASS SiTU contract with MXFP8 x MXFP4."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
import sys
|
||||
|
||||
import torch
|
||||
|
||||
from flashinfer import mxfp8_quantize
|
||||
from flashinfer.fused_moe import cutlass_fused_moe
|
||||
from flashinfer.tllm_enums import ActivationType
|
||||
|
||||
|
||||
def load_official_helpers(flashinfer_source: Path):
|
||||
sys.path.insert(0, str(flashinfer_source))
|
||||
return importlib.import_module(
|
||||
"tests.moe.test_trtllm_cutlass_fused_moe"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--flashinfer-source",
|
||||
type=Path,
|
||||
default=Path("/data/hzy/src/flashinfer-kimi-sm120"),
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
helpers = load_official_helpers(args.flashinfer_source)
|
||||
torch.manual_seed(42)
|
||||
|
||||
tokens, hidden, experts, top_k, intermediate = 1, 128, 2, 2, 128
|
||||
dtype = torch.bfloat16
|
||||
x = torch.randn(tokens, hidden, dtype=dtype, device="cuda")
|
||||
w31 = torch.randn(
|
||||
experts, 2 * intermediate, hidden, dtype=dtype, device="cuda"
|
||||
) / 10
|
||||
w2 = torch.randn(
|
||||
experts, hidden, intermediate, dtype=dtype, device="cuda"
|
||||
) / 10
|
||||
|
||||
x_q, x_sf = mxfp8_quantize(x, True, 32)
|
||||
w31_q, w31_sf = helpers.quant_mxfp4_batches(w31, experts)
|
||||
w2_q, w2_sf = helpers.quant_mxfp4_batches(w2, experts)
|
||||
logits = torch.randn(tokens, experts, dtype=dtype, device="cuda")
|
||||
weights, ids = helpers.compute_routing(logits, top_k)
|
||||
|
||||
neutral_global_scale = torch.ones(experts, dtype=torch.float32, device="cuda")
|
||||
situ_gate_beta = torch.full(
|
||||
(experts,), 4.0, dtype=torch.float32, device="cuda"
|
||||
)
|
||||
situ_linear_beta = torch.full(
|
||||
(experts,), 25.0, dtype=torch.float32, device="cuda"
|
||||
)
|
||||
output = torch.empty_like(x)
|
||||
|
||||
cutlass_fused_moe(
|
||||
input=x_q,
|
||||
token_selected_experts=ids.to(torch.int32),
|
||||
token_final_scales=weights,
|
||||
fc1_expert_weights=w31_q.contiguous().view(torch.int64),
|
||||
fc2_expert_weights=w2_q.contiguous().view(torch.int64),
|
||||
output_dtype=dtype,
|
||||
quant_scales=[
|
||||
w31_sf.view(torch.int32),
|
||||
neutral_global_scale,
|
||||
w2_sf.view(torch.int32),
|
||||
neutral_global_scale,
|
||||
],
|
||||
input_sf=x_sf,
|
||||
swiglu_alpha=situ_gate_beta,
|
||||
swiglu_beta=situ_linear_beta,
|
||||
swiglu_limit=None,
|
||||
activation_type=ActivationType.Situ,
|
||||
use_mxfp8_act_scaling=True,
|
||||
output=output,
|
||||
)
|
||||
|
||||
raise SystemExit("unexpected success: SM120 CUTLASS SiTU gap is no longer present")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -0,0 +1,81 @@
|
||||
{
|
||||
"arguments": {
|
||||
"activation": "swiglu",
|
||||
"benchmark_iterations": 0,
|
||||
"output": "results/phase2_control_swiglu_real_gpu6.json",
|
||||
"require_situ": false,
|
||||
"seed": 20260817,
|
||||
"suite": "real",
|
||||
"warmup": 3
|
||||
},
|
||||
"compute_capability": [
|
||||
12,
|
||||
0
|
||||
],
|
||||
"created_at": "2026-08-17T12:02:07.099079+08:00",
|
||||
"cuda_version": "13.0",
|
||||
"cuda_visible_devices": "6",
|
||||
"device": "NVIDIA RTX 6000D",
|
||||
"flashinfer_version": "0.6.18",
|
||||
"logical_cuda_device": 0,
|
||||
"results": [
|
||||
{
|
||||
"activation": "swiglu",
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_decode_m1_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 1,
|
||||
"top_k": 16
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 13.75742506980896,
|
||||
"metrics": {
|
||||
"cosine": 0.9995965957641602,
|
||||
"finite_elements": 3584,
|
||||
"max_abs": 0.01025390625,
|
||||
"nrmse": 0.028520602732896805,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "swiglu",
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m128_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 128,
|
||||
"top_k": 16
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 12.48392629623413,
|
||||
"metrics": {
|
||||
"cosine": 0.9996211528778076,
|
||||
"finite_elements": 458752,
|
||||
"max_abs": 0.013671875,
|
||||
"nrmse": 0.027537915855646133,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
}
|
||||
],
|
||||
"schema_version": 1,
|
||||
"torch_version": "2.11.0+cu130"
|
||||
}
|
||||
@ -0,0 +1,214 @@
|
||||
{
|
||||
"arguments": {
|
||||
"activation": "swiglu",
|
||||
"benchmark_iterations": 0,
|
||||
"output": "results/phase2_control_swiglu_smoke_gpu6.json",
|
||||
"require_situ": false,
|
||||
"seed": 20260817,
|
||||
"suite": "smoke",
|
||||
"warmup": 3
|
||||
},
|
||||
"compute_capability": [
|
||||
12,
|
||||
0
|
||||
],
|
||||
"created_at": "2026-08-17T12:01:32.407549+08:00",
|
||||
"cuda_version": "13.0",
|
||||
"cuda_visible_devices": "6",
|
||||
"device": "NVIDIA RTX 6000D",
|
||||
"flashinfer_version": "0.6.18",
|
||||
"logical_cuda_device": 0,
|
||||
"results": [
|
||||
{
|
||||
"activation": "swiglu",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "decode_m1_uniform",
|
||||
"routing": "uniform",
|
||||
"tokens": 1,
|
||||
"top_k": 2
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 1.2227470874786377,
|
||||
"metrics": {
|
||||
"cosine": 0.9998119473457336,
|
||||
"finite_elements": 128,
|
||||
"max_abs": 0.03125,
|
||||
"nrmse": 0.019757770001888275,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "swiglu",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "decode_m2_skewed",
|
||||
"routing": "skewed",
|
||||
"tokens": 2,
|
||||
"top_k": 2
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.03631448745727539,
|
||||
"metrics": {
|
||||
"cosine": 0.9997045397758484,
|
||||
"finite_elements": 256,
|
||||
"max_abs": 0.03125,
|
||||
"nrmse": 0.024355093017220497,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "swiglu",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 8,
|
||||
"name": "batch_m8_empty_experts",
|
||||
"routing": "skewed",
|
||||
"tokens": 8,
|
||||
"top_k": 2
|
||||
},
|
||||
"cuda_graph": {
|
||||
"replay_matches_eager": {
|
||||
"cosine": 0.9999999403953552,
|
||||
"finite_elements": 1024,
|
||||
"max_abs": 0.0,
|
||||
"nrmse": 0.0,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"replays_bitwise_equal": true
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.08502435684204102,
|
||||
"metrics": {
|
||||
"cosine": 0.9996250867843628,
|
||||
"finite_elements": 1024,
|
||||
"max_abs": 0.0546875,
|
||||
"nrmse": 0.02746359072625637,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "swiglu",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 8,
|
||||
"name": "batch_m32_masked_slot",
|
||||
"routing": "masked",
|
||||
"tokens": 32,
|
||||
"top_k": 4
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.013071775436401367,
|
||||
"metrics": {
|
||||
"cosine": 0.9996730089187622,
|
||||
"finite_elements": 4096,
|
||||
"max_abs": 0.0390625,
|
||||
"nrmse": 0.02586047165095806,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "swiglu",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 2,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "ep2_global_ids",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 8,
|
||||
"top_k": 4
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.0670325756072998,
|
||||
"metrics": {
|
||||
"cosine": 0.9996550679206848,
|
||||
"finite_elements": 1024,
|
||||
"max_abs": 0.01953125,
|
||||
"nrmse": 0.02657792530953884,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "swiglu",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": true,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "nonfinite_activation",
|
||||
"routing": "uniform",
|
||||
"tokens": 2,
|
||||
"top_k": 2
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.009414196014404297,
|
||||
"metrics": {
|
||||
"candidate_inf_elements": 68,
|
||||
"candidate_nan_elements": 60,
|
||||
"candidate_nonfinite_rows": [
|
||||
false,
|
||||
true
|
||||
],
|
||||
"contract": "same_affected_rows_and_close_unaffected_rows",
|
||||
"cosine": 0.9998903870582581,
|
||||
"finite_elements": 128,
|
||||
"max_abs": 0.025146484375,
|
||||
"nrmse": 0.015044113621115685,
|
||||
"reference_nonfinite_rows": [
|
||||
false,
|
||||
true
|
||||
],
|
||||
"same_nonfinite_mask": false
|
||||
},
|
||||
"status": "PASS"
|
||||
}
|
||||
],
|
||||
"schema_version": 1,
|
||||
"torch_version": "2.11.0+cu130"
|
||||
}
|
||||
@ -0,0 +1,61 @@
|
||||
{
|
||||
"arguments": {
|
||||
"activation": "situ",
|
||||
"benchmark_iterations": 0,
|
||||
"output": "results/phase2_target_situ_real_gpu6.json",
|
||||
"require_situ": false,
|
||||
"seed": 20260817,
|
||||
"suite": "real",
|
||||
"warmup": 3
|
||||
},
|
||||
"compute_capability": [
|
||||
12,
|
||||
0
|
||||
],
|
||||
"created_at": "2026-08-17T12:04:05.231369+08:00",
|
||||
"cuda_version": "13.0",
|
||||
"cuda_visible_devices": "6",
|
||||
"device": "NVIDIA RTX 6000D",
|
||||
"flashinfer_version": "0.6.18",
|
||||
"logical_cuda_device": 0,
|
||||
"results": [
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_decode_m1_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 1,
|
||||
"top_k": 16
|
||||
},
|
||||
"elapsed_s": 15.5287344455719,
|
||||
"error": "Check failed: fc1_expert_weights.size(1) == fc2_expert_weights.size(2) * mInnerDimMultiplier (6144 vs. 3072) : fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.",
|
||||
"status": "EXPECTED_MISSING_SITU"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m128_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 128,
|
||||
"top_k": 16
|
||||
},
|
||||
"elapsed_s": 11.09748888015747,
|
||||
"error": "Check failed: fc1_expert_weights.size(1) == fc2_expert_weights.size(2) * mInnerDimMultiplier (6144 vs. 3072) : fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.",
|
||||
"status": "EXPECTED_MISSING_SITU"
|
||||
}
|
||||
],
|
||||
"schema_version": 1,
|
||||
"torch_version": "2.11.0+cu130"
|
||||
}
|
||||
@ -0,0 +1,133 @@
|
||||
{
|
||||
"arguments": {
|
||||
"activation": "situ",
|
||||
"benchmark_iterations": 0,
|
||||
"output": "results/phase2_target_situ_smoke_gpu6.json",
|
||||
"require_situ": false,
|
||||
"seed": 20260817,
|
||||
"suite": "smoke",
|
||||
"warmup": 3
|
||||
},
|
||||
"compute_capability": [
|
||||
12,
|
||||
0
|
||||
],
|
||||
"created_at": "2026-08-17T12:00:47.947426+08:00",
|
||||
"cuda_version": "13.0",
|
||||
"cuda_visible_devices": "6",
|
||||
"device": "NVIDIA RTX 6000D",
|
||||
"flashinfer_version": "0.6.18",
|
||||
"logical_cuda_device": 0,
|
||||
"results": [
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "decode_m1_uniform",
|
||||
"routing": "uniform",
|
||||
"tokens": 1,
|
||||
"top_k": 2
|
||||
},
|
||||
"elapsed_s": 2.840806484222412,
|
||||
"error": "Check failed: fc1_expert_weights.size(1) == fc2_expert_weights.size(2) * mInnerDimMultiplier (256 vs. 128) : fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.",
|
||||
"status": "EXPECTED_MISSING_SITU"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "decode_m2_skewed",
|
||||
"routing": "skewed",
|
||||
"tokens": 2,
|
||||
"top_k": 2
|
||||
},
|
||||
"elapsed_s": 0.03530144691467285,
|
||||
"error": "Check failed: fc1_expert_weights.size(1) == fc2_expert_weights.size(2) * mInnerDimMultiplier (256 vs. 128) : fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.",
|
||||
"status": "EXPECTED_MISSING_SITU"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 8,
|
||||
"name": "batch_m8_empty_experts",
|
||||
"routing": "skewed",
|
||||
"tokens": 8,
|
||||
"top_k": 2
|
||||
},
|
||||
"elapsed_s": 0.016995668411254883,
|
||||
"error": "Check failed: fc1_expert_weights.size(1) == fc2_expert_weights.size(2) * mInnerDimMultiplier (256 vs. 128) : fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.",
|
||||
"status": "EXPECTED_MISSING_SITU"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 8,
|
||||
"name": "batch_m32_masked_slot",
|
||||
"routing": "masked",
|
||||
"tokens": 32,
|
||||
"top_k": 4
|
||||
},
|
||||
"elapsed_s": 0.01120138168334961,
|
||||
"error": "Check failed: fc1_expert_weights.size(1) == fc2_expert_weights.size(2) * mInnerDimMultiplier (256 vs. 128) : fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.",
|
||||
"status": "EXPECTED_MISSING_SITU"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 2,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "ep2_global_ids",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 8,
|
||||
"top_k": 4
|
||||
},
|
||||
"elapsed_s": 0.007741689682006836,
|
||||
"error": "Check failed: fc1_expert_weights.size(1) == fc2_expert_weights.size(2) * mInnerDimMultiplier (256 vs. 128) : fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.",
|
||||
"status": "EXPECTED_MISSING_SITU"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": true,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "nonfinite_activation",
|
||||
"routing": "uniform",
|
||||
"tokens": 2,
|
||||
"top_k": 2
|
||||
},
|
||||
"elapsed_s": 0.007404804229736328,
|
||||
"error": "Check failed: fc1_expert_weights.size(1) == fc2_expert_weights.size(2) * mInnerDimMultiplier (256 vs. 128) : fc1_expert_weights inter size must be equal to fc2_expert_weights inter size.",
|
||||
"status": "EXPECTED_MISSING_SITU"
|
||||
}
|
||||
],
|
||||
"schema_version": 1,
|
||||
"torch_version": "2.11.0+cu130"
|
||||
}
|
||||
@ -0,0 +1,296 @@
|
||||
{
|
||||
"arguments": {
|
||||
"activation": "situ",
|
||||
"benchmark_iterations": 3,
|
||||
"output": "results/phase3_target_situ_all_gpu6_v3.json",
|
||||
"require_situ": true,
|
||||
"seed": 20260817,
|
||||
"suite": "all",
|
||||
"warmup": 2
|
||||
},
|
||||
"compute_capability": [
|
||||
12,
|
||||
0
|
||||
],
|
||||
"created_at": "2026-08-17T14:54:33.526558+08:00",
|
||||
"cuda_version": "13.0",
|
||||
"cuda_visible_devices": "6",
|
||||
"device": "NVIDIA RTX 6000D",
|
||||
"flashinfer_version": "0.6.18",
|
||||
"logical_cuda_device": 0,
|
||||
"results": [
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "decode_m1_uniform",
|
||||
"routing": "uniform",
|
||||
"tokens": 1,
|
||||
"top_k": 2
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 1.2600979804992676,
|
||||
"latency_ms": {
|
||||
"candidate": 0.24848000208536783,
|
||||
"reference": 0.7109653155008951
|
||||
},
|
||||
"metrics": {
|
||||
"cosine": 0.9996331930160522,
|
||||
"finite_elements": 128,
|
||||
"max_abs": 0.046875,
|
||||
"nrmse": 0.038191329687833786,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "decode_m2_skewed",
|
||||
"routing": "skewed",
|
||||
"tokens": 2,
|
||||
"top_k": 2
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.041979312896728516,
|
||||
"latency_ms": {
|
||||
"candidate": 0.24086399873097739,
|
||||
"reference": 0.6774506568908691
|
||||
},
|
||||
"metrics": {
|
||||
"cosine": 0.9996317625045776,
|
||||
"finite_elements": 256,
|
||||
"max_abs": 0.046875,
|
||||
"nrmse": 0.03233937919139862,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 8,
|
||||
"name": "batch_m8_empty_experts",
|
||||
"routing": "skewed",
|
||||
"tokens": 8,
|
||||
"top_k": 2
|
||||
},
|
||||
"cuda_graph": {
|
||||
"replay_matches_eager": {
|
||||
"cosine": 0.9999998807907104,
|
||||
"finite_elements": 1024,
|
||||
"max_abs": 0.0,
|
||||
"nrmse": 0.0,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"replays_bitwise_equal": true
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.0908958911895752,
|
||||
"latency_ms": {
|
||||
"candidate": 0.23886932929356894,
|
||||
"reference": 0.7901653448740641
|
||||
},
|
||||
"metrics": {
|
||||
"cosine": 0.9996166229248047,
|
||||
"finite_elements": 1024,
|
||||
"max_abs": 0.046875,
|
||||
"nrmse": 0.02830333448946476,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 8,
|
||||
"name": "batch_m32_masked_slot",
|
||||
"routing": "masked",
|
||||
"tokens": 32,
|
||||
"top_k": 4
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.02089524269104004,
|
||||
"latency_ms": {
|
||||
"candidate": 0.22908800840377808,
|
||||
"reference": 1.0701546669006348
|
||||
},
|
||||
"metrics": {
|
||||
"cosine": 0.9996225237846375,
|
||||
"finite_elements": 4096,
|
||||
"max_abs": 0.03125,
|
||||
"nrmse": 0.027592986822128296,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 2,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "ep2_global_ids",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 8,
|
||||
"top_k": 4
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.07790017127990723,
|
||||
"latency_ms": {
|
||||
"candidate": 0.23056000471115112,
|
||||
"reference": 1.2077759901682537
|
||||
},
|
||||
"metrics": {
|
||||
"cosine": 0.9995924830436707,
|
||||
"finite_elements": 1024,
|
||||
"max_abs": 0.0166015625,
|
||||
"nrmse": 0.02855238877236843,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 0,
|
||||
"ep_size": 1,
|
||||
"hidden": 128,
|
||||
"inject_nonfinite": true,
|
||||
"intermediate": 128,
|
||||
"local_experts": 4,
|
||||
"name": "nonfinite_activation",
|
||||
"routing": "uniform",
|
||||
"tokens": 2,
|
||||
"top_k": 2
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 0.008931636810302734,
|
||||
"metrics": {
|
||||
"candidate_inf_elements": 0,
|
||||
"candidate_nan_elements": 0,
|
||||
"candidate_nonfinite_rows": [
|
||||
false,
|
||||
false
|
||||
],
|
||||
"contract": "mxfp8_satfinite_input_and_finite_candidate",
|
||||
"input_payload_finite": true,
|
||||
"max_scale_code": 254,
|
||||
"reference_nonfinite_rows_diagnostic_only": [
|
||||
false,
|
||||
true
|
||||
],
|
||||
"scale_has_nan_code": false
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_decode_m1_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 1,
|
||||
"top_k": 16
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 12.66703748703003,
|
||||
"latency_ms": {
|
||||
"candidate": 0.295904000600179,
|
||||
"reference": 2.942485491434733
|
||||
},
|
||||
"metrics": {
|
||||
"cosine": 0.9996223449707031,
|
||||
"finite_elements": 3584,
|
||||
"max_abs": 0.009765625,
|
||||
"nrmse": 0.027483593672513962,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
},
|
||||
{
|
||||
"activation": "situ",
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m128_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 128,
|
||||
"top_k": 16
|
||||
},
|
||||
"determinism": {
|
||||
"bitwise_equal": true,
|
||||
"runs": 3
|
||||
},
|
||||
"elapsed_s": 12.947205066680908,
|
||||
"latency_ms": {
|
||||
"candidate": 0.6200213432312012,
|
||||
"reference": 8.530954360961914
|
||||
},
|
||||
"metrics": {
|
||||
"cosine": 0.9996151924133301,
|
||||
"finite_elements": 458752,
|
||||
"max_abs": 0.013671875,
|
||||
"nrmse": 0.02774309180676937,
|
||||
"same_nonfinite_mask": true
|
||||
},
|
||||
"status": "PASS"
|
||||
}
|
||||
],
|
||||
"schema_version": 1,
|
||||
"torch_version": "2.11.0+cu130"
|
||||
}
|
||||
@ -0,0 +1,293 @@
|
||||
{
|
||||
"compute_capability": [
|
||||
12,
|
||||
0
|
||||
],
|
||||
"created_at": "2026-08-17T09:47:32.426483+00:00",
|
||||
"cuda_visible_devices": null,
|
||||
"device": "NVIDIA RTX 6000D",
|
||||
"iterations": 100,
|
||||
"rounds": 5,
|
||||
"shapes": {
|
||||
"decode": {
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_decode_m1_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 1,
|
||||
"top_k": 16
|
||||
},
|
||||
"correctness": {
|
||||
"flashinfer_vs_marlin": {
|
||||
"cosine": 0.9996320605278015,
|
||||
"finite": true,
|
||||
"max_abs": 0.00978851318359375,
|
||||
"mean_abs": 0.002149919280782342,
|
||||
"nrmse": 0.02712450921535492,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.007080078125,
|
||||
"reference_max_abs": 0.37890625,
|
||||
"reference_rms": 0.09998540580272675
|
||||
},
|
||||
"flashinfer_vs_reference": {
|
||||
"cosine": 0.9996223449707031,
|
||||
"finite": true,
|
||||
"max_abs": 0.009765625,
|
||||
"mean_abs": 0.00218405993655324,
|
||||
"nrmse": 0.027483593672513962,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.0072021484375,
|
||||
"reference_max_abs": 0.37890625,
|
||||
"reference_rms": 0.10006484389305115
|
||||
},
|
||||
"marlin_vs_reference": {
|
||||
"cosine": 0.9999805688858032,
|
||||
"finite": true,
|
||||
"max_abs": 0.00390625,
|
||||
"mean_abs": 0.00045798509381711483,
|
||||
"nrmse": 0.0062673031352460384,
|
||||
"p95_abs": 0.0009765625,
|
||||
"p99_abs": 0.001953125,
|
||||
"reference_max_abs": 0.37890625,
|
||||
"reference_rms": 0.10006484389305115
|
||||
}
|
||||
},
|
||||
"rounds": [
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.3265369415283203,
|
||||
"wall_ms": 0.32803289126604795
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.30863103866577146,
|
||||
"wall_ms": 0.3124843700788915
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 1
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.3199043273925781,
|
||||
"wall_ms": 0.32471023965626955
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.3098236846923828,
|
||||
"wall_ms": 0.3145282808691263
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 2
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.32463966369628905,
|
||||
"wall_ms": 0.3285454004071653
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.30631423950195313,
|
||||
"wall_ms": 0.3102407115511596
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 3
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.31916448593139646,
|
||||
"wall_ms": 0.32354930881410837
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.3083350372314453,
|
||||
"wall_ms": 0.3130973596125841
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 4
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.3248342514038086,
|
||||
"wall_ms": 0.326665451284498
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.31008224487304686,
|
||||
"wall_ms": 0.31198404962196946
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 5
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.32463966369628905,
|
||||
"wall_ms": 0.326665451284498
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.30863103866577146,
|
||||
"wall_ms": 0.3124843700788915
|
||||
},
|
||||
"speedup_marlin_over_flashinfer": {
|
||||
"gpu_ms": 0.9506880186843276,
|
||||
"wall_ms": 0.9565883654061232
|
||||
}
|
||||
}
|
||||
},
|
||||
"prefill": {
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m128_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 128,
|
||||
"top_k": 16
|
||||
},
|
||||
"correctness": {
|
||||
"flashinfer_vs_marlin": {
|
||||
"cosine": 0.9996392726898193,
|
||||
"finite": true,
|
||||
"max_abs": 0.0126953125,
|
||||
"mean_abs": 0.0021022639703005552,
|
||||
"nrmse": 0.02686966024339199,
|
||||
"p95_abs": 0.005126953125,
|
||||
"p99_abs": 0.0068359375,
|
||||
"reference_max_abs": 0.462890625,
|
||||
"reference_rms": 0.09850681573152542
|
||||
},
|
||||
"flashinfer_vs_reference": {
|
||||
"cosine": 0.9996151924133301,
|
||||
"finite": true,
|
||||
"max_abs": 0.013671875,
|
||||
"mean_abs": 0.00217337254434824,
|
||||
"nrmse": 0.02774309180676937,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.00689697265625,
|
||||
"reference_max_abs": 0.46484375,
|
||||
"reference_rms": 0.09856347739696503
|
||||
},
|
||||
"marlin_vs_reference": {
|
||||
"cosine": 0.9999709129333496,
|
||||
"finite": true,
|
||||
"max_abs": 0.00390625,
|
||||
"mean_abs": 0.000566127069760114,
|
||||
"nrmse": 0.007645560894161463,
|
||||
"p95_abs": 0.00146484375,
|
||||
"p99_abs": 0.001953125,
|
||||
"reference_max_abs": 0.46484375,
|
||||
"reference_rms": 0.09856347739696503
|
||||
}
|
||||
},
|
||||
"rounds": [
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9432563018798829,
|
||||
"wall_ms": 0.9479056694544852
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8844432067871094,
|
||||
"wall_ms": 0.8977009798400104
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 1
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9416944122314453,
|
||||
"wall_ms": 0.9463393897749484
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.88412353515625,
|
||||
"wall_ms": 0.8888133312575519
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 2
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9426467132568359,
|
||||
"wall_ms": 0.9474158892408013
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8948384094238281,
|
||||
"wall_ms": 0.8993784594349563
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 3
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9381990051269531,
|
||||
"wall_ms": 0.9428730909712613
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.89110595703125,
|
||||
"wall_ms": 0.8956197183579206
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 4
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9436812591552735,
|
||||
"wall_ms": 0.9483181592077017
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8909331512451172,
|
||||
"wall_ms": 0.895569019485265
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 5
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9426467132568359,
|
||||
"wall_ms": 0.9474158892408013
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8909331512451172,
|
||||
"wall_ms": 0.8956197183579206
|
||||
},
|
||||
"speedup_marlin_over_flashinfer": {
|
||||
"gpu_ms": 0.945140038909117,
|
||||
"wall_ms": 0.9453290033752898
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,173 @@
|
||||
{
|
||||
"compute_capability": [
|
||||
12,
|
||||
0
|
||||
],
|
||||
"created_at": "2026-08-18T02:17:06.360422+00:00",
|
||||
"cuda_visible_devices": null,
|
||||
"device": "NVIDIA RTX 6000D",
|
||||
"iterations": 1,
|
||||
"rounds": 1,
|
||||
"shapes": {
|
||||
"decode": {
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_decode_m1_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 1,
|
||||
"top_k": 16
|
||||
},
|
||||
"correctness": {
|
||||
"flashinfer_vs_marlin": {
|
||||
"cosine": 0.9996320605278015,
|
||||
"finite": true,
|
||||
"max_abs": 0.00978851318359375,
|
||||
"mean_abs": 0.002149919280782342,
|
||||
"nrmse": 0.02712450921535492,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.007080078125,
|
||||
"reference_max_abs": 0.37890625,
|
||||
"reference_rms": 0.09998540580272675
|
||||
},
|
||||
"flashinfer_vs_reference": {
|
||||
"cosine": 0.9996223449707031,
|
||||
"finite": true,
|
||||
"max_abs": 0.009765625,
|
||||
"mean_abs": 0.00218405993655324,
|
||||
"nrmse": 0.027483593672513962,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.0072021484375,
|
||||
"reference_max_abs": 0.37890625,
|
||||
"reference_rms": 0.10006484389305115
|
||||
},
|
||||
"marlin_vs_reference": {
|
||||
"cosine": 0.9999805688858032,
|
||||
"finite": true,
|
||||
"max_abs": 0.00390625,
|
||||
"mean_abs": 0.00045798509381711483,
|
||||
"nrmse": 0.0062673031352460384,
|
||||
"p95_abs": 0.0009765625,
|
||||
"p99_abs": 0.001953125,
|
||||
"reference_max_abs": 0.37890625,
|
||||
"reference_rms": 0.10006484389305115
|
||||
}
|
||||
},
|
||||
"rounds": [
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.5889919996261597,
|
||||
"wall_ms": 0.6109171081334352
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.5362880229949951,
|
||||
"wall_ms": 0.6202820222824812
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 1
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.5889919996261597,
|
||||
"wall_ms": 0.6109171081334352
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.5362880229949951,
|
||||
"wall_ms": 0.6202820222824812
|
||||
},
|
||||
"speedup_marlin_over_flashinfer": {
|
||||
"gpu_ms": 0.9105183488661707,
|
||||
"wall_ms": 1.015329271392741
|
||||
}
|
||||
}
|
||||
},
|
||||
"prefill": {
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m128_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 128,
|
||||
"top_k": 16
|
||||
},
|
||||
"correctness": {
|
||||
"flashinfer_vs_marlin": {
|
||||
"cosine": 0.9996392726898193,
|
||||
"finite": true,
|
||||
"max_abs": 0.0126953125,
|
||||
"mean_abs": 0.0021022639703005552,
|
||||
"nrmse": 0.02686966024339199,
|
||||
"p95_abs": 0.005126953125,
|
||||
"p99_abs": 0.0068359375,
|
||||
"reference_max_abs": 0.462890625,
|
||||
"reference_rms": 0.09850681573152542
|
||||
},
|
||||
"flashinfer_vs_reference": {
|
||||
"cosine": 0.9996151924133301,
|
||||
"finite": true,
|
||||
"max_abs": 0.013671875,
|
||||
"mean_abs": 0.00217337254434824,
|
||||
"nrmse": 0.02774309180676937,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.00689697265625,
|
||||
"reference_max_abs": 0.46484375,
|
||||
"reference_rms": 0.09856347739696503
|
||||
},
|
||||
"marlin_vs_reference": {
|
||||
"cosine": 0.9999709129333496,
|
||||
"finite": true,
|
||||
"max_abs": 0.00390625,
|
||||
"mean_abs": 0.000566127069760114,
|
||||
"nrmse": 0.007645560894161463,
|
||||
"p95_abs": 0.00146484375,
|
||||
"p99_abs": 0.001953125,
|
||||
"reference_max_abs": 0.46484375,
|
||||
"reference_rms": 0.09856347739696503
|
||||
}
|
||||
},
|
||||
"rounds": [
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9108160138130188,
|
||||
"wall_ms": 0.9314471390098333
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.813152015209198,
|
||||
"wall_ms": 0.8787710685282946
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 1
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9108160138130188,
|
||||
"wall_ms": 0.9314471390098333
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.813152015209198,
|
||||
"wall_ms": 0.8787710685282946
|
||||
},
|
||||
"speedup_marlin_over_flashinfer": {
|
||||
"gpu_ms": 0.8927730769741712,
|
||||
"wall_ms": 0.943447063955196
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,573 @@
|
||||
{
|
||||
"compute_capability": [
|
||||
12,
|
||||
0
|
||||
],
|
||||
"created_at": "2026-08-18T02:53:56.080340+00:00",
|
||||
"cuda_visible_devices": "0",
|
||||
"device": "NVIDIA RTX 6000D",
|
||||
"iterations": 30,
|
||||
"rounds": 5,
|
||||
"shapes": {
|
||||
"prefill": {
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m128_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 128,
|
||||
"top_k": 16
|
||||
},
|
||||
"correctness": {
|
||||
"flashinfer_vs_marlin": {
|
||||
"cosine": 0.9996392726898193,
|
||||
"finite": true,
|
||||
"max_abs": 0.0126953125,
|
||||
"mean_abs": 0.0021022639703005552,
|
||||
"nrmse": 0.02686966024339199,
|
||||
"p95_abs": 0.005126953125,
|
||||
"p99_abs": 0.0068359375,
|
||||
"reference_max_abs": 0.462890625,
|
||||
"reference_rms": 0.09850681573152542
|
||||
},
|
||||
"flashinfer_vs_reference": {
|
||||
"cosine": 0.9996151924133301,
|
||||
"finite": true,
|
||||
"max_abs": 0.013671875,
|
||||
"mean_abs": 0.00217337254434824,
|
||||
"nrmse": 0.02774309180676937,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.00689697265625,
|
||||
"reference_max_abs": 0.46484375,
|
||||
"reference_rms": 0.09856347739696503
|
||||
},
|
||||
"marlin_vs_reference": {
|
||||
"cosine": 0.9999709129333496,
|
||||
"finite": true,
|
||||
"max_abs": 0.00390625,
|
||||
"mean_abs": 0.000566127069760114,
|
||||
"nrmse": 0.007645560894161463,
|
||||
"p95_abs": 0.00146484375,
|
||||
"p99_abs": 0.001953125,
|
||||
"reference_max_abs": 0.46484375,
|
||||
"reference_rms": 0.09856347739696503
|
||||
}
|
||||
},
|
||||
"rounds": [
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9099744160970052,
|
||||
"wall_ms": 0.9106991346925497
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8695210774739583,
|
||||
"wall_ms": 0.871785699079434
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 1
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.910638936360677,
|
||||
"wall_ms": 0.91134087027361
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8706239700317383,
|
||||
"wall_ms": 0.8713232663770517
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 2
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.928394635518392,
|
||||
"wall_ms": 0.9290497982874513
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8668832143147787,
|
||||
"wall_ms": 0.8675735288610061
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 3
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9538229624430339,
|
||||
"wall_ms": 0.9544791265701255
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8697290420532227,
|
||||
"wall_ms": 0.8704239269718528
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 4
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9259029388427734,
|
||||
"wall_ms": 0.9265734348446131
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8654357274373372,
|
||||
"wall_ms": 0.8660859350735942
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 5
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"flashinfer": {
|
||||
"gpu_ms": 0.9259029388427734,
|
||||
"wall_ms": 0.9265734348446131
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 0.8695210774739583,
|
||||
"wall_ms": 0.8704239269718528
|
||||
},
|
||||
"speedup_marlin_over_flashinfer": {
|
||||
"gpu_ms": 0.9391060779661385,
|
||||
"wall_ms": 0.9394009090254388
|
||||
}
|
||||
}
|
||||
},
|
||||
"prefill_m2048": {
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m2048_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 2048,
|
||||
"top_k": 16
|
||||
},
|
||||
"correctness": {
|
||||
"flashinfer_vs_marlin": {
|
||||
"cosine": 0.9996377229690552,
|
||||
"finite": true,
|
||||
"max_abs": 0.013671875,
|
||||
"mean_abs": 0.0021010376513004303,
|
||||
"nrmse": 0.026925193145871162,
|
||||
"p95_abs": 0.005126953125,
|
||||
"p99_abs": 0.0068359375,
|
||||
"reference_max_abs": 0.5546875,
|
||||
"reference_rms": 0.09821589291095734
|
||||
},
|
||||
"flashinfer_vs_reference": {
|
||||
"cosine": 0.9996229410171509,
|
||||
"finite": true,
|
||||
"max_abs": 0.013671875,
|
||||
"mean_abs": 0.002144214231520891,
|
||||
"nrmse": 0.027460230514407158,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.0068359375,
|
||||
"reference_max_abs": 0.55859375,
|
||||
"reference_rms": 0.09827093780040741
|
||||
},
|
||||
"marlin_vs_reference": {
|
||||
"cosine": 0.9999799132347107,
|
||||
"finite": true,
|
||||
"max_abs": 0.005859375,
|
||||
"mean_abs": 0.00045405104174278677,
|
||||
"nrmse": 0.006360496394336224,
|
||||
"p95_abs": 0.0009765625,
|
||||
"p99_abs": 0.001953125,
|
||||
"reference_max_abs": 0.55859375,
|
||||
"reference_rms": 0.09827093780040741
|
||||
}
|
||||
},
|
||||
"rounds": [
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 4.343468729654948,
|
||||
"wall_ms": 4.344448462749521
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 9.073473103841145,
|
||||
"wall_ms": 9.075486830746135
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 1
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 4.388905843098958,
|
||||
"wall_ms": 4.389951967944701
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 9.065022786458334,
|
||||
"wall_ms": 9.065993401842812
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 2
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 4.355183919270833,
|
||||
"wall_ms": 4.35632496761779
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 9.069696044921875,
|
||||
"wall_ms": 9.070580204327902
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 3
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 4.35906016031901,
|
||||
"wall_ms": 4.360084631480277
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 9.070074462890625,
|
||||
"wall_ms": 9.071069831649462
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 4
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 4.34208984375,
|
||||
"wall_ms": 4.343056639966865
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 9.072443644205729,
|
||||
"wall_ms": 9.073465705538789
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 5
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"flashinfer": {
|
||||
"gpu_ms": 4.355183919270833,
|
||||
"wall_ms": 4.35632496761779
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 9.070074462890625,
|
||||
"wall_ms": 9.071069831649462
|
||||
},
|
||||
"speedup_marlin_over_flashinfer": {
|
||||
"gpu_ms": 2.0825927517681464,
|
||||
"wall_ms": 2.082275748269046
|
||||
}
|
||||
}
|
||||
},
|
||||
"prefill_m512": {
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m512_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 512,
|
||||
"top_k": 16
|
||||
},
|
||||
"correctness": {
|
||||
"flashinfer_vs_marlin": {
|
||||
"cosine": 0.9996380805969238,
|
||||
"finite": true,
|
||||
"max_abs": 0.01318359375,
|
||||
"mean_abs": 0.0021038446575403214,
|
||||
"nrmse": 0.026911206543445587,
|
||||
"p95_abs": 0.005126953125,
|
||||
"p99_abs": 0.0068359375,
|
||||
"reference_max_abs": 0.5,
|
||||
"reference_rms": 0.0984053835272789
|
||||
},
|
||||
"flashinfer_vs_reference": {
|
||||
"cosine": 0.9996137619018555,
|
||||
"finite": true,
|
||||
"max_abs": 0.012939453125,
|
||||
"mean_abs": 0.002174302702769637,
|
||||
"nrmse": 0.027790850028395653,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.0069580078125,
|
||||
"reference_max_abs": 0.498046875,
|
||||
"reference_rms": 0.09845462441444397
|
||||
},
|
||||
"marlin_vs_reference": {
|
||||
"cosine": 0.9999691247940063,
|
||||
"finite": true,
|
||||
"max_abs": 0.005859375,
|
||||
"mean_abs": 0.0005840405938215554,
|
||||
"nrmse": 0.007868881337344646,
|
||||
"p95_abs": 0.00146484375,
|
||||
"p99_abs": 0.001953125,
|
||||
"reference_max_abs": 0.498046875,
|
||||
"reference_rms": 0.09845462441444397
|
||||
}
|
||||
},
|
||||
"rounds": [
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 1.81865603129069,
|
||||
"wall_ms": 1.8193970046316583
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 2.5967445373535156,
|
||||
"wall_ms": 2.5989100647469363
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 1
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 1.7714890797932943,
|
||||
"wall_ms": 1.7721137730404735
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 2.600511932373047,
|
||||
"wall_ms": 2.6012318286423883
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 2
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 1.7895263671875,
|
||||
"wall_ms": 1.7902158355961244
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 2.6133761088053387,
|
||||
"wall_ms": 2.6140614645555615
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 3
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 1.7714154561360678,
|
||||
"wall_ms": 1.7720839319129784
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 2.6104395548502604,
|
||||
"wall_ms": 2.611126432505747
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 4
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 1.7820064544677734,
|
||||
"wall_ms": 1.7826298717409372
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 2.6109214782714845,
|
||||
"wall_ms": 2.6116128312423825
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 5
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"flashinfer": {
|
||||
"gpu_ms": 1.7820064544677734,
|
||||
"wall_ms": 1.7826298717409372
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 2.6104395548502604,
|
||||
"wall_ms": 2.611126432505747
|
||||
},
|
||||
"speedup_marlin_over_flashinfer": {
|
||||
"gpu_ms": 1.4648878225471482,
|
||||
"wall_ms": 1.4647608423366596
|
||||
}
|
||||
}
|
||||
},
|
||||
"prefill_m8192": {
|
||||
"case": {
|
||||
"ep_rank": 7,
|
||||
"ep_size": 32,
|
||||
"hidden": 3584,
|
||||
"inject_nonfinite": false,
|
||||
"intermediate": 3072,
|
||||
"local_experts": 28,
|
||||
"name": "kimi_prefill_m8192_ep32",
|
||||
"routing": "ep_mixed",
|
||||
"tokens": 8192,
|
||||
"top_k": 16
|
||||
},
|
||||
"correctness": {
|
||||
"flashinfer_vs_marlin": {
|
||||
"cosine": 0.9996379613876343,
|
||||
"finite": true,
|
||||
"max_abs": 0.0146484375,
|
||||
"mean_abs": 0.0021007475443184376,
|
||||
"nrmse": 0.0269180778414011,
|
||||
"p95_abs": 0.005126953125,
|
||||
"p99_abs": 0.0068359375,
|
||||
"reference_max_abs": 0.5703125,
|
||||
"reference_rms": 0.09823954105377197
|
||||
},
|
||||
"flashinfer_vs_reference": {
|
||||
"cosine": 0.9996234178543091,
|
||||
"finite": true,
|
||||
"max_abs": 0.0146484375,
|
||||
"mean_abs": 0.0021431385539472103,
|
||||
"nrmse": 0.027442343533039093,
|
||||
"p95_abs": 0.00537109375,
|
||||
"p99_abs": 0.0068359375,
|
||||
"reference_max_abs": 0.5703125,
|
||||
"reference_rms": 0.09829506278038025
|
||||
},
|
||||
"marlin_vs_reference": {
|
||||
"cosine": 0.9999799132347107,
|
||||
"finite": true,
|
||||
"max_abs": 0.005859375,
|
||||
"mean_abs": 0.0004543001123238355,
|
||||
"nrmse": 0.006363501772284508,
|
||||
"p95_abs": 0.0009765625,
|
||||
"p99_abs": 0.001953125,
|
||||
"reference_max_abs": 0.5703125,
|
||||
"reference_rms": 0.09829506278038025
|
||||
}
|
||||
},
|
||||
"rounds": [
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 16.167036946614584,
|
||||
"wall_ms": 16.168367435845237
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 34.38465983072917,
|
||||
"wall_ms": 34.38796966802329
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 1
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 16.178935750325522,
|
||||
"wall_ms": 16.179970768280327
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 34.393994140625,
|
||||
"wall_ms": 34.39523999889692
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 2
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 16.147456868489584,
|
||||
"wall_ms": 16.148571266482275
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 34.362894694010414,
|
||||
"wall_ms": 34.36435852975895
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 3
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 16.169014485677085,
|
||||
"wall_ms": 16.169893369078636
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 34.378727213541666,
|
||||
"wall_ms": 34.37981543441614
|
||||
},
|
||||
"order": [
|
||||
"flashinfer",
|
||||
"marlin"
|
||||
],
|
||||
"round": 4
|
||||
},
|
||||
{
|
||||
"flashinfer": {
|
||||
"gpu_ms": 16.138995361328124,
|
||||
"wall_ms": 16.140034631825984
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 34.34004313151042,
|
||||
"wall_ms": 34.341279161162674
|
||||
},
|
||||
"order": [
|
||||
"marlin",
|
||||
"flashinfer"
|
||||
],
|
||||
"round": 5
|
||||
}
|
||||
],
|
||||
"summary": {
|
||||
"flashinfer": {
|
||||
"gpu_ms": 16.167036946614584,
|
||||
"wall_ms": 16.168367435845237
|
||||
},
|
||||
"marlin": {
|
||||
"gpu_ms": 34.378727213541666,
|
||||
"wall_ms": 34.37981543441614
|
||||
},
|
||||
"speedup_marlin_over_flashinfer": {
|
||||
"gpu_ms": 2.126470504586844,
|
||||
"wall_ms": 2.126362823632778
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Binary file not shown.
File diff suppressed because one or more lines are too long
Binary file not shown.
@ -0,0 +1,4 @@
|
||||
Time (%),Total Time (ns),Num Calls,Avg (ns),Med (ns),Min (ns),Max (ns),StdDev (ns),Name
|
||||
87.0,139383,10,13938.3,11355.5,9370,21685,4468.6,cudaLaunchKernelExC
|
||||
7.2,11500,12,958.3,916.0,880,1324,122.6,cudaGetDriverEntryPointByVersion
|
||||
5.8,9349,12,779.1,705.5,565,1516,256.7,cuTensorMapEncodeTiled
|
||||
|
File diff suppressed because one or more lines are too long
@ -0,0 +1,4 @@
|
||||
Time (%),Total Time (ns),Num Calls,Avg (ns),Med (ns),Min (ns),Max (ns),StdDev (ns),Name
|
||||
85.2,127306,10,12730.6,10292.0,8707,19421,4222.6,cudaLaunchKernelExC
|
||||
7.8,11727,12,977.3,641.0,518,3906,959.1,cuTensorMapEncodeTiled
|
||||
7.0,10386,12,865.5,818.5,790,1158,114.1,cudaGetDriverEntryPointByVersion
|
||||
|
File diff suppressed because one or more lines are too long
@ -0,0 +1,5 @@
|
||||
Time (%),Total Time (ns),Num Calls,Avg (ns),Med (ns),Min (ns),Max (ns),StdDev (ns),Name
|
||||
61.2,71109,5,14221.8,14981.0,12385,15556,1523.4,cudaLaunchKernel
|
||||
24.7,28632,2,14316.0,14316.0,13516,15116,1131.4,cuLaunchKernelEx
|
||||
11.3,13082,1,13082.0,13082.0,13082,13082,0.0,cudaLaunchKernelExC
|
||||
2.8,3307,5,661.4,627.0,593,747,67.7,cuKernelGetName
|
||||
|
@ -0,0 +1,7 @@
|
||||
Time (%),Total Time (ns),Instances,Avg (ns),Med (ns),Min (ns),Max (ns),StdDev (ns),Name
|
||||
94.4,109374,2,54687.0,54687.0,38783,70591,22491.7,"void sglang::device::marlin_moe::Marlin<__nv_bfloat16, (long)562949953487106, (long)2814749767106568, (int)128, (int)1, (int)8, (int)4, (bool)1, (int)4, (int)2, (bool)0, (bool)1, (bool)0>(const int4 *, const int4 *, int4 *, int4 *, const int4 *, const int4 *, const unsigned short *, const int4 *, const int *, const int *, const int *, const int *, const float *, int, bool, bool, int, int, int, int, int *, bool, bool, bool, int)"
|
||||
1.6,1888,1,1888.0,1888.0,1888,1888,0.0,"void sglang::topk_sum_kernel<(int)256, (bool)1>(sglang::TopkSumParams)"
|
||||
1.2,1344,1,1344.0,1344.0,1344,1344,0.0,_moe_align_small_numel_kernel
|
||||
1.1,1312,1,1312.0,1312.0,1312,1312,0.0,_situ_and_mul_kernel
|
||||
1.1,1248,2,624.0,624.0,608,640,22.6,"void at::native::vectorized_elementwise_kernel<(int)4, at::native::FillFunctor<c10::BFloat16>, std::array<char *, (unsigned long)1>>(int, T2, T3)"
|
||||
0.6,704,1,704.0,704.0,704,704,0.0,"void at::native::vectorized_elementwise_kernel<(int)4, at::native::FillFunctor<int>, std::array<char *, (unsigned long)1>>(int, T2, T3)"
|
||||
|
@ -0,0 +1,5 @@
|
||||
Time (%),Total Time (ns),Num Calls,Avg (ns),Med (ns),Min (ns),Max (ns),StdDev (ns),Name
|
||||
75.4,96744,7,13820.6,14257.0,10984,15791,1900.7,cudaLaunchKernel
|
||||
11.6,14830,1,14830.0,14830.0,14830,14830,0.0,cuLaunchKernelEx
|
||||
9.7,12416,1,12416.0,12416.0,12416,12416,0.0,cudaLaunchKernelExC
|
||||
3.3,4295,7,613.6,582.0,475,798,114.7,cuKernelGetName
|
||||
|
@ -0,0 +1,8 @@
|
||||
Time (%),Total Time (ns),Instances,Avg (ns),Med (ns),Min (ns),Max (ns),StdDev (ns),Name
|
||||
94.1,709879,3,236626.3,181534.0,180702,347643,96144.2,"void sglang::device::marlin_moe::Marlin<__nv_bfloat16, (long)562949953487106, (long)2814749767106568, (int)256, (int)4, (int)16, (int)4, (bool)0, (int)4, (int)2, (bool)0, (bool)1, (bool)0>(const int4 *, const int4 *, int4 *, int4 *, const int4 *, const int4 *, const unsigned short *, const int4 *, const int *, const int *, const int *, const int *, const float *, int, bool, bool, int, int, int, int, int *, bool, bool, bool, int)"
|
||||
2.2,16671,1,16671.0,16671.0,16671,16671,0.0,_situ_and_mul_kernel
|
||||
1.7,13120,2,6560.0,6560.0,4896,8224,2353.3,"void at::native::vectorized_elementwise_kernel<(int)4, at::native::FillFunctor<c10::BFloat16>, std::array<char *, (unsigned long)1>>(int, T2, T3)"
|
||||
1.1,8448,1,8448.0,8448.0,8448,8448,0.0,"void sglang::topk_sum_kernel<(int)256, (bool)1>(sglang::TopkSumParams)"
|
||||
0.5,4064,1,4064.0,4064.0,4064,4064,0.0,"void moe_align_block_size_kernel<int>(const T1 *, int *, int *, int *, int, int, unsigned long, int *, bool, bool, int, int)"
|
||||
0.2,1696,1,1696.0,1696.0,1696,1696,0.0,"void count_and_sort_expert_tokens_kernel<int>(const T1 *, int *, int *, unsigned long, bool)"
|
||||
0.1,768,1,768.0,768.0,768,768,0.0,"void at::native::vectorized_elementwise_kernel<(int)4, at::native::FillFunctor<int>, std::array<char *, (unsigned long)1>>(int, T2, T3)"
|
||||
|
@ -0,0 +1,9 @@
|
||||
Time (%),Total Time (ns),Num Calls,Avg (ns),Med (ns),Min (ns),Max (ns),StdDev (ns),Name
|
||||
45.0,1657471,110,15067.9,13500.5,8707,37211,5252.2,cudaLaunchKernelExC
|
||||
27.8,1024642,60,17077.4,15240.0,10648,99929,11361.2,cudaLaunchKernel
|
||||
11.0,404869,1,404869.0,404869.0,404869,404869,0.0,cudaDeviceSynchronize
|
||||
7.7,283828,15,18921.9,16005.0,12865,41329,8082.9,cuLaunchKernelEx
|
||||
3.3,121690,120,1014.1,915.0,790,4516,414.6,cudaGetDriverEntryPointByVersion
|
||||
2.9,105105,120,875.9,712.0,518,3906,507.8,cuTensorMapEncodeTiled
|
||||
1.2,43805,60,730.1,650.0,468,2774,328.4,cuKernelGetName
|
||||
1.1,40880,1,40880.0,40880.0,40880,40880,0.0,cuProfilerStart
|
||||
|
File diff suppressed because one or more lines are too long
@ -0,0 +1,21 @@
|
||||
Range,Style,Total Proj Time (ns),Total Range Time (ns),Range Instances,Proj Avg (ns),Proj Med (ns),Proj Min (ns),Proj Max (ns),Proj StdDev (ns),Total GPU Ops,Avg GPU Ops,Avg Range Lvl,Avg Num Child
|
||||
:marlin_decode_iteration_1,PushPop,1646632,2030238,1,1646632.0,1646632.0,1646632,1646632,0.0,8,8.0,0.0,0.0
|
||||
:flashinfer_prefill_iteration_1,PushPop,1499497,1236905,1,1499497.0,1499497.0,1499497,1499497,0.0,10,10.0,0.0,0.0
|
||||
:marlin_prefill_iteration_1,PushPop,1400491,1514107,1,1400491.0,1400491.0,1400491,1400491,0.0,9,9.0,0.0,0.0
|
||||
:marlin_decode_iteration_2,PushPop,1317581,1499161,1,1317581.0,1317581.0,1317581,1317581,0.0,8,8.0,0.0,0.0
|
||||
:flashinfer_decode_iteration_1,PushPop,1304108,1433573,1,1304108.0,1304108.0,1304108,1304108,0.0,10,10.0,0.0,0.0
|
||||
:marlin_prefill_iteration_2,PushPop,1252429,1302983,1,1252429.0,1252429.0,1252429,1252429,0.0,9,9.0,0.0,0.0
|
||||
:marlin_prefill_iteration_3,PushPop,1181582,1220545,1,1181582.0,1181582.0,1181582,1181582,0.0,9,9.0,0.0,0.0
|
||||
:marlin_decode_iteration_3,PushPop,1171887,1334668,1,1171887.0,1171887.0,1171887,1171887,0.0,8,8.0,0.0,0.0
|
||||
:marlin_prefill_iteration_4,PushPop,1161231,1184045,1,1161231.0,1161231.0,1161231,1161231,0.0,9,9.0,0.0,0.0
|
||||
:marlin_prefill_iteration_5,PushPop,1156238,1170223,1,1156238.0,1156238.0,1156238,1156238,0.0,9,9.0,0.0,0.0
|
||||
:marlin_decode_iteration_5,PushPop,1147087,1295194,1,1147087.0,1147087.0,1147087,1147087,0.0,8,8.0,0.0,0.0
|
||||
:flashinfer_prefill_iteration_2,PushPop,1132207,1082480,1,1132207.0,1132207.0,1132207,1132207,0.0,10,10.0,0.0,0.0
|
||||
:marlin_decode_iteration_4,PushPop,1126735,1274663,1,1126735.0,1126735.0,1126735,1126735,0.0,8,8.0,0.0,0.0
|
||||
:flashinfer_decode_iteration_2,PushPop,1081231,1167266,1,1081231.0,1081231.0,1081231,1081231,0.0,10,10.0,0.0,0.0
|
||||
:flashinfer_prefill_iteration_3,PushPop,1067216,1023953,1,1067216.0,1067216.0,1067216,1067216,0.0,10,10.0,0.0,0.0
|
||||
:flashinfer_prefill_iteration_4,PushPop,1003985,967242,1,1003985.0,1003985.0,1003985,1003985,0.0,10,10.0,0.0,0.0
|
||||
:flashinfer_decode_iteration_3,PushPop,991185,1054475,1,991185.0,991185.0,991185,991185,0.0,10,10.0,0.0,0.0
|
||||
:flashinfer_prefill_iteration_5,PushPop,955122,904533,1,955122.0,955122.0,955122,955122,0.0,10,10.0,0.0,0.0
|
||||
:flashinfer_decode_iteration_4,PushPop,924402,977971,1,924402.0,924402.0,924402,924402,0.0,10,10.0,0.0,0.0
|
||||
:flashinfer_decode_iteration_5,PushPop,900691,947759,1,900691.0,900691.0,900691,900691,0.0,10,10.0,0.0,0.0
|
||||
|
File diff suppressed because one or more lines are too long
@ -0,0 +1,21 @@
|
||||
Time (%),Total Time (ns),Instances,Avg (ns),Med (ns),Min (ns),Max (ns),StdDev (ns),Style,Range
|
||||
8.2,2030238,1,2030238.0,2030238.0,2030238,2030238,0.0,PushPop,:marlin_decode_iteration_1
|
||||
6.1,1514107,1,1514107.0,1514107.0,1514107,1514107,0.0,PushPop,:marlin_prefill_iteration_1
|
||||
6.1,1499161,1,1499161.0,1499161.0,1499161,1499161,0.0,PushPop,:marlin_decode_iteration_2
|
||||
5.8,1433573,1,1433573.0,1433573.0,1433573,1433573,0.0,PushPop,:flashinfer_decode_iteration_1
|
||||
5.4,1334668,1,1334668.0,1334668.0,1334668,1334668,0.0,PushPop,:marlin_decode_iteration_3
|
||||
5.3,1302983,1,1302983.0,1302983.0,1302983,1302983,0.0,PushPop,:marlin_prefill_iteration_2
|
||||
5.3,1295194,1,1295194.0,1295194.0,1295194,1295194,0.0,PushPop,:marlin_decode_iteration_5
|
||||
5.2,1274663,1,1274663.0,1274663.0,1274663,1274663,0.0,PushPop,:marlin_decode_iteration_4
|
||||
5.0,1236905,1,1236905.0,1236905.0,1236905,1236905,0.0,PushPop,:flashinfer_prefill_iteration_1
|
||||
5.0,1220545,1,1220545.0,1220545.0,1220545,1220545,0.0,PushPop,:marlin_prefill_iteration_3
|
||||
4.8,1184045,1,1184045.0,1184045.0,1184045,1184045,0.0,PushPop,:marlin_prefill_iteration_4
|
||||
4.8,1170223,1,1170223.0,1170223.0,1170223,1170223,0.0,PushPop,:marlin_prefill_iteration_5
|
||||
4.7,1167266,1,1167266.0,1167266.0,1167266,1167266,0.0,PushPop,:flashinfer_decode_iteration_2
|
||||
4.4,1082480,1,1082480.0,1082480.0,1082480,1082480,0.0,PushPop,:flashinfer_prefill_iteration_2
|
||||
4.3,1054475,1,1054475.0,1054475.0,1054475,1054475,0.0,PushPop,:flashinfer_decode_iteration_3
|
||||
4.2,1023953,1,1023953.0,1023953.0,1023953,1023953,0.0,PushPop,:flashinfer_prefill_iteration_3
|
||||
4.0,977971,1,977971.0,977971.0,977971,977971,0.0,PushPop,:flashinfer_decode_iteration_4
|
||||
3.9,967242,1,967242.0,967242.0,967242,967242,0.0,PushPop,:flashinfer_prefill_iteration_4
|
||||
3.8,947759,1,947759.0,947759.0,947759,947759,0.0,PushPop,:flashinfer_decode_iteration_5
|
||||
3.7,904533,1,904533.0,904533.0,904533,904533,0.0,PushPop,:flashinfer_prefill_iteration_5
|
||||
|
@ -0,0 +1,13 @@
|
||||
backend,chunked_prefill_size,concurrency,repeat,completed,duration,request_throughput,input_throughput,output_throughput,total_throughput,mean_ttft_ms,median_ttft_ms,p95_ttft_ms,p99_ttft_ms,mean_e2e_latency_ms,median_e2e_latency_ms,p95_e2e_latency_ms,p99_e2e_latency_ms
|
||||
flashinfer_mxfp4,4096,16,1,40,221.47698600799777,0.18060567249436724,2959.043338147713,0.18060567249436724,2959.2239438202073,72603.08281157631,87635.37671801168,89080.70434568217,89951.34060245706,72603.12579314923,87635.42024500202,89080.75052604545,89951.384754295
|
||||
flashinfer_mxfp4,4096,16,2,40,221.06144716404378,0.1809451648541732,2964.6055809707736,0.1809451648541732,2964.786526135628,72440.78481566976,87606.67828551959,88998.53118621977,89664.52184239868,72440.8288814302,87606.71954799909,88998.57463886729,89664.56696634414
|
||||
flashinfer_mxfp4,4096,16,3,40,221.3745526280254,0.18068924149205084,2960.412532605761,0.18068924149205084,2960.593221847253,72566.08219580376,87594.9859394459,89038.73863747576,89889.50585184386,72566.12714612857,87595.03182361368,89038.78330618609,89889.55271399347
|
||||
flashinfer_mxfp4,4096,8,1,40,221.64001683914103,0.1804728251262978,2956.8667668692633,0.1804728251262978,2957.0472396943896,40754.60876591387,44480.76946451329,44945.41641417891,45654.79073660448,40754.653389629675,44480.81373295281,44945.46504281461,45654.83455775771
|
||||
flashinfer_mxfp4,4096,8,2,40,221.6125555918552,0.18049518852022162,2957.233168715311,0.18049518852022162,2957.4136639038315,40748.24627056951,44506.1850650236,44862.61514539365,45574.393972705584,40748.29224284622,44506.22835196555,44862.659902649466,45574.439538051374
|
||||
flashinfer_mxfp4,4096,8,3,40,221.44526491500437,0.1806315434893263,2959.467208529122,0.1806315434893263,2959.647840072612,40712.37499640556,44453.15464294981,44811.94024076685,45526.09015815426,40712.42369815009,44453.20301700849,44811.989083641674,45526.13830527756
|
||||
marlin,4096,16,1,40,243.22425961703993,0.16445727931490292,2694.4680642953695,0.16445727931490292,2694.6325215746842,79749.63234093157,96212.83750247676,97888.6785025592,98836.45602966193,79749.6743179392,96212.87480182946,97888.72231170535,98836.50125178043
|
||||
marlin,4096,16,2,40,243.39303705701604,0.1643432387534973,2692.5996237372997,0.1643432387534973,2692.763966976053,79811.2635479949,96251.52806309052,97955.36642154912,98938.37275844067,79811.30552649265,96251.56629201956,97955.41314255679,98938.41525124619
|
||||
marlin,4096,16,3,40,243.62333092605695,0.16418788729286585,2690.054345406314,0.16418788729286585,2690.218533293607,79882.47500148718,96372.09082697518,98021.89738572342,99003.82948873332,79882.51752386568,96372.13580147363,98021.94035288412,99003.87401685584
|
||||
marlin,4096,8,1,40,243.21893303701654,0.16446088098705797,2694.5270740919577,0.16446088098705797,2694.6915349729447,44725.21210900741,48850.54129792843,49209.063365869224,50020.41739182314,44725.258720549755,48850.58849258348,49209.11144018173,50020.46818676172
|
||||
marlin,4096,8,2,40,243.2142291299533,0.16446406175778208,2694.5791878395016,0.16446406175778208,2694.743651901259,44724.67264174484,48868.09575161897,49182.01875761151,50006.30366436206,44724.714723881334,48868.13376657665,49182.06452928716,50006.34664747398
|
||||
marlin,4096,8,3,40,243.25217532995157,0.1644384061344705,2694.1588461071647,0.1644384061344705,2694.323284513299,44730.88195841992,48857.85997705534,49201.6794851399,50016.33869816782,44730.92947611003,48857.90269903373,49201.72802127199,50016.38562648789
|
||||
|
@ -0,0 +1,109 @@
|
||||
{
|
||||
"run_id": "kimi3-ep4-moe-chunk4k-20260819-150632",
|
||||
"summary": [
|
||||
{
|
||||
"backend": "flashinfer_mxfp4",
|
||||
"chunked_prefill_size": 4096,
|
||||
"concurrency": 8,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.18049518852022162,
|
||||
"median_input_throughput": 2957.233168715311,
|
||||
"median_total_throughput": 2957.4136639038315,
|
||||
"median_mean_ttft_ms": 40748.24627056951,
|
||||
"median_median_ttft_ms": 44480.76946451329,
|
||||
"median_p95_ttft_ms": 44862.61514539365,
|
||||
"median_p99_ttft_ms": 45574.393972705584,
|
||||
"median_mean_e2e_latency_ms": 40748.29224284622,
|
||||
"median_median_e2e_latency_ms": 44480.81373295281,
|
||||
"median_p95_e2e_latency_ms": 44862.659902649466
|
||||
},
|
||||
{
|
||||
"backend": "flashinfer_mxfp4",
|
||||
"chunked_prefill_size": 4096,
|
||||
"concurrency": 16,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.18068924149205084,
|
||||
"median_input_throughput": 2960.412532605761,
|
||||
"median_total_throughput": 2960.593221847253,
|
||||
"median_mean_ttft_ms": 72566.08219580376,
|
||||
"median_median_ttft_ms": 87606.67828551959,
|
||||
"median_p95_ttft_ms": 89038.73863747576,
|
||||
"median_p99_ttft_ms": 89889.50585184386,
|
||||
"median_mean_e2e_latency_ms": 72566.12714612857,
|
||||
"median_median_e2e_latency_ms": 87606.71954799909,
|
||||
"median_p95_e2e_latency_ms": 89038.78330618609
|
||||
},
|
||||
{
|
||||
"backend": "marlin",
|
||||
"chunked_prefill_size": 4096,
|
||||
"concurrency": 8,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.16446088098705797,
|
||||
"median_input_throughput": 2694.5270740919577,
|
||||
"median_total_throughput": 2694.6915349729447,
|
||||
"median_mean_ttft_ms": 44725.21210900741,
|
||||
"median_median_ttft_ms": 48857.85997705534,
|
||||
"median_p95_ttft_ms": 49201.6794851399,
|
||||
"median_p99_ttft_ms": 50016.33869816782,
|
||||
"median_mean_e2e_latency_ms": 44725.258720549755,
|
||||
"median_median_e2e_latency_ms": 48857.90269903373,
|
||||
"median_p95_e2e_latency_ms": 49201.72802127199
|
||||
},
|
||||
{
|
||||
"backend": "marlin",
|
||||
"chunked_prefill_size": 4096,
|
||||
"concurrency": 16,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.1643432387534973,
|
||||
"median_input_throughput": 2692.5996237372997,
|
||||
"median_total_throughput": 2692.763966976053,
|
||||
"median_mean_ttft_ms": 79811.2635479949,
|
||||
"median_median_ttft_ms": 96251.52806309052,
|
||||
"median_p95_ttft_ms": 97955.36642154912,
|
||||
"median_p99_ttft_ms": 98938.37275844067,
|
||||
"median_mean_e2e_latency_ms": 79811.30552649265,
|
||||
"median_median_e2e_latency_ms": 96251.56629201956,
|
||||
"median_p95_e2e_latency_ms": 97955.41314255679
|
||||
}
|
||||
],
|
||||
"comparisons": [
|
||||
{
|
||||
"chunked_prefill_size": 4096,
|
||||
"concurrency": 8,
|
||||
"flashinfer_input_throughput_change_pct": 9.749617925508659,
|
||||
"flashinfer_total_throughput_change_pct": 9.749617925508659,
|
||||
"flashinfer_median_ttft_ms_change_pct": -8.958825692729121,
|
||||
"flashinfer_p95_ttft_ms_change_pct": -8.818935420805618,
|
||||
"flashinfer_median_e2e_latency_ms_change_pct": -8.958814693794636
|
||||
},
|
||||
{
|
||||
"chunked_prefill_size": 4096,
|
||||
"concurrency": 16,
|
||||
"flashinfer_input_throughput_change_pct": 9.946258125697117,
|
||||
"flashinfer_total_throughput_change_pct": 9.946258125697138,
|
||||
"flashinfer_median_ttft_ms_change_pct": -8.981519516141546,
|
||||
"flashinfer_p95_ttft_ms_change_pct": -9.102745576695416,
|
||||
"flashinfer_median_e2e_latency_ms_change_pct": -8.981512797197189
|
||||
}
|
||||
]
|
||||
}
|
||||
@ -0,0 +1,25 @@
|
||||
backend,chunked_prefill_size,concurrency,repeat,completed,duration,request_throughput,input_throughput,output_throughput,total_throughput,mean_ttft_ms,median_ttft_ms,p95_ttft_ms,p99_ttft_ms,mean_e2e_latency_ms,median_e2e_latency_ms,p95_e2e_latency_ms,p99_e2e_latency_ms
|
||||
flashinfer_mxfp4,16384,16,1,40,216.50833019288257,0.1847503971988739,3026.95050770635,0.1847503971988739,3027.135258103549,72316.50786341052,84522.44199160486,84848.25946678175,91778.32209909335,72316.55178854708,84522.48549158685,84848.3054406708,91778.3653298742
|
||||
flashinfer_mxfp4,16384,16,2,40,216.24680145387538,0.1849738342073552,3030.6112996533075,0.1849738342073552,3030.796273487515,72215.42132570758,84481.09910136554,84765.86677560116,91561.01098880405,72215.46733585419,84481.14997893572,84765.91166217111,91561.05765777174
|
||||
flashinfer_mxfp4,16384,16,3,40,216.48059168993495,0.18477406998818618,3027.3383626864425,0.18477406998818618,3027.5231367564306,72312.7285073162,84540.76406499371,84853.01167063879,91778.99107307894,72312.77419594117,84540.81119748298,84853.05667960782,91779.0367784421
|
||||
flashinfer_mxfp4,16384,8,1,40,216.53949848609045,0.18472380456985968,3026.514814072581,0.18472380456985968,3026.6995378771508,40339.51893080957,42272.18415902462,42653.667874005616,49590.42487496044,40339.55953667755,42272.22537947819,42653.713031823274,49590.473274949465
|
||||
flashinfer_mxfp4,16384,8,2,40,216.4690438960679,0.18478392697666732,3027.4998595857173,0.18478392697666732,3027.684643512694,40325.29723601765,42272.719990462065,42599.55736350965,49530.63742398982,40325.33475726959,42272.74997299537,42599.60809347685,49530.707661700435
|
||||
flashinfer_mxfp4,16384,8,3,40,216.46678952896036,0.18478585138644807,3027.531389115565,0.18478585138644807,3027.7161749669513,40328.56577691855,42263.43939546496,42628.12673794104,49566.40897058882,40328.609929577215,42263.489623437636,42628.16974768647,49566.45701683592
|
||||
flashinfer_mxfp4,8192,16,1,40,201.0955079239793,0.1989104600741321,3258.9489778545803,0.1989104600741321,3259.147888314654,66464.03402339201,78428.61395492218,80900.29283260228,82248.80441807676,66464.07695321832,78428.6568548996,80900.33666997915,82248.85000464972
|
||||
flashinfer_mxfp4,8192,16,2,40,201.02179174800403,0.1989834020091862,3260.144058518507,0.1989834020091862,3260.343041920516,66441.70982562937,78389.55687300768,80875.61333409976,82222.40279711317,66441.75368165597,78389.59934201557,80875.65720523708,82222.44497414678
|
||||
flashinfer_mxfp4,8192,16,3,40,200.99849134706892,0.19900646881439046,3260.5219850549734,0.19900646881439046,3260.7209915237877,66426.99547338416,78381.78810151294,80856.50636442006,82182.5564461993,66427.0412968297,78381.83325249702,80856.55405281577,82182.60281610303
|
||||
flashinfer_mxfp4,8192,8,1,40,201.29355681990273,0.19871475586169896,3255.7425600380757,0.19871475586169896,3255.9412747939377,37166.8226861686,39200.42008545715,41681.43365625292,43097.31771401828,37166.86569513404,39200.45927807223,41681.47955433233,43097.35930965515
|
||||
flashinfer_mxfp4,8192,8,2,40,201.1566553460434,0.1988499954485188,3257.958325428532,0.1988499954485188,3258.157175423981,37137.24017466302,39192.07067997195,41617.94732188573,43050.63308592187,37137.28631501435,39192.11111799814,41617.99226633739,43050.676568434574
|
||||
flashinfer_mxfp4,8192,8,3,40,201.02496831398457,0.198980257703726,3260.0925422178466,0.198980257703726,3260.29152247555,37112.54617974046,39172.589769004844,41604.221663821954,43020.16220120946,37112.590136541985,39172.621404984966,41604.265929351095,43020.20668042591
|
||||
marlin,16384,16,1,40,232.999846650986,0.17167393272973525,2812.7057138439823,0.17167393272973525,2812.877387776712,77838.70012289844,90956.05235511903,91350.91751561266,98844.17919690021,77838.74413453159,90956.10008307267,91350.963561167,98844.22676438
|
||||
marlin,16384,16,2,40,233.0310456538573,0.17165094842948833,2812.329139068737,0.17165094842948833,2812.5007900171663,77849.1269534803,90990.29910599347,91359.9561718176,98847.1738615795,77849.17135268915,90990.34174648114,91360.00232317018,98847.21923549892
|
||||
marlin,16384,16,3,40,233.0339416109491,0.1716488152905216,2812.2941897199057,0.1716488152905216,2812.4658385351963,77849.66164607904,90972.34335646499,91361.77453174023,98851.06839299435,77849.70525740064,90972.38598857075,91361.8199098506,98851.1173977633
|
||||
marlin,16384,8,1,40,233.02748973597772,0.17165356776284363,2812.37205422643,0.17165356776284363,2812.543707794193,43413.6843872373,45488.352137384936,45877.185943827484,53396.961687700816,43413.72910180944,45488.39849594515,45877.235028997515,53397.00680534588
|
||||
marlin,16384,8,2,40,232.74746951600537,0.17186008545304213,2815.755640062642,0.17186008545304213,2815.927500148095,43359.47861029417,45474.66833144426,45731.592441652894,53163.52174516069,43359.524680202594,45474.7117330553,45731.640505709205,53163.57911798171
|
||||
marlin,16384,8,3,40,232.98255807999521,0.1716866718677966,2812.9144318819795,0.1716866718677966,2813.086118553847,43401.67168469634,45470.05465638358,45861.754696990814,53373.858650114846,43401.71520576696,45470.09659849573,45861.829353275236,53373.90170200961
|
||||
marlin,8192,16,1,40,223.29927809908986,0.17913179272460458,2934.8952919999215,0.17913179272460458,2935.074423792646,73815.6138905033,87082.98856834881,89829.93190474808,91355.88917817222,73815.65661169589,87083.03174446337,89829.9770831829,91355.93388687819
|
||||
marlin,8192,16,2,40,223.2892873159144,0.17913980773921837,2935.0266099993537,0.17913980773921837,2935.205749807093,73812.15487574227,87086.07016585302,89856.62238476798,91390.91669618152,73812.19822114217,87086.11152344383,89856.66677479167,91390.96168792807
|
||||
marlin,8192,16,3,40,223.33232814306393,0.1791052837383063,2934.4609687684106,0.1791052837383063,2934.640074052149,73832.87409407203,87089.25215096679,89883.80955031607,91406.37799005723,73832.91684102733,87089.2927980749,89883.8516859687,91406.42066749977
|
||||
marlin,8192,8,1,40,223.2894644900225,0.17913966559666036,2935.0242811356834,0.17913966559666036,2935.20342080128,41227.18363396125,43507.58656300604,46237.23470707191,47776.81458188454,41227.227372379275,43507.62654503342,46237.28018605616,47776.86155497795
|
||||
marlin,8192,8,2,40,223.2271411139518,0.17918968007380884,2935.843718329284,0.17918968007380884,2936.0229080093577,41215.50619005575,43508.31371441018,46201.26801881706,47751.084343302064,41215.55230935919,43508.363836910576,46201.314437552355,47751.12928038696
|
||||
marlin,8192,8,3,40,223.2904916331172,0.17913884154871654,2935.0107799341718,0.17913884154871654,2935.189918775721,41230.31144623528,43496.72723049298,46221.18599918904,47788.806532013696,41230.35565857426,43496.76744209137,46221.24284039019,47788.8507199334
|
||||
|
@ -0,0 +1,211 @@
|
||||
{
|
||||
"run_id": "kimi3-ep4-moe-full-20260818-151349",
|
||||
"summary": [
|
||||
{
|
||||
"backend": "flashinfer_mxfp4",
|
||||
"chunked_prefill_size": 8192,
|
||||
"concurrency": 8,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.1988499954485188,
|
||||
"median_input_throughput": 3257.958325428532,
|
||||
"median_total_throughput": 3258.157175423981,
|
||||
"median_mean_ttft_ms": 37137.24017466302,
|
||||
"median_median_ttft_ms": 39192.07067997195,
|
||||
"median_p95_ttft_ms": 41617.94732188573,
|
||||
"median_p99_ttft_ms": 43050.63308592187,
|
||||
"median_mean_e2e_latency_ms": 37137.28631501435,
|
||||
"median_median_e2e_latency_ms": 39192.11111799814,
|
||||
"median_p95_e2e_latency_ms": 41617.99226633739
|
||||
},
|
||||
{
|
||||
"backend": "flashinfer_mxfp4",
|
||||
"chunked_prefill_size": 8192,
|
||||
"concurrency": 16,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.1989834020091862,
|
||||
"median_input_throughput": 3260.144058518507,
|
||||
"median_total_throughput": 3260.343041920516,
|
||||
"median_mean_ttft_ms": 66441.70982562937,
|
||||
"median_median_ttft_ms": 78389.55687300768,
|
||||
"median_p95_ttft_ms": 80875.61333409976,
|
||||
"median_p99_ttft_ms": 82222.40279711317,
|
||||
"median_mean_e2e_latency_ms": 66441.75368165597,
|
||||
"median_median_e2e_latency_ms": 78389.59934201557,
|
||||
"median_p95_e2e_latency_ms": 80875.65720523708
|
||||
},
|
||||
{
|
||||
"backend": "flashinfer_mxfp4",
|
||||
"chunked_prefill_size": 16384,
|
||||
"concurrency": 8,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.18478392697666732,
|
||||
"median_input_throughput": 3027.4998595857173,
|
||||
"median_total_throughput": 3027.684643512694,
|
||||
"median_mean_ttft_ms": 40328.56577691855,
|
||||
"median_median_ttft_ms": 42272.18415902462,
|
||||
"median_p95_ttft_ms": 42628.12673794104,
|
||||
"median_p99_ttft_ms": 49566.40897058882,
|
||||
"median_mean_e2e_latency_ms": 40328.609929577215,
|
||||
"median_median_e2e_latency_ms": 42272.22537947819,
|
||||
"median_p95_e2e_latency_ms": 42628.16974768647
|
||||
},
|
||||
{
|
||||
"backend": "flashinfer_mxfp4",
|
||||
"chunked_prefill_size": 16384,
|
||||
"concurrency": 16,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.18477406998818618,
|
||||
"median_input_throughput": 3027.3383626864425,
|
||||
"median_total_throughput": 3027.5231367564306,
|
||||
"median_mean_ttft_ms": 72312.7285073162,
|
||||
"median_median_ttft_ms": 84522.44199160486,
|
||||
"median_p95_ttft_ms": 84848.25946678175,
|
||||
"median_p99_ttft_ms": 91778.32209909335,
|
||||
"median_mean_e2e_latency_ms": 72312.77419594117,
|
||||
"median_median_e2e_latency_ms": 84522.48549158685,
|
||||
"median_p95_e2e_latency_ms": 84848.3054406708
|
||||
},
|
||||
{
|
||||
"backend": "marlin",
|
||||
"chunked_prefill_size": 8192,
|
||||
"concurrency": 8,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.17913966559666036,
|
||||
"median_input_throughput": 2935.0242811356834,
|
||||
"median_total_throughput": 2935.20342080128,
|
||||
"median_mean_ttft_ms": 41227.18363396125,
|
||||
"median_median_ttft_ms": 43507.58656300604,
|
||||
"median_p95_ttft_ms": 46221.18599918904,
|
||||
"median_p99_ttft_ms": 47776.81458188454,
|
||||
"median_mean_e2e_latency_ms": 41227.227372379275,
|
||||
"median_median_e2e_latency_ms": 43507.62654503342,
|
||||
"median_p95_e2e_latency_ms": 46221.24284039019
|
||||
},
|
||||
{
|
||||
"backend": "marlin",
|
||||
"chunked_prefill_size": 8192,
|
||||
"concurrency": 16,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.17913179272460458,
|
||||
"median_input_throughput": 2934.8952919999215,
|
||||
"median_total_throughput": 2935.074423792646,
|
||||
"median_mean_ttft_ms": 73815.6138905033,
|
||||
"median_median_ttft_ms": 87086.07016585302,
|
||||
"median_p95_ttft_ms": 89856.62238476798,
|
||||
"median_p99_ttft_ms": 91390.91669618152,
|
||||
"median_mean_e2e_latency_ms": 73815.65661169589,
|
||||
"median_median_e2e_latency_ms": 87086.11152344383,
|
||||
"median_p95_e2e_latency_ms": 89856.66677479167
|
||||
},
|
||||
{
|
||||
"backend": "marlin",
|
||||
"chunked_prefill_size": 16384,
|
||||
"concurrency": 8,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.1716866718677966,
|
||||
"median_input_throughput": 2812.9144318819795,
|
||||
"median_total_throughput": 2813.086118553847,
|
||||
"median_mean_ttft_ms": 43401.67168469634,
|
||||
"median_median_ttft_ms": 45474.66833144426,
|
||||
"median_p95_ttft_ms": 45861.754696990814,
|
||||
"median_p99_ttft_ms": 53373.858650114846,
|
||||
"median_mean_e2e_latency_ms": 43401.71520576696,
|
||||
"median_median_e2e_latency_ms": 45474.7117330553,
|
||||
"median_p95_e2e_latency_ms": 45861.829353275236
|
||||
},
|
||||
{
|
||||
"backend": "marlin",
|
||||
"chunked_prefill_size": 16384,
|
||||
"concurrency": 16,
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.17165094842948833,
|
||||
"median_input_throughput": 2812.329139068737,
|
||||
"median_total_throughput": 2812.5007900171663,
|
||||
"median_mean_ttft_ms": 77849.1269534803,
|
||||
"median_median_ttft_ms": 90972.34335646499,
|
||||
"median_p95_ttft_ms": 91359.9561718176,
|
||||
"median_p99_ttft_ms": 98847.1738615795,
|
||||
"median_mean_e2e_latency_ms": 77849.17135268915,
|
||||
"median_median_e2e_latency_ms": 90972.38598857075,
|
||||
"median_p95_e2e_latency_ms": 91360.00232317018
|
||||
}
|
||||
],
|
||||
"comparisons": [
|
||||
{
|
||||
"chunked_prefill_size": 8192,
|
||||
"concurrency": 8,
|
||||
"flashinfer_input_throughput_change_pct": 11.002772493857949,
|
||||
"flashinfer_total_throughput_change_pct": 11.002772493857949,
|
||||
"flashinfer_median_ttft_ms_change_pct": -9.918996257778002,
|
||||
"flashinfer_p95_ttft_ms_change_pct": -9.959153097854456,
|
||||
"flashinfer_median_e2e_latency_ms_change_pct": -9.918986094468796
|
||||
},
|
||||
{
|
||||
"chunked_prefill_size": 8192,
|
||||
"concurrency": 16,
|
||||
"flashinfer_input_throughput_change_pct": 11.082125055880665,
|
||||
"flashinfer_total_throughput_change_pct": 11.082125055880665,
|
||||
"flashinfer_median_ttft_ms_change_pct": -9.986112906786438,
|
||||
"flashinfer_p95_ttft_ms_change_pct": -9.994821541601407,
|
||||
"flashinfer_median_e2e_latency_ms_change_pct": -9.98610688810826
|
||||
},
|
||||
{
|
||||
"chunked_prefill_size": 16384,
|
||||
"concurrency": 8,
|
||||
"flashinfer_input_throughput_change_pct": 7.628579997727458,
|
||||
"flashinfer_total_throughput_change_pct": 7.62857999772748,
|
||||
"flashinfer_median_ttft_ms_change_pct": -7.042347508899216,
|
||||
"flashinfer_p95_ttft_ms_change_pct": -7.050816045775821,
|
||||
"flashinfer_median_e2e_latency_ms_change_pct": -7.0423455840166405
|
||||
},
|
||||
{
|
||||
"chunked_prefill_size": 16384,
|
||||
"concurrency": 16,
|
||||
"flashinfer_input_throughput_change_pct": 7.645236847665093,
|
||||
"flashinfer_total_throughput_change_pct": 7.645236847665093,
|
||||
"flashinfer_median_ttft_ms_change_pct": -7.089958471869751,
|
||||
"flashinfer_p95_ttft_ms_change_pct": -7.127517325850641,
|
||||
"flashinfer_median_e2e_latency_ms_change_pct": -7.089954195324976
|
||||
}
|
||||
]
|
||||
}
|
||||
@ -0,0 +1,7 @@
|
||||
backend,repeat,completed,request_throughput,input_throughput,total_throughput,median_ttft_ms,p95_ttft_ms,median_e2e_latency_ms
|
||||
flashinfer_mxfp4,1,40,0.19844076757687357,3251.2535359794965,3251.451976747074,39266.71834103763,41694.44455524208,39266.759283607826
|
||||
flashinfer_mxfp4,2,40,0.19842530057784252,3251.000124667372,3251.19854996795,39268.08554097079,41711.827712808736,39268.12894491013
|
||||
flashinfer_mxfp4,3,40,0.198540460122689,3252.8868986501366,3253.0854391102594,39249.804597580805,41681.47474176949,39249.838569085114
|
||||
marlin,1,40,0.17882371078176418,2929.8476774484243,2930.026501159206,43587.397116585635,46338.065504981205,43587.43864367716
|
||||
marlin,2,40,0.17884277000539248,2930.1599437683503,2930.3387865383556,43590.378262568265,46287.321725080255,43590.42251249775
|
||||
marlin,3,40,0.178834593712505,2930.025983385682,2930.2048179793946,43578.72363855131,46280.03058874747,43578.76205956563
|
||||
|
@ -0,0 +1,48 @@
|
||||
{
|
||||
"run_id": "kimi3-sm120-pr-representative-20260819-115530",
|
||||
"shape": {
|
||||
"input_len": 16384,
|
||||
"output_len": 1,
|
||||
"concurrency": 8,
|
||||
"chunked_prefill_size": 8192,
|
||||
"tp": 32,
|
||||
"ep": 4
|
||||
},
|
||||
"summary": [
|
||||
{
|
||||
"backend": "marlin",
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.178834593712505,
|
||||
"median_input_throughput": 2930.025983385682,
|
||||
"median_total_throughput": 2930.2048179793946,
|
||||
"median_median_ttft_ms": 43587.397116585635,
|
||||
"median_p95_ttft_ms": 46287.321725080255,
|
||||
"median_median_e2e_latency_ms": 43587.43864367716
|
||||
},
|
||||
{
|
||||
"backend": "flashinfer_mxfp4",
|
||||
"repeats": 3,
|
||||
"completed_each": [
|
||||
40,
|
||||
40,
|
||||
40
|
||||
],
|
||||
"median_request_throughput": 0.19844076757687357,
|
||||
"median_input_throughput": 3251.2535359794965,
|
||||
"median_total_throughput": 3251.451976747074,
|
||||
"median_median_ttft_ms": 39266.71834103763,
|
||||
"median_p95_ttft_ms": 41694.44455524208,
|
||||
"median_median_e2e_latency_ms": 39266.759283607826
|
||||
}
|
||||
],
|
||||
"comparison": {
|
||||
"input_throughput_change_pct": 10.963300476353854,
|
||||
"median_ttft_change_pct": -9.912679034242956,
|
||||
"p95_ttft_change_pct": -9.922538178201778
|
||||
}
|
||||
}
|
||||
@ -0,0 +1,424 @@
|
||||
#!/usr/bin/env bash
|
||||
# Validate the exact SGLang Draft with one four-node Kimi-K3 serving point.
|
||||
set -Eeuo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
# shellcheck source=/dev/null
|
||||
source "${REPO_ROOT}/scripts/common/lib.sh"
|
||||
|
||||
ACTION="${1:-all}"
|
||||
RUN_ID="${RUN_ID:-kimi3-sm120-pr-representative-$(date '+%Y%m%d-%H%M%S')}"
|
||||
RESULT_ROOT="${RESULT_ROOT:-${SCRIPT_DIR}/results/${RUN_ID}}"
|
||||
|
||||
MODEL_PATH="${MODEL_PATH:-/data/hf_models/Kimi-K3}"
|
||||
SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-kimi-k3}"
|
||||
BASE_IMAGE="${BASE_IMAGE:-lmsysorg/sglang:kimi-k3-fiv617situ-warm}"
|
||||
PR_IMAGE="${PR_IMAGE:-local/sglang:kimi-k3-sm120-pr-300c87a-fi-b525c51}"
|
||||
SGLANG_SOURCE="${SGLANG_SOURCE:-/data/hzy/src/sglang-kimi-sm120-draft-wt}"
|
||||
FLASHINFER_SOURCE="${FLASHINFER_SOURCE:-/data/hzy/src/flashinfer-pr4460-kimi-test-wt}"
|
||||
ARTIFACT_DIR="${ARTIFACT_DIR:-/data/hzy/artifacts/sglang-pr-kimi-sm120-300c87a}"
|
||||
FLASHINFER_WHEEL="${FLASHINFER_WHEEL:-${ARTIFACT_DIR}/flashinfer_python-0.6.18-py3-none-any.whl}"
|
||||
BUILD_CONTEXT="${BUILD_CONTEXT:-/tmp/kimi3_sm120_pr_validation_context}"
|
||||
FLASHINFER_CACHE="${FLASHINFER_CACHE:-/data/hzy/cache/flashinfer-pr-b525c51}"
|
||||
|
||||
HEAD_HOST="${HEAD_HOST:-174.1.60.1}"
|
||||
NODE_SSH_USER="${NODE_SSH_USER:-user}"
|
||||
NODE_HOSTS=(174.1.60.1 174.1.60.2 174.1.60.3 174.1.60.4)
|
||||
SSH_OPTS=(-o BatchMode=yes -o StrictHostKeyChecking=no -o ConnectTimeout=10)
|
||||
DIST_PORT="${DIST_PORT:-20000}"
|
||||
PORT="${PORT:-30000}"
|
||||
CONTAINER_PREFIX="kimi3_sm120_pr_validation"
|
||||
|
||||
INPUT_LEN=16384
|
||||
OUTPUT_LEN=1
|
||||
CONCURRENCY=8
|
||||
CHUNKED_PREFILL_SIZE=8192
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-40}"
|
||||
REPEATS="${REPEATS:-3}"
|
||||
WARMUP_REQUESTS="${WARMUP_REQUESTS:-2}"
|
||||
HEALTH_WAIT_S="${HEALTH_WAIT_S:-2400}"
|
||||
BACKENDS=(marlin flashinfer_mxfp4)
|
||||
|
||||
mkdir -p "${RESULT_ROOT}"/{build,service,raw,bench,gpu}
|
||||
log_init "${RESULT_ROOT}/orchestrator.log"
|
||||
|
||||
require_password() {
|
||||
if [[ -z "${SUDO_PASSWORD:-}" && -n "${SUDO_PASSWORD_FILE:-}" ]]; then
|
||||
[[ -r "${SUDO_PASSWORD_FILE}" ]] || {
|
||||
echo "ERROR: cannot read SUDO_PASSWORD_FILE=${SUDO_PASSWORD_FILE}" >&2
|
||||
exit 2
|
||||
}
|
||||
IFS= read -r SUDO_PASSWORD <"${SUDO_PASSWORD_FILE}"
|
||||
fi
|
||||
[[ -n "${SUDO_PASSWORD:-}" ]] || {
|
||||
echo "ERROR: set SUDO_PASSWORD or SUDO_PASSWORD_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
}
|
||||
|
||||
is_head() { [[ "$1" == "${HEAD_HOST}" ]]; }
|
||||
|
||||
sudo_host() {
|
||||
local host="$1"
|
||||
shift
|
||||
require_password
|
||||
if is_head "$host"; then
|
||||
printf '%s\n' "${SUDO_PASSWORD}" | sudo -S -p '' -- "$@"
|
||||
else
|
||||
local remote_cmd
|
||||
printf -v remote_cmd '%q ' "$@"
|
||||
printf '%s\n' "${SUDO_PASSWORD}" | \
|
||||
ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \
|
||||
"sudo -S -p '' -- ${remote_cmd}"
|
||||
fi
|
||||
}
|
||||
|
||||
check_inputs() {
|
||||
local path
|
||||
for path in "${MODEL_PATH}" "${SGLANG_SOURCE}/python/sglang" \
|
||||
"${FLASHINFER_SOURCE}" "${FLASHINFER_WHEEL}" \
|
||||
"${SCRIPT_DIR}/Dockerfile.pr_validation"; do
|
||||
[[ -e "$path" ]] || { echo "ERROR: missing ${path}" >&2; exit 2; }
|
||||
done
|
||||
|
||||
SGLANG_COMMIT="$(git -C "${SGLANG_SOURCE}" rev-parse HEAD)"
|
||||
FLASHINFER_COMMIT="$(git -C "${FLASHINFER_SOURCE}" rev-parse HEAD~1)"
|
||||
[[ "${SGLANG_COMMIT}" == 300c87a431ac40d3e7817246376b7fe20932db09 ]] || {
|
||||
echo "ERROR: unexpected SGLang commit ${SGLANG_COMMIT}" >&2
|
||||
exit 2
|
||||
}
|
||||
[[ "${FLASHINFER_COMMIT}" == b525c51* ]] || {
|
||||
echo "ERROR: unexpected FlashInfer implementation commit ${FLASHINFER_COMMIT}" >&2
|
||||
exit 2
|
||||
}
|
||||
export SGLANG_COMMIT FLASHINFER_COMMIT
|
||||
}
|
||||
|
||||
prepare_build_context() {
|
||||
log "preparing exact Draft build context"
|
||||
rm -rf "${BUILD_CONTEXT}"
|
||||
mkdir -p "${BUILD_CONTEXT}/sglang"
|
||||
cp "${SCRIPT_DIR}/Dockerfile.pr_validation" "${BUILD_CONTEXT}/Dockerfile"
|
||||
cp "${FLASHINFER_WHEEL}" "${BUILD_CONTEXT}/"
|
||||
cp "${SGLANG_SOURCE}/test/registered/unit/layers/quantization/test_mxfp4_sm120_cutlass.py" \
|
||||
"${BUILD_CONTEXT}/test_mxfp4_sm120_cutlass.py"
|
||||
tar -C "${SGLANG_SOURCE}/python/sglang" -cf - . | \
|
||||
tar -C "${BUILD_CONTEXT}/sglang" -xf -
|
||||
{
|
||||
printf 'sglang=%s\n' "${SGLANG_COMMIT}"
|
||||
printf 'flashinfer_implementation=%s\n' "${FLASHINFER_COMMIT}"
|
||||
sha256sum "${FLASHINFER_WHEEL}"
|
||||
} >"${BUILD_CONTEXT}/source_identity.txt"
|
||||
cp "${BUILD_CONTEXT}/source_identity.txt" "${RESULT_ROOT}/build/"
|
||||
du -sh "${BUILD_CONTEXT}" | tee "${RESULT_ROOT}/build/context_size.txt"
|
||||
}
|
||||
|
||||
stage_context() {
|
||||
local host="$1"
|
||||
is_head "$host" && return
|
||||
ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \
|
||||
"rm -rf '${BUILD_CONTEXT}' && mkdir -p '${BUILD_CONTEXT}'"
|
||||
tar -C "${BUILD_CONTEXT}" -cf - . | \
|
||||
ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \
|
||||
"tar -C '${BUILD_CONTEXT}' -xf -"
|
||||
}
|
||||
|
||||
build_one_node() {
|
||||
local host="$1"
|
||||
sudo_host "$host" docker build \
|
||||
--build-arg "BASE_IMAGE=${BASE_IMAGE}" \
|
||||
--build-arg "SGLANG_COMMIT=${SGLANG_COMMIT}" \
|
||||
--build-arg "FLASHINFER_COMMIT=${FLASHINFER_COMMIT}" \
|
||||
--tag "${PR_IMAGE}" "${BUILD_CONTEXT}" \
|
||||
>"${RESULT_ROOT}/build/${host}.log" 2>&1
|
||||
sudo_host "$host" docker image inspect "${PR_IMAGE}" \
|
||||
--format '{{.Id}} {{.Size}} {{json .Config.Labels}}' \
|
||||
>"${RESULT_ROOT}/build/${host}.image.txt"
|
||||
sudo_host "$host" docker run --rm --entrypoint python3 "${PR_IMAGE}" -c \
|
||||
"import importlib.metadata as m; print('sglang=' + m.version('sglang')); print('sglang-kernel=' + m.version('sglang-kernel')); print('flashinfer-python=' + m.version('flashinfer-python')); print(open('/opt/pr-build/source_identity.txt').read(), end='')" \
|
||||
>"${RESULT_ROOT}/build/${host}.packages.txt"
|
||||
}
|
||||
|
||||
build_all_nodes() {
|
||||
check_inputs
|
||||
prepare_build_context
|
||||
local host pid rc=0
|
||||
for host in "${NODE_HOSTS[@]}"; do stage_context "$host"; done
|
||||
local -a pids=()
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
build_one_node "$host" &
|
||||
pids+=("$!")
|
||||
done
|
||||
for pid in "${pids[@]}"; do wait "$pid" || rc=1; done
|
||||
(( rc == 0 )) || { log "ERROR: image build failed"; return 1; }
|
||||
log "exact Draft image built on all nodes"
|
||||
}
|
||||
|
||||
prewarm_one_node() {
|
||||
local host="$1"
|
||||
sudo_host "$host" mkdir -p "${FLASHINFER_CACHE}"
|
||||
sudo_host "$host" docker run --rm --gpus device=0 \
|
||||
-v "${FLASHINFER_CACHE}:/root/.cache/flashinfer" \
|
||||
-e FLASHINFER_WORKSPACE_BASE=/root \
|
||||
-e FLASHINFER_DISABLE_VERSION_CHECK=1 \
|
||||
--entrypoint python3 "${PR_IMAGE}" -m pytest -q -s \
|
||||
/opt/pr-tests/test_mxfp4_sm120_cutlass.py \
|
||||
-k kimi_k3_sm120_situ_layout_and_noncontiguous_input \
|
||||
>"${RESULT_ROOT}/build/${host}.prewarm.log" 2>&1
|
||||
}
|
||||
|
||||
prewarm_all_nodes() {
|
||||
local host pid rc=0
|
||||
local -a pids=()
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
prewarm_one_node "$host" &
|
||||
pids+=("$!")
|
||||
done
|
||||
for pid in "${pids[@]}"; do wait "$pid" || rc=1; done
|
||||
(( rc == 0 )) || { log "ERROR: FlashInfer prewarm failed"; return 1; }
|
||||
log "exact #4460 kernel prewarmed on all nodes"
|
||||
}
|
||||
|
||||
container_name() { printf '%s_node%s' "${CONTAINER_PREFIX}" "$1"; }
|
||||
|
||||
stop_service() {
|
||||
local rank host
|
||||
for rank in 0 1 2 3; do
|
||||
host="${NODE_HOSTS[$rank]}"
|
||||
sudo_host "$host" docker rm -f "$(container_name "$rank")" \
|
||||
>/dev/null 2>&1 || true
|
||||
done
|
||||
}
|
||||
|
||||
collect_service_logs() {
|
||||
local label="$1" rank host
|
||||
for rank in 0 1 2 3; do
|
||||
host="${NODE_HOSTS[$rank]}"
|
||||
sudo_host "$host" docker logs "$(container_name "$rank")" \
|
||||
>"${RESULT_ROOT}/service/${label}_node${rank}.log" 2>&1 || true
|
||||
done
|
||||
}
|
||||
|
||||
collect_gpu() {
|
||||
local label="$1" rank host
|
||||
for rank in 0 1 2 3; do
|
||||
host="${NODE_HOSTS[$rank]}"
|
||||
sudo_host "$host" nvidia-smi \
|
||||
--query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu,power.draw \
|
||||
--format=csv,noheader,nounits \
|
||||
>"${RESULT_ROOT}/gpu/${label}_node${rank}.csv" 2>&1 || true
|
||||
done
|
||||
}
|
||||
|
||||
verify_service_logs() {
|
||||
local label="$1"
|
||||
local pattern='CUDA out of memory|torch\.OutOfMemoryError|Traceback|EngineDeadError|NCCL[^[:cntrl:]]*(error|failed)|connection refused|Terminated'
|
||||
if grep -Ein "${pattern}" "${RESULT_ROOT}/service/${label}_node"*.log \
|
||||
>"${RESULT_ROOT}/service/${label}_fatal_scan.txt"; then
|
||||
log "ERROR: fatal pattern found in service logs label=${label}"
|
||||
return 1
|
||||
fi
|
||||
: >"${RESULT_ROOT}/service/${label}_fatal_scan.txt"
|
||||
}
|
||||
|
||||
start_node() {
|
||||
local rank="$1" backend="$2"
|
||||
local host="${NODE_HOSTS[$rank]}" name bootstrap
|
||||
name="$(container_name "$rank")"
|
||||
bootstrap="export SGLANG_HOST_IP=174.1.60.$((rank + 1)); exec python3 -m sglang.launch_server --model-path ${MODEL_PATH} --served-model-name ${SERVED_MODEL_NAME} --tp-size 32 --ep-size 4 --nnodes 4 --node-rank ${rank} --dist-init-addr ${HEAD_HOST}:${DIST_PORT} --trust-remote-code --moe-runner-backend ${backend} --chunked-prefill-size ${CHUNKED_PREFILL_SIZE} --mem-fraction-static 0.88 --cuda-graph-max-bs-decode 16 --mamba-radix-cache-strategy extra_buffer_lazy --disable-radix-cache --dist-timeout 3600 --mamba-full-memory-ratio 0.36 --host 0.0.0.0 --port ${PORT}"
|
||||
local -a cmd=(
|
||||
docker run -d --name "$name"
|
||||
--gpus all --network host --ipc=host --ulimit memlock=-1
|
||||
--device /dev/infiniband --shm-size 32g --entrypoint bash
|
||||
-v "${MODEL_PATH}:${MODEL_PATH}:ro"
|
||||
-v "${FLASHINFER_CACHE}:/root/.cache/flashinfer"
|
||||
-e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
|
||||
-e NCCL_SOCKET_IFNAME=bond0 -e GLOO_SOCKET_IFNAME=bond0
|
||||
-e NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3
|
||||
-e NCCL_IB_GID_INDEX=3 -e NCCL_IB_TIMEOUT=22 -e NCCL_IB_RETRY_CNT=7
|
||||
-e NCCL_CUMEM_ENABLE=1 -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
-e SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK=0 -e SGLANG_MOE_FUSED_GATE_RADIX=1
|
||||
-e SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK=1
|
||||
-e FLASHINFER_WORKSPACE_BASE=/root
|
||||
-e FLASHINFER_DISABLE_VERSION_CHECK=1
|
||||
"${PR_IMAGE}" -lc "$bootstrap"
|
||||
)
|
||||
printf '%q ' "${cmd[@]}" \
|
||||
>"${RESULT_ROOT}/service/${backend}_node${rank}.cmd.txt"
|
||||
printf '\n' >>"${RESULT_ROOT}/service/${backend}_node${rank}.cmd.txt"
|
||||
sudo_host "$host" "${cmd[@]}" >/dev/null
|
||||
}
|
||||
|
||||
wait_health() {
|
||||
local backend="$1" i
|
||||
for ((i = 1; i <= HEALTH_WAIT_S; i++)); do
|
||||
if curl --fail --silent --max-time 5 \
|
||||
"http://${HEAD_HOST}:${PORT}/health" >/dev/null 2>&1; then
|
||||
log "service healthy backend=${backend} wait_s=${i}"
|
||||
return 0
|
||||
fi
|
||||
if (( i % 30 == 0 )); then
|
||||
log "waiting backend=${backend} elapsed_s=${i}"
|
||||
collect_service_logs "${backend}_starting"
|
||||
if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|EngineDeadError' \
|
||||
"${RESULT_ROOT}/service/${backend}_starting_node"*.log; then
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
start_service() {
|
||||
local backend="$1"
|
||||
stop_service
|
||||
log "starting TP32/EP4 backend=${backend} chunk=8192"
|
||||
start_node 1 "$backend"
|
||||
start_node 2 "$backend"
|
||||
start_node 3 "$backend"
|
||||
sleep 5
|
||||
start_node 0 "$backend"
|
||||
wait_health "$backend" || {
|
||||
collect_service_logs "${backend}_startup_failed"
|
||||
return 1
|
||||
}
|
||||
collect_service_logs "${backend}_healthy"
|
||||
collect_gpu "${backend}_healthy"
|
||||
}
|
||||
|
||||
run_bench() {
|
||||
local backend="$1" repeat="$2"
|
||||
local stem="${backend}_chunk8192_c8_r${repeat}"
|
||||
local output="${RESULT_ROOT}/raw/${stem}.jsonl"
|
||||
rm -f "$output"
|
||||
log "bench backend=${backend} repeat=${repeat}/${REPEATS}"
|
||||
sudo_host "${HEAD_HOST}" docker run --rm --network host \
|
||||
-v "${MODEL_PATH}:${MODEL_PATH}:ro" \
|
||||
-v "${RESULT_ROOT}:${RESULT_ROOT}" \
|
||||
-e PYTHONUNBUFFERED=1 --entrypoint python3 "${PR_IMAGE}" \
|
||||
-m sglang.benchmark.serving \
|
||||
--backend sglang --host "${HEAD_HOST}" --port "${PORT}" \
|
||||
--tokenizer "${MODEL_PATH}" --dataset-name random-ids \
|
||||
--random-input-len "${INPUT_LEN}" --random-output-len "${OUTPUT_LEN}" \
|
||||
--random-range-ratio 1.0 --num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency "${CONCURRENCY}" --request-rate 10000 \
|
||||
--warmup-requests "${WARMUP_REQUESTS}" --output-file "$output" \
|
||||
--output-details --disable-tqdm \
|
||||
>"${RESULT_ROOT}/bench/${stem}.log" 2>&1
|
||||
python3 - "$output" "${NUM_PROMPTS}" "${OUTPUT_LEN}" <<'PY'
|
||||
import json, sys
|
||||
path, expected, expected_output_len = sys.argv[1], int(sys.argv[2]), int(sys.argv[3])
|
||||
rows = [json.loads(x) for x in open(path, encoding="utf-8") if x.strip()]
|
||||
assert len(rows) == 1, (path, len(rows))
|
||||
assert rows[0].get("completed") == expected, rows[0].get("completed")
|
||||
errors = rows[0].get("errors") or []
|
||||
assert all(not error for error in errors), errors
|
||||
output_lens = rows[0].get("output_lens") or []
|
||||
assert len(output_lens) == expected, (len(output_lens), expected)
|
||||
assert all(length == expected_output_len for length in output_lens), output_lens
|
||||
PY
|
||||
}
|
||||
|
||||
run_ab() {
|
||||
local backend repeat
|
||||
collect_gpu before
|
||||
for backend in "${BACKENDS[@]}"; do
|
||||
start_service "$backend"
|
||||
for ((repeat = 1; repeat <= REPEATS; repeat++)); do
|
||||
run_bench "$backend" "$repeat"
|
||||
done
|
||||
collect_service_logs "${backend}_completed"
|
||||
verify_service_logs "${backend}_completed"
|
||||
collect_gpu "${backend}_completed"
|
||||
stop_service
|
||||
sleep 5
|
||||
done
|
||||
collect_gpu after
|
||||
}
|
||||
|
||||
summarize() {
|
||||
python3 - "${RESULT_ROOT}" <<'PY'
|
||||
import csv, json, re, statistics, sys
|
||||
from pathlib import Path
|
||||
|
||||
root = Path(sys.argv[1])
|
||||
pat = re.compile(r"(.+)_chunk8192_c8_r(\d+)\.jsonl$")
|
||||
rows = []
|
||||
for path in sorted((root / "raw").glob("*.jsonl")):
|
||||
m = pat.match(path.name)
|
||||
if not m:
|
||||
continue
|
||||
data = next(json.loads(x) for x in path.read_text().splitlines() if x.strip())
|
||||
# The benchmark JSON reports the framework as backend="sglang". Keep the
|
||||
# MoE implementation parsed from the filename as the comparison backend.
|
||||
rows.append({**data, "backend": m.group(1), "repeat": int(m.group(2))})
|
||||
|
||||
metrics = ["request_throughput", "input_throughput", "total_throughput",
|
||||
"median_ttft_ms", "p95_ttft_ms", "median_e2e_latency_ms"]
|
||||
summary = []
|
||||
for backend in ("marlin", "flashinfer_mxfp4"):
|
||||
group = [x for x in rows if x["backend"] == backend]
|
||||
if len(group) != 3:
|
||||
raise SystemExit(f"expected 3 repeats for {backend}, got {len(group)}")
|
||||
item = {"backend": backend, "repeats": len(group),
|
||||
"completed_each": [x.get("completed") for x in group]}
|
||||
for metric in metrics:
|
||||
item[f"median_{metric}"] = statistics.median(float(x[metric]) for x in group)
|
||||
summary.append(item)
|
||||
idx = {x["backend"]: x for x in summary}
|
||||
base, cand = idx["marlin"], idx["flashinfer_mxfp4"]
|
||||
comparison = {
|
||||
"input_throughput_change_pct":
|
||||
(cand["median_input_throughput"] / base["median_input_throughput"] - 1) * 100,
|
||||
"median_ttft_change_pct":
|
||||
(cand["median_median_ttft_ms"] / base["median_median_ttft_ms"] - 1) * 100,
|
||||
"p95_ttft_change_pct":
|
||||
(cand["median_p95_ttft_ms"] / base["median_p95_ttft_ms"] - 1) * 100,
|
||||
}
|
||||
payload = {
|
||||
"run_id": root.name,
|
||||
"shape": {"input_len": 16384, "output_len": 1,
|
||||
"concurrency": 8, "chunked_prefill_size": 8192,
|
||||
"tp": 32, "ep": 4},
|
||||
"summary": summary,
|
||||
"comparison": comparison,
|
||||
}
|
||||
(root / "summary.json").write_text(json.dumps(payload, indent=2))
|
||||
with (root / "results.csv").open("w", newline="") as f:
|
||||
fields = ["backend", "repeat", "completed", *metrics]
|
||||
w = csv.DictWriter(f, fieldnames=fields, extrasaction="ignore")
|
||||
w.writeheader(); w.writerows(rows)
|
||||
print(json.dumps(payload, indent=2))
|
||||
PY
|
||||
}
|
||||
|
||||
cleanup() {
|
||||
collect_service_logs cleanup 2>/dev/null || true
|
||||
stop_service 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
case "${ACTION}" in
|
||||
build)
|
||||
require_password; check_inputs; build_all_nodes; prewarm_all_nodes
|
||||
;;
|
||||
run)
|
||||
require_password; check_inputs; run_ab; summarize
|
||||
;;
|
||||
all)
|
||||
require_password; check_inputs; build_all_nodes; prewarm_all_nodes
|
||||
run_ab; summarize
|
||||
;;
|
||||
summarize)
|
||||
summarize
|
||||
;;
|
||||
stop)
|
||||
require_password; stop_service
|
||||
;;
|
||||
*)
|
||||
echo "Usage: $0 {all|build|run|summarize|stop}" >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
@ -0,0 +1,682 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Validate Kimi-K3 SiTU semantics for FlashInfer SM120 MXFP8 x MXFP4 MoE.
|
||||
|
||||
The control and target paths share the same generated case, quantized tensors,
|
||||
routing metadata, CUTLASS ABI wrapper, synchronization, and timing method. The
|
||||
only intentional difference is the activation contract: supported SwiGLU is a
|
||||
harness control, while SiTU is the Kimi-K3 target.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime
|
||||
import importlib.metadata
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
import sys
|
||||
import time
|
||||
from typing import Callable, Literal
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from flashinfer import (
|
||||
mxfp4_dequantize,
|
||||
mxfp4_quantize,
|
||||
mxfp8_dequantize_host,
|
||||
mxfp8_quantize,
|
||||
)
|
||||
from flashinfer.fused_moe import cutlass_fused_moe
|
||||
from flashinfer.tllm_enums import ActivationType
|
||||
|
||||
|
||||
ActivationName = Literal["swiglu", "situ"]
|
||||
RoutingName = Literal["uniform", "skewed", "masked", "ep_mixed"]
|
||||
|
||||
KNOWN_MISSING_SITU_ERROR = (
|
||||
"fc1_expert_weights inter size must be equal to "
|
||||
"fc2_expert_weights inter size"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Case:
|
||||
name: str
|
||||
tokens: int
|
||||
hidden: int
|
||||
intermediate: int
|
||||
local_experts: int
|
||||
top_k: int
|
||||
routing: RoutingName
|
||||
ep_size: int = 1
|
||||
ep_rank: int = 0
|
||||
inject_nonfinite: bool = False
|
||||
|
||||
@property
|
||||
def total_experts(self) -> int:
|
||||
return self.local_experts * self.ep_size
|
||||
|
||||
@property
|
||||
def local_expert_start(self) -> int:
|
||||
return self.local_experts * self.ep_rank
|
||||
|
||||
|
||||
@dataclass
|
||||
class QuantizedCase:
|
||||
case: Case
|
||||
input_q: torch.Tensor
|
||||
input_sf: torch.Tensor
|
||||
input_dq: torch.Tensor
|
||||
fc1_q: torch.Tensor
|
||||
fc1_sf: torch.Tensor
|
||||
fc1_dq: torch.Tensor
|
||||
fc2_q: torch.Tensor
|
||||
fc2_sf: torch.Tensor
|
||||
fc2_dq: torch.Tensor
|
||||
topk_ids: torch.Tensor
|
||||
topk_weights: torch.Tensor
|
||||
expert_scale: torch.Tensor
|
||||
situ_gate_beta: torch.Tensor
|
||||
situ_linear_beta: torch.Tensor
|
||||
|
||||
|
||||
def quantize_mxfp4_experts(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
packed = []
|
||||
scales = []
|
||||
for expert_weight in weight:
|
||||
expert_packed, expert_scale = mxfp4_quantize(expert_weight)
|
||||
packed.append(expert_packed)
|
||||
scales.append(expert_scale)
|
||||
return torch.stack(packed), torch.stack(scales)
|
||||
|
||||
|
||||
def dequantize_mxfp4_experts(
|
||||
packed: torch.Tensor, scales: torch.Tensor, dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
outputs = []
|
||||
flat_scales = scales.reshape(scales.shape[0], -1)
|
||||
for expert in range(packed.shape[0]):
|
||||
outputs.append(
|
||||
mxfp4_dequantize(
|
||||
packed[expert].cpu().view(torch.uint8),
|
||||
flat_scales[expert].cpu().view(torch.uint8),
|
||||
)
|
||||
)
|
||||
return torch.stack(outputs).to(device=packed.device, dtype=dtype)
|
||||
|
||||
|
||||
def dequantize_mxfp8_input(
|
||||
packed: torch.Tensor, scales: torch.Tensor, dtype: torch.dtype
|
||||
) -> torch.Tensor:
|
||||
return (
|
||||
mxfp8_dequantize_host(
|
||||
packed.cpu().view(torch.uint8),
|
||||
scales.cpu().view(torch.uint8).reshape(-1),
|
||||
True,
|
||||
)
|
||||
.to(device=packed.device, dtype=dtype)
|
||||
.reshape(packed.shape)
|
||||
)
|
||||
|
||||
|
||||
def make_routing(case: Case, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
ids = torch.empty(case.tokens, case.top_k, dtype=torch.int32, device=device)
|
||||
local_start = case.local_expert_start
|
||||
local_end = local_start + case.local_experts
|
||||
|
||||
if case.routing == "uniform":
|
||||
for token in range(case.tokens):
|
||||
for slot in range(case.top_k):
|
||||
ids[token, slot] = (token * case.top_k + slot) % case.total_experts
|
||||
elif case.routing == "skewed":
|
||||
chosen = torch.arange(case.top_k, device=device) % case.total_experts
|
||||
ids.copy_(chosen.expand(case.tokens, -1))
|
||||
elif case.routing == "masked":
|
||||
chosen = torch.arange(case.top_k, device=device) % case.total_experts
|
||||
ids.copy_(chosen.expand(case.tokens, -1))
|
||||
ids[:, -1] = -1
|
||||
elif case.routing == "ep_mixed":
|
||||
remote_start = local_end % case.total_experts
|
||||
for token in range(case.tokens):
|
||||
for slot in range(case.top_k):
|
||||
if slot % 2 == 0:
|
||||
ids[token, slot] = local_start + (token + slot) % case.local_experts
|
||||
else:
|
||||
ids[token, slot] = remote_start + (token + slot) % case.local_experts
|
||||
else:
|
||||
raise ValueError(f"unsupported routing mode: {case.routing}")
|
||||
|
||||
logits = torch.linspace(
|
||||
0.25,
|
||||
1.25,
|
||||
case.top_k,
|
||||
dtype=torch.float32,
|
||||
device=device,
|
||||
).expand(case.tokens, -1)
|
||||
weights = torch.softmax(logits, dim=-1)
|
||||
weights = torch.where(ids >= 0, weights, torch.zeros_like(weights))
|
||||
weights /= weights.sum(dim=-1, keepdim=True)
|
||||
return ids, weights.contiguous()
|
||||
|
||||
|
||||
def prepare_case(case: Case, seed: int, device: torch.device) -> QuantizedCase:
|
||||
generator = torch.Generator(device=device).manual_seed(seed)
|
||||
dtype = torch.bfloat16
|
||||
x = torch.randn(
|
||||
case.tokens,
|
||||
case.hidden,
|
||||
generator=generator,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
if case.inject_nonfinite:
|
||||
x[0, 0] = float("nan")
|
||||
if case.hidden > 1:
|
||||
x[-1, 1] = float("inf")
|
||||
|
||||
std_fc1 = 1.0 / math.sqrt(case.hidden)
|
||||
std_fc2 = 1.0 / math.sqrt(case.intermediate)
|
||||
gate = torch.randn(
|
||||
case.local_experts,
|
||||
case.intermediate,
|
||||
case.hidden,
|
||||
generator=generator,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
) * std_fc1
|
||||
up = torch.randn(
|
||||
case.local_experts,
|
||||
case.intermediate,
|
||||
case.hidden,
|
||||
generator=generator,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
) * std_fc1
|
||||
down = torch.randn(
|
||||
case.local_experts,
|
||||
case.hidden,
|
||||
case.intermediate,
|
||||
generator=generator,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
) * std_fc2
|
||||
|
||||
# CUTLASS consumes [up; gate]. Kimi checkpoints load as [gate; up], so
|
||||
# this is also the oracle layout required from the SGLang weight processor.
|
||||
fc1 = torch.cat([up, gate], dim=1).contiguous()
|
||||
input_q, input_sf = mxfp8_quantize(x, True, 32)
|
||||
fc1_q, fc1_sf = quantize_mxfp4_experts(fc1)
|
||||
fc2_q, fc2_sf = quantize_mxfp4_experts(down)
|
||||
input_dq = dequantize_mxfp8_input(input_q, input_sf, dtype)
|
||||
fc1_dq = dequantize_mxfp4_experts(fc1_q, fc1_sf, dtype)
|
||||
fc2_dq = dequantize_mxfp4_experts(fc2_q, fc2_sf, dtype)
|
||||
topk_ids, topk_weights = make_routing(case, device)
|
||||
|
||||
return QuantizedCase(
|
||||
case=case,
|
||||
input_q=input_q,
|
||||
input_sf=input_sf,
|
||||
input_dq=input_dq,
|
||||
fc1_q=fc1_q,
|
||||
fc1_sf=fc1_sf,
|
||||
fc1_dq=fc1_dq,
|
||||
fc2_q=fc2_q,
|
||||
fc2_sf=fc2_sf,
|
||||
fc2_dq=fc2_dq,
|
||||
topk_ids=topk_ids,
|
||||
topk_weights=topk_weights,
|
||||
expert_scale=torch.ones(
|
||||
case.local_experts, dtype=torch.float32, device=device
|
||||
),
|
||||
situ_gate_beta=torch.full(
|
||||
(case.local_experts,), 4.0, dtype=torch.float32, device=device
|
||||
),
|
||||
situ_linear_beta=torch.full(
|
||||
(case.local_experts,), 25.0, dtype=torch.float32, device=device
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def apply_activation(
|
||||
gate: torch.Tensor, up: torch.Tensor, activation: ActivationName
|
||||
) -> torch.Tensor:
|
||||
if activation == "swiglu":
|
||||
return F.silu(gate) * up
|
||||
if activation == "situ":
|
||||
gate = 4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)
|
||||
up = 25.0 * torch.tanh(up / 25.0)
|
||||
return gate * up
|
||||
raise ValueError(f"unsupported activation: {activation}")
|
||||
|
||||
|
||||
def run_reference(data: QuantizedCase, activation: ActivationName) -> torch.Tensor:
|
||||
case = data.case
|
||||
output = torch.zeros(
|
||||
case.tokens,
|
||||
case.hidden,
|
||||
dtype=torch.bfloat16,
|
||||
device=data.input_q.device,
|
||||
)
|
||||
local_start = case.local_expert_start
|
||||
for local_expert in range(case.local_experts):
|
||||
global_expert = local_start + local_expert
|
||||
selected = data.topk_ids == global_expert
|
||||
if not bool(selected.any()):
|
||||
continue
|
||||
token_ids, slots = torch.where(selected)
|
||||
up_weight, gate_weight = torch.chunk(data.fc1_dq[local_expert], 2, dim=0)
|
||||
expert_input = data.input_dq[token_ids]
|
||||
gate = expert_input @ gate_weight.transpose(0, 1)
|
||||
up = expert_input @ up_weight.transpose(0, 1)
|
||||
hidden = apply_activation(gate, up, activation)
|
||||
expert_output = hidden @ data.fc2_dq[local_expert].transpose(0, 1)
|
||||
output[token_ids] += (
|
||||
data.topk_weights[token_ids, slots, None].to(expert_output.dtype)
|
||||
* expert_output
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def run_candidate(
|
||||
data: QuantizedCase,
|
||||
activation: ActivationName,
|
||||
*,
|
||||
use_fused_finalize: bool,
|
||||
output: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
case = data.case
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
case.tokens,
|
||||
case.hidden,
|
||||
dtype=torch.bfloat16,
|
||||
device=data.input_q.device,
|
||||
)
|
||||
activation_type = (
|
||||
ActivationType.Swiglu if activation == "swiglu" else ActivationType.Situ
|
||||
)
|
||||
alpha = data.situ_gate_beta if activation == "situ" else None
|
||||
beta = data.situ_linear_beta if activation == "situ" else None
|
||||
cutlass_fused_moe(
|
||||
input=data.input_q,
|
||||
token_selected_experts=data.topk_ids,
|
||||
token_final_scales=data.topk_weights,
|
||||
fc1_expert_weights=data.fc1_q.contiguous().view(torch.int64),
|
||||
fc2_expert_weights=data.fc2_q.contiguous().view(torch.int64),
|
||||
output_dtype=torch.bfloat16,
|
||||
quant_scales=[
|
||||
data.fc1_sf.view(torch.int32),
|
||||
data.expert_scale,
|
||||
data.fc2_sf.view(torch.int32),
|
||||
data.expert_scale,
|
||||
],
|
||||
input_sf=data.input_sf,
|
||||
swiglu_alpha=alpha,
|
||||
swiglu_beta=beta,
|
||||
swiglu_limit=None,
|
||||
tp_size=1,
|
||||
tp_rank=0,
|
||||
ep_size=case.ep_size,
|
||||
ep_rank=case.ep_rank,
|
||||
use_mxfp8_act_scaling=True,
|
||||
activation_type=activation_type,
|
||||
tune_max_num_tokens=1 << (case.tokens - 1).bit_length(),
|
||||
output=output,
|
||||
use_fused_finalize=use_fused_finalize,
|
||||
)
|
||||
return output
|
||||
|
||||
|
||||
def tensor_metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict:
|
||||
ref_nan = torch.isnan(reference)
|
||||
got_nan = torch.isnan(candidate)
|
||||
ref_inf = torch.isinf(reference)
|
||||
got_inf = torch.isinf(candidate)
|
||||
same_nonfinite = bool(torch.equal(ref_nan, got_nan) and torch.equal(ref_inf, got_inf))
|
||||
finite = torch.isfinite(reference) & torch.isfinite(candidate)
|
||||
if not bool(finite.any()):
|
||||
return {
|
||||
"same_nonfinite_mask": same_nonfinite,
|
||||
"finite_elements": 0,
|
||||
"max_abs": None,
|
||||
"nrmse": None,
|
||||
"cosine": None,
|
||||
}
|
||||
ref = reference[finite].float()
|
||||
got = candidate[finite].float()
|
||||
diff = got - ref
|
||||
denominator = torch.sqrt(torch.mean(ref.square())).clamp_min(1e-12)
|
||||
if bool(torch.count_nonzero(ref)) and bool(torch.count_nonzero(got)):
|
||||
cosine = float(F.cosine_similarity(ref, got, dim=0).item())
|
||||
else:
|
||||
cosine = 1.0 if torch.equal(ref, got) else 0.0
|
||||
return {
|
||||
"same_nonfinite_mask": same_nonfinite,
|
||||
"finite_elements": int(finite.sum().item()),
|
||||
"max_abs": float(diff.abs().max().item()),
|
||||
"nrmse": float((torch.sqrt(torch.mean(diff.square())) / denominator).item()),
|
||||
"cosine": cosine,
|
||||
}
|
||||
|
||||
|
||||
def assert_matches(reference: torch.Tensor, candidate: torch.Tensor) -> dict:
|
||||
metrics = tensor_metrics(reference, candidate)
|
||||
if not metrics["same_nonfinite_mask"]:
|
||||
raise AssertionError(f"NaN/Inf masks differ: {metrics}")
|
||||
torch.testing.assert_close(
|
||||
candidate,
|
||||
reference,
|
||||
rtol=0.2,
|
||||
atol=0.2,
|
||||
equal_nan=True,
|
||||
)
|
||||
return metrics
|
||||
|
||||
|
||||
def assert_nonfinite_contract(
|
||||
data: QuantizedCase,
|
||||
reference: torch.Tensor,
|
||||
candidate: torch.Tensor,
|
||||
) -> dict:
|
||||
"""Validate the defined MXFP8 boundary for non-finite BF16 inputs.
|
||||
|
||||
``mxfp8_quantize`` uses saturating conversion: the generated E4M3 payload
|
||||
is finite and E8M0 codes 0x00..0xFE denote finite powers of two; only 0xFF
|
||||
denotes NaN. A BF16 dequantized reference can overflow for the 0xFE scale,
|
||||
so its NaN/Inf rows are diagnostic only and are not a valid Tensor Core
|
||||
oracle. The kernel contract here is finite, deterministic output from the
|
||||
finite quantized representation.
|
||||
"""
|
||||
input_payload_finite = bool(torch.isfinite(data.input_q.float()).all())
|
||||
scale_has_nan_code = bool((data.input_sf == 0xFF).any())
|
||||
candidate_finite = bool(torch.isfinite(candidate).all())
|
||||
if not input_payload_finite or scale_has_nan_code or not candidate_finite:
|
||||
raise AssertionError(
|
||||
"MXFP8 saturation contract failed: "
|
||||
f"input_payload_finite={input_payload_finite} "
|
||||
f"scale_has_nan_code={scale_has_nan_code} "
|
||||
f"candidate_finite={candidate_finite}"
|
||||
)
|
||||
|
||||
reference_rows = (~torch.isfinite(reference)).any(dim=1)
|
||||
candidate_rows = (~torch.isfinite(candidate)).any(dim=1)
|
||||
return {
|
||||
"contract": "mxfp8_satfinite_input_and_finite_candidate",
|
||||
"input_payload_finite": input_payload_finite,
|
||||
"scale_has_nan_code": scale_has_nan_code,
|
||||
"max_scale_code": int(data.input_sf.max().item()),
|
||||
"reference_nonfinite_rows_diagnostic_only": reference_rows.tolist(),
|
||||
"candidate_nonfinite_rows": candidate_rows.tolist(),
|
||||
"candidate_nan_elements": int(torch.isnan(candidate).sum().item()),
|
||||
"candidate_inf_elements": int(torch.isinf(candidate).sum().item()),
|
||||
}
|
||||
|
||||
|
||||
def timed_ms(fn: Callable[[], torch.Tensor], warmup: int, iterations: int) -> float:
|
||||
for _ in range(warmup):
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
fn()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return float(start.elapsed_time(end) / iterations)
|
||||
|
||||
|
||||
def validate_determinism(data: QuantizedCase, activation: ActivationName) -> dict:
|
||||
outputs = [
|
||||
run_candidate(data, activation, use_fused_finalize=False).clone()
|
||||
for _ in range(3)
|
||||
]
|
||||
exact = True
|
||||
for item in outputs[1:]:
|
||||
try:
|
||||
torch.testing.assert_close(
|
||||
outputs[0], item, rtol=0.0, atol=0.0, equal_nan=True
|
||||
)
|
||||
except AssertionError:
|
||||
exact = False
|
||||
break
|
||||
if not exact:
|
||||
raise AssertionError("non-fused finalize is not bitwise deterministic")
|
||||
return {"runs": len(outputs), "bitwise_equal": exact}
|
||||
|
||||
|
||||
def validate_cuda_graph(data: QuantizedCase, activation: ActivationName) -> dict:
|
||||
eager = run_candidate(data, activation, use_fused_finalize=False).clone()
|
||||
static_output = torch.empty_like(eager)
|
||||
# Warm module loading and tactic selection before capture.
|
||||
run_candidate(
|
||||
data,
|
||||
activation,
|
||||
use_fused_finalize=False,
|
||||
output=static_output,
|
||||
)
|
||||
torch.cuda.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
run_candidate(
|
||||
data,
|
||||
activation,
|
||||
use_fused_finalize=False,
|
||||
output=static_output,
|
||||
)
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
first = static_output.clone()
|
||||
graph.replay()
|
||||
torch.cuda.synchronize()
|
||||
second = static_output.clone()
|
||||
metrics = assert_matches(eager, first)
|
||||
if not torch.equal(first, second):
|
||||
raise AssertionError("CUDA Graph replay is not bitwise deterministic")
|
||||
return {"replay_matches_eager": metrics, "replays_bitwise_equal": True}
|
||||
|
||||
|
||||
def smoke_cases() -> list[Case]:
|
||||
return [
|
||||
Case("decode_m1_uniform", 1, 128, 128, 4, 2, "uniform"),
|
||||
Case("decode_m2_skewed", 2, 128, 128, 4, 2, "skewed"),
|
||||
Case("batch_m8_empty_experts", 8, 128, 128, 8, 2, "skewed"),
|
||||
Case("batch_m32_masked_slot", 32, 128, 128, 8, 4, "masked"),
|
||||
Case("ep2_global_ids", 8, 128, 128, 4, 4, "ep_mixed", ep_size=2),
|
||||
Case(
|
||||
"nonfinite_activation",
|
||||
2,
|
||||
128,
|
||||
128,
|
||||
4,
|
||||
2,
|
||||
"uniform",
|
||||
inject_nonfinite=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def real_cases() -> list[Case]:
|
||||
# One TP32/EP32 rank owns 28 of Kimi-K3's 896 experts. Global expert IDs
|
||||
# intentionally mix local and remote routes, matching the per-rank kernel ABI.
|
||||
return [
|
||||
Case(
|
||||
"kimi_decode_m1_ep32",
|
||||
1,
|
||||
3584,
|
||||
3072,
|
||||
28,
|
||||
16,
|
||||
"ep_mixed",
|
||||
ep_size=32,
|
||||
ep_rank=7,
|
||||
),
|
||||
Case(
|
||||
"kimi_prefill_m128_ep32",
|
||||
128,
|
||||
3584,
|
||||
3072,
|
||||
28,
|
||||
16,
|
||||
"ep_mixed",
|
||||
ep_size=32,
|
||||
ep_rank=7,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def run_one(
|
||||
case: Case,
|
||||
activation: ActivationName,
|
||||
args: argparse.Namespace,
|
||||
) -> dict:
|
||||
started = time.time()
|
||||
data = prepare_case(case, args.seed, torch.device("cuda"))
|
||||
reference = run_reference(data, activation)
|
||||
result = {
|
||||
"case": asdict(case),
|
||||
"activation": activation,
|
||||
"status": "UNKNOWN",
|
||||
}
|
||||
try:
|
||||
candidate = run_candidate(
|
||||
data,
|
||||
activation,
|
||||
use_fused_finalize=False,
|
||||
)
|
||||
result["metrics"] = (
|
||||
assert_nonfinite_contract(data, reference, candidate)
|
||||
if case.inject_nonfinite
|
||||
else assert_matches(reference, candidate)
|
||||
)
|
||||
result["determinism"] = validate_determinism(data, activation)
|
||||
if case.name == "batch_m8_empty_experts":
|
||||
result["cuda_graph"] = validate_cuda_graph(data, activation)
|
||||
if args.benchmark_iterations > 0 and not case.inject_nonfinite:
|
||||
result["latency_ms"] = {
|
||||
"reference": timed_ms(
|
||||
lambda: run_reference(data, activation),
|
||||
args.warmup,
|
||||
args.benchmark_iterations,
|
||||
),
|
||||
"candidate": timed_ms(
|
||||
lambda: run_candidate(
|
||||
data,
|
||||
activation,
|
||||
use_fused_finalize=False,
|
||||
),
|
||||
args.warmup,
|
||||
args.benchmark_iterations,
|
||||
),
|
||||
}
|
||||
result["status"] = "PASS"
|
||||
except RuntimeError as error:
|
||||
if activation == "situ" and KNOWN_MISSING_SITU_ERROR in str(error):
|
||||
result["status"] = "EXPECTED_MISSING_SITU"
|
||||
result["error"] = str(error)
|
||||
if args.require_situ:
|
||||
raise
|
||||
else:
|
||||
result["status"] = "ERROR"
|
||||
result["error"] = repr(error)
|
||||
raise
|
||||
finally:
|
||||
result["elapsed_s"] = time.time() - started
|
||||
return result
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--suite",
|
||||
choices=("smoke", "real", "all"),
|
||||
default="smoke",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--activation",
|
||||
choices=("swiglu", "situ", "both"),
|
||||
default="both",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--require-situ",
|
||||
action="store_true",
|
||||
help="Fail instead of recording the known missing-SiTU boundary.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=20260817)
|
||||
parser.add_argument("--warmup", type=int, default=3)
|
||||
parser.add_argument("--benchmark-iterations", type=int, default=0)
|
||||
parser.add_argument("--output", type=Path)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA is required")
|
||||
major, minor = torch.cuda.get_device_capability()
|
||||
if major != 12:
|
||||
raise SystemExit(f"SM120/SM121 is required, got SM{major}{minor}")
|
||||
|
||||
cases = []
|
||||
if args.suite in ("smoke", "all"):
|
||||
cases.extend(smoke_cases())
|
||||
if args.suite in ("real", "all"):
|
||||
cases.extend(real_cases())
|
||||
activations: list[ActivationName] = (
|
||||
["swiglu", "situ"] if args.activation == "both" else [args.activation]
|
||||
)
|
||||
|
||||
report = {
|
||||
"schema_version": 1,
|
||||
"created_at": datetime.now().astimezone().isoformat(),
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"flashinfer_version": importlib.metadata.version("flashinfer-python"),
|
||||
"device": torch.cuda.get_device_name(),
|
||||
"cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
|
||||
"logical_cuda_device": torch.cuda.current_device(),
|
||||
"compute_capability": [major, minor],
|
||||
"arguments": vars(args) | {"output": str(args.output) if args.output else None},
|
||||
"results": [],
|
||||
}
|
||||
failed = False
|
||||
for activation in activations:
|
||||
for case in cases:
|
||||
print(f"[RUN] activation={activation} case={case.name}", flush=True)
|
||||
try:
|
||||
result = run_one(case, activation, args)
|
||||
except Exception as error:
|
||||
failed = True
|
||||
result = {
|
||||
"case": asdict(case),
|
||||
"activation": activation,
|
||||
"status": "ERROR",
|
||||
"error": repr(error),
|
||||
}
|
||||
report["results"].append(result)
|
||||
print(
|
||||
f"[{result['status']}] activation={activation} case={case.name}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
output = args.output
|
||||
if output is None:
|
||||
stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
|
||||
output = Path("results") / f"correctness_{args.suite}_{stamp}.json"
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(report, indent=2, sort_keys=True), encoding="utf-8")
|
||||
print(f"report={output.resolve()}")
|
||||
|
||||
statuses = [item["status"] for item in report["results"]]
|
||||
counts = {status: statuses.count(status) for status in sorted(set(statuses))}
|
||||
print("summary=" + json.dumps(counts))
|
||||
if failed:
|
||||
raise SystemExit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@ -0,0 +1,39 @@
|
||||
ARG BASE_IMAGE=lmsysorg/sglang:kimi-k3-fiv617situ-warm
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ENV FLASHINFER_DISABLE_JIT=1 \
|
||||
FLASHINFER_DISABLE_VERSION_CHECK=1 \
|
||||
PYTHONUNBUFFERED=1
|
||||
|
||||
RUN rm -rf \
|
||||
/usr/local/lib/python3.12/dist-packages/flashinfer \
|
||||
/usr/local/lib/python3.12/dist-packages/flashinfer_python-*.dist-info
|
||||
|
||||
COPY flashinfer_python-0.6.18-py3-none-any.whl /tmp/flashinfer_python-0.6.18-py3-none-any.whl
|
||||
RUN python3 -m pip install --no-deps --force-reinstall \
|
||||
/tmp/flashinfer_python-0.6.18-py3-none-any.whl && \
|
||||
rm -f /tmp/flashinfer_python-0.6.18-py3-none-any.whl
|
||||
|
||||
# The wheel contains Python/C++ sources. These three libraries are the exact
|
||||
# SM120 AOT artifacts that passed the Phase 3/4 correctness gate.
|
||||
COPY aot/ /usr/local/lib/python3.12/dist-packages/flashinfer/data/aot/
|
||||
|
||||
# Apply only the SM120 backend-selection, weight-layout, and SiTU activation
|
||||
# changes against this Kimi image's exact SGLang source. Replacing the newer
|
||||
# Python tree would break the image's intentional sglang-kernel 0.4.5 pin.
|
||||
COPY sglang.patch /tmp/sglang.patch
|
||||
RUN patch --dry-run --batch --forward -p1 -d /sgl-workspace/sglang \
|
||||
< /tmp/sglang.patch && \
|
||||
patch --batch --forward -p1 -d /sgl-workspace/sglang \
|
||||
< /tmp/sglang.patch && \
|
||||
rm -f /tmp/sglang.patch
|
||||
|
||||
# Existing Kimi image compatibility fix: tcgen05 attention-residual is valid on
|
||||
# SM100/SM103, not consumer Blackwell SM120. Bake the fallback into the image.
|
||||
RUN python3 -c "from pathlib import Path; p=Path('/sgl-workspace/sglang/python/sglang/srt/layers/attn_residual.py'); s=p.read_text(); old=' _FAST_SUPPORTED = major >= 10'; new=' _FAST_SUPPORTED = major in (10, 11)'; assert old in s; p.write_text(s.replace(old,new))" && \
|
||||
python3 -m compileall -q /sgl-workspace/sglang/python/sglang/srt/layers && \
|
||||
python3 -c "from flashinfer.fused_moe import core; assert core.CUTLASS_FUSED_MOE_SUPPORTS_SITU is True"
|
||||
|
||||
LABEL ai.meta-stone.purpose="Kimi-K3 SM120 FlashInfer MXFP4 Phase 5 acceptance" \
|
||||
ai.meta-stone.sglang.patch="kimi-image-sm120-flashinfer-mxfp4" \
|
||||
ai.meta-stone.flashinfer.patch="7d1a4ee"
|
||||
@ -0,0 +1,390 @@
|
||||
# Kimi-K3 EP32 与 EP4：实现方式及 TTFT 差异分析
|
||||
|
||||
## 1. 先给结论
|
||||
|
||||
在当前四机 32 卡、`TP=32`、`DP=1`、`moe_a2a_backend=none` 的 Kimi-K3 部署里：
|
||||
|
||||
- `EP32` 不是 32 份模型副本，而是 896 个专家分散到 32 张卡；每张卡持有 28 个完整专家。
|
||||
- `EP4` 也不是只用 4 张卡跑 MoE，而是把 32 张卡分解成 `EP4 × MoE-TP8`。每个节点持有 224 个专家，每个专家的中间维再由节点内 8 张卡切分。
|
||||
- 两种布局每卡承载的专家参数量和理论 MoE FLOPs 接近，但计算颗粒度不同。
|
||||
- 当前没有启用 A2A。每张卡仍看到同一批 token，只计算自己负责的专家或专家分片，随后在 TP32 上汇总结果。
|
||||
- 同后端 Marlin、同请求口径下，EP4 相比 EP32 的 Input TPS 提高约 **15.9%**，TTFT P50 降低约 **13.9%**。
|
||||
- 这次收益不是“把 AllReduce 从 32 卡缩成 4 卡”。Kimi-K3 当前 `A2A=none` 路径仍在完整 TP32 组上归约。更合理的解释是：EP4 把每个专家的计算切到 8 卡，改善了路由偏斜、热点专家关键路径和 grouped MoE kernel 的并行形状。
|
||||
|
||||
## 2. 审计对象
|
||||
|
||||
### 2.1 模型参数
|
||||
|
||||
来源：`/data/hf_models/Kimi-K3/config.json` 的 `text_config`。
|
||||
|
||||
| 参数 | 数值 |
|
||||
|---|---:|
|
||||
| hidden size | 7168 |
|
||||
| Latent MoE hidden size | 3584 |
|
||||
| 单专家 intermediate size | 3072 |
|
||||
| routed experts | 896 |
|
||||
| 每 token 激活专家数 | 16 |
|
||||
| Transformer 层数 | 93 |
|
||||
|
||||
### 2.2 源码版本
|
||||
|
||||
审计仓库：
|
||||
|
||||
```text
|
||||
/data/hzy/src/sglang-kimi-sm120
|
||||
commit fb929bbccb7a34640f3d4904174767eb7ae95509
|
||||
```
|
||||
|
||||
实验镜像：
|
||||
|
||||
```text
|
||||
local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5
|
||||
```
|
||||
|
||||
EP32 和 EP4 的 Marlin 实验使用同一镜像、同一模型、同一请求、同一 chunk 和同一服务参数。启动命令中唯一关键差异是：
|
||||
|
||||
```text
|
||||
EP32: --tp-size 32 --ep-size 32
|
||||
EP4 : --tp-size 32 --ep-size 4
|
||||
```
|
||||
|
||||
## 3. EP、MoE-TP 和外层 TP 的关系
|
||||
|
||||
SGLang 不是把 `TP` 和 `EP` 相乘申请 GPU，而是在外层 TP 组内部进一步分解 MoE：
|
||||
|
||||
```python
|
||||
moe_tp_size = tensor_model_parallel_size // moe_ep_size // moe_dp_size
|
||||
```
|
||||
|
||||
源码：
|
||||
|
||||
```text
|
||||
python/sglang/srt/distributed/parallel_state.py:2526-2528
|
||||
python/sglang/srt/model_executor/model_runner_components/moe_ep_setup.py:123-137
|
||||
```
|
||||
|
||||
本实验 `TP=32`、`MoE-DP=1`，因此：
|
||||
|
||||
```text
|
||||
TP32 = EP32 × MoE-TP1
|
||||
TP32 = EP4 × MoE-TP8
|
||||
```
|
||||
|
||||
| 配置 | EP size | 派生 MoE-TP size | 含义 |
|
||||
|---|---:|---:|---|
|
||||
| EP32 | 32 | 1 | 每个专家只落在 1 张 GPU 上 |
|
||||
| EP4 | 4 | 8 | 每个专家落在一个 8 GPU MoE-TP 组上 |
|
||||
|
||||
所以 `EP4` 的准确叫法是“较小 EP + 较大 MoE-TP”的混合专家布局。
|
||||
|
||||
## 4. 32 张卡具体怎么分组
|
||||
|
||||
假设全局 rank `0-7 / 8-15 / 16-23 / 24-31` 分别属于四个节点。
|
||||
|
||||
### 4.1 EP32
|
||||
|
||||
```text
|
||||
MoE-EP group: [0, 1, 2, ..., 31]
|
||||
MoE-TP group: 每个 rank 自己，size=1
|
||||
```
|
||||
|
||||
896 个专家平均切到 32 个 EP rank：
|
||||
|
||||
```text
|
||||
每 rank 专家数 = 896 / 32 = 28
|
||||
```
|
||||
|
||||
例如：
|
||||
|
||||
```text
|
||||
rank 0 负责专家 0-27
|
||||
rank 1 负责专家 28-55
|
||||
...
|
||||
rank 31 负责专家 868-895
|
||||
```
|
||||
|
||||
每个专家在其所属 GPU 上保留完整的 intermediate `3072`。
|
||||
|
||||
### 4.2 EP4
|
||||
|
||||
SGLang 构造的 MoE-TP 组是连续 8 rank：
|
||||
|
||||
```text
|
||||
[0-7], [8-15], [16-23], [24-31]
|
||||
```
|
||||
|
||||
每组正好对应一个 8 卡节点。MoE-EP 组则按 MoE-TP rank 跨节点取 rank：
|
||||
|
||||
```text
|
||||
[0,8,16,24], [1,9,17,25], ..., [7,15,23,31]
|
||||
```
|
||||
|
||||
源码：
|
||||
|
||||
```text
|
||||
python/sglang/srt/distributed/parallel_state.py:2559-2615
|
||||
```
|
||||
|
||||
896 个专家先切为 4 份：
|
||||
|
||||
```text
|
||||
每个 EP rank 专家数 = 896 / 4 = 224
|
||||
```
|
||||
|
||||
每个专家的 intermediate 再切给 MoE-TP8：
|
||||
|
||||
```text
|
||||
每卡专家分片宽度 = 3072 / 8 = 384
|
||||
```
|
||||
|
||||
因此：
|
||||
|
||||
```text
|
||||
节点 0 的 8 张卡共同持有专家 0-223
|
||||
节点 1 的 8 张卡共同持有专家 224-447
|
||||
节点 2 的 8 张卡共同持有专家 448-671
|
||||
节点 3 的 8 张卡共同持有专家 672-895
|
||||
```
|
||||
|
||||
节点内 8 张卡不是各存一份 224 个完整专家，而是各存这些专家的 `1/8` intermediate 分片。
|
||||
|
||||
## 5. 为什么 EP4 没有让每卡专家权重暴涨 8 倍
|
||||
|
||||
SGLang 先按 EP 切专家数量，再按 MoE-TP 切专家中间维：
|
||||
|
||||
```python
|
||||
self._num_local_routed = self._num_global_routed // storage_ep_size
|
||||
self.intermediate_size_per_partition = intermediate_size // self.moe_tp_size
|
||||
```
|
||||
|
||||
源码：
|
||||
|
||||
```text
|
||||
python/sglang/srt/layers/moe/fused_moe_triton/layer.py:291-309
|
||||
```
|
||||
|
||||
权重加载时，`w1/w3` 沿输出维切分，`w2` 沿输入维切分：
|
||||
|
||||
```text
|
||||
python/sglang/srt/layers/moe/fused_moe_triton/layer.py:618-682
|
||||
python/sglang/srt/layers/moe/fused_moe_triton/layer.py:702-770
|
||||
```
|
||||
|
||||
只看与 intermediate 成正比的部分，每卡工作量满足：
|
||||
|
||||
```text
|
||||
EP32: 28 × 3072 = 86016
|
||||
EP4 : 224 × 384 = 86016
|
||||
```
|
||||
|
||||
真实服务也印证了每卡权重显存接近：
|
||||
|
||||
| 配置 | 后端 | 每卡 weight memory |
|
||||
|---|---|---:|
|
||||
| EP32 | Marlin | 59.63 GiB |
|
||||
| EP4 | Marlin | 59.28 GiB |
|
||||
| EP4 | FlashInfer MXFP4 | 59.65 GiB |
|
||||
|
||||
小幅差异来自权重布局、padding、workspace 和后端转换，不是模型参数量发生了数量级变化。
|
||||
|
||||
## 6. 当前 `A2A=none` 时，一个 token 怎么执行
|
||||
|
||||
### 6.1 路由
|
||||
|
||||
Kimi-K3 的 gate 为每个 token 计算 896 个专家分数，再选 top-16：
|
||||
|
||||
```text
|
||||
hidden states
|
||||
-> gate / grouped top-k
|
||||
-> 16 个全局 expert id + routing weight
|
||||
```
|
||||
|
||||
模型入口：
|
||||
|
||||
```text
|
||||
python/sglang/srt/models/kimi_k3.py:386-469
|
||||
python/sglang/srt/models/kimi_k3.py:1020-1059
|
||||
```
|
||||
|
||||
### 6.2 没有 token A2A 搬运
|
||||
|
||||
`moe_a2a_backend=none` 会创建 `StandardDispatcher`：
|
||||
|
||||
```text
|
||||
python/sglang/srt/layers/moe/fused_moe_triton/layer.py:137-178
|
||||
```
|
||||
|
||||
StandardDispatcher 不会把 token 发到专家所在 GPU。它保留 hidden states，只处理专家 ID：
|
||||
|
||||
```python
|
||||
hidden_states = hidden_states
|
||||
```
|
||||
|
||||
Marlin 路径把不属于本 EP rank 的全局专家映射为 `-1`，本地 kernel 只计算自己负责的专家：
|
||||
|
||||
```text
|
||||
python/sglang/srt/layers/moe/token_dispatcher/standard.py:182-246
|
||||
python/sglang/srt/layers/moe/fused_moe_triton/layer.py:903-911
|
||||
```
|
||||
|
||||
FlashInfer 路径不在 dispatcher 中改 ID，而是把 `moe_ep_size/rank` 与 `moe_tp_size/rank` 传给 CUTLASS runner，由 runner 解释全局专家 ID：
|
||||
|
||||
```text
|
||||
python/sglang/srt/layers/quantization/mxfp4.py:1378-1401
|
||||
```
|
||||
|
||||
两者实现位置不同，但语义一致：没有 A2A token dispatch，每个 rank 从同一批 token 中计算自己的专家贡献。
|
||||
|
||||
### 6.3 汇总部分结果
|
||||
|
||||
Kimi-K3 的 routed expert 输出位于 3584 维 latent 空间。A2A 未启用时，源码要求在 RMSNorm 前完成归约：
|
||||
|
||||
```python
|
||||
def _routed_needs_reduce(self):
|
||||
return self.tp_size > 1 and get_moe_a2a_backend().is_none()
|
||||
|
||||
return self._latent_norm(tensor_model_parallel_all_reduce(latent))
|
||||
```
|
||||
|
||||
源码：
|
||||
|
||||
```text
|
||||
python/sglang/srt/models/kimi_k3.py:699-702
|
||||
python/sglang/srt/models/kimi_k3.py:955-960
|
||||
```
|
||||
|
||||
这里调用的是完整的 `get_tp_group()`，即 TP32。某些优化路径会把这次归约与相邻算子融合，但数学上的通信域仍是 TP32。
|
||||
|
||||
因此，本次 EP4 收益不能解释成“MoE AllReduce 从 32 rank 缩成 4 rank”。
|
||||
|
||||
## 7. EP4 为什么更快
|
||||
|
||||
### 7.1 同口径实验结果
|
||||
|
||||
唯一能隔离 EP 变量的比较是 Marlin 对 Marlin：
|
||||
|
||||
```text
|
||||
模型 Kimi-K3
|
||||
4 节点 / 32 GPU
|
||||
TP32 / PP1 / DP1
|
||||
ISL=16K / OSL=1 / C=8
|
||||
chunked_prefill_size=8K
|
||||
moe_a2a_backend=none
|
||||
moe_runner_backend=marlin
|
||||
```
|
||||
|
||||
| 配置 | repeats | Input TPS | TTFT P50 | TTFT P95 |
|
||||
|---|---:|---:|---:|---:|
|
||||
| EP32 | 2 | 2531.43 | 50.52 s | 53.64 s |
|
||||
| EP4 | 3 | 2935.02 | 43.51 s | 46.22 s |
|
||||
| EP4 vs EP32 | | **+15.9%** | **-13.9%** | **-13.8%** |
|
||||
|
||||
结果路径：
|
||||
|
||||
```text
|
||||
EP32:
|
||||
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/results/kimi3-moe-prefill-20260818-130900/raw/
|
||||
|
||||
EP4:
|
||||
/data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill/results/kimi3-ep4-moe-full-20260818-151349/raw/
|
||||
```
|
||||
|
||||
EP4 + FlashInfer 的同场景结果为 3257.96 Input TPS、39.19 s TTFT P50。它相对 EP4 + Marlin 还包含约 11% 的 MoE backend 收益，不能全部归因于 EP。
|
||||
|
||||
### 7.2 主要机制：把单专家关键路径从 1 卡摊到 8 卡
|
||||
|
||||
Kimi-K3 每个 token 激活 16 个专家。路由并不保证所有专家收到完全一样数量的 token，总会出现热点专家和冷门专家。
|
||||
|
||||
EP32 中：
|
||||
|
||||
```text
|
||||
一个专家只属于一张 GPU
|
||||
热点专家的完整 3584 -> 3072 -> 3584 计算由这张 GPU 独自完成
|
||||
其他 GPU 即使更空，也不能帮助它完成该专家
|
||||
```
|
||||
|
||||
EP4 中：
|
||||
|
||||
```text
|
||||
一个专家属于一个节点的 8 卡 MoE-TP 组
|
||||
同一专家的 intermediate 3072 被切成 8 份，每卡计算 384
|
||||
热点专家的矩阵计算由 8 张 GPU 共同承担
|
||||
```
|
||||
|
||||
这不会减少总 FLOPs，但会降低单个热门专家压在单卡上的关键路径，并让节点内 8 卡对同一专家协同计算。
|
||||
|
||||
### 7.3 grouped MoE kernel 的执行形状改变
|
||||
|
||||
每卡的理论 FLOPs 接近，但 kernel 看到的形状不同：
|
||||
|
||||
| 配置 | 本地专家槽位 | 单专家 intermediate 分片 | 平均每 token 本 rank 参与的专家分片 |
|
||||
|---|---:|---:|---:|
|
||||
| EP32 | 28 | 3072 | 16/32 = 0.5 |
|
||||
| EP4 | 224 | 384 | 16/4 = 4 |
|
||||
|
||||
EP4 让更多专家分片同时出现在一个 rank 的 grouped MoE 调度中，并把单个专家的大矩阵切成更多可并行的小分片。在本次 SM120 + Marlin 实测中，这个形状明显更合适。
|
||||
|
||||
这属于“同样总计算，不同并行颗粒度”的收益，不是模型数学变化。
|
||||
|
||||
### 7.4 通信没有消失，但也没有因 MoE-TP8 再增加一次完整归约
|
||||
|
||||
当前 Kimi `A2A=none` 路径最终直接在 TP32 汇总 latent 输出。这一次 TP32 AllReduce 同时合并了：
|
||||
|
||||
- 不同 EP rank 持有的不同专家贡献；
|
||||
- EP4 中同一专家的 8 个 MoE-TP 分片贡献。
|
||||
|
||||
所以 EP4 没有额外增加一轮独立的“EP4 AllReduce + TP8 AllReduce”。这使它可以获得更好的计算分摊，而不在当前实现中付出第二次完整 collective 的代价。
|
||||
|
||||
## 8. EP4 与真正 A2A EP 的区别
|
||||
|
||||
当前实验：
|
||||
|
||||
```text
|
||||
--ep-size 4
|
||||
--moe-a2a-backend none
|
||||
```
|
||||
|
||||
它只是改变专家权重的 EP/MoE-TP 布局，不会按路由结果搬 token。
|
||||
|
||||
真正的 DeepEP 或 FlashInfer A2A 会执行：
|
||||
|
||||
```text
|
||||
dispatch: token 按 top-k 发往专家所在 rank
|
||||
compute : rank 只计算收到的 token
|
||||
combine : 专家结果返回 token 原所属 rank 并加权合并
|
||||
```
|
||||
|
||||
SGLang 的 dispatcher 选择源码：
|
||||
|
||||
```text
|
||||
Standard/A2A 分支：python/sglang/srt/layers/moe/fused_moe_triton/layer.py:137-178
|
||||
DeepEP dispatcher：python/sglang/srt/layers/moe/token_dispatcher/
|
||||
```
|
||||
|
||||
Kimi-K3 源码也明确说明，A2A 后 combine 已返回完整 routed sum，因此不能再做 TP AllReduce：
|
||||
|
||||
```text
|
||||
python/sglang/srt/models/kimi_k3.py:506-516
|
||||
python/sglang/srt/models/kimi_k3.py:699-709
|
||||
```
|
||||
|
||||
所以“EP4 比 EP32 快”不能直接推出“DeepEP EP4 一定更快”。A2A 会引入完全不同的 token dispatch/combine、buffer 和共享专家约束，需要单独实验。
|
||||
|
||||
## 9. 证据边界
|
||||
|
||||
已经由源码和实验确认：
|
||||
|
||||
1. `TP32/EP32` 派生 `MoE-TP1`，`TP32/EP4` 派生 `MoE-TP8`。
|
||||
2. EP32 每卡 28 个完整专家；EP4 每个节点 224 个专家、每卡持有其 1/8 intermediate 分片。
|
||||
3. 两种 Marlin 实验只改变 `--ep-size`，EP4 的 TTFT P50 低 13.9%。
|
||||
4. 当前 `A2A=none` 不搬 token，Kimi routed latent 仍在完整 TP32 上归约。
|
||||
5. 每卡权重显存基本相同，收益不是因为 EP4 少加载了大量权重。
|
||||
|
||||
目前还没有由成对 Nsight Trace 直接确认：
|
||||
|
||||
1. 13.9% TTFT 收益中，Marlin grouped MoE kernel 本身占多少；
|
||||
2. 路由负载偏斜改善占多少；
|
||||
3. padding、排序和 expert mapping 开销分别变化多少。
|
||||
|
||||
若需要把原因从“源码约束下的机制解释”升级为“算子级定量归因”，下一步只需补一组相同 16K、C8、chunk8K 的 EP32/EP4 Nsight 对照，比较每层 Marlin MoE kernel、top-k/sort 与 TP32 AllReduce 时间。现有结果已经足够说明 EP4 是当前部署更好的基线，但不应虚构尚未采集的算子时间占比。
|
||||
|
||||
@ -0,0 +1,135 @@
|
||||
# Kimi-K3 / Pro6000D / SGLang TP32 EP4 MoE Backend Prefill
|
||||
|
||||
This experiment is the real-serving acceptance gate for Kimi-K3 FlashInfer
|
||||
MXFP4 MoE on NVIDIA SM120. It does not modify the existing baseline experiment.
|
||||
|
||||
## Fixed matrix
|
||||
|
||||
| Dimension | Values |
|
||||
|---|---|
|
||||
| Nodes | 601-604, 8 GPUs per node |
|
||||
| Parallelism | TP32, EP4, DP1 |
|
||||
| MoE runner | `marlin`, `flashinfer_mxfp4` |
|
||||
| Input / output | 16K / 1 token |
|
||||
| Concurrency | 8, 16 |
|
||||
| Chunked prefill | 8K, 16K |
|
||||
| Requests | 40 per repeat |
|
||||
| Repeats | 3 |
|
||||
|
||||
The eight configurations are identical except for MoE runner, chunk size, and
|
||||
benchmark concurrency. The directory retains its historical `tp32ep32` name so
|
||||
existing evidence paths remain stable; the production comparison and the
|
||||
default `EP_SIZE` are EP4.
|
||||
The client uses SGLang `random-ids`, which produces the exact synthetic token
|
||||
length without depending on a ShareGPT file or text truncation. Radix cache is
|
||||
disabled. `OSL=1` makes TTFT, E2E latency and input TPS the primary metrics;
|
||||
TPOT is not meaningful for this matrix.
|
||||
|
||||
The defaults above can be narrowed for a capacity-first run with `EP_SIZE`,
|
||||
`BACKENDS_CSV`, `CHUNKED_PREFILL_SIZES_CSV`, and `CONCURRENCIES_CSV`. For
|
||||
example, the TP32/EP4 maximum-pressure probe uses Marlin, 16K chunk, and C=16
|
||||
before spending time on the complete backend comparison.
|
||||
|
||||
## Accepted EP4 result
|
||||
|
||||
Run ID: `kimi3-ep4-moe-full-20260818-151349`.
|
||||
|
||||
All 24 repeat files completed 40/40 requests with empty error arrays. No OOM,
|
||||
Traceback, NCCL error, or EngineDeadError was found. Values below are medians
|
||||
across three repeats.
|
||||
|
||||
| Chunk | C | Backend | Input TPS | TTFT P50 | TTFT P95 |
|
||||
|---:|---:|---|---:|---:|---:|
|
||||
| 8K | 8 | Marlin | 2935.02 | 43.51 s | 46.22 s |
|
||||
| 8K | 8 | FlashInfer MXFP4 | 3257.96 | 39.19 s | 41.62 s |
|
||||
| 8K | 16 | Marlin | 2934.90 | 87.09 s | 89.86 s |
|
||||
| 8K | 16 | FlashInfer MXFP4 | 3260.14 | 78.39 s | 80.88 s |
|
||||
| 16K | 8 | Marlin | 2812.91 | 45.47 s | 45.86 s |
|
||||
| 16K | 8 | FlashInfer MXFP4 | 3027.50 | 42.27 s | 42.63 s |
|
||||
| 16K | 16 | Marlin | 2812.33 | 90.97 s | 91.36 s |
|
||||
| 16K | 16 | FlashInfer MXFP4 | 3027.34 | 84.52 s | 84.85 s |
|
||||
|
||||
FlashInfer improves Input TPS by 11.00%-11.08% with an 8K chunk and by
|
||||
7.63%-7.65% with a 16K chunk. It reduces median TTFT by 9.92%-9.99% and
|
||||
7.04%-7.09%, respectively. The accepted Prefill setting is therefore
|
||||
`flashinfer_mxfp4`, EP4, and an 8K chunk for this 16K/1 workload.
|
||||
|
||||
The only matching existing EP32 evidence is Marlin/8K/C8 with two repeats.
|
||||
Its median Input TPS is 2531.43 and TTFT P50/P95 are 50.52/53.64 seconds.
|
||||
Against that evidence, EP4 raises Input TPS by 15.94% and lowers TTFT P50 by
|
||||
13.88%. EP32 was not rerun.
|
||||
|
||||
## Unique entrypoint
|
||||
|
||||
Run on 601 only. The sudo password is supplied at runtime and is never written
|
||||
to Git or result files.
|
||||
|
||||
```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 build-image
|
||||
bash run_moe_backend_prefill.sh run
|
||||
```
|
||||
|
||||
For an unattended run, use the same entrypoint under tmux:
|
||||
|
||||
```bash
|
||||
umask 077
|
||||
printf '%s\n' "$SUDO_PASSWORD" > /tmp/kimi3-phase5-sudo
|
||||
unset SUDO_PASSWORD
|
||||
|
||||
tmux new-session -d -s kimi3-moe-prefill \
|
||||
"cd /data/hzy/sskj/experiments/pro6000/kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill && \
|
||||
SUDO_PASSWORD_FILE=/tmp/kimi3-phase5-sudo RUN_ID=\"${RUN_ID}\" \
|
||||
bash run_moe_backend_prefill.sh all 2>&1 | tee /data/hzy/${RUN_ID}.log"
|
||||
```
|
||||
|
||||
Delete `/tmp/kimi3-phase5-sudo` after the run. The file is mode `0600` and is
|
||||
never copied to result directories.
|
||||
|
||||
## Outputs
|
||||
|
||||
```text
|
||||
results/<RUN_ID>/
|
||||
build/ per-node image build and smoke evidence
|
||||
service/ exact docker commands and all four server logs
|
||||
raw/ one benchmark JSONL per repeat
|
||||
bench/ human-readable benchmark logs
|
||||
gpu/ per-node GPU snapshots
|
||||
results.csv flat repeat-level metrics
|
||||
summary.json median metrics and FlashInfer-vs-Marlin deltas
|
||||
orchestrator.log lifecycle and progress
|
||||
```
|
||||
|
||||
The run aborts on a failed service start, any failed request, a missing result,
|
||||
or an image capability smoke failure. Containers are removed on exit.
|
||||
|
||||
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`.
|
||||
@ -0,0 +1,34 @@
|
||||
#!/usr/bin/env bash
|
||||
# Kimi-K3 real-serving MoE backend acceptance on 601-604.
|
||||
|
||||
EXPERIMENT="kimi3_pro6000_sglang_tp32ep32_moe_backend_prefill"
|
||||
MODEL_PATH="${MODEL_PATH:-/data/hf_models/Kimi-K3}"
|
||||
SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-kimi-k3}"
|
||||
PORT="${PORT:-30000}"
|
||||
|
||||
HEAD_HOST="${HEAD_HOST:-174.1.60.1}"
|
||||
NODE_SSH_USER="${NODE_SSH_USER:-user}"
|
||||
NODE_HOSTS=(174.1.60.1 174.1.60.2 174.1.60.3 174.1.60.4)
|
||||
DIST_PORT="${DIST_PORT:-20000}"
|
||||
EP_SIZE="${EP_SIZE:-4}"
|
||||
|
||||
BASE_IMAGE="${BASE_IMAGE:-lmsysorg/sglang:kimi-k3-fiv617situ-warm}"
|
||||
CANDIDATE_IMAGE="${CANDIDATE_IMAGE:-local/sglang:kimi-k3-sm120-flashinfer-mxfp4-phase5}"
|
||||
FLASHINFER_TARGET="${FLASHINFER_TARGET:-/tmp/hzy-fi-situ-phase3-v3}"
|
||||
FLASHINFER_WHEEL="${FLASHINFER_WHEEL:-/data/hzy/artifacts/flashinfer-kimi-sm120-situ-phase3-v3/flashinfer_python-0.6.18-py3-none-any.whl}"
|
||||
|
||||
BENCH_DATASET_NAME="${BENCH_DATASET_NAME:-random-ids}"
|
||||
BENCH_MODULE="${BENCH_MODULE:-sglang.benchmark.serving}"
|
||||
RESULT_BASE="${RESULT_BASE:-${SCRIPT_DIR}/results}"
|
||||
|
||||
INPUT_LEN="${INPUT_LEN:-16384}"
|
||||
OUTPUT_LEN="${OUTPUT_LEN:-1}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-40}"
|
||||
REPEATS="${REPEATS:-3}"
|
||||
WARMUP_REQUESTS="${WARMUP_REQUESTS:-2}"
|
||||
HEALTH_WAIT_S="${HEALTH_WAIT_S:-2400}"
|
||||
|
||||
IFS=',' read -r -a BACKENDS <<<"${BACKENDS_CSV:-marlin,flashinfer_mxfp4}"
|
||||
IFS=',' read -r -a CHUNKED_PREFILL_SIZES \
|
||||
<<<"${CHUNKED_PREFILL_SIZES_CSV:-8192,16384}"
|
||||
IFS=',' read -r -a CONCURRENCIES <<<"${CONCURRENCIES_CSV:-8,16}"
|
||||
@ -0,0 +1,309 @@
|
||||
--- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
|
||||
+++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
|
||||
@@ -114,17 +114,21 @@
|
||||
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,
|
||||
+ if runner_config.activation == "situ" and runner_config.is_gated:
|
||||
+ activation = ActivationType.Situ
|
||||
+ else:
|
||||
+ activation = ActivationType(
|
||||
+ get_activation_type(
|
||||
+ runner_config.activation,
|
||||
+ is_gated=runner_config.is_gated,
|
||||
+ )
|
||||
)
|
||||
- )
|
||||
supported = {
|
||||
ActivationType.Swiglu,
|
||||
ActivationType.Geglu,
|
||||
ActivationType.Relu2,
|
||||
ActivationType.Identity,
|
||||
+ ActivationType.Situ,
|
||||
}
|
||||
assert activation in supported, (
|
||||
f"Activation {runner_config.activation!r} "
|
||||
@@ -306,7 +310,7 @@
|
||||
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
|
||||
@@ -343,3 +347,4 @@
|
||||
+ x = x.contiguous()
|
||||
x, input_sf = mxfp8_quantize(
|
||||
x,
|
||||
is_sf_swizzled_layout=True,
|
||||
@@ -384,7 +389,7 @@
|
||||
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,
|
||||
)
|
||||
--- a/python/sglang/srt/layers/quantization/mxfp4.py
|
||||
+++ b/python/sglang/srt/layers/quantization/mxfp4.py
|
||||
@@ -340,8 +340,7 @@
|
||||
self.flashinfer_mxfp4_moe_precision = (
|
||||
get_server_args().flashinfer_mxfp4_moe_precision
|
||||
)
|
||||
- # When `flashinfer_mxfp4` is enabled, dispatch to one of two FlashInfer
|
||||
- # entry points depending on the GPU:
|
||||
+ # When `flashinfer_mxfp4` is enabled, dispatch by GPU architecture:
|
||||
# - SM100 (Blackwell) -> trtllm_fp4_block_scale_moe (existing)
|
||||
# - SM90 (Hopper) -> cutlass_fused_moe(use_w4_group_scaling=True)
|
||||
# (FlashInfer PR #3084, post-0.6.10)
|
||||
@@ -349,6 +348,8 @@
|
||||
if self.use_flashinfer:
|
||||
if is_sm100_supported():
|
||||
self._fi_kernel = "trtllm_sm100"
|
||||
+ elif is_sm120_supported():
|
||||
+ self._fi_kernel = "cutlass_sm120"
|
||||
elif is_sm90_supported():
|
||||
if not _FI_HAS_SM90_CUTLASS_MXFP4:
|
||||
raise RuntimeError(
|
||||
@@ -360,7 +361,8 @@
|
||||
self._fi_kernel = "cutlass_sm90"
|
||||
else:
|
||||
raise NotImplementedError(
|
||||
- "moe_runner_backend=flashinfer_mxfp4 requires SM90 or SM100."
|
||||
+ "moe_runner_backend=flashinfer_mxfp4 requires SM90, SM100, "
|
||||
+ "or SM120."
|
||||
)
|
||||
|
||||
def create_weights(
|
||||
@@ -412,7 +414,7 @@
|
||||
intermediate_size_per_partition_after_pad = round_up(
|
||||
intermediate_size_per_partition, triton_kernels_padding_alignment
|
||||
)
|
||||
- elif self._fi_kernel == "cutlass_sm90":
|
||||
+ elif self._fi_kernel in ("cutlass_sm90", "cutlass_sm120"):
|
||||
# cutlass mixed-input GEMM contraction dim K must be % 128 == 0
|
||||
# (interleave factor for MXFP4 group_size=32 is 4). The kernel
|
||||
# also expects ``fc1_expert_weights`` in halved ``[up; gate]``
|
||||
@@ -609,6 +611,9 @@
|
||||
if self._fi_kernel == "cutlass_sm90":
|
||||
self._process_weights_for_sm90_cutlass(layer)
|
||||
return
|
||||
+ if self._fi_kernel == "cutlass_sm120":
|
||||
+ self._process_weights_for_sm120_cutlass(layer)
|
||||
+ return
|
||||
if self.use_flashinfer:
|
||||
# Per-expert buffers are local (create_weights uses num_local_experts);
|
||||
# the global self.num_experts here breaks EP>1. Mirrors the SM90 path.
|
||||
@@ -1130,7 +1135,136 @@
|
||||
layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
+
|
||||
+ def _process_weights_for_sm120_cutlass(self, layer):
|
||||
+ """Prepare Kimi/GPT-OSS MXFP4 experts for FlashInfer CUTLASS SM120."""
|
||||
+ from flashinfer import block_scale_interleave
|
||||
+
|
||||
+ sf_block_size = 32
|
||||
+ n_un = layer.w13_weight.shape[1] // 2
|
||||
+ k_un = layer.w13_weight.shape[2] * 2
|
||||
+ n_pad = self._padded_intermediate
|
||||
+ k_pad = self._padded_hidden
|
||||
+ num_experts = layer.num_local_experts
|
||||
+ device = layer.w13_weight.device
|
||||
+ gate_up_interleaved = getattr(
|
||||
+ layer.moe_runner_config, "gate_up_interleaved", True
|
||||
+ )
|
||||
+
|
||||
+ 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(unpadded, last_pad, last_un):
|
||||
+ gate_rows, up_rows = split_gate_up(unpadded)
|
||||
+ out = torch.zeros(
|
||||
+ num_experts,
|
||||
+ 2 * n_pad,
|
||||
+ last_pad,
|
||||
+ dtype=unpadded.dtype,
|
||||
+ device=device,
|
||||
+ )
|
||||
+ out[:, :n_un, :last_un] = up_rows
|
||||
+ out[:, n_pad : n_pad + n_un, :last_un] = gate_rows
|
||||
+ return out
|
||||
+
|
||||
+ w13_padded = stack_up_gate(
|
||||
+ layer.w13_weight.data, k_pad // 2, k_un // 2
|
||||
+ )
|
||||
+ w13_scale_padded = stack_up_gate(
|
||||
+ layer.w13_weight_scale.data,
|
||||
+ k_pad // sf_block_size,
|
||||
+ k_un // sf_block_size,
|
||||
+ )
|
||||
|
||||
+ bias_dtype = layer.w13_weight_bias.dtype
|
||||
+ w13_bias_padded = torch.zeros(
|
||||
+ num_experts, 2 * n_pad, dtype=bias_dtype, device=device
|
||||
+ )
|
||||
+ 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(unpadded, last_pad, last_un):
|
||||
+ out = torch.zeros(
|
||||
+ num_experts,
|
||||
+ k_pad,
|
||||
+ last_pad,
|
||||
+ dtype=unpadded.dtype,
|
||||
+ device=device,
|
||||
+ )
|
||||
+ out[:, :k_un, :last_un] = unpadded[:, :k_un, :]
|
||||
+ return out
|
||||
+
|
||||
+ w2_padded = pad_w2(layer.w2_weight.data, n_pad // 2, n_un // 2)
|
||||
+ w2_scale_padded = pad_w2(
|
||||
+ layer.w2_weight_scale.data,
|
||||
+ n_pad // sf_block_size,
|
||||
+ n_un // sf_block_size,
|
||||
+ )
|
||||
+ w2_bias_padded = torch.zeros(
|
||||
+ num_experts, k_pad, dtype=bias_dtype, device=device
|
||||
+ )
|
||||
+ w2_bias_padded[:, :k_un] = layer.w2_weight_bias.data
|
||||
+
|
||||
+ layer.w13_weight = Parameter(w13_padded, requires_grad=False)
|
||||
+ layer.w2_weight = Parameter(w2_padded, requires_grad=False)
|
||||
+ layer.w13_weight_scale = Parameter(
|
||||
+ block_scale_interleave(w13_scale_padded).reshape_as(w13_scale_padded),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
+ layer.w2_weight_scale = Parameter(
|
||||
+ block_scale_interleave(w2_scale_padded).reshape_as(w2_scale_padded),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
+ layer.w13_weight_bias = Parameter(w13_bias_padded, requires_grad=False)
|
||||
+ layer.w2_weight_bias = Parameter(w2_bias_padded, requires_grad=False)
|
||||
+
|
||||
+ activation = layer.moe_runner_config.activation
|
||||
+ alpha = layer.moe_runner_config.gemm1_alpha
|
||||
+ # Older Kimi image baselines predate the generic gemm1_beta field.
|
||||
+ beta = getattr(layer.moe_runner_config, "gemm1_beta", None)
|
||||
+ limit = layer.moe_runner_config.gemm1_clamp_limit
|
||||
+ if activation == "situ":
|
||||
+ alpha = 4.0 if alpha is None else alpha
|
||||
+ beta = 25.0 if limit is None else limit
|
||||
+ limit = None
|
||||
+ else:
|
||||
+ 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 = Parameter(
|
||||
+ torch.full((num_experts,), alpha, dtype=torch.float32, device=device),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
+ layer.swiglu_beta = Parameter(
|
||||
+ torch.full((num_experts,), beta, dtype=torch.float32, device=device),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
+ layer.swiglu_limit = (
|
||||
+ None
|
||||
+ if limit is None
|
||||
+ else Parameter(
|
||||
+ torch.full(
|
||||
+ (num_experts,), limit, dtype=torch.float32, device=device
|
||||
+ ),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
+ )
|
||||
+ layer.mxfp4_weight_global_scale = Parameter(
|
||||
+ torch.ones(num_experts, dtype=torch.float32, device=device),
|
||||
+ requires_grad=False,
|
||||
+ )
|
||||
+ layer._mxfp4_backend = "flashinfer_cutlass_sm120"
|
||||
+ torch.cuda.empty_cache()
|
||||
+
|
||||
def create_moe_runner(
|
||||
self, layer: torch.nn.Module, moe_runner_config: MoeRunnerConfig
|
||||
):
|
||||
@@ -1154,10 +1288,23 @@
|
||||
or moe_runner_backend.is_deep_gemm()
|
||||
):
|
||||
self.runner = MoeRunner(moe_runner_backend, moe_runner_config)
|
||||
- elif (
|
||||
- moe_runner_backend.is_flashinfer_mxfp4()
|
||||
- and self._fi_kernel == "cutlass_sm90"
|
||||
+ elif moe_runner_backend.is_flashinfer_mxfp4() and self._fi_kernel in (
|
||||
+ "cutlass_sm90",
|
||||
+ "cutlass_sm120",
|
||||
):
|
||||
+ if (
|
||||
+ self._fi_kernel == "cutlass_sm120"
|
||||
+ and moe_runner_config.activation == "situ"
|
||||
+ ):
|
||||
+ from flashinfer.fused_moe import core as flashinfer_moe_core
|
||||
+
|
||||
+ if not getattr(
|
||||
+ flashinfer_moe_core, "CUTLASS_FUSED_MOE_SUPPORTS_SITU", False
|
||||
+ ):
|
||||
+ raise RuntimeError(
|
||||
+ "Kimi-K3 FlashInfer MXFP4 MoE on SM120 requires a "
|
||||
+ "FlashInfer build with CUTLASS SiTU support."
|
||||
+ )
|
||||
# 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
|
||||
@@ -1215,8 +1362,33 @@
|
||||
quant_info = build_marlin_moe_quant_info(layer)
|
||||
return self.runner.run(
|
||||
dispatch_output._replace(hidden_states=x_padded), quant_info
|
||||
+ )
|
||||
+
|
||||
+ def _apply_sm120_cutlass(self, layer, dispatch_output):
|
||||
+ """SM120 MXFP8 x MXFP4 MoE via FlashInfer CUTLASS."""
|
||||
+ from sglang.srt.layers.moe.moe_runner.flashinfer_cutlass import (
|
||||
+ FlashInferCutlassMxfp4MoeQuantInfo,
|
||||
)
|
||||
|
||||
+ quant_info = FlashInferCutlassMxfp4MoeQuantInfo(
|
||||
+ w13_weight=layer.w13_weight,
|
||||
+ w2_weight=layer.w2_weight,
|
||||
+ 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 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,
|
||||
+ moe_tp_size=layer.moe_tp_size,
|
||||
+ moe_tp_rank=layer.moe_tp_rank,
|
||||
+ moe_ep_size=layer.moe_ep_size,
|
||||
+ moe_ep_rank=layer.moe_ep_rank,
|
||||
+ padded_hidden=self._padded_hidden,
|
||||
+ )
|
||||
+ return self.runner.run(dispatch_output, quant_info)
|
||||
+
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
@@ -1290,6 +1462,8 @@
|
||||
|
||||
if self._fi_kernel == "cutlass_sm90":
|
||||
return self._apply_sm90_cutlass(layer, dispatch_output)
|
||||
+ if self._fi_kernel == "cutlass_sm120":
|
||||
+ return self._apply_sm120_cutlass(layer, dispatch_output)
|
||||
if self.use_flashinfer:
|
||||
# When bf16 mode is enabled, we don't need to quantize the input,
|
||||
# TRT-LLM automatically handles quantization in the kernel implementation and pipelines it with GEMM operations,
|
||||
@ -0,0 +1,432 @@
|
||||
#!/usr/bin/env bash
|
||||
# Compare Marlin and FlashInfer MXFP4 on real four-node Kimi-K3 prefill.
|
||||
set -Eeuo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
# shellcheck source=/dev/null
|
||||
source "${SCRIPT_DIR}/../../../scripts/common/lib.sh"
|
||||
# shellcheck source=/dev/null
|
||||
source "${SCRIPT_DIR}/../../../scripts/common/platform.sh"
|
||||
# shellcheck source=/dev/null
|
||||
source "${SCRIPT_DIR}/config.env"
|
||||
|
||||
ACTION="${1:-all}"
|
||||
RUN_ID="${RUN_ID:-kimi3-moe-prefill-$(date '+%Y%m%d-%H%M%S')}"
|
||||
RESULT_ROOT="${RESULT_BASE}/${RUN_ID}"
|
||||
BUILD_CONTEXT="/tmp/${EXPERIMENT}_image_context"
|
||||
CONTAINER_PREFIX="${EXPERIMENT}"
|
||||
SSH_OPTS=(-o BatchMode=yes -o StrictHostKeyChecking=no -o ConnectTimeout=10)
|
||||
|
||||
mkdir -p "${RESULT_ROOT}"/{build,service,raw,bench,gpu}
|
||||
log_init "${RESULT_ROOT}/orchestrator.log"
|
||||
|
||||
usage() {
|
||||
cat <<'EOF'
|
||||
Usage: SUDO_PASSWORD=... bash run_moe_backend_prefill.sh ACTION
|
||||
|
||||
Actions:
|
||||
build-image Build the same candidate image on 601-604.
|
||||
run Run the fixed 8-case matrix with an existing image.
|
||||
all Build the image, run the matrix, and summarize it (default).
|
||||
summarize Rebuild CSV summaries from an existing RUN_ID.
|
||||
stop Remove this experiment's containers from all four nodes.
|
||||
EOF
|
||||
}
|
||||
|
||||
require_password() {
|
||||
if [[ -z "${SUDO_PASSWORD:-}" && -n "${SUDO_PASSWORD_FILE:-}" ]]; then
|
||||
[[ -r "$SUDO_PASSWORD_FILE" ]] || {
|
||||
echo "ERROR: cannot read SUDO_PASSWORD_FILE=$SUDO_PASSWORD_FILE" >&2
|
||||
exit 2
|
||||
}
|
||||
IFS= read -r SUDO_PASSWORD <"$SUDO_PASSWORD_FILE"
|
||||
fi
|
||||
if [[ -z "${SUDO_PASSWORD:-}" ]]; then
|
||||
echo "ERROR: set SUDO_PASSWORD or SUDO_PASSWORD_FILE at runtime" >&2
|
||||
exit 2
|
||||
fi
|
||||
}
|
||||
|
||||
is_head() {
|
||||
[[ "$1" == "$HEAD_HOST" ]]
|
||||
}
|
||||
|
||||
run_host() {
|
||||
local host="$1"
|
||||
shift
|
||||
if is_head "$host"; then
|
||||
"$@"
|
||||
else
|
||||
ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" "$@"
|
||||
fi
|
||||
}
|
||||
|
||||
sudo_host() {
|
||||
local host="$1"
|
||||
shift
|
||||
require_password
|
||||
if is_head "$host"; then
|
||||
printf '%s\n' "$SUDO_PASSWORD" | sudo -S -p '' -- "$@"
|
||||
return
|
||||
fi
|
||||
local remote_cmd
|
||||
printf -v remote_cmd '%q ' "$@"
|
||||
printf '%s\n' "$SUDO_PASSWORD" | \
|
||||
ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \
|
||||
"sudo -S -p '' -- ${remote_cmd}"
|
||||
}
|
||||
|
||||
check_inputs() {
|
||||
local path
|
||||
for path in "$MODEL_PATH" "$FLASHINFER_TARGET/flashinfer" \
|
||||
"$FLASHINFER_WHEEL" \
|
||||
"$SCRIPT_DIR/patches/sglang-kimi-image-sm120-flashinfer-mxfp4.patch"; do
|
||||
[[ -e "$path" ]] || { echo "ERROR: required path missing: $path" >&2; exit 2; }
|
||||
done
|
||||
[[ -d "$FLASHINFER_TARGET/flashinfer/data/aot/fused_moe_120" ]] || {
|
||||
echo "ERROR: fused_moe_120 AOT artifact is missing" >&2
|
||||
exit 2
|
||||
}
|
||||
}
|
||||
|
||||
prepare_build_context() {
|
||||
log "preparing candidate image context at ${BUILD_CONTEXT}"
|
||||
rm -rf "$BUILD_CONTEXT"
|
||||
mkdir -p "$BUILD_CONTEXT/aot"
|
||||
cp "$SCRIPT_DIR/Dockerfile" "$BUILD_CONTEXT/Dockerfile"
|
||||
cp "$SCRIPT_DIR/patches/sglang-kimi-image-sm120-flashinfer-mxfp4.patch" \
|
||||
"$BUILD_CONTEXT/sglang.patch"
|
||||
cp "$FLASHINFER_WHEEL" "$BUILD_CONTEXT/flashinfer_python-0.6.18-py3-none-any.whl"
|
||||
cp -a "$FLASHINFER_TARGET/flashinfer/data/aot/." "$BUILD_CONTEXT/aot/"
|
||||
find "$BUILD_CONTEXT" -type d -name __pycache__ -prune -exec rm -rf {} +
|
||||
find "$BUILD_CONTEXT" -type f -name '*.pyc' -delete
|
||||
du -sh "$BUILD_CONTEXT" | tee "$RESULT_ROOT/build/context_size.txt"
|
||||
}
|
||||
|
||||
stage_context() {
|
||||
local host="$1"
|
||||
if is_head "$host"; then
|
||||
return
|
||||
fi
|
||||
log "staging image context on ${host}"
|
||||
ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \
|
||||
"rm -rf '$BUILD_CONTEXT' && mkdir -p '$BUILD_CONTEXT'"
|
||||
tar -C "$BUILD_CONTEXT" -cf - . | \
|
||||
ssh "${SSH_OPTS[@]}" "${NODE_SSH_USER}@${host}" \
|
||||
"tar -C '$BUILD_CONTEXT' -xf -"
|
||||
}
|
||||
|
||||
build_one_node() {
|
||||
local host="$1"
|
||||
local log_path="$RESULT_ROOT/build/${host}.log"
|
||||
log "building ${CANDIDATE_IMAGE} on ${host}"
|
||||
sudo_host "$host" docker build \
|
||||
--build-arg "BASE_IMAGE=${BASE_IMAGE}" \
|
||||
--tag "$CANDIDATE_IMAGE" "$BUILD_CONTEXT" >"$log_path" 2>&1
|
||||
sudo_host "$host" docker image inspect "$CANDIDATE_IMAGE" \
|
||||
--format '{{.Id}} {{.Size}}' >"$RESULT_ROOT/build/${host}.image.txt"
|
||||
}
|
||||
|
||||
smoke_image() {
|
||||
local host="$1"
|
||||
sudo_host "$host" docker run --rm --gpus device=0 \
|
||||
--entrypoint python3 "$CANDIDATE_IMAGE" -c \
|
||||
"import torch; from flashinfer.fused_moe import core; import sglang; assert torch.cuda.get_device_capability()[0] == 12; assert core.CUTLASS_FUSED_MOE_SUPPORTS_SITU is True; print(sglang.__file__, core.ActivationType.Situ)" \
|
||||
>"$RESULT_ROOT/build/${host}.smoke.txt" 2>&1
|
||||
}
|
||||
|
||||
build_image_all_nodes() {
|
||||
require_password
|
||||
check_inputs
|
||||
prepare_build_context
|
||||
local host
|
||||
for host in "${NODE_HOSTS[@]}"; do stage_context "$host"; done
|
||||
local -a pids=()
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
build_one_node "$host" &
|
||||
pids+=("$!")
|
||||
done
|
||||
local rc=0 pid
|
||||
for pid in "${pids[@]}"; do wait "$pid" || rc=1; done
|
||||
(( rc == 0 )) || { log "ERROR: candidate image build failed"; return 1; }
|
||||
for host in "${NODE_HOSTS[@]}"; do smoke_image "$host"; done
|
||||
log "candidate image built and smoke-tested on all four nodes"
|
||||
}
|
||||
|
||||
container_name() {
|
||||
local rank="$1"
|
||||
printf '%s_node%s' "$CONTAINER_PREFIX" "$rank"
|
||||
}
|
||||
|
||||
collect_gpu_snapshot() {
|
||||
local label="$1" host rank=0
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
sudo_host "$host" nvidia-smi \
|
||||
--query-gpu=timestamp,index,memory.used,memory.total,utilization.gpu,power.draw \
|
||||
--format=csv,noheader,nounits \
|
||||
>"$RESULT_ROOT/gpu/${label}_node${rank}.csv" 2>&1 || true
|
||||
rank=$((rank + 1))
|
||||
done
|
||||
}
|
||||
|
||||
collect_service_logs() {
|
||||
local label="$1" host rank=0 name
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
name="$(container_name "$rank")"
|
||||
sudo_host "$host" docker logs "$name" \
|
||||
>"$RESULT_ROOT/service/${label}_node${rank}.log" 2>&1 || true
|
||||
rank=$((rank + 1))
|
||||
done
|
||||
}
|
||||
|
||||
stop_service() {
|
||||
local host rank=0 name
|
||||
for host in "${NODE_HOSTS[@]}"; do
|
||||
name="$(container_name "$rank")"
|
||||
sudo_host "$host" docker rm -f "$name" >/dev/null 2>&1 || true
|
||||
rank=$((rank + 1))
|
||||
done
|
||||
}
|
||||
|
||||
start_node() {
|
||||
local rank="$1" backend="$2" chunk="$3"
|
||||
local host="${NODE_HOSTS[$rank]}" name bootstrap
|
||||
name="$(container_name "$rank")"
|
||||
bootstrap="export SGLANG_HOST_IP=174.1.60.$((rank + 1)); exec python3 -m sglang.launch_server --model-path ${MODEL_PATH} --served-model-name ${SERVED_MODEL_NAME} --tp-size 32 --ep-size ${EP_SIZE} --nnodes 4 --node-rank ${rank} --dist-init-addr ${HEAD_HOST}:${DIST_PORT} --trust-remote-code --moe-runner-backend ${backend} --chunked-prefill-size ${chunk} --mem-fraction-static 0.88 --cuda-graph-max-bs-decode 16 --mamba-radix-cache-strategy extra_buffer_lazy --disable-radix-cache --dist-timeout 3600 --mamba-full-memory-ratio 0.36 --host 0.0.0.0 --port ${PORT}"
|
||||
local -a cmd=(
|
||||
docker run -d --name "$name"
|
||||
--gpus all --network host --ipc=host --ulimit memlock=-1
|
||||
--device /dev/infiniband --shm-size 32g --entrypoint bash
|
||||
-v "${MODEL_PATH}:${MODEL_PATH}:ro"
|
||||
-e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
|
||||
-e NCCL_SOCKET_IFNAME=bond0 -e GLOO_SOCKET_IFNAME=bond0
|
||||
-e NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3
|
||||
-e NCCL_IB_GID_INDEX=3 -e NCCL_IB_TIMEOUT=22 -e NCCL_IB_RETRY_CNT=7
|
||||
-e NCCL_CUMEM_ENABLE=1 -e PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
|
||||
-e SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK=0 -e SGLANG_MOE_FUSED_GATE_RADIX=1
|
||||
-e FLASHINFER_DISABLE_JIT=1 -e FLASHINFER_DISABLE_VERSION_CHECK=1
|
||||
"$CANDIDATE_IMAGE" -lc "$bootstrap"
|
||||
)
|
||||
printf '%q ' "${cmd[@]}" >"$RESULT_ROOT/service/${backend}_chunk${chunk}_node${rank}.cmd.txt"
|
||||
printf '\n' >>"$RESULT_ROOT/service/${backend}_chunk${chunk}_node${rank}.cmd.txt"
|
||||
sudo_host "$host" "${cmd[@]}" >/dev/null
|
||||
}
|
||||
|
||||
wait_health() {
|
||||
local backend="$1" chunk="$2" i
|
||||
for ((i = 1; i <= HEALTH_WAIT_S; i++)); do
|
||||
if curl --fail --silent --max-time 5 "http://${HEAD_HOST}:${PORT}/health" >/dev/null 2>&1; then
|
||||
log "service healthy backend=${backend} chunk=${chunk} wait_s=${i}"
|
||||
return 0
|
||||
fi
|
||||
if (( i % 30 == 0 )); then
|
||||
log "waiting for service backend=${backend} chunk=${chunk} elapsed_s=${i}"
|
||||
collect_service_logs "${backend}_chunk${chunk}_starting"
|
||||
if grep -Eiq 'Traceback|CUDA out of memory|NCCL.*(error|failed)|RuntimeError' \
|
||||
"$RESULT_ROOT/service/${backend}_chunk${chunk}_starting_node"*.log; then
|
||||
log "startup log contains an error; stopping early"
|
||||
return 1
|
||||
fi
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
return 1
|
||||
}
|
||||
|
||||
start_service() {
|
||||
local backend="$1" chunk="$2"
|
||||
stop_service
|
||||
log "starting TP32/EP${EP_SIZE} backend=${backend} chunk=${chunk} on 601-604"
|
||||
start_node 1 "$backend" "$chunk"
|
||||
start_node 2 "$backend" "$chunk"
|
||||
start_node 3 "$backend" "$chunk"
|
||||
sleep 5
|
||||
start_node 0 "$backend" "$chunk"
|
||||
if ! wait_health "$backend" "$chunk"; then
|
||||
collect_service_logs "${backend}_chunk${chunk}_startup_failed"
|
||||
stop_service
|
||||
return 1
|
||||
fi
|
||||
collect_gpu_snapshot "${backend}_chunk${chunk}_healthy"
|
||||
collect_service_logs "${backend}_chunk${chunk}_healthy"
|
||||
}
|
||||
|
||||
run_one_bench() {
|
||||
local backend="$1" chunk="$2" concurrency="$3" repeat="$4"
|
||||
local stem="${backend}_chunk${chunk}_c${concurrency}_r${repeat}"
|
||||
local output_file="${RESULT_ROOT}/raw/${stem}.jsonl"
|
||||
local bench_log="${RESULT_ROOT}/bench/${stem}.log"
|
||||
local -a dataset_args=(--dataset-name "$BENCH_DATASET_NAME")
|
||||
rm -f "$output_file"
|
||||
log "bench backend=${backend} chunk=${chunk} C=${concurrency} repeat=${repeat}/${REPEATS}"
|
||||
sudo_host "$HEAD_HOST" docker run --rm --network host \
|
||||
-v "${MODEL_PATH}:${MODEL_PATH}:ro" \
|
||||
-v "${RESULT_ROOT}:${RESULT_ROOT}" \
|
||||
-e PYTHONUNBUFFERED=1 --entrypoint python3 "$CANDIDATE_IMAGE" \
|
||||
-m "$BENCH_MODULE" \
|
||||
--backend sglang --host "$HEAD_HOST" --port "$PORT" \
|
||||
--tokenizer "$MODEL_PATH" \
|
||||
"${dataset_args[@]}" \
|
||||
--random-input-len "$INPUT_LEN" --random-output-len "$OUTPUT_LEN" \
|
||||
--random-range-ratio 1.0 --num-prompts "$NUM_PROMPTS" \
|
||||
--max-concurrency "$concurrency" --request-rate 10000 \
|
||||
--warmup-requests "$WARMUP_REQUESTS" --output-file "$output_file" \
|
||||
--output-details --disable-tqdm >"$bench_log" 2>&1
|
||||
python3 - "$output_file" "$NUM_PROMPTS" <<'PY'
|
||||
import json, sys
|
||||
path, expected = sys.argv[1], int(sys.argv[2])
|
||||
with open(path, encoding="utf-8") as f:
|
||||
rows = [json.loads(line) for line in f if line.strip()]
|
||||
assert len(rows) == 1, (path, len(rows))
|
||||
assert rows[0].get("completed") == expected, rows[0].get("completed")
|
||||
assert not any(rows[0].get("errors", [])), "benchmark contains request errors"
|
||||
PY
|
||||
}
|
||||
|
||||
run_matrix() {
|
||||
require_password
|
||||
check_inputs
|
||||
local backend chunk concurrency repeat label
|
||||
collect_gpu_snapshot before_all
|
||||
for chunk in "${CHUNKED_PREFILL_SIZES[@]}"; do
|
||||
for backend in "${BACKENDS[@]}"; do
|
||||
label="${backend}_chunk${chunk}"
|
||||
if ! start_service "$backend" "$chunk"; then
|
||||
log "ERROR: service failed for ${label}; preserving evidence and aborting"
|
||||
return 1
|
||||
fi
|
||||
for concurrency in "${CONCURRENCIES[@]}"; do
|
||||
for ((repeat = 1; repeat <= REPEATS; repeat++)); do
|
||||
run_one_bench "$backend" "$chunk" "$concurrency" "$repeat"
|
||||
done
|
||||
done
|
||||
collect_service_logs "${label}_completed"
|
||||
collect_gpu_snapshot "${label}_completed"
|
||||
stop_service
|
||||
sleep 5
|
||||
done
|
||||
done
|
||||
collect_gpu_snapshot after_all
|
||||
}
|
||||
|
||||
summarize() {
|
||||
python3 - "$RESULT_ROOT" <<'PY'
|
||||
import csv, json, re, statistics, sys
|
||||
from pathlib import Path
|
||||
|
||||
root = Path(sys.argv[1])
|
||||
pattern = re.compile(r"(.+)_chunk(\d+)_c(\d+)_r(\d+)\.jsonl$")
|
||||
fields = [
|
||||
"backend", "chunked_prefill_size", "concurrency", "repeat", "completed",
|
||||
"duration", "request_throughput", "input_throughput", "output_throughput",
|
||||
"total_throughput", "mean_ttft_ms", "median_ttft_ms", "p95_ttft_ms",
|
||||
"p99_ttft_ms", "mean_e2e_latency_ms", "median_e2e_latency_ms",
|
||||
"p95_e2e_latency_ms", "p99_e2e_latency_ms",
|
||||
]
|
||||
rows = []
|
||||
for path in sorted((root / "raw").glob("*.jsonl")):
|
||||
match = pattern.match(path.name)
|
||||
if not match:
|
||||
continue
|
||||
with path.open(encoding="utf-8") as f:
|
||||
data = next(json.loads(line) for line in f if line.strip())
|
||||
backend, chunk, concurrency, repeat = match.groups()
|
||||
row = {key: data.get(key) for key in fields}
|
||||
row.update(
|
||||
backend=backend,
|
||||
chunked_prefill_size=int(chunk),
|
||||
concurrency=int(concurrency),
|
||||
repeat=int(repeat),
|
||||
)
|
||||
rows.append(row)
|
||||
|
||||
with (root / "results.csv").open("w", newline="", encoding="utf-8") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=fields)
|
||||
writer.writeheader()
|
||||
writer.writerows(rows)
|
||||
|
||||
metric_names = [
|
||||
"request_throughput", "input_throughput", "total_throughput",
|
||||
"mean_ttft_ms", "median_ttft_ms", "p95_ttft_ms", "p99_ttft_ms",
|
||||
"mean_e2e_latency_ms", "median_e2e_latency_ms", "p95_e2e_latency_ms",
|
||||
]
|
||||
groups = {}
|
||||
for row in rows:
|
||||
key = (row["backend"], row["chunked_prefill_size"], row["concurrency"])
|
||||
groups.setdefault(key, []).append(row)
|
||||
|
||||
summary = []
|
||||
for (backend, chunk, concurrency), group in sorted(groups.items()):
|
||||
item = {
|
||||
"backend": backend,
|
||||
"chunked_prefill_size": chunk,
|
||||
"concurrency": concurrency,
|
||||
"repeats": len(group),
|
||||
"completed_each": [row["completed"] for row in group],
|
||||
}
|
||||
for metric in metric_names:
|
||||
values = [float(row[metric]) for row in group if row.get(metric) is not None]
|
||||
item[f"median_{metric}"] = statistics.median(values) if values else None
|
||||
summary.append(item)
|
||||
|
||||
comparisons = []
|
||||
index = {(x["backend"], x["chunked_prefill_size"], x["concurrency"]): x for x in summary}
|
||||
for chunk in sorted({x["chunked_prefill_size"] for x in summary}):
|
||||
for concurrency in sorted({x["concurrency"] for x in summary}):
|
||||
base = index.get(("marlin", chunk, concurrency))
|
||||
cand = index.get(("flashinfer_mxfp4", chunk, concurrency))
|
||||
if not base or not cand:
|
||||
continue
|
||||
comp = {"chunked_prefill_size": chunk, "concurrency": concurrency}
|
||||
for metric in ("input_throughput", "total_throughput"):
|
||||
b = base[f"median_{metric}"]
|
||||
c = cand[f"median_{metric}"]
|
||||
comp[f"flashinfer_{metric}_change_pct"] = (c / b - 1.0) * 100.0
|
||||
for metric in ("median_ttft_ms", "p95_ttft_ms", "median_e2e_latency_ms"):
|
||||
b = base[f"median_{metric}"]
|
||||
c = cand[f"median_{metric}"]
|
||||
comp[f"flashinfer_{metric}_change_pct"] = (c / b - 1.0) * 100.0
|
||||
comparisons.append(comp)
|
||||
|
||||
payload = {"run_id": root.name, "summary": summary, "comparisons": comparisons}
|
||||
(root / "summary.json").write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||||
print(json.dumps(payload, indent=2))
|
||||
PY
|
||||
}
|
||||
|
||||
cleanup() {
|
||||
collect_service_logs cleanup 2>/dev/null || true
|
||||
stop_service 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT INT TERM
|
||||
|
||||
case "$ACTION" in
|
||||
build-image)
|
||||
build_image_all_nodes
|
||||
;;
|
||||
run)
|
||||
run_matrix
|
||||
summarize
|
||||
;;
|
||||
all)
|
||||
build_image_all_nodes
|
||||
run_matrix
|
||||
summarize
|
||||
;;
|
||||
summarize)
|
||||
summarize
|
||||
;;
|
||||
stop)
|
||||
stop_service
|
||||
;;
|
||||
-h|--help|help)
|
||||
usage
|
||||
;;
|
||||
*)
|
||||
usage >&2
|
||||
exit 2
|
||||
;;
|
||||
esac
|
||||
|
||||
trap - EXIT INT TERM
|
||||
log "action=${ACTION} complete run_id=${RUN_ID} result_root=${RESULT_ROOT}"
|
||||
@ -0,0 +1,273 @@
|
||||
# Kimi K3 推理优化 - 优化 Prefill - TP Reduce Scatter
|
||||
|
||||
## 结论
|
||||
|
||||
**不推荐在当前 Kimi-K3 TP32 部署上继续这条优化路线。**
|
||||
|
||||
SGLang 的通用 `--enable-attn-tp-input-scattered` 优化适用于后续 Attention
|
||||
只依赖 QKV latent 的 MLA 模型。Kimi-K3 的 MLA 额外包含
|
||||
`g_proj(hidden_states)` 输出门控，而且 93 层中只有 24 层为 MLA、其余 69
|
||||
层为 KDA。为了保证门控和 KDA 正确，仍需恢复完整的 7168 维 hidden。
|
||||
这样不仅无法消除原有 hidden All-Reduce，还会增加一次 2112 维 latent
|
||||
All-Gather，与“降低跨节点通信”的目标相反。
|
||||
|
||||
因此本任务在源码审计和最小原型验证后停止，不进入四节点性能测试，不形成
|
||||
SGLang PR。后续 Prefill 通信优化应优先转向 MoE A2A 和 Pipeline
|
||||
Parallelism。
|
||||
|
||||
## 目标与固定测试口径
|
||||
|
||||
原计划比较以下两组：
|
||||
|
||||
```text
|
||||
模型：Kimi-K3
|
||||
节点：601-604，32 GPU
|
||||
并行：TP32 / EP4 / DP1
|
||||
MoE runner：flashinfer_mxfp4
|
||||
请求：16K -> 1
|
||||
并发：C=8、C=16
|
||||
Chunked Prefill：8K
|
||||
变量：--enable-attn-tp-input-scattered 关闭 / 开启
|
||||
```
|
||||
|
||||
源码审计已证明当前开关在 Kimi-K3 中没有完整接线，而且直接移植不能降低
|
||||
通信，因此没有消耗四节点资源运行该矩阵。
|
||||
|
||||
## DeepSeek 路径为什么成立
|
||||
|
||||
通用优化把原路径：
|
||||
|
||||
```text
|
||||
embed / MLP partial output
|
||||
-> All-Reduce hidden
|
||||
-> RMSNorm(full hidden)
|
||||
-> qkv_a(full hidden)
|
||||
```
|
||||
|
||||
改为：
|
||||
|
||||
```text
|
||||
embed / MLP partial output
|
||||
-> Reduce-Scatter hidden by token
|
||||
-> RMSNorm(local tokens)
|
||||
-> qkv_a(local tokens)
|
||||
-> All-Gather QKV latent
|
||||
```
|
||||
|
||||
Kimi-K3 的张量宽度为：
|
||||
|
||||
```text
|
||||
hidden = 7168
|
||||
QKV latent = q_lora_rank + kv_lora_rank + qk_rope_head_dim
|
||||
= 1536 + 512 + 64
|
||||
= 2112
|
||||
```
|
||||
|
||||
如果后续只需要 latent，将 7168 维 hidden 的 All-Reduce 改成 7168 维
|
||||
Reduce-Scatter，再 All-Gather 2112 维 latent，可以同时减少 RMSNorm/QKV
|
||||
重复计算和通信量。
|
||||
|
||||
## Kimi-K3 为什么不能直接复用
|
||||
|
||||
### MLA 输出门控仍依赖完整 hidden
|
||||
|
||||
Kimi-K3 MLA 有一条独立门控支路：
|
||||
|
||||
```text
|
||||
full hidden
|
||||
|-- qkv_a -> QKV latent -> Attention --|
|
||||
`-- g_proj -> sigmoid(gate) ------------|-> elementwise gate
|
||||
```
|
||||
|
||||
相关实现位于：
|
||||
|
||||
```text
|
||||
python/sglang/srt/models/kimi_k3.py
|
||||
KimiK3MLAAttention.__init__
|
||||
KimiK3MLAAttention.forward
|
||||
_gated_o_proj_forward
|
||||
```
|
||||
|
||||
`g_proj` 是按 Attention Head 切分的 `ColumnParallelLinear`。TP32 时，每个
|
||||
rank 只持有自己的 3 个 head 的权重，但必须为全部 token 计算门控：
|
||||
|
||||
```text
|
||||
hidden_all: [T, 7168]
|
||||
W_gate_rank: [384, 7168]
|
||||
gate_rank: [T, 384]
|
||||
```
|
||||
|
||||
Reduce-Scatter 后每个 rank 只有 `[T/32, 7168]`，只能得到
|
||||
`[T/32, 384]`。把各 rank 的 gate 直接 All-Gather 也不正确，因为每段
|
||||
token 使用了不同 rank、不同 head 的权重；本 rank 需要的是自己的
|
||||
`W_gate_rank` 对全部 token 的结果。
|
||||
|
||||
QKV latent 的 All-Gather 不能解决这个问题，因为 `g_proj` 的输入是原始
|
||||
hidden，不是 QKV latent。
|
||||
|
||||
### KDA 层要求完整 hidden
|
||||
|
||||
Kimi-K3 共 93 层，其中 24 层为 MLA、69 层为 KDA。KDA 的宽 Q/K/V/G
|
||||
投影直接消费完整 hidden，不能复用 MLA 的 2112 维 latent。即使只在 MLA
|
||||
层建立 scattered window，也需要在频繁的 MLA/KDA 边界恢复完整 hidden，
|
||||
进一步削弱收益并增加实现复杂度。
|
||||
|
||||
## 通信量判断
|
||||
|
||||
忽略常数、数据类型和不同 collective 算法，令：
|
||||
|
||||
```text
|
||||
H = 7168 # hidden width
|
||||
L = 2112 # QKV latent width
|
||||
```
|
||||
|
||||
Ring collective 的每 rank 主量级可写为：
|
||||
|
||||
| 路径 | 相对通信项 |
|
||||
|---|---:|
|
||||
| 原始 hidden All-Reduce | `2H = 14336` |
|
||||
| 理想 DeepSeek 路径：hidden RS + latent AG | `H + L = 9280` |
|
||||
| K3 若恢复 gate hidden：hidden RS + hidden AG + latent AG | `2H + L = 16448` |
|
||||
| K3 最小原型：保留 hidden AR + latent AG | `2H + L = 16448` |
|
||||
|
||||
理想路径比原始路径少约 35.3% 通信；K3 为门控保留或恢复完整 hidden 后，
|
||||
反而比原始路径多约 14.7%。实际延迟还会叠加跨四节点 collective 的固定
|
||||
开销，因此没有理由期待它改善当前以通信为目标的 Prefill。
|
||||
|
||||
## Prefill 通信归因
|
||||
|
||||
复用既有 `16K -> 256、C=8、Chunk=8K、TP32、无 MoE A2A` 的
|
||||
PyTorch Trace，并用 `record_shapes + with_stack` 原始事件重新统计。虽然该
|
||||
Trace 使用 EP32/Marlin，而当前验收配置为 EP4/FlashInfer，但两者均为
|
||||
`TP32 + moe_a2a_backend=none`；MoE runner 和 EP 值会改变专家计算，不改变
|
||||
下列 TP32 collective 的调用位置与张量形状。
|
||||
|
||||
原始 Trace：
|
||||
|
||||
```text
|
||||
/data/yy/sskj/experiments/pro6000/kimi3_pro6000_decode_profile/
|
||||
trace/1786591959.655009/
|
||||
decode_c8_i16384_o256-1786591959.6615255-TP-0-EP-0-EXTEND.trace.json.gz
|
||||
```
|
||||
|
||||
Trace 捕获了 9 个 Prefill engine step。每个 step 的主 AllReduce 构成为：
|
||||
|
||||
| 通信来源 | 每 step 次数 | `M=8192` 单次输入 | 每 step 输入量 | 占比 |
|
||||
|---|---:|---:|---:|---:|
|
||||
| Embedding + 93 层 Attention `o_proj` + 首层 dense MLP | 95 | `[8192,7168]`，112 MiB | 10.39 GiB | 40.8% |
|
||||
| 92 个 MoE 层的 routed latent + shared expert | 92 | `8192x(3584+7168)`，168 MiB | 15.09 GiB | 59.2% |
|
||||
| 合计 | 187 | - | 25.48 GiB | 100% |
|
||||
|
||||
这也解释了旧资料里的约 187 次 collective：它不是“全部来自 MoE”，而是
|
||||
95 次完整 hidden 归约和 92 次 MoE 尾部归约之和。Kimi-K3 配置有 93 层，
|
||||
`first_k_dense_replace=1`，因此只有后 92 层走 MoE。
|
||||
|
||||
当前 SM120 四机部署不会启用 `k3_ar_fusion`：该路径只在 SM100/SM103 且
|
||||
CustomAllReduceV2 multicast 可用时自动开启；跨节点日志也明确显示
|
||||
`CustomAllreduce is disabled because this process group spans across nodes`。
|
||||
所以以上标准 NCCL 调用仍是当前 EP4/FlashInfer 配置的真实结构。
|
||||
|
||||
### 当前 EP4 + FlashInfer 真机时间线
|
||||
|
||||
为消除旧 Trace 使用 EP32/Marlin 的不确定性，又对当前正式配置做了三步
|
||||
Prefill Nsight 捕获：
|
||||
|
||||
```text
|
||||
TP32 / EP4 / DP1
|
||||
MoE runner: flashinfer_mxfp4
|
||||
MoE A2A: none
|
||||
16K -> 1, C=8, Chunk=8K
|
||||
```
|
||||
|
||||
原始结果位于 601：
|
||||
|
||||
```text
|
||||
/data/hzy/sskj/experiments/pro6000/
|
||||
kimi3_pro6000_sglang_prefill_communication_profile/results/
|
||||
kimi3-prefill-comm-20260820-143749/
|
||||
```
|
||||
|
||||
四个节点均生成 `.nsys-rep` 和 `.sqlite`，每份报告覆盖本节点 8 张 GPU。
|
||||
`nsys_analysis.json` 和 `nsys_device_metrics.csv` 是离线分析结果。为避免
|
||||
`cudaProfilerStart` 边界开销污染结论，统计排除了出现单次 AllReduce
|
||||
大于 50 ms 的 5 条触发 lane，保留 27/32 条稳定 lane 取中位数。
|
||||
|
||||
| 指标 | 稳定 rank 中位数 | 解释 |
|
||||
|---|---:|---|
|
||||
| 捕获窗口 | 7274.23 ms | 三个 8K Prefill step |
|
||||
| GPU busy | 99.37% | 几乎没有 GPU 空洞 |
|
||||
| NCCL 时间 | 4161.36 ms，57.27% | 当前 Prefill 的第一瓶颈 |
|
||||
| 非 NCCL 计算 | 3062.48 ms，42.10% | Attention、MoE GEMM、KDA 等 |
|
||||
| NCCL 与计算重叠 | 0 ms，0% | Collective 完整暴露在关键路径上 |
|
||||
| AllReduce 次数 | 557 | 捕获边界略少于理论 561 次 |
|
||||
| AllReduce P50 / P95 | 7.10 / 9.14 ms | 全部选择 `RING_LL` |
|
||||
|
||||
AllReduce 时长呈现无重叠的双峰：约 `6.1 ms` 和 `8.7 ms`，与源码 Trace
|
||||
中的 112 MiB hidden 消息和 168 MiB MoE 消息相符。按两个峰的累计时间：
|
||||
|
||||
| 通信类别 | 三步累计时间 | 占 AllReduce 时间 | 占完整窗口 |
|
||||
|---|---:|---:|---:|
|
||||
| hidden / Attention 路径 | 1719.78 ms | 41.27% | 23.64% |
|
||||
| MoE 尾部归约 | 2440.52 ms | 58.73% | 33.55% |
|
||||
|
||||
稳定 rank 的 AllReduce 平均值约为 `7.34–7.56 ms`，未发现固定慢节点。
|
||||
这说明问题不是某张卡掉速，而是 TP32 每层同步通信本身占据关键路径。
|
||||
|
||||
带 Nsight 的完整 benchmark 为 40/40 请求成功，但 Input TPS 为 2573，低于
|
||||
无 Profiler 基线 3257.96；该值只用于证明 workload 完整成功，不能作为性能
|
||||
回归结果。性能比较仍使用无 Profiler 的正式基线。
|
||||
|
||||
### 优化优先级
|
||||
|
||||
1. **MoE A2A / SP-MoE 优先验证。** 它瞄准占输入字节 59.2% 的 MoE 尾部
|
||||
collective，并可避免每个 TP rank 对同一批 token 的重复路由与专家计算。
|
||||
但 A2A 会引入 dispatch/combine，必须以四机实测判断净收益。
|
||||
2. **Pipeline Parallelism 第二。** 用 TP16/PP2 或 TP8/PP4 缩小 TP
|
||||
collective 通信域，同时只增加少数 stage-boundary P2P；需验证流水线空泡
|
||||
对 TTFT 的影响。
|
||||
3. **通信量化作为独立高风险项。** 它能同时压缩两类大消息，但需要准确性与
|
||||
backend 支持验证。
|
||||
4. **不优先做 NCCL 算法或 launch fusion。** Prefill 单消息为 112/168 MiB，
|
||||
首要矛盾是字节量和 32-rank 跨节点通信域，不是小消息启动延迟。
|
||||
|
||||
因此，通用 TP Reduce Scatter 即使成功，也只触及字节口径约 40.8%、时间
|
||||
口径约 41.3% 的 hidden 通信；它无法处理占完整 Prefill 时间约 33.6% 的
|
||||
MoE 尾部归约，而且 K3 gate/KDA 又迫使完整 hidden 存在。它不如先处理
|
||||
MoE 通信，再用 PP 缩小剩余 TP collective 的通信域。
|
||||
|
||||
## FlashInfer MoE 的关系
|
||||
|
||||
当前 `flashinfer_mxfp4` 是 MoE runner，且实验配置保持
|
||||
`moe_a2a_backend=none`，因此它不违反通用 input-scattered 的静态约束。
|
||||
但它也无法解决 MLA gate 或 KDA 对完整 hidden 的依赖。
|
||||
|
||||
换言之：**FlashInfer MoE 与该优化不冲突，但不会使该优化在 Kimi-K3 上
|
||||
成立。**
|
||||
|
||||
## 原型与验证记录
|
||||
|
||||
为验证最小可行边界，建立过仅分摊 MLA qkv latent GEMM 的研究分支：
|
||||
|
||||
```text
|
||||
601 worktree: /data/hzy/src/sglang-kimi-tp-rs-wt
|
||||
branch: hzy/kimi-k3-attn-tp-latent-scattered
|
||||
commit: f5577d09
|
||||
```
|
||||
|
||||
原型保留完整 hidden 给 K3 gate，只把 replicated 的
|
||||
`7168 -> 2112` qkv latent GEMM 按 token 分给 TP ranks，再 All-Gather
|
||||
latent。CPU 边界单测为 `3 passed`，证明 token 分片与 fallback 逻辑正确。
|
||||
|
||||
该原型没有进入 GPU/四节点服务测试，因为通信公式已经表明它保留原
|
||||
All-Reduce 并新增 latent All-Gather。分支仅作为否决证据保留，不应合并、
|
||||
发布或提交上游。
|
||||
|
||||
## 后续建议
|
||||
|
||||
1. 保留已验收的 `flashinfer_mxfp4 + EP4 + chunked_prefill_size=8192`。
|
||||
2. 下一步测试 MoE A2A，目标是减少或替换 MoE 的跨 rank collective，而非
|
||||
在 MLA gate 前后增加 collective。
|
||||
3. 再测试 Prefill Pipeline Parallelism，用更小 TP 通信域换取流水线开销，
|
||||
继续使用固定口径 `16K -> 1、C=8/16、Chunk=8K`。
|
||||
4. 只有出现 gate-aware 的实现（例如低成本复制/重排 gate 权重，且通信模型
|
||||
明确优于原始 All-Reduce）时，才重新开启 TP Reduce Scatter 方向。
|
||||
Loading…
x
Reference in New Issue
Block a user