perf(Spec Decode): skip dead-position compute in draft catch-up step(decode)

python/tokenspeed/runtime/distributed/comm_manager.py

@@ -60,17 +60,21 @@ def get_num_tokens(self, ctx: ForwardContext):
         return sum(scattered), max(scattered)
 
     def scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
-        if ctx.global_num_tokens is not None:
+        # Under draft first-step reduce, comm operates on bs / global_bs since
+        # the midlayer pruned activations to one row per request.
+        global_counts = (
+            ctx.global_bs if ctx.draft_first_step_reduce else ctx.global_num_tokens
+        )
+        if global_counts is not None:
             scattered = []
             for attn_dp_rank in range(self.mapping.attn.dp_size):
-                num_tokens = ctx.global_num_tokens[
-                    attn_dp_rank * self.mapping.attn.tp_size
-                ]
+                num_tokens = global_counts[attn_dp_rank * self.mapping.attn.tp_size]
                 scattered.extend(
                     self._scatter_count(num_tokens, self.mapping.attn.tp_size)
                 )
             return scattered
-        return self._scatter_count(ctx.input_num_tokens, self.mapping.attn.tp_size)
+        num_tokens = ctx.bs if ctx.draft_first_step_reduce else ctx.input_num_tokens
+        return self._scatter_count(num_tokens, self.mapping.attn.tp_size)
 
     def attn_tp_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
         start = self.mapping.attn.tp_size * self.mapping.attn.dp_rank
@@ -84,15 +88,17 @@ def dense_tp_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
 
     def moe_tp_ep_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]:
         tp_ep_size = self.mapping.moe.tp_ep_size
-        if ctx.global_num_tokens is not None:
-            # DP path: global_num_tokens has one entry per world rank.
-            # MoE group = tp_ep_size contiguous ranks starting at dp_rank * tp_ep_size.
+        # Under draft first-step reduce, the midlayer pruned activations to bs
+        # rows before pre_moe_comm; MoE collectives must size accordingly.
+        global_counts = (
+            ctx.global_bs if ctx.draft_first_step_reduce else ctx.global_num_tokens
+        )
+        if global_counts is not None:
             start = self.mapping.moe.dp_rank * tp_ep_size
-            return list(ctx.global_num_tokens[start : start + tp_ep_size])
-        # No global_num_tokens (e.g., drafter context or non-DP).
-        # This rank has all input_num_tokens; other ranks in the EP group have 0.
+            return list(global_counts[start : start + tp_ep_size])
+        num_tokens = ctx.bs if ctx.draft_first_step_reduce else ctx.input_num_tokens
         result = [0] * tp_ep_size
-        result[self.mapping.moe.tp_ep_rank] = ctx.input_num_tokens
+        result[self.mapping.moe.tp_ep_rank] = num_tokens
         return result
 
     # ---- Communication patterns ----

python/tokenspeed/runtime/execution/context.py

@@ -50,6 +50,8 @@ class ForwardContext:
     forward_mode: ForwardMode | None
     req_to_page: torch.Tensor | None = None
     capture_hidden_mode: CaptureHiddenMode | None = CaptureHiddenMode.NULL
+    # Spec decode draft head's first step prunes to one live row per request.
+    draft_first_step_reduce: bool = False
 
     # --- dp attention ---
     global_num_tokens: list[int] | None = None

python/tokenspeed/runtime/execution/drafter/eagle.py

@@ -212,6 +212,16 @@ def _run_first_step(
             draft_input, bs, draft_input.input_num_tokens
         )
 
