sskj/experiments/pro6000/kimi3_pro6000_sglang_tp_reduce_scatter_prefill

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。

目标与固定测试口径

原计划比较以下两组：

模型：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 路径为什么成立

通用优化把原路径：

embed / MLP partial output
  -> All-Reduce hidden
  -> RMSNorm(full hidden)
  -> qkv_a(full hidden)

改为：

embed / MLP partial output
  -> Reduce-Scatter hidden by token
  -> RMSNorm(local tokens)
  -> qkv_a(local tokens)
  -> All-Gather QKV latent

Kimi-K3 的张量宽度为：

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 有一条独立门控支路：

full hidden
  |-- qkv_a -> QKV latent -> Attention --|
  `-- g_proj -> sigmoid(gate) ------------|-> elementwise gate

相关实现位于：

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 计算门控：

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 算法，令：

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：

/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 捕获：

TP32 / EP4 / DP1
MoE runner: flashinfer_mxfp4
MoE A2A: none
16K -> 1, C=8, Chunk=8K

原始结果位于 601：

/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 的研究分支：

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 方向。