[Attention][TokenSpeed MLA] Also warm up prefill kernel from decode impl

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) <[email protected]>

Signed-off-by: Yongye Zhu <[email protected]>

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):

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