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, 18:15 UTC
[FlashInfer] Avoid FlashInfer block_size 16 + head_size 256 on blackwell (#27994)
Signed-off-by: Chen Zhang <zhangch99@outlook.com>
Chen Zhang committed
Nov 5, 2025 at 09:25 UTC
c765f0b443c2de886a5771efcbc154ee84c2c19d
2 files changed
+21
vllm/model_executor/models/config.py
+12
@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING
7
import vllm.envs as envs
8
from vllm.logger import init_logger
9
from vllm.model_executor.models import ModelRegistry
10
+from vllm.platforms import current_platform
11
from vllm.utils.math_utils import cdiv, round_up
12
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
13
from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, MLAAttentionSpec
@@ -356,6 +357,17 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
357
).page_size_bytes
358
else:
359
kernel_block_alignment_size = 16
360
+ if (
361
+ current_platform.is_device_capability(100)
362
+ and model_config.get_head_size() == 256
363
+ and (
364
+ envs.VLLM_ATTENTION_BACKEND is None
365
+ or envs.VLLM_ATTENTION_BACKEND == "FLASHINFER"
366
+ )
367
+ ):
368
+ # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that`
369
+ # head size 256 and block size 16 is not supported on blackwell.
370
+ kernel_block_alignment_size = 32
371
attn_page_size_1_token = FullAttentionSpec(
372
block_size=1,
373
num_kv_heads=model_config.get_num_kv_heads(parallel_config),
vllm/v1/attention/backends/flashinfer.py
+9
@@ -402,6 +402,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
402
)
403
self.paged_kv_last_page_len_np = self.paged_kv_last_page_len_cpu.numpy()
404
405
+ if self.head_dim == 256 and current_platform.is_device_capability(100):
406
+ # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that
407
+ # head size 256 and block size 16 is not supported on blackwell.
408
+ assert kv_cache_spec.block_size != 16, (
409
+ "There is a bug in FlashInfer "
410
+ "block_size 16 head size 256 support. Please avoid this combination by "
411
+ "passing --block-size 32 or --block-size 64."
412
+ )
413
+
414
def _get_workspace_buffer(self):
415
if self._workspace_buffer is None:
416
buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE