diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py index 81208140d..d51618813 100644 --- a/python/tokenspeed/runtime/execution/drafter/eagle.py +++ b/python/tokenspeed/runtime/execution/drafter/eagle.py @@ -36,6 +36,7 @@ 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 @@ -226,9 +227,16 @@ 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. + + # 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 ) diff --git a/python/tokenspeed/runtime/models/base/causal_lm.py b/python/tokenspeed/runtime/models/base/causal_lm.py index 24b1d05fa..1cd84880c 100644 --- a/python/tokenspeed/runtime/models/base/causal_lm.py +++ b/python/tokenspeed/runtime/models/base/causal_lm.py @@ -44,6 +44,10 @@ 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_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py index a64745e47..9c160f53b 100644 --- a/python/tokenspeed/runtime/models/llama_eagle3.py +++ b/python/tokenspeed/runtime/models/llama_eagle3.py @@ -583,6 +583,11 @@ 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,