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
[Bugfix][Attention] Fix FlashInfer MLA block size logic (#24692)
Signed-off-by: Matthew Bonanni <mbonanni001@gmail.com>
Matthew Bonanni committed
Sep 11, 2025 at 18:39 UTC
d4fd2768ef1cf3db52956a3aa881ef3a5ff72a51
1 file changed
+11
-2
vllm/platforms/cuda.py
+11
-2
@@ -146,6 +146,7 @@ class CudaPlatformBase(Platform):
146
# required block_size.
147
use_flashmla = False
148
use_cutlass_mla = False
149
+ use_flashinfer_mla = False
150
151
if envs.VLLM_ATTENTION_BACKEND is None:
152
# Default case
@@ -164,6 +165,8 @@ class CudaPlatformBase(Platform):
165
use_flashmla = (envs.VLLM_ATTENTION_BACKEND == "FLASHMLA")
166
use_cutlass_mla = (
167
envs.VLLM_ATTENTION_BACKEND == "CUTLASS_MLA")
168
+ use_flashinfer_mla = (
169
+ envs.VLLM_ATTENTION_BACKEND == "FLASHINFER_MLA")
170
171
from vllm.attention.ops.flashmla import is_flashmla_supported
172
if use_flashmla and is_flashmla_supported()[0] \
@@ -176,6 +179,11 @@ class CudaPlatformBase(Platform):
179
cache_config.block_size = 128
180
logger.info("Forcing kv cache block size to 128 for "
181
"CUTLASS_MLA backend.")
182
+ if use_flashinfer_mla and cache_config.block_size not in [32, 64]:
183
+ cache_config.block_size = 64
184
+ logger.info(
185
+ "Forcing kv cache block size to 64 for FlashInferMLA "
186
+ "backend.")
187
188
# lazy import to avoid circular import
189
from vllm.config import CUDAGraphMode
@@ -228,8 +236,9 @@ class CudaPlatformBase(Platform):
236
use_cutlassmla = selected_backend == _Backend.CUTLASS_MLA or (
237
selected_backend is None and cls.is_device_capability(100)
238
and block_size == 128)
231
- use_flashinfermla = (selected_backend == _Backend.FLASHINFER_MLA
232
- and cls.has_device_capability(100))
239
+ use_flashinfermla = selected_backend == _Backend.FLASHINFER_MLA or (
240
+ selected_backend is None and cls.is_device_capability(100)
241
+ and block_size in [32, 64])
242
use_flashmla = selected_backend in [
243
_Backend.FLASHMLA, _Backend.FLASHMLA_VLLM_V1
244
] or (selected_backend is None and is_flashmla_supported()[0])