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 13, 20:45 UTC
[Bug] Batch invariant: Fix flash attn MLA `RuntimeError: scheduler_metadata must have shape (metadata_size)` (#27884)
Wentao Ye committed
Nov 4, 2025 at 01:05 UTC
7e4be741044bfead91afc418100ff9a4d804bf7f
2 files changed
+5
-3
vllm/model_executor/layers/batch_invariant.py
+2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import contextlib
4
+import functools
5
import os
6
from collections import namedtuple
7
from collections.abc import Callable
@@ -846,6 +847,7 @@ def get_batch_invariant_attention_block_size() -> AttentionBlockSize:
847
return AttentionBlockSize(block_m=16, block_n=16)
848
849
850
+@functools.cache
851
def vllm_is_batch_invariant():
852
env_key = "VLLM_BATCH_INVARIANT"
853
is_overridden = False
vllm/v1/attention/backends/mla/flashattn_mla.py
+3
-3
@@ -163,6 +163,9 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
163
# we only set num_splits when using cuda graphs.
164
max_num_splits = self.max_num_splits
165
166
+ if vllm_is_batch_invariant():
167
+ max_num_splits = 1
168
+
169
scheduler_metadata = self._schedule_decode(
170
num_reqs=seq_lens_cpu.numel(),
171
cu_query_lens=query_start_loc_device,
@@ -188,9 +191,6 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
191
self.scheduler_metadata[n:] = 0
192
scheduler_metadata = self.scheduler_metadata[:n]
193
191
- if vllm_is_batch_invariant():
192
- max_num_splits = 1
193
-
194
metadata = FlashAttnMLADecodeMetadata(
195
block_table=block_table_tensor,
196
seq_lens=seq_lens_device,