+        draft_first_step_reduce = forward_mode.is_decode()
+
+        if draft_first_step_reduce and self.attn_backend.support_kv_cache_prewrite:
+            # Trim seq_lens by rejected-draft count so the sliced decode
+            # query does not attend to dead positions.
+            correction = (self.spec_num_tokens - draft_input.accept_lengths).to(
+                self.draft_seq_lens_buf.dtype
+            )
+            self.draft_seq_lens_buf[:bs].sub_(correction)
+
         ctx = ForwardContext(
             attn_backend=self.attn_backend,
             token_to_kv_pool=self.token_to_kv_pool,
@@ -225,6 +235,7 @@ def _run_first_step(
             global_num_tokens=draft_input.global_num_tokens,
             global_bs=draft_input.global_bs,
             all_decode_or_idle=draft_input.all_decode_or_idle,
+            draft_first_step_reduce=draft_first_step_reduce,
         )
 
         return self.draft_model_runner.forward(

python/tokenspeed/runtime/execution/model_executor.py

@@ -882,19 +882,22 @@ def execute_idle_forward(
         # NCCL collectives. Idle ranks must match those collectives:
         # 1 first-step forward + (spec_num_steps - 1) multi-step decode forwards.
         if self.drafter is not None:
-            draft_ctx = ForwardContext(
-                attn_backend=self.drafter.attn_backend,
-                token_to_kv_pool=self.drafter.token_to_kv_pool,
-                req_to_page=self.drafter.req_to_page,
-                bs=0,
-                num_extends=0,
-                input_num_tokens=0,
-                forward_mode=ForwardMode.IDLE,
-                global_num_tokens=global_num_tokens,
-                global_bs=global_bs,
-                all_decode_or_idle=all_decode_or_idle,
-            )
-            for _ in range(self.drafter.spec_num_steps):
+            for step_idx in range(self.drafter.spec_num_steps):
+                # Mirror active rank's catch-up step: when all non-idle ranks
+                # are decoding, step 0 sizes collectives from bs/global_bs.
+                draft_ctx = ForwardContext(
+                    attn_backend=self.drafter.attn_backend,
+                    token_to_kv_pool=self.drafter.token_to_kv_pool,
+                    req_to_page=self.drafter.req_to_page,
+                    bs=0,
+                    num_extends=0,
+                    input_num_tokens=0,
+                    forward_mode=ForwardMode.IDLE,
+                    global_num_tokens=global_num_tokens,
+                    global_bs=global_bs,
+                    all_decode_or_idle=all_decode_or_idle,
+                    draft_first_step_reduce=(step_idx == 0 and all_decode_or_idle),
+                )
                 self.drafter.draft_model_runner.forward(
                     draft_ctx,
                     input_ids=empty,

python/tokenspeed/runtime/layers/logits_processor.py

@@ -210,7 +210,9 @@ def forward(
         # Get the last hidden states and last logits for the next token prediction
         if not logits_metadata.extend_return_logprob:
             gather_ids = logits_metadata.gather_ids
-            if gather_ids is None:
+            # Shapes align iff midlayer already pruned to one row per request
+            # (draft first-step reduce). Other paths emit [N, H] with N > bs.
+            if gather_ids is None or gather_ids.shape[0] == hidden_states.shape[0]:
                 pruned_states = hidden_states
                 if aux_hidden_states is not None:
                     aux_pruned_states = list(aux_hidden_states)

python/tokenspeed/runtime/models/deepseek_v3.py

@@ -675,6 +675,10 @@ def forward(
                 attn_output[num_prefill_tokens:],
             )
 
+        if ctx.draft_first_step_reduce:
+            # KV already written; drop dead-position rows so o_proj / MLP /
+            # post-norms only run on one live row per request.
+            attn_output = attn_output.index_select(0, ctx.gather_ids)
         output, _ = self.o_proj(attn_output)
         return output
 
@@ -1153,6 +1157,9 @@ def forward(
                 out_cache_loc=out_cache_loc,
                 comm_manager=self.comm_manager,
             )
+            if ctx.draft_first_step_reduce:
+                # Gather residual to self_attn's [bs, H].
+                residual = residual.index_select(0, ctx.gather_ids)
             hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                 hidden_states, residual, ctx
             )
@@ -1697,6 +1704,9 @@ def forward(
                 comm_manager=self.comm_manager,
             )
 
+            if ctx.draft_first_step_reduce:
+                # Gather residual to self_attn's [bs, H].
+                residual = residual.index_select(0, ctx.gather_ids)
             hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                 hidden_states, residual, ctx
             )

python/tokenspeed/runtime/models/llama_eagle3.py

@@ -35,6 +35,7 @@
 from tokenspeed.runtime.configs.utils import get_rope_theta
 from tokenspeed.runtime.distributed.mapping import Mapping
 from tokenspeed.runtime.execution.context import ForwardContext
+from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
 from tokenspeed.runtime.layers.activation import SiluAndMul
 from tokenspeed.runtime.layers.common import concat
 from tokenspeed.runtime.layers.layernorm import RMSNorm
@@ -185,17 +186,37 @@ def forward(
                 output_q_rope=q_rope,
                 enable_pdl=pdl_enabled(),
             )
-            attn_output = self.attn(
-                q_rope,
-                None,
-                None,
-                save_kv_cache=False,
-                ctx=ctx,
-                out_cache_loc=out_cache_loc,
-            )
+            if ctx.draft_first_step_reduce:
+                # KV already written via fused_set_kv_buffer_arg above; slice Q
+                # to one query per request and route attn as decode.
+                q_rope = q_rope.index_select(0, ctx.gather_ids)
+                attn_output = ctx.attn_backend.forward(
+                    q_rope,
+                    None,
+                    None,
+                    self.attn,
+                    out_cache_loc,
+                    ctx.token_to_kv_pool,
+                    ForwardMode.DECODE,
+                    ctx.bs,
+                    save_kv_cache=False,
+                )
+            else:
+                attn_output = self.attn(
+                    q_rope,
+                    None,
+                    None,
+                    save_kv_cache=False,
+                    ctx=ctx,
+                    out_cache_loc=out_cache_loc,
+                )
         else:
             q, k = self.rotary_emb(positions, q, k)
             attn_output = self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc)
+            if ctx.draft_first_step_reduce:
+                # KV written by self.attn above; slice attn_output so o_proj
+                # and the rest of the layer only run on the live rows.
+                attn_output = attn_output.index_select(0, ctx.gather_ids)
 
         output, _ = self.o_proj(attn_output)
         return output
@@ -361,6 +382,9 @@ def forward_low_latency(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
+        if ctx.draft_first_step_reduce and not ctx.forward_mode.is_idle():
+            # Gather residual to self_attn's [bs, H]; idle has no gather_ids.
+            residual = residual.index_select(0, ctx.gather_ids)
 
         # Fused post-attn allreduce + norm (uses attn tp group)
         block_scale = None
@@ -426,6 +450,9 @@ def forward(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
+        if ctx.draft_first_step_reduce and not ctx.forward_mode.is_idle():
+            # Gather residual to self_attn's [bs, H]; idle has no gather_ids.
+            residual = residual.index_select(0, ctx.gather_ids)
         hidden_states, residual = self.comm_manager.post_attn_comm(
             hidden_states, residual, ctx
         )
@@ -524,7 +551,9 @@ def forward(
             fuse_embed_reduce=fuse_embed_reduce,
         )
 
-        if midlayer.comm_manager.should_fuse(hidden_states.shape[0]):
+        # Decide on pre-slice token count so this matches the path midlayer
+        # actually took; under draft reduce, hidden_states.shape[0] shrinks.
+        if midlayer.comm_manager.should_fuse(input_ids.shape[0]):
             hidden_states_to_logits, hidden_states_to_aux = hidden_states, residual
         else:
             hidden_states_to_logits, hidden_states_to_aux = self.norm(

python/tokenspeed/runtime/models/qwen3_5.py

@@ -747,6 +747,10 @@ def self_attention(
         if self.attn_output_gate:
             sigmoid_mul(attn_output, gate)
 
+        if ctx.draft_first_step_reduce:
+            # Slice attn_output to [bs, H] so o_proj runs on live rows only.
+            attn_output = attn_output.index_select(0, ctx.gather_ids)
+
         output, _ = self.o_proj(attn_output)
         return output
 
@@ -774,6 +778,9 @@ def forward(
                 ctx=ctx,
                 out_cache_loc=out_cache_loc,
             )
+            if ctx.draft_first_step_reduce:
+                # Gather residual to self_attention's [bs, H].
+                residual = residual.index_select(0, ctx.gather_ids)
             hidden_states, residual = self.comm_manager.post_attn_reduce_norm(
                 hidden_states, residual, ctx
             )