xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2025-12-14 14:55:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Attention] Fix FlashMLA metadata builder arguments for q_len > 1 (#27368)

Browse Source 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
This commit is contained in:
Matthew Bonanni 2025-10-23 15:58:15 -04:00 committed by GitHub
parent ca76486a16
commit dbfbf9f324
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
vllm/v1/attention/backends/mla/flashmla.py
View File

@ -120,9 +120,13 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):
num_decode_tokens: int, num_decode_tokens: int,
dcp_tot_seq_lens_device: torch.Tensor | None, dcp_tot_seq_lens_device: torch.Tensor | None,
) -> FlashMLADecodeMetadata: ) -> FlashMLADecodeMetadata:
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
# we use the max but all should be the same due to uniform length requirement
max_query_len = query_lens_cpu.max().item()
num_q_tokens_per_head_k = max_query_len * self.num_q_heads // 1
tile_scheduler_metadata, num_splits = get_mla_metadata( tile_scheduler_metadata, num_splits = get_mla_metadata(
seq_lens_device, seq_lens_device,
self.num_q_heads, num_q_tokens_per_head_k,
1, # MQA for the decode path 1, # MQA for the decode path
is_fp8_kvcache=self.is_fp8_kvcache, is_fp8_kvcache=self.is_fp8_kvcache,
) )