From d21ac4da917c9ab8515d8e2b8853ac6a2791d938 Mon Sep 17 00:00:00 2001 From: zhangqi-chen Date: Mon, 29 Jun 2026 11:43:24 +0800 Subject: [PATCH 1/2] feat(dsv4): single-token decode config (B=8, S=1, T=8) Switch the decode config from MTP (B=4, S=2) to single-token decode (B=8, S=1), keeping T = B*S = 8 so the MoE pipeline width is unchanged. The kernel-side enablement for small/single-token decode already lives on main: qkv_proj_rope pads T to MATMUL_T_TILE=16, decode_sparse_attn_swa zero-pads the MTP overlay up to ATTN_K_TILE, and #647 drives the MoE pipeline natively at T = MOE_TOKENS (prefill entries override to PREFILL_TOKENS). So this is now a pure config flip. Validated a2a3 ep2: decode_fwd (T=8) and prefill_fwd (T=128) run end-to-end. --- models/deepseek/v4/config.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/models/deepseek/v4/config.py b/models/deepseek/v4/config.py index 56e7c723..be1cbf0d 100644 --- a/models/deepseek/v4/config.py +++ b/models/deepseek/v4/config.py @@ -239,8 +239,8 @@ def mix_hc(self) -> int: # Deployment constants -DECODE_BATCH = 4 # B: requests per decode step -DECODE_SEQ = 2 # S: 2 tokens per step (MTP) +DECODE_BATCH = 8 # B: requests per decode step +DECODE_SEQ = 1 # S: 1 token per step (single-token decode) DECODE_TOKENS = DECODE_BATCH * DECODE_SEQ PREFILL_BATCH = 1 # B: prefill batch for the current kernel programs PREFILL_SEQ = 128 # S: prefill sequence for the current kernel programs From d184688e1b5c13d050da004755f88a66955cef40 Mon Sep 17 00:00:00 2001 From: zhangqi-chen Date: Wed, 1 Jul 2026 01:52:19 +0800 Subject: [PATCH 2/2] feat(dsv4): size prefill RECV_MAX for real-weight routing; auto-scope MoE ring-heap config: PREFILL_RECV_MAX=1024 -- real gate routing skews past the RECV_SAFETY=4 uniform bound, overflowing the default recv depth and deadlocking dispatch/combine on real-weight ep8 prefill. moe: mark `moe` @pl.jit.inline(auto_scope=False) and use a plain pl.scope() around expert_routed + combine + hc_post so the compiler places AUTO runtime scopes across the whole MoE instead of one hand-placed MANUAL scope. The auto placement recycles the large recv buffers per stage and spreads the MoE over two rings, dropping the worst-ring MoE ring-heap high-water from 126% (unscoped) / ~90% (single MANUAL scope) to ~71% (keeps the larger RECV_MAX within the 1G ring). --- models/deepseek/v4/config.py | 1 + models/deepseek/v4/moe.py | 33 +++++++++++++++++---------------- 2 files changed, 18 insertions(+), 16 deletions(-) diff --git a/models/deepseek/v4/config.py b/models/deepseek/v4/config.py index be1cbf0d..7cd8f195 100644 --- a/models/deepseek/v4/config.py +++ b/models/deepseek/v4/config.py @@ -268,6 +268,7 @@ def mix_hc(self) -> int: # bake different depths; default to decode, prefill overrides to PREFILL_RECV_MAX. DECODE_RECV_MAX = max(16, DECODE_TOKENS * FLASH.num_experts_per_tok * RECV_SAFETY // (FLASH.n_routed_experts // EP_WORLD_SIZE)) PREFILL_RECV_MAX = max(16, PREFILL_TOKENS * FLASH.num_experts_per_tok * RECV_SAFETY // (FLASH.n_routed_experts // EP_WORLD_SIZE)) +PREFILL_RECV_MAX = 1024 # real routing skews past the RECV_SAFETY=4 uniform bound; size for the observed peak RECV_MAX = DECODE_RECV_MAX # When True, gate.py's N_EXPERTS uses the full global expert space diff --git a/models/deepseek/v4/moe.py b/models/deepseek/v4/moe.py index 84bd4b10..c8bdda38 100644 --- a/models/deepseek/v4/moe.py +++ b/models/deepseek/v4/moe.py @@ -78,7 +78,7 @@ def _parse_ep_argv(): assert N_EXPERTS_GLOBAL == N_RANKS * N_LOCAL -@pl.jit.inline +@pl.jit.inline(auto_scope=False) def moe( # model inputs x_hc: pl.Tensor[[T, HC_MULT, D], pl.BF16], @@ -168,23 +168,24 @@ def moe( num_tokens, my_rank, moe_epoch, ) - recv_y = pl.create_tensor([N_LOCAL, RECV_MAX, D], dtype=pl.BF16) - expert_routed( - recv_x_out, recv_scale_out, recv_w_out, recv_count_out, - routed_w1, routed_w1_scale, routed_w3, routed_w3_scale, - routed_w2, routed_w2_scale, - recv_y, - ) + with pl.scope(): + recv_y = pl.create_tensor([N_LOCAL, RECV_MAX, D], dtype=pl.BF16) + expert_routed( + recv_x_out, recv_scale_out, recv_w_out, recv_count_out, + routed_w1, routed_w1_scale, routed_w3, routed_w3_scale, + routed_w2, routed_w2_scale, + recv_y, + ) - ffn_out = pl.create_tensor([T, D], dtype=pl.BF16) - combine( - recv_y, recv_r_route_out, sh, - ffn_out, - pub_counts, routed_y_buf, combine_done, - num_tokens, my_rank, moe_epoch, - ) + ffn_out = pl.create_tensor([T, D], dtype=pl.BF16) + combine( + recv_y, recv_r_route_out, sh, + ffn_out, + pub_counts, routed_y_buf, combine_done, + num_tokens, my_rank, moe_epoch, + ) - hc_post(ffn_out, x_hc, post_ffn, comb_ffn, x_next) + hc_post(ffn_out, x_hc, post_ffn, comb_ffn, x_next) return x_next