refactor(spec-decode): simplify Llama Eagle3 attention path for #217 (1/3)

python/tokenspeed/runtime/execution/context.py

@@ -50,7 +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.
+    # Legacy draft first-step flag; Qwen / DeepSeek NextN still set this until
+    # their attention subclasses own trim.  Llama Eagle3 uses accept_lengths.
     draft_first_step_reduce: bool = False
     # Normalized explicit decode input overrides for this forward, if any.
     decode_input_ids: list[int] | None = None
@@ -62,3 +63,9 @@ class ForwardContext:
 
     # --- logits processor ---
     gather_ids: torch.Tensor | None = None
+
+    # --- spec-decode draft (drafter-owned buffers plumbed per forward) ---
+    # draft_seq_lens_buf: mutable per-request seq_lens alias the draft backend reads.
+    draft_seq_lens_buf: torch.Tensor | None = None
+    # accept_lengths: per-request accepted verify width for cache_seqlens correction.
+    accept_lengths: torch.Tensor | None = None

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

@@ -36,7 +36,7 @@
     CaptureHiddenMode,
     ForwardMode,
 )
-from tokenspeed.runtime.models.base import BaseCausalLM
+from tokenspeed.runtime.models.llama_eagle3 import LlamaForCausalLMEagle3
 from tokenspeed.runtime.multimodal.inputs import maybe_substitute_mm_pad
 from tokenspeed.runtime.utils import get_colorful_logger
 from tokenspeed.runtime.utils.nvtx import nvtx_range
@@ -218,21 +218,12 @@ def _run_first_step(
             draft_input, bs, draft_input.input_num_tokens
         )
         input_ids = maybe_substitute_mm_pad(input_ids, self.mm_pad_substitute_id)
