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 6, 17:00 UTC
[FlashInfer] Revert block_size 16 + head_size 256 workaround on Blackwell (#36987)
Signed-off-by: Vadim Gimpelson <vadim.gimpelson@gmail.com>
Vadim Gimpelson committed
Mar 16, 2026 at 13:04 UTC
8374387bd8ef747ea331d84e911cdaaed4eb7124
2 files changed
-21
vllm/model_executor/models/config.py
-12
@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING
6
7
from vllm.logger import init_logger
8
from vllm.model_executor.models import ModelRegistry
9
-from vllm.platforms import current_platform
9
from vllm.utils.math_utils import cdiv, round_up
10
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
11
from vllm.v1.attention.backends.registry import AttentionBackendEnum
@@ -148,17 +147,6 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
147
).page_size_bytes
148
else:
149
kernel_block_alignment_size = 16
151
- if (
152
- current_platform.is_device_capability_family(100)
153
- and model_config.get_head_size() == 256
154
- and (
155
- attention_config.backend is None
156
- or attention_config.backend == AttentionBackendEnum.FLASHINFER
157
- )
158
- ):
159
- # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that`
160
- # head size 256 and block size 16 is not supported on blackwell.
161
- kernel_block_alignment_size = 32
150
attn_page_size_1_token = FullAttentionSpec(
151
block_size=1,
152
num_kv_heads=model_config.get_num_kv_heads(parallel_config),
vllm/v1/attention/backends/flashinfer.py
-9
@@ -630,15 +630,6 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
630
self.paged_kv_indices = self._make_buffer(max_num_pages)
631
self.paged_kv_last_page_len = self._make_buffer(max_num_reqs)
632
633
- if self.head_dim == 256 and current_platform.is_device_capability_family(100):
634
- # https://github.com/flashinfer-ai/flashinfer/issues/1993 reports that
635
- # head size 256 and block size 16 is not supported on blackwell.
636
- assert kv_cache_spec.block_size != 16, (
637
- "There is a bug in FlashInfer "
638
- "block_size 16 head size 256 support. Please avoid this combination by "
639
- "passing --block-size 32 or --block-size 64."
640
- )
641
-
633
def _make_buffer(
634
self, *size: int | torch.SymInt, dtype: torch.dtype = torch.int32
635
) -> CpuGpuBuffer: