Fix(spec decode): catch up trim bug

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

@@ -226,13 +226,6 @@ def _run_first_step(
         )
         input_ids = maybe_substitute_mm_pad(input_ids, self.mm_pad_substitute_id)
         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)
 
         draft_first_mode = (
             ForwardMode.DRAFT_EXTEND

python/tokenspeed/runtime/models/llama_eagle3.py

@@ -35,7 +35,6 @@
 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
@@ -187,37 +186,22 @@ def forward(
                 output_q_rope=q_rope,
                 enable_pdl=pdl_enabled(),
             )
-            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,
-                )
+            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)
+
+        if ctx.draft_first_step_reduce:
+            # KV written 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