From 73cd7e25aee261b27f5f9fea06fdf7f5a806a09b Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Tue, 5 May 2026 22:57:37 +0000 Subject: [PATCH] [Attention][TokenSpeed MLA] Warm up BF16 prefill compile, drop seq_lens computation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pre-JIT both BF16 and FP8 prefill kernels at backend init since the dtype isn't visible from `__init__` — depends on `use_prefill_query_quantization`. Move the per-forward `seq_lens` computation into `prepare_metadata` and document the cuda-graph padding interaction with `query_start_loc`. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Yongye Zhu --- .../backends/mla/prefill/tokenspeed_mla.py | 28 +++++++++++++------ 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py index f3a335ed94f..3fa2b2cb1ca 100644 --- a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py @@ -64,24 +64,34 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): layer_names=layer_names, ) - # Pre-JIT the kernel for the FP8 prefill shape so the first forward - # pass doesn't pay the compile cost. warmup_compile_prefill is + # 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, so calling it once per layer instantiation is fine. + # process-wide. Whether prefill runs FP8 or BF16 depends on + # `use_prefill_query_quantization`, which we can't see here, so warm up + # both. from tokenspeed_mla import warmup_compile_prefill - warmup_compile_prefill( - q_dtype=torch.float8_e4m3fn, - d_qk=qk_nope_head_dim + qk_rope_head_dim, - d_v=v_head_dim, - enable_pdl=False, - ) + for q_dtype in (torch.bfloat16, torch.float8_e4m3fn): + warmup_compile_prefill( + q_dtype=q_dtype, + d_qk=qk_nope_head_dim + qk_rope_head_dim, + d_v=v_head_dim, + enable_pdl=False, + ) def prepare_metadata( self, prefill_metadata: "MLACommonPrefillMetadata", ) -> None: super().prepare_metadata(prefill_metadata) + # Kernel signature requires `seq_lens` but the implementation never reads + # it (per-batch lengths are derived from `cum_seq_lens` diffs); compute + # for parity with trtllm_ragged. cuda-graph padding in + # `query_start_loc` is saturated to `total_num_tokens` + # (gpu_model_runner.py:1905), so trailing diffs are 0 and padded batches + # are kernel no-ops — same reason trtllm passes the padded length as + # batch_size directly. self._query_seq_lens = ( prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1] )