From d7882e7b54fabd1c1412a71b6bb4b743180ef4bc Mon Sep 17 00:00:00 2001 From: rjzhb Date: Tue, 9 Jun 2026 00:24:52 +0000 Subject: [PATCH 1/8] refactor(spec-decode): wrap Eagle3 attention via base llama._attn Signed-off-by: rjzhb --- .../tokenspeed/runtime/execution/context.py | 9 +- .../runtime/execution/drafter/eagle.py | 17 +- .../runtime/models/base/causal_lm.py | 4 - python/tokenspeed/runtime/models/llama.py | 69 +++++- .../tokenspeed/runtime/models/llama_eagle3.py | 227 ++++++------------ 5 files changed, 143 insertions(+), 183 deletions(-) diff --git a/python/tokenspeed/runtime/execution/context.py b/python/tokenspeed/runtime/execution/context.py index 194a9709d..67a889e4d 100644 --- a/python/tokenspeed/runtime/execution/context.py +++ b/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 diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py index 0aab7a1d9..6466c8290 100644 --- a/python/tokenspeed/runtime/execution/drafter/eagle.py +++ b/python/tokenspeed/runtime/execution/drafter/eagle.py @@ -36,7 +36,6 @@ CaptureHiddenMode, ForwardMode, ) -from tokenspeed.runtime.models.base import BaseCausalLM 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 @@ -221,20 +220,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() - # 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) - draft_first_mode = ( ForwardMode.DRAFT_EXTEND if forward_mode.is_target_verify() @@ -255,6 +240,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( diff --git a/python/tokenspeed/runtime/models/base/causal_lm.py b/python/tokenspeed/runtime/models/base/causal_lm.py index b87e40e96..b83ce4be0 100644 --- a/python/tokenspeed/runtime/models/base/causal_lm.py +++ b/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, diff --git a/python/tokenspeed/runtime/models/llama.py b/python/tokenspeed/runtime/models/llama.py index 1baf785bf..0fc39e1f1 100644 --- a/python/tokenspeed/runtime/models/llama.py +++ b/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,72 @@ 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, fused rope writes KV directly + into the cache so the attention call can run with ``save_kv_cache=False`` + (saves one kernel launch). 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): + q_rope = self._fused_rope_kv_write(positions, q, k, v, ctx, out_cache_loc) + 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 _fused_rope_kv_write( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + """Fused RoPE that writes KV into cache and returns the rope'd Q.""" + n = q.shape[0] + fused_kv_arg = 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, + ) + 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): diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py index 9c160f53b..b633bfbb8 100644 --- a/python/tokenspeed/runtime/models/llama_eagle3.py +++ b/python/tokenspeed/runtime/models/llama_eagle3.py @@ -32,7 +32,6 @@ from torch import nn from transformers import LlamaConfig -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 @@ -42,12 +41,9 @@ from tokenspeed.runtime.layers.linear import ( ColumnParallelLinear, MergedColumnParallelLinear, - QKVParallelLinear, RowParallelLinear, ) -from tokenspeed.runtime.layers.paged_attention import PagedAttention from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig -from tokenspeed.runtime.layers.rotary_embedding import get_rope from tokenspeed.runtime.layers.vocab_parallel_embedding import ParallelLMHead from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader from tokenspeed.runtime.models.base import ( @@ -55,9 +51,8 @@ BaseDecoderLayer, BaseTransformerModel, ) -from tokenspeed.runtime.models.utils import create_fused_set_kv_buffer_arg +from tokenspeed.runtime.models.llama import LlamaAttention as BaseLlamaAttention from tokenspeed.runtime.utils import add_prefix, get_colorful_logger -from tokenspeed.runtime.utils.pdl import pdl_enabled logger = get_colorful_logger(__name__) @@ -67,160 +62,72 @@ # --------------------------------------------------------------------------- -class LlamaAttention(nn.Module): +class LlamaAttention(BaseLlamaAttention): + """Eagle3 draft head attention. - def __init__( - self, - config: LlamaConfig, - mapping: Mapping, - hidden_size: int, - num_heads: int, - num_kv_heads: int, - layer_id: int = 0, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - qkv_input_size: int | None = None, - ) -> None: + Inherits ``__init__`` (with ``qkv_input_size=2*hidden_size`` to accommodate + the [embed || hidden] concat) and ``forward`` (= empty-shape handling + + qkv projection + o_proj scaffolding) from base. - super().__init__() - self.hidden_size = hidden_size - - self.attn_tp_size = mapping.attn.tp_size - self.attn_tp_rank = mapping.attn.tp_rank - attn_tp_group = mapping.attn.tp_group - - self.total_num_heads = num_heads - assert self.total_num_heads % self.attn_tp_size == 0 - self.num_heads = self.total_num_heads // self.attn_tp_size - self.total_num_kv_heads = num_kv_heads - if self.total_num_kv_heads >= self.attn_tp_size: - assert self.total_num_kv_heads % self.attn_tp_size == 0 - else: - assert self.attn_tp_size % self.total_num_kv_heads == 0 - self.num_kv_heads = max(1, self.total_num_kv_heads // self.attn_tp_size) - self.head_dim = getattr( - config, "head_dim", self.hidden_size // self.total_num_heads - ) - self.q_size = self.num_heads * self.head_dim - self.kv_size = self.num_kv_heads * self.head_dim - - rope_theta = get_rope_theta(config) - rope_scaling = getattr(config, "rope_scaling", None) - max_position_embeddings = getattr(config, "max_position_embeddings", 8192) - - self.qkv_proj = QKVParallelLinear( - qkv_input_size or hidden_size, - self.head_dim, - self.total_num_heads, - self.total_num_kv_heads, - bias=False, - quant_config=quant_config, - tp_rank=self.attn_tp_rank, - tp_size=self.attn_tp_size, - tp_group=attn_tp_group, - prefix=add_prefix("qkv_proj", prefix), - ) - self.o_proj = RowParallelLinear( - self.total_num_heads * self.head_dim, - hidden_size, - bias=False, - quant_config=quant_config, - reduce_results=False, - tp_rank=self.attn_tp_rank, - tp_size=self.attn_tp_size, - tp_group=attn_tp_group, - prefix=add_prefix("o_proj", prefix), - ) - self.rotary_emb = get_rope( - self.head_dim, - rotary_dim=self.head_dim, - max_position=max_position_embeddings, - base=rope_theta, - rope_scaling=rope_scaling, - ) - self.attn = PagedAttention( - self.num_heads, - self.head_dim, - self.head_dim**-0.5, - num_kv_heads=self.num_kv_heads, - layer_id=layer_id, - ) + Overrides ``_attn`` to add the draft-first-step dispatch B: trim + cache_seqlens, slice q to one live row per request, and route the attention + call as ``DECODE`` (fused KV pre-write path) or post-slice attn_output + (non-fused fallback). Inactive draft steps delegate to base. + """ - def forward( + def _attn( self, positions: torch.Tensor, - hidden_states: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, ctx: ForwardContext, out_cache_loc: torch.Tensor, ) -> torch.Tensor: - - if hidden_states.shape[0] == 0: - # Under DP attention the caller concatenates [embeds, hidden_states] - # to width 2*H before attention. Peers with N>0 return an H-wide - # tensor from o_proj; idle ranks must match that invariant so the - # subsequent dense-TP RSAG agrees on the last dim. - 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) - - fused_kv_arg = None + # Active draft first step (= target verify was pure decode + drafter set + # up gather_ids). Other forwards (multi-step decode without gather_ids, + # or first step after a target verify with extends) delegate to base. + if ctx.gather_ids is None or not ctx.forward_mode.is_decode(): + return super()._attn(positions, q, k, v, ctx, out_cache_loc) + + # Active dispatch B: correction + q-slice + DECODE (fused) or post-slice (non-fused). + self._apply_correction(ctx) if ctx.attn_backend.support_kv_cache_prewrite(ctx.forward_mode): - n = q.shape[0] - v_3d = v.view(n, self.num_kv_heads, self.head_dim) - fused_kv_arg = create_fused_set_kv_buffer_arg( - value=v_3d, - layer=self.attn, - out_cache_loc=out_cache_loc, - token_to_kv_pool=ctx.token_to_kv_pool, + q_rope = self._fused_rope_kv_write(positions, q, k, v, ctx, out_cache_loc) + q_rope = q_rope.index_select(0, ctx.gather_ids) + return 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, ) + q, k = self.rotary_emb(positions, q, k) + return self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc).index_select( + 0, ctx.gather_ids + ) - if fused_kv_arg is not None: - n = q.shape[0] - q_rope = torch.empty((n, self.q_size), dtype=q.dtype, device=q.device) - q, k = self.rotary_emb( - positions, - q, - k, - fused_set_kv_buffer_arg=fused_kv_arg, - 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, - ) - 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) + def _apply_correction(self, ctx: ForwardContext) -> None: + """Trim cache_seqlens by ``spec_num_tokens - accept_lengths``. - output, _ = self.o_proj(attn_output) - return output + Idempotent across draft layers via ``ctx.accept_lengths = None``. + No-op when drafter did not populate ``accept_lengths`` (e.g. backends + without KV pre-write). + """ + if ctx.accept_lengths is None: + return + seq_lens_buf = ctx.draft_seq_lens_buf + if seq_lens_buf is None: + return + correction = (ctx.attn_backend.spec_num_tokens - ctx.accept_lengths).to( + seq_lens_buf.dtype + ) + seq_lens_buf[: ctx.bs].sub_(correction) + ctx.accept_lengths = None # --------------------------------------------------------------------------- @@ -348,6 +255,17 @@ def resolve_mlp(self, prefix: str) -> nn.Module: prefix=f"{prefix}.mlp", ) + def _maybe_narrow_residual( + self, + residual: torch.Tensor, + ctx: ForwardContext, + ) -> torch.Tensor: + """Wrapper: align residual with attn output narrowed to [bs, H].""" + 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. + return residual.index_select(0, ctx.gather_ids) + return residual + def forward_low_latency( self, positions: torch.Tensor, @@ -383,9 +301,7 @@ 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) + residual = self._maybe_narrow_residual(residual, ctx) # Fused post-attn allreduce + norm (uses attn tp group) block_scale = None @@ -451,9 +367,7 @@ 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) + residual = self._maybe_narrow_residual(residual, ctx) hidden_states, residual = self.comm_manager.post_attn_comm( hidden_states, residual, ctx ) @@ -583,11 +497,6 @@ class LlamaForCausalLMEagle3(BaseCausalLM): model_cls = Eagle3LlamaModel - # Eagle3 catch-up pre-slices q to active row before attn; trim must pair. - # TODO: remove together with the base flag once Qwen NextN / DeepSeek V3 - # NextN also pre-slice. - pre_attention_trim: bool = True - def __init__( self, config: LlamaConfig, From 6fbab67283600e17691ca30d047f55560a061328 Mon Sep 17 00:00:00 2001 From: rjzhb Date: Tue, 9 Jun 2026 01:13:11 +0000 Subject: [PATCH 2/8] update Signed-off-by: rjzhb --- python/tokenspeed/runtime/models/llama_eagle3.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py index b633bfbb8..2db130705 100644 --- a/python/tokenspeed/runtime/models/llama_eagle3.py +++ b/python/tokenspeed/runtime/models/llama_eagle3.py @@ -93,8 +93,9 @@ def _attn( # Active dispatch B: correction + q-slice + DECODE (fused) or post-slice (non-fused). self._apply_correction(ctx) if ctx.attn_backend.support_kv_cache_prewrite(ctx.forward_mode): - q_rope = self._fused_rope_kv_write(positions, q, k, v, ctx, out_cache_loc) - q_rope = q_rope.index_select(0, ctx.gather_ids) + q_rope = self._fused_rope_kv_write( + positions, q, k, v, ctx, out_cache_loc + ).index_select(0, ctx.gather_ids) return ctx.attn_backend.forward( q_rope, None, From 399b793e959079297cf835023c6d3ddab6c38b9d Mon Sep 17 00:00:00 2001 From: rjzhb Date: Tue, 9 Jun 2026 03:54:14 +0000 Subject: [PATCH 3/8] feat(spec-decode): extend Llama Eagle3 dispatch B to prefill catch-up Signed-off-by: rjzhb --- .../runtime/execution/drafter/eagle.py | 6 +++- .../tokenspeed/runtime/models/llama_eagle3.py | 36 ++++++++----------- 2 files changed, 20 insertions(+), 22 deletions(-) diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py index 6466c8290..9264b1788 100644 --- a/python/tokenspeed/runtime/execution/drafter/eagle.py +++ b/python/tokenspeed/runtime/execution/drafter/eagle.py @@ -218,7 +218,11 @@ 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() + # Llama Eagle3 narrows for prefill catch-up too; Qwen/DeepSeek do not. + draft_first_step_reduce = forward_mode.is_decode() or ( + isinstance(self.draft_model_runner.model, LlamaForCausalLMEagle3) + and forward_mode.is_target_verify() + ) draft_first_mode = ( ForwardMode.DRAFT_EXTEND diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py index 2db130705..e6494ef13 100644 --- a/python/tokenspeed/runtime/models/llama_eagle3.py +++ b/python/tokenspeed/runtime/models/llama_eagle3.py @@ -84,10 +84,10 @@ def _attn( ctx: ForwardContext, out_cache_loc: torch.Tensor, ) -> torch.Tensor: - # Active draft first step (= target verify was pure decode + drafter set - # up gather_ids). Other forwards (multi-step decode without gather_ids, - # or first step after a target verify with extends) delegate to base. - if ctx.gather_ids is None or not ctx.forward_mode.is_decode(): + # Active draft first step (drafter set up gather_ids + accept_lengths). + # Covers both decode catch-up and prefill catch-up; multi-step decode + # delegates to base. + if ctx.accept_lengths is None: return super()._attn(positions, q, k, v, ctx, out_cache_loc) # Active dispatch B: correction + q-slice + DECODE (fused) or post-slice (non-fused). @@ -113,22 +113,17 @@ def _attn( ) def _apply_correction(self, ctx: ForwardContext) -> None: - """Trim cache_seqlens by ``spec_num_tokens - accept_lengths``. - - Idempotent across draft layers via ``ctx.accept_lengths = None``. - No-op when drafter did not populate ``accept_lengths`` (e.g. backends - without KV pre-write). - """ - if ctx.accept_lengths is None: - return + """Trim decode rows' cache_seqlens by ``spec_num_tokens - accept_lengths``.""" seq_lens_buf = ctx.draft_seq_lens_buf - if seq_lens_buf is None: + if seq_lens_buf is None or ctx.accept_lengths is None: return - correction = (ctx.attn_backend.spec_num_tokens - ctx.accept_lengths).to( - seq_lens_buf.dtype - ) - seq_lens_buf[: ctx.bs].sub_(correction) - ctx.accept_lengths = None + num_extends = ctx.num_extends + if num_extends >= ctx.bs: + return + correction = ( + ctx.attn_backend.spec_num_tokens - ctx.accept_lengths[num_extends:] + ).to(seq_lens_buf.dtype) + seq_lens_buf[num_extends : ctx.bs].sub_(correction) # --------------------------------------------------------------------------- @@ -261,9 +256,8 @@ def _maybe_narrow_residual( residual: torch.Tensor, ctx: ForwardContext, ) -> torch.Tensor: - """Wrapper: align residual with attn output narrowed to [bs, H].""" - 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. + """Align residual with attn output narrowed to [bs, H].""" + if ctx.accept_lengths is not None and not ctx.forward_mode.is_idle(): return residual.index_select(0, ctx.gather_ids) return residual From 90ec73a083e66905147e5992382f1829f74c2380 Mon Sep 17 00:00:00 2001 From: rjzhb Date: Tue, 9 Jun 2026 04:04:48 +0000 Subject: [PATCH 4/8] update Signed-off-by: rjzhb --- python/tokenspeed/runtime/execution/drafter/eagle.py | 1 + 1 file changed, 1 insertion(+) diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py index 9264b1788..99e8e8b24 100644 --- a/python/tokenspeed/runtime/execution/drafter/eagle.py +++ b/python/tokenspeed/runtime/execution/drafter/eagle.py @@ -27,6 +27,7 @@ from tokenspeed_kernel.ops.sampling.cute_dsl import argmax as cute_argmax from typing_extensions import override +from python.tokenspeed.runtime.models.llama_eagle3 import LlamaForCausalLMEagle3 from tokenspeed.runtime.execution.cache_loc_kernel import ( compute_out_cache_loc_uniform, ) From 2ca7ebe69370ce25e6f24e8f0aadcbc6e1b3ce38 Mon Sep 17 00:00:00 2001 From: rjzhb Date: Tue, 9 Jun 2026 19:51:17 +0000 Subject: [PATCH 5/8] fix(spec-decode): correct LlamaForCausalLMEagle3 import path Signed-off-by: rjzhb --- python/tokenspeed/runtime/execution/drafter/eagle.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py index 7a5494b0a..db79965a1 100644 --- a/python/tokenspeed/runtime/execution/drafter/eagle.py +++ b/python/tokenspeed/runtime/execution/drafter/eagle.py @@ -27,7 +27,6 @@ from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax from typing_extensions import override -from python.tokenspeed.runtime.models.llama_eagle3 import LlamaForCausalLMEagle3 from tokenspeed.runtime.execution.cache_loc_kernel import ( compute_out_cache_loc_uniform, ) @@ -37,6 +36,7 @@ CaptureHiddenMode, ForwardMode, ) +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 From cc698bf429e110c1fc36f5ac192da54d3d14c46e Mon Sep 17 00:00:00 2001 From: rjzhb Date: Tue, 9 Jun 2026 20:04:56 +0000 Subject: [PATCH 6/8] fix(spec-decode): cover EXTEND/MIXED catch-up in dispatch B flag broaden Signed-off-by: rjzhb --- python/tokenspeed/runtime/execution/drafter/eagle.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py index db79965a1..ef96a782b 100644 --- a/python/tokenspeed/runtime/execution/drafter/eagle.py +++ b/python/tokenspeed/runtime/execution/drafter/eagle.py @@ -219,10 +219,11 @@ 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) - # Llama Eagle3 narrows for prefill catch-up too; Qwen/DeepSeek do not. + # 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 forward_mode.is_target_verify() + and not forward_mode.is_idle() ) draft_first_mode = ( From 5ac5d7976fb8a3710899b3ff942e7c4683375c76 Mon Sep 17 00:00:00 2001 From: rjzhb Date: Wed, 10 Jun 2026 19:11:01 +0000 Subject: [PATCH 7/8] fix(spec-decode): fall back when fused KV prewrite arg is None Signed-off-by: rjzhb --- python/tokenspeed/runtime/models/llama.py | 57 ++++++++++++------- .../tokenspeed/runtime/models/llama_eagle3.py | 30 +++++----- 2 files changed, 52 insertions(+), 35 deletions(-) diff --git a/python/tokenspeed/runtime/models/llama.py b/python/tokenspeed/runtime/models/llama.py index 0fc39e1f1..3f99c7eb6 100644 --- a/python/tokenspeed/runtime/models/llama.py +++ b/python/tokenspeed/runtime/models/llama.py @@ -222,41 +222,56 @@ def _attn( ) -> torch.Tensor: """RoPE + attention (pre-o_proj), with optional fused KV pre-write. - When the backend supports KV pre-write, fused rope writes KV directly - into the cache so the attention call can run with ``save_kv_cache=False`` - (saves one kernel launch). Subclasses (e.g. Eagle3 draft head) override - this hook to insert spec-decode behaviour around the same scaffolding. + 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): - q_rope = self._fused_rope_kv_write(positions, q, k, v, ctx, out_cache_loc) - return self.attn( - q_rope, - None, - None, - save_kv_cache=False, - ctx=ctx, - out_cache_loc=out_cache_loc, - ) + 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 _fused_rope_kv_write( + def _build_fused_kv_arg( self, - positions: torch.Tensor, - q: torch.Tensor, - k: torch.Tensor, v: torch.Tensor, ctx: ForwardContext, out_cache_loc: torch.Tensor, - ) -> torch.Tensor: - """Fused RoPE that writes KV into cache and returns the rope'd Q.""" - n = q.shape[0] - fused_kv_arg = create_fused_set_kv_buffer_arg( + ): + """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, diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py index e6494ef13..ee8942975 100644 --- a/python/tokenspeed/runtime/models/llama_eagle3.py +++ b/python/tokenspeed/runtime/models/llama_eagle3.py @@ -93,20 +93,22 @@ def _attn( # Active dispatch B: correction + q-slice + DECODE (fused) or post-slice (non-fused). self._apply_correction(ctx) if ctx.attn_backend.support_kv_cache_prewrite(ctx.forward_mode): - q_rope = self._fused_rope_kv_write( - positions, q, k, v, ctx, out_cache_loc - ).index_select(0, ctx.gather_ids) - return 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, - ) + 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 + ).index_select(0, ctx.gather_ids) + return 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, + ) q, k = self.rotary_emb(positions, q, k) return self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc).index_select( 0, ctx.gather_ids From 0c8345f58a7bdf1239ae59c2586e2e44f3c7c6ab Mon Sep 17 00:00:00 2001 From: rjzhb Date: Wed, 10 Jun 2026 20:45:44 +0000 Subject: [PATCH 8/8] fix(spec-decode): only trim cache_seqlens on the sliced prewrite path Signed-off-by: rjzhb --- .../tokenspeed/runtime/models/llama_eagle3.py | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py index ee8942975..258f7c703 100644 --- a/python/tokenspeed/runtime/models/llama_eagle3.py +++ b/python/tokenspeed/runtime/models/llama_eagle3.py @@ -65,14 +65,13 @@ class LlamaAttention(BaseLlamaAttention): """Eagle3 draft head attention. - Inherits ``__init__`` (with ``qkv_input_size=2*hidden_size`` to accommodate - the [embed || hidden] concat) and ``forward`` (= empty-shape handling + - qkv projection + o_proj scaffolding) from base. - - Overrides ``_attn`` to add the draft-first-step dispatch B: trim - cache_seqlens, slice q to one live row per request, and route the attention - call as ``DECODE`` (fused KV pre-write path) or post-slice attn_output - (non-fused fallback). Inactive draft steps delegate to base. + Inherits ``__init__`` (with ``qkv_input_size=2*hidden_size`` for the + [embed || hidden] concat) and ``forward`` (= qkv_proj + o_proj scaffolding) + from base. Overrides ``_attn`` so the draft's first step skips dead + catch-up rows: on backends that support fused KV pre-write, q is sliced + to one live row per request and dispatched as DECODE; otherwise the + fallback runs the full N-row attn and post-slices the output. Inactive + draft steps delegate to base. """ def _attn( @@ -90,11 +89,13 @@ def _attn( if ctx.accept_lengths is None: return super()._attn(positions, q, k, v, ctx, out_cache_loc) - # Active dispatch B: correction + q-slice + DECODE (fused) or post-slice (non-fused). - self._apply_correction(ctx) 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: + # Trim only on the sliced single-token decode path; the + # post-slice fallback below still runs full N-row attn and + # needs the original seq_lens. + self._apply_correction(ctx) q_rope = self._fused_rope_kv_write( positions, q, k, fused_kv_arg ).index_select(0, ctx.gather_ids)