From 58c8a5eaa505f7b6799eaf535874f8167bd5b3fd Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 6 May 2026 07:45:16 +0000 Subject: [PATCH] [Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The prefill backend may be paired with flash_attn / trtllm in production — in that case the prefill backend's __init__ never runs and the prefill kernel's first call pays a 1.5–2 minute JIT cost. Add the same idempotent `warmup_compile_prefill` invocation to TokenspeedMLAImpl.__init__ (the decode-side backend, always present when tokenspeed is selected). The function dedupes by config key, so the double call is a no-op when both backends are tokenspeed. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Yongye Zhu --- .../backends/mla/prefill/tokenspeed_mla.py | 8 ++------ vllm/v1/attention/backends/mla/tokenspeed_mla.py | 13 +++++++++++++ 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py index 135e76795b3..d6e4fca172a 100644 --- a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py @@ -81,12 +81,8 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): vllm_config=vllm_config, ) - # Pre-JIT the kernel for both BF16 and FP8 prefill shapes so the first - # forward pass doesn't pay the compile cost. warmup_compile_prefill is - # idempotent: each (q_dtype, d_qk, d_v) is compiled at most once - # process-wide. Whether prefill runs FP8 or BF16 depends on - # `use_prefill_query_quantization`, which we can't see here, so warm up - # both. + # Pre-JIT BF16 and FP8 prefill kernels. Idempotent — also called from + # TokenspeedMLAImpl.__init__; second call is a no-op. from tokenspeed_mla import warmup_compile_prefill for q_dtype in (torch.bfloat16, torch.float8_e4m3fn): diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index 23aac770f50..6c8dedd77f2 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -187,6 +187,19 @@ class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]): self.softmax_scale: float | None = None self.output_scale: float | None = None + # Pre-JIT BF16 and FP8 prefill kernels here too — decode impl always + # runs when tokenspeed is selected, prefill backend may not (user can + # pair with flash_attn / trtllm). Idempotent. + from tokenspeed_mla import warmup_compile_prefill + + for q_dtype in (torch.bfloat16, torch.float8_e4m3fn): + warmup_compile_prefill( + q_dtype=q_dtype, + d_qk=self.qk_nope_head_dim + self.qk_rope_head_dim, + d_v=self.v_head_dim, + enable_pdl=False, + ) + def forward_mqa( self, q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],