fix: size FlashInfer prefill workspace to batch head footprint (#48428)
Signed-off-by: Joe Rowell <joerowell4@gmail.com>
Joe Rowell committed
Jul 15, 2026 at 02:18 UTC
7e950521b39fd201979a888eda82ff0f7ecc4731
1 file changed
+19
vllm/v1/attention/backends/flashinfer.py
+19
@@ -82,6 +82,7 @@ from vllm.v1.kv_cache_interface import (
82
from vllm.v1.utils import CpuGpuBuffer
83
84
FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024
85
+FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM = 16
86
87
FP8_DTYPE = current_platform.fp8_dtype()
88
FP4_DTYPE = torch.uint8
@@ -648,6 +649,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
649
self.disable_split_kv = False
650
651
self.compilation_config = vllm_config.compilation_config
652
+ self.max_num_batched_tokens = (
653
+ vllm_config.scheduler_config.max_num_batched_tokens
654
+ )
655
max_num_pages_per_req = cdiv(
656
self.model_config.max_model_len, self.kv_cache_spec.block_size
657
)
@@ -933,6 +937,21 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
937
buffer_size = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
938
if envs.VLLM_BATCH_INVARIANT:
939
buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT
940
+ else:
941
+ # FlashInfer prefill temp buffers (batch_prefill_tmp_v, ...)
942
+ # scale with the prefill chunk and query-head footprint, NOT
943
+ # context length. The fixed ~394 MiB default is too small for
944
+ # wide-head models at the default 8192-token chunk on some
945
+ # archs (e.g. sm_120), where FlashInfer hard-errors instead of
946
+ # growing. Size to the batch's head footprint; never shrink
947
+ # below the configured default.
948
+ est = (
949
+ self.max_num_batched_tokens
950
+ * self.num_qo_heads
951
+ * self.head_dim
952
+ * FLASHINFER_PREFILL_WORKSPACE_BYTES_PER_ELEM
953
+ )
954
+ buffer_size = max(buffer_size, est)
955
self._workspace_buffer = torch.zeros(
956
buffer_size, dtype=torch.uint8, device=self.device
957
)