A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Attention] Fix FlashMLA metadata builder arguments for q_len > 1 (#27368)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Oct 23, 2025 at 15:58 UTC
dbfbf9f32445ffadcbb731adef00b4a393612be9
1 file changed
+5
-1
vllm/v1/attention/backends/mla/flashmla.py
+5
-1
@@ -120,9 +120,13 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
120
num_decode_tokens: int,
121
dcp_tot_seq_lens_device: torch.Tensor | None,
122
) -> FlashMLADecodeMetadata:
123
+ query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
124
+ # we use the max but all should be the same due to uniform length requirement
125
+ max_query_len = query_lens_cpu.max().item()
126
+ num_q_tokens_per_head_k = max_query_len * self.num_q_heads // 1
127
tile_scheduler_metadata, num_splits = get_mla_metadata(
128
seq_lens_device,
125
- self.num_q_heads,
129
+ num_q_tokens_per_head_k,
130
1, # MQA for the decode path
131
is_fp8_kvcache=self.is_fp8_kvcache,
132
)