[abmfy]

Updates kimik3.1-fp4-b300-vllm-agentic-dspark (one B300 node, aggregated) to a per-concurrency TP8 / TP8 x DCP8 / TEP recipe on inferactinc/public:kimi-k31-mc-f819987-tep-shared-tp-x86_64: vllm/vllm-openai:kimi-k31-mc plus one opt-in vLLM change that is not in upstream vLLM yet.

Draft until the blockers below are settled (upstream image and engine-patch rule, vLLM recipe coverage).

Image

inferactinc/public:kimi-k31-mc-f819987-tep-shared-tp-x86_64@sha256:acf503a6058e75a63f0902a7ebb8200f6757d5d5ed92ca50edb67cbbb691ace9 (public), built by Inferact's image pipeline from:

• vllm/vllm-openai:kimi-k31-mc@sha256:f81998764e8594d71b40788c5285ae112acf9230129d5d0d238b5df83da5acf6 unchanged (vLLM 0.1.dev24218+g78e9a0c28, the day-0-mc image [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871 uses). Its wide-MoE grid is sized by SM count, so latency_specialized decode graphs wider than 64 rows run on the 148-SM B300.
• Plus VLLM_KIMI_K3_TEP_SHARED_TP (default off): vllm/envs.py and the Kimi K3 TEP hybrid decode runtime are copied over the package behind sha256 guards on the base files (diff below). With the variable set, a TEP decode step (TP8 x DCP8 attention + EP8 MegaMoE) runs each rank's TP-8 shard of the BF16 shared expert and folds the routed rows into one all-reduce, instead of every rank streaming the whole replicated shared expert (24 GB per step and rank). Prefill is unchanged. Without the variable the code is that of kimi-k31-mc.

Recipe

Conc	Topology	DSpark draft tokens	max-num-seqs	Point settings
1	TP8	7	2	TOKENSPEED_MLA draft attention
4, 8	TP8 x DCP8	7	8, 16	ReplaySSM; a2a DCP backend with direct A2A, Q and KV gathers
16, 24	TP8 x DCP8	3	96	as c4-c8
32-70	TP8 x DCP8 attention, EP8 MegaMoE (deep_gemm_mega_moe, flashinfer_nvlink_one_sided)	4	96, 96, 96, 112, 140	as c4-c8, plus VLLM_KIMI_K3_TEP_SHARED_TP=1 and gpu-memory-utilization 0.961

Every point runs latency_specialized decode with explicit CUDA graph capture sizes up to 4096 tokens, so prefill and mixed steps of up to 4096 tokens replay piecewise graphs instead of running eager; larger steps (max-num-batched-tokens is 16384) still run eager and are GPU-bound. gpu-memory-utilization is 0.95 (0.961 at c32-c70) and the Simple CPU KV offload stays at 281.125 GB per rank. The speculative config pins the DSpark draft at 6a2ef20a, the #4851 golden-AL revision. The per-point max-num-seqs and capture-size lists are the ones measured; c2 was not measured and is left out.

Results

tok/s/GPU @ P90 interactivity (tok/s/user) on our own HGX B300 nodes (8 GPUs, not the b300-dsxe cluster): 3600 s per point, 10 warmup requests per lane, the #4851 golden acceptance lengths, Kimi parsers off as in the synthetic runs, InferenceX scoring. c32-c70 ran on this PR's image; c1-c24 ran on kimi-k31-mc itself, which has the same code with the variable off. SGLang, #4833 and #4871 are their published or PR-sweep values.

Conc	This recipe	SGLang #4803	vLLM #4833 (current)	vLLM #4871 sweep
1	2,713 @ 300.0	2,370 @ 241.0	2,558 @ 284.9	2,756 @ 280.1
4	4,036 @ 223.9	3,425 @ 189.0	4,064 @ 213.7	4,082 @ 231.4
8	6,986 @ 140.1	5,651 @ 119.9	5,538 @ 94.6	6,992 @ 139.0
16	12,275 @ 92.4	6,647 @ 93.2	9,515 @ 58.8	11,753 @ 83.4
24	17,803 @ 66.1	15,286 @ 50.9	14,204 @ 45.3
32	20,219 @ 51.2	18,953 @ 43.3	17,550 @ 38.9	18,605 @ 44.3
40	22,560 @ 48.6	19,627 @ 39.3	19,296 @ 36.8
48	25,238 @ 42.0	22,611 @ 35.0	21,873 @ 31.4	23,499 @ 36.9
56	26,355 @ 35.8	22,877 @ 34.4	21,480 @ 31.5	24,679 @ 33.5
70	27,823 @ 32.1	21,333 @ 32.8	19,464 @ 28.6

Every SGLang #4803 point is dominated by a point of this curve (its c16 by our c8, its c70 by our c56). Against #4871's sweep, this curve is ahead at c16-c56 by 4-9% tok/s/GPU at 7-16% higher P90 interactivity and level with it at c8 (-0.1% throughput, +0.8% P90); #4871 is slightly ahead at c4 (+1% throughput, +3% P90), and c1 is a trade-off (+1.6% throughput for #4871, +7% P90 for this recipe). The two sweeps ran on different B300 clusters.

VLLM_KIMI_K3_TEP_SHARED_TP, same recipe: variable off on kimi-k31-mc itself (gpu-memory-utilization 0.95) vs on with this PR's image (0.961). Only the c48 pair ran on the same node:

Conc	Off	On	Change (tput / P90 / P50)
32	19,184 @ 48.4	20,219 @ 51.2	+5.4% / +5.8% / +9.5%
40	21,653 @ 46.3	22,560 @ 48.6	+4.2% / +5.0% / +7.5%
48	24,553 @ 40.7	25,238 @ 42.0	+2.8% / +3.3% / +6.9%
56	26,000 @ 35.9	26,355 @ 35.8	+1.4% / -0.3% / +0.8%
70	26,412 @ 32.1	27,823 @ 32.1	+5.3% / 0.0% / +1.0%

Single-request decode goes from 158.8 to 190.8 tok/s/user (8K-token prompt). The "on" runs also hold about 5% more KV tokens (for example 8.24M vs 7.82M at c32, 8.85M vs 8.45M at c70), so part of the gain at c56-c70, where P90 interactivity does not move and TTFT drops, may come from KV capacity rather than the faster decode step. Runs of equivalent code on different nodes differed by up to 2.5% throughput, so the c56 change is within noise.

Run notes

• Nodes for the main table: c1, c4 and c8 on one node; c16, c24 and c56 on a second; c32, c40 and c70 on a third; c48 on a fourth. The variable-off runs: c32-c56 on the fourth node, c70 on the second.
• A first c48 run with the variable on, on the node that had just produced c1, c4 and c8, had about 3.9x lower interactivity than the rerun for the whole hour (P50 15.1 vs 58.3 tok/s/user), and the node's GPUs stopped responding afterwards (nvidia-smi hung). It was discarded; the c48 value above is the rerun on another node. The c1, c4 and c8 runs from that node had normal warm-probe decode speed.
• The preliminary curve in the first version of this description ran on vllm/vllm-openai:kimi-k31 with the changes mounted as Python files, an earlier version of the shared-expert change at c32, c40 and c70, and none at c48 and c56. It is within about 2% of the numbers above at c1-c40 except c8 (2.5% higher throughput and 4.9% higher P90 there), and 1-4% lower in throughput at c48-c70 (at c48 also 4.9% lower P90).
• Our runs loaded an earlier local snapshot of the draft (5e654fbe, same config.json as 6a2ef20a); with synthetic acceptance the draft weights do not affect timing. The harness also enabled iteration logging.

Differences from #4871

Both PRs move this config to kimi-k31-mc and use latency_specialized decode at every point, so they conflict on the same files; this is not a duplicate of it. None of the differences below was ablated against #4871 on B300 except the shared-expert variable, so the gap is not attributed to individual settings:

• Graph coverage: [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871 captures decode graphs only (FULL_DECODE_ONLY), so every prefill and mixed step runs eager; this recipe captures steps up to 4096 tokens. On GB300 at c4 (5 draft tokens, 1200 s runs, graphs up to 16384 tokens vs decode sizes only), graphing those steps gave +14% P50 and +18% P90 interactivity at equal throughput and cut the decode stall a 1K-token turn causes from 238 ms to 86 ms.
• TEP shared expert: VLLM_KIMI_K3_TEP_SHARED_TP at c32-c70 (table above). With the variable off, c32, c48 and c56 are still 3-5% ahead of [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871 in throughput.
• Topology and depth: c1 runs TP8 without DCP and 7 draft tokens ([AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871: TP8 x DCP8 with 4); c4-c8 draft 7 tokens ([AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871: 4); c16 matches [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871's topology and depth (TP8 x DCP8, TP MoE, 3 draft tokens). TokenSpeed draft attention only at c1 here ([AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871 uses it at every point). Three more points (c24, c40, c70).
• Other settings: max-num-seqs 96 at c16 and c32 ([AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871: 20 and 64); gpu-memory-utilization 0.95 at c1-c24 ([AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871: 0.9); max-num-batched-tokens 16384 and 2249 GB of host DRAM at the TEP points ([AgentX] B300 Kimi-K3.1-Pro FP4 vLLM aggregated configs / [AgentX] B300 Kimi-K3.1-Pro FP4 vLLM 聚合式配置 #4871: 8192 and 2999 GB).

Blockers

• Engine change and image. docs/PR_REVIEW_CHECKLIST.md requires vLLM submissions on the upstream vllm Docker Hub repositories and an engine that runs as shipped; copying patched files over kimi-k31-mc counts as an engine patch unless docs/waiver/4878.md is filed (what is patched, why, upstream PR, removal plan). VLLM_KIMI_K3_TEP_SHARED_TP has no upstream vLLM PR yet. Until it lands, or a waiver is accepted and the image is re-tagged under vllm/vllm-openai, c32-c70 can run with the variable off on kimi-k31-mc itself (the "Off" column above), which needs no patch.
• vLLM recipe. The checklist also asks for every production-useful flag and env var to be documented in the official vLLM recipes. The published Kimi-K3.1-Pro recipe has no B300 AgentX preset and does not cover latency_specialized with EP8 MegaMoE, --use-replayssm, --dcp-comm-backend a2a with the direct gathers, TOKENSPEED_MLA draft attention, explicit capture sizes up to 4096, or VLLM_KIMI_K3_TEP_SHARED_TP.

Accuracy

VLLM_KIMI_K3_TEP_SHARED_TP sums BF16 shared-expert partials across the TP group, the precision of the existing VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT, not bitwise equal to the replicated path. GSM8K on GB300 with the InferenceX lm-eval recipe (TEP c70, real acceptance, both Kimi parsers): strict-match 0.9818 ± 0.0037 with it vs 0.9795 ± 0.0039 without (official #4833 B300 c70: 0.9773 ± 0.0041). Both arms ran on the kimi-k31 image with our development overlay, which also forced the stock DSA indexer and carried a Kimi tool-parser fix, and with an earlier version of the change that differs only in GEMM row padding and memory layout. GSM8K has not run on this PR's image yet; this PR's sweep evaluates it at c70.

The vLLM change (diff against the image's files)

diff --git c/vllm/envs.py i/vllm/envs.py
index 064565cf1e..7f01d49cda 100755
--- c/vllm/envs.py
+++ i/vllm/envs.py
@@ -222,6 +222,7 @@ if TYPE_CHECKING:
     VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False
     VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False
     VLLM_KIMI_K3_GEMM_AR: bool = True
+    VLLM_KIMI_K3_TEP_SHARED_TP: bool = False
     VLLM_ENABLE_GEMM_RS: bool = False
     VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True
     VLLM_USE_FLASHINFER_MOE_INT4: bool = False
@@ -1723,6 +1724,17 @@ environment_variables: dict[str, Callable[[], Any]] = {
     # Use the SM100 BF16 GEMM-AR kernel for eligible Kimi-K3 row-parallel
     # attention projections. All TP ranks must belong to one NVLink domain.
     "VLLM_KIMI_K3_GEMM_AR": lambda: bool(int(os.getenv("VLLM_KIMI_K3_GEMM_AR", "1"))),
+    # Kimi-K3 latency_specialized decode under TEP (TP-G attention groups with
+    # MegaMoE expert parallelism over every rank). Its decode steps run each MoE
+    # layer on the rank's 1/G row shard with the BF16 shared expert replicated, so
+    # every rank streams the whole shared expert. Shard it across the TP group
+    # instead: each rank computes its shard's partial for all rows, adds its routed
+    # rows, and one all-reduce sums both. Costs a copy of the gate/up shard (about
+    # 1.9 GiB per rank on Kimi-K3.1 at TP8) and changes the shared expert's
+    # summation order.
+    "VLLM_KIMI_K3_TEP_SHARED_TP": lambda: bool(
+        int(os.getenv("VLLM_KIMI_K3_TEP_SHARED_TP", "0"))
+    ),
     # Fuse eligible sequence-parallel row-parallel projections with their TP
     # reduce-scatter using the SM100 BF16/MXFP8 GEMM-RS kernel (Kimi-K3
     # attention/shared-expert projections, DeepSeek-V4.1 ``wo_b``). All TP
diff --git c/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/hybrid.py i/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/hybrid.py
index 8ed4bf33e8..6c3283f0a0 100644
--- c/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/hybrid.py
+++ i/vllm/models/kimi_k3/nvidia/specialized_decode/runtime/hybrid.py
@@ -12,7 +12,7 @@ the o-projection, and the MoE runs on the rank's 1/G row shard with the shared
 expert replicated.
 
 The hybrid step keeps every row on every TP rank instead (the TP8 decode's
-layout) for attention, and runs the MLPs exactly as stock does:
+layout) for attention, and runs the MLPs as stock does (or TP-sharded, below):
 
   AttnRes            stock, on the group's rows (as the TP8 wide step)
   attention          the TP-G specialized input projection and core (KDA, DSA; the
@@ -20,15 +20,24 @@ layout) for attention, and runs the MLPs exactly as stock does:
                      o_proj partial, one all-reduce
   MoE                the stock SP MoE module on this rank's row shard (stock
                      ``sp_shard``, zero padding marked as padding for MegaMoE),
-                     all-gathered back to the group rows
+                     all-gathered back to the group rows. VLLM_KIMI_K3_TEP_SHARED_TP:
+                     the replicated BF16 shared expert's TP-G shard on all group
+                     rows, this rank's routed rows added into its partial by the
+                     stock up-projection epilogue, one all-reduce
   dense MLP          the stock (replicated) dense MLP on all rows
 
 Precision is the stock layer's at every stage (BF16 GEMMs with FP32 accumulation,
-BF16 all-reduces); only the attention's summation order differs.
+BF16 all-reduces); only the attention's summation order differs, and under
+VLLM_KIMI_K3_TEP_SHARED_TP the shared expert's (TP-G BF16 partials, all-reduced:
+the precision of VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT, not bitwise stock).
 """
 
 from typing import Any
 
+import torch
+import torch.nn.functional as F
+
+from vllm import envs
 from vllm.distributed import tensor_model_parallel_all_reduce as all_reduce
 from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.models.common.ops.sequence_parallel import sp_padding_mask, sp_shard
@@ -138,6 +147,44 @@ def run_stock_moe(module, x):
     return sp_all_gather(module(sp_shard(x)))[: x.shape[0]]
 
 
+def shard_mlp(mlp, tp, rank):
+    """This rank's TP shard of a replicated BF16 MLP (gate and up rows copied, down
+    columns a strided view); None if stock already shards it or it is quantized."""
+    if mlp.shard_sequence_parallel:
+        return None
+    w13, w2 = mlp.gate_up_proj.weight, mlp.down_proj.weight
+    if w13.dtype != torch.bfloat16 or w2.dtype != torch.bfloat16:
+        return None
+    inter = w2.shape[1]
+    rows = slice(rank * inter // tp, (rank + 1) * inter // tp)
+    return torch.cat((w13[:inter][rows], w13[inter:][rows])), w2[:, rows], mlp.act_fn
+
+
+def padded_shard(x, rank, tp):
+    """The group rows ``x`` zero-padded to c * tp rows (``sp_shard``'s padding, made
+    once for the layer's all-rows GEMMs too) and this rank's ``sp_shard`` rows [c]."""
+    pad = -x.shape[0] % tp
+    x = F.pad(x, (0, 0, 0, pad)) if pad else x
+    c = x.shape[0] // tp
+    return x, x[rank * c : (rank + 1) * c]
+
+
+def run_split_moe(module, x, w13, w2, act, rank, tp):
+    """A MegaMoE layer with the shared expert TP-sharded: this rank's shared partial
+    on all group rows, its routed rows (stock, on its ``sp_shard`` rows) added by
+    the stock up-projection epilogue, one all-reduce. The shared GEMMs run on the
+    padded rows, so the SP padding rows' partials are zeros (SiTU(0) = 0) and the
+    all-reduce never sees uninitialized memory; ``[:b]`` drops them."""
+    b = x.shape[0]
+    x, shard = padded_shard(x, rank, tp)
+    c = shard.shape[0]
+    latent, router, topk_ids = module._maybe_overlap_router_and_down_proj(shard)
+    routed = module.experts(latent, router, topk_ids)
+    out = torch.mm(act(F.linear(x, w13)), w2.t())
+    module.routed_output_transform(routed, residual=out[rank * c : (rank + 1) * c])
+    return all_reduce(out)[:b]
+
+
 class HybridStep:
     """Builds the hybrid providers on a ``SpecializedDecode`` and runs its steps."""
 
@@ -156,7 +203,7 @@ class HybridStep:
         caps = wide_caps(runtime.dcp > 1, recoverssm)
         # Every verify step whose MoE shard fits MegaMoE's buffer is a hybrid step
         # (its rows are the KV capacity's point: past the cores' ATTENTION_MAX rows
-        # the attention layers run the stock module too; the MLPs are always stock).
+        # the attention layers run the stock module too; the MLPs run at any row count).
         tokens = min(
             (
                 d.mlp.experts.max_num_tokens
@@ -177,8 +224,16 @@ class HybridStep:
             (m for m in runtime.wide if m <= ATTENTION_MAX), default=0
         )
         runtime.attention, runtime.feed_forward, runtime.wide_moe = {}, {}, {}
+        self.shared: dict[int, tuple[torch.Tensor, torch.Tensor, Any]] = {}
         if not runtime.wide:
             return
+        if envs.VLLM_KIMI_K3_TEP_SHARED_TP:
+            for layer in decoders:
+                index, mlp = layer.layer_idx, layer.mlp
+                if layer.is_moe_layer:
+                    shard = shard_mlp(mlp.shared_experts, runtime.tp, runtime.rank)
+                    if shard is not None:
+                        self.shared[index] = shard
         for layer in decoders:
             kind = attention_kind(layer.self_attn)
             if (
@@ -190,7 +245,8 @@ class HybridStep:
 
     def run(self, step, input_ids, inputs_embeds):
         """A hybrid step: the stock decoder loop over the group's rows with the
-        specialized attention and the stock MLPs (module docstring)."""
+        specialized attention and the stock (or TP-sharded) MLPs (module
+        docstring)."""
         from ...ops import attn_res
 
         runtime = self.runtime
@@ -223,7 +279,11 @@ class HybridStep:
                 hidden, prefix_sum, residual = layer._post_attn_norm(
                     hidden, residual, prefix_sum
                 )
-                if layer.is_moe_layer:
+                if index in self.shared:
+                    hidden = run_split_moe(
+                        layer.mlp, hidden, *self.shared[index], runtime.rank, runtime.tp
+                    )
+                elif layer.is_moe_layer:
                     hidden = run_stock_moe(layer.mlp, hidden)
                 else:
                     hidden = run_dense(layer.mlp, hidden)

AI assistance: the recipe search, the vLLM change, the runs and this description were prepared with Claude Code (claude-opus-5-5). The submitter will review the change and the measurements before this leaves draft.

中文摘要

把 kimik3.1-fp4-b300-vllm-agentic-dspark 改为按并发选择拓扑：c1 用 TP8，c4–c24 用 TP8 x DCP8，c32–c70 用 TEP（TP8 x DCP8 注意力 + EP8 MegaMoE）。镜像为 inferactinc/public:kimi-k31-mc-f819987-tep-shared-tp-x86_64，即 vllm/vllm-openai:kimi-k31-mc（与 #4871 相同）加上一项默认关闭、尚未进入上游 vLLM 的改动 VLLM_KIMI_K3_TEP_SHARED_TP：TEP decode 时把 BF16 shared expert 按 TP 切分，并和 routed 输出合并成一次 all-reduce。所有点都用 latency_specialized decode，CUDA graph 捕获到 4096 个 token，不超过 4096 token 的 prefill 和混合 step 也走 graph。

在我们自己的 B300 节点上实测（c1–c24 用 kimi-k31-mc 本身，c32–c70 用本 PR 的镜像）：SGLang #4803 的每个点都被这条曲线支配；相对 #4871，c16–c56 吞吐高 4–9%，P90 交互性高 7–16%，c8 持平，c4 #4871 略好，c1 互有取舍。VLLM_KIMI_K3_TEP_SHARED_TP 在 c32–c48 带来 3–5% 的吞吐和 3–6% 的 P90 提升；c56–c70 的 P90 不变，吞吐变化部分可能来自开启后约多 5% 的 KV 容量。由于审查清单要求使用上游镜像且不打引擎补丁，该改动需要上游 PR 或 waiver；在此之前 c32–c70 可以关闭该变量，直接用 kimi-k31-mc。

🤖 Generated with Claude Code