-        draft_first_step_reduce = forward_mode.is_decode()
-
-        # TODO: remove the isinstance/flag gate together with pre_attention_trim
-        # once Qwen NextN and DeepSeek V3 NextN also pre-slice q.
-        draft_model = self.draft_model_runner.model
-        if (
-            draft_first_step_reduce
-            and self.attn_backend.support_kv_cache_prewrite
-            and isinstance(draft_model, BaseCausalLM)
-            and draft_model.pre_attention_trim
-        ):
-            correction = (self.spec_num_tokens - draft_input.accept_lengths).to(
-                self.draft_seq_lens_buf.dtype
-            )
-            self.draft_seq_lens_buf[:bs].sub_(correction)
+        # Llama Eagle3 narrows for any non-idle catch-up (EXTEND/MIXED/
+        # TARGET_VERIFY/DECODE); Qwen/DeepSeek keep is_decode() only.
+        draft_first_step_reduce = forward_mode.is_decode() or (
+            isinstance(self.draft_model_runner.model, LlamaForCausalLMEagle3)
+            and not forward_mode.is_idle()
+        )
 
         draft_first_mode = (
             ForwardMode.DRAFT_EXTEND
@@ -254,6 +245,8 @@ def _run_first_step(
             global_bs=draft_input.global_bs,
             all_decode_or_idle=draft_input.all_decode_or_idle,
             draft_first_step_reduce=draft_first_step_reduce,
+            draft_seq_lens_buf=self.draft_seq_lens_buf,
+            accept_lengths=draft_input.accept_lengths,
         )
 
         return self.draft_model_runner.forward(

python/tokenspeed/runtime/models/base/causal_lm.py

@@ -44,10 +44,6 @@ class BaseCausalLM(nn.Module):
 
     model_cls: type[BaseTransformerModel]
 
-    # TODO: temporary; remove in the follow-up refactoring that extends
-    # pre-attn q-slice to Qwen NextN and DeepSeek V3 NextN.
-    pre_attention_trim: bool = False
-
     def __init__(
         self,
         config: PretrainedConfig,

python/tokenspeed/runtime/models/llama.py

@@ -52,7 +52,9 @@
     BaseDecoderLayer,
     BaseTransformerModel,
 )
+from tokenspeed.runtime.models.utils import create_fused_set_kv_buffer_arg
 from tokenspeed.runtime.utils import add_prefix
+from tokenspeed.runtime.utils.pdl import pdl_enabled
 
 
 class LlamaMLP(nn.Module):
@@ -119,6 +121,7 @@ def __init__(
         layer_id: int = 0,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
+        qkv_input_size: int | None = None,
     ) -> None:
         super().__init__()
         self.hidden_size = hidden_size
@@ -151,7 +154,7 @@ def __init__(
         attention_bias = getattr(config, "attention_bias", False)
 
         self.qkv_proj = QKVParallelLinear(
-            hidden_size,
+            qkv_input_size or hidden_size,
             self.head_dim,
             self.total_num_heads,
             self.total_num_kv_heads,
@@ -199,14 +202,87 @@ def forward(
         # the batch row is empty (e.g. idle ranks under DP attention). Matches
         # the short-circuit ``LlamaMLP.forward`` already has.
         if hidden_states.shape[0] == 0:
-            return hidden_states
+            return hidden_states.new_zeros(
+                (0, self.hidden_size), dtype=hidden_states.dtype
+            )
         qkv, _ = self.qkv_proj(hidden_states)
         q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-        q, k = self.rotary_emb(positions, q, k)
-        attn_output = self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc)
+        attn_output = self._attn(positions, q, k, v, ctx, out_cache_loc)
         output, _ = self.o_proj(attn_output)
         return output
 
+    def _attn(
+        self,
+        positions: torch.Tensor,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        ctx: ForwardContext,
+        out_cache_loc: torch.Tensor,
+    ) -> torch.Tensor:
+        """RoPE + attention (pre-o_proj), with optional fused KV pre-write.
+
+        When the backend supports KV pre-write *and* ``create_fused_set_kv_buffer_arg``
+        accepts the layer's scales, fused rope writes KV directly into the cache
+        so the attention call can run with ``save_kv_cache=False`` (saves one
+        kernel launch). Otherwise we fall back to plain RoPE + ``self.attn(q, k, v)``
+        so the backend writes KV the normal way — without this fallback, layers
+        with non-trivial k/v scales silently lose their KV writes. Subclasses
+        (e.g. Eagle3 draft head) override this hook to insert spec-decode
+        behaviour around the same scaffolding.
+        """
+        if ctx.attn_backend.support_kv_cache_prewrite(ctx.forward_mode):
+            fused_kv_arg = self._build_fused_kv_arg(v, ctx, out_cache_loc)
+            if fused_kv_arg is not None:
+                q_rope = self._fused_rope_kv_write(positions, q, k, fused_kv_arg)
+                return self.attn(
+                    q_rope,
+                    None,
+                    None,
+                    save_kv_cache=False,
+                    ctx=ctx,
+                    out_cache_loc=out_cache_loc,
+                )
+        q, k = self.rotary_emb(positions, q, k)
+        return self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc)
+
+    def _build_fused_kv_arg(
+        self,
+        v: torch.Tensor,
+        ctx: ForwardContext,
+        out_cache_loc: torch.Tensor,
+    ):
+        """Try to build the fused RoPE+KV-write descriptor; returns ``None`` if
+        the helper rejects the layer (e.g. non-trivial k/v scales)."""
+        n = v.shape[0]
+        return create_fused_set_kv_buffer_arg(
+            value=v.view(n, self.num_kv_heads, self.head_dim),
+            layer=self.attn,
+            out_cache_loc=out_cache_loc,
+            token_to_kv_pool=ctx.token_to_kv_pool,
+        )
+
+    def _fused_rope_kv_write(
+        self,
+        positions: torch.Tensor,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        fused_kv_arg,
+    ) -> torch.Tensor:
+        """Fused RoPE that writes KV into cache (via ``fused_kv_arg``) and
+        returns the rope'd Q."""
+        n = q.shape[0]
+        q_rope = torch.empty((n, self.q_size), dtype=q.dtype, device=q.device)
+        self.rotary_emb(
+            positions,
+            q,
+            k,
+            fused_set_kv_buffer_arg=fused_kv_arg,
+            output_q_rope=q_rope,
+            enable_pdl=pdl_enabled(),
+        )
+        return q_rope
+
 
 class LlamaDecoderLayer(BaseDecoderLayer):