[Docs] Reject Kimi-K3 Prefill TP Reduce Scatter path
This commit is contained in:
parent
08a35066d7
commit
2a3b12fa78
@ -1,5 +1,7 @@
|
|||||||
# sskj — 多平台大模型推理性能基准测试项目
|
# 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 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 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 分支保持不变。
|
||||||
@ -70,6 +72,7 @@
|
|||||||
| `experiments/pro6000/dsv4_pro6000_sglang_tp_dp_matrix/` | RTX 6000D + SGLang |
|
| `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/` | 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_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) |
|
| `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`)。
|
每个目录内:`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,173 @@
|
|||||||
|
# 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。
|
||||||
|
|
||||||
|
## 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