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

python/tokenspeed/runtime/execution/context.py

@@ -58,3 +58,8 @@ class ForwardContext:
 
     # --- logits processor ---
     gather_ids: torch.Tensor | None = None
+
+    # When True, the draft head's first step prunes attn/MLP/post-norms to the
+    # one live position per request using ctx.gather_ids. LogitsProcessor
+    # bypasses its own slicing (gather_ids set to None on LogitsMetadata).
+    draft_reduce_to_last: bool = False

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

@@ -36,6 +36,7 @@
     CaptureHiddenMode,
     ForwardMode,
 )
+from tokenspeed.runtime.models.llama_eagle3 import LlamaForCausalLMEagle3
 from tokenspeed.runtime.utils import get_colorful_logger
 from tokenspeed.runtime.utils.nvtx import nvtx_range
 
@@ -212,6 +213,16 @@ def _run_first_step(
             draft_input, bs, draft_input.input_num_tokens
         )
 
+        # Catch-up step with multi-token decode input: every position except
+        # the live one per request is purely a KV-cache write. The midlayer
+        # slices Q after KV write via ctx.gather_ids and runs attn/MLP/post-
+        # norms on just the [bs] live positions. Only LlamaForCausalLMEagle3
+        # implements the midlayer slice today; MLA / NextN drafts fall back
+        # to the full path until a follow-up PR adds matching slice code.
+        draft_reduce_to_last = forward_mode.is_decode() and isinstance(
+            self.draft_model_runner.model, LlamaForCausalLMEagle3
+        )
+
         ctx = ForwardContext(
             attn_backend=self.attn_backend,
             token_to_kv_pool=self.token_to_kv_pool,
@@ -225,6 +236,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_reduce_to_last=draft_reduce_to_last,
         )
 
         return self.draft_model_runner.forward(

python/tokenspeed/runtime/layers/logits_processor.py

@@ -111,10 +111,14 @@ def from_forward_context(
         ctx: ForwardContext,
         input_lengths: torch.Tensor,
     ):
+        # When the midlayer already pruned to one row per request (EAGLE draft
+        # first-step reduce), drop gather_ids so LogitsProcessor passes through
+        # instead of re-indexing.
+        gather_ids = None if ctx.draft_reduce_to_last else ctx.gather_ids
         return cls(
             forward_mode=ctx.forward_mode,
             capture_hidden_mode=ctx.capture_hidden_mode,
-            gather_ids=ctx.gather_ids,
+            gather_ids=gather_ids,
             extend_seq_lens=input_lengths,
         )
 

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,42 @@ 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_reduce_to_last:
+                # KV cache for all bs*spec_num_tokens positions already written
+                # by fused_set_kv_buffer_arg above. Reduce Q to one query per
+                # request (live position) and route to decode kernel: bs queries
+                # against the full cache (which now includes the just-written
+                # spec_num_tokens new tokens per request). DRAFT_EXTEND init
+                # populates forward_decode_metadata precisely for this case.
+                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_reduce_to_last:
+                # Non-prewrite backend: KV was written by self.attn above;
+                # slice attn_output to one row per request so o_proj and the
+                # rest of the layer only see the live positions.
+                attn_output = attn_output.index_select(0, ctx.gather_ids)
 
         output, _ = self.o_proj(attn_output)
         return output
@@ -361,6 +387,11 @@ def forward_low_latency(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
+        if ctx.draft_reduce_to_last:
+            # self_attn returned [bs, H]; gather residual to match before the
+            # fused allreduce+norm so shapes line up. Everything downstream
+            # (post-norm, MLP, final norm) now runs on [bs, H].
+            residual = residual.index_select(0, ctx.gather_ids)
 
         # Fused post-attn allreduce + norm (uses attn tp group)
         block_scale = None
@@ -426,6 +457,11 @@ def forward(
             ctx=ctx,
             out_cache_loc=out_cache_loc,
         )
+        if ctx.draft_reduce_to_last:
+            # self_attn returned [bs, H]; gather residual to match before the
+            # post-attn norm+comm so shapes line up. Everything downstream
+            # (post-norm, MLP, final norm) now runs on [bs, H].
+            residual = residual.index_select(0, ctx.gather_ids)
         hidden_states, residual = self.comm_manager.post_attn_comm(
             hidden_states, residual, ctx
         )