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],