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, 10:45 UTC
[Feature] Allow configuring FlashInfer workspace size (#28269)
Signed-off-by: Max Hu <hyoung2991@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Max Hu committed
Nov 11, 2025 at 18:32 UTC
412e153df557bbae541363ac4abde879a6d84488
3 files changed
+16
-12
vllm/envs.py
+6
@@ -159,6 +159,7 @@ if TYPE_CHECKING:
159
VLLM_USE_FLASHINFER_MOE_FP8: bool = False
160
VLLM_USE_FLASHINFER_MOE_FP4: bool = False
161
VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency"] = "latency"
162
+ VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024
163
VLLM_XGRAMMAR_CACHE_MB: int = 0
164
VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
165
VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
@@ -1237,6 +1238,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
1238
"VLLM_FLASHINFER_MOE_BACKEND": env_with_choices(
1239
"VLLM_FLASHINFER_MOE_BACKEND", "latency", ["throughput", "latency"]
1240
),
1241
+ # Control the workspace buffer size for the FlashInfer backend.
1242
+ "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int(
1243
+ os.getenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(394 * 1024 * 1024))
1244
+ ),
1245
# Control the maximum number of tokens per expert supported by the
1246
# NVFP4 MoE CUTLASS Kernel. This value is used to create a buffer for
1247
# the blockscale tensor of activations NVFP4 Quantization.
@@ -1583,6 +1588,7 @@ def compute_hash() -> str:
1588
"VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8",
1589
"VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS",
1590
"VLLM_USE_FLASHINFER_MOE_MXFP4_BF16",
1591
+ "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE",
1592
"VLLM_USE_CUDNN_PREFILL",
1593
"VLLM_USE_TRTLLM_RAGGED_DEEPSEEK_PREFILL",
1594
"VLLM_USE_TRTLLM_ATTENTION",
vllm/v1/attention/backends/flashinfer.py
+3
-3
@@ -16,6 +16,7 @@ from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache
16
from flashinfer.prefill import trtllm_batch_context_with_kv_cache
17
from flashinfer.utils import FP4Tensor
18
19
+from vllm import envs
20
from vllm.attention.backends.abstract import (
21
AttentionBackend,
22
AttentionImpl,
@@ -55,7 +56,6 @@ from vllm.v1.attention.backends.utils import (
56
)
57
from vllm.v1.kv_cache_interface import AttentionSpec
58
58
-FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024
59
FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024
60
61
FP8_DTYPE = current_platform.fp8_dtype()
@@ -70,7 +70,7 @@ def _get_trtllm_gen_workspace_buffer():
70
global trtllm_gen_workspace_buffer
71
if trtllm_gen_workspace_buffer is None:
72
trtllm_gen_workspace_buffer = torch.zeros(
73
- FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device="cuda"
73
+ envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device="cuda"
74
)
75
return trtllm_gen_workspace_buffer
76
@@ -414,7 +414,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
414
415
def _get_workspace_buffer(self):
416
if self._workspace_buffer is None:
417
- buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE
417
+ buffer_size = envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE
418
if vllm_is_batch_invariant():
419
buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT
420
self._workspace_buffer = torch.zeros(
vllm/v1/attention/backends/mla/common.py
+7
-9
@@ -196,8 +196,8 @@ from typing import ClassVar, Generic, TypeVar
196
import torch
197
from tqdm import tqdm
198
199
-import vllm.envs as envs
199
from vllm import _custom_ops as ops
200
+from vllm import envs
201
from vllm._aiter_ops import rocm_aiter_ops
202
from vllm.attention.backends.abstract import (
203
AttentionBackend,
@@ -453,12 +453,6 @@ def use_trtllm_ragged_deepseek_prefill() -> bool:
453
)
454
455
456
-# Currently 394MB, this can be tuned based on GEMM sizes used.
457
-# Chosen to be the same as sglang:
458
-# https://github.com/sgl-project/sglang/blob/766392c6bda2558b61ce6d1c1bfd8081a549e1f1/python/sglang/global_config.py#L37
459
-FLASHINFER_WORKSPACE_BUFFER_SIZE = 394 * 1024 * 1024
460
-
461
-
456
class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
457
"""
458
NOTE: Please read the comment at the top of the file before trying to
@@ -590,7 +584,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
584
585
if self._use_fi_prefill:
586
self._workspace_buffer = torch.empty(
593
- FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device=device
587
+ envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
588
+ dtype=torch.uint8,
589
+ device=device,
590
)
591
592
self._fi_prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
@@ -602,7 +598,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
598
599
if self._use_trtllm_ragged_prefill:
600
self._workspace_buffer = torch.empty(
605
- FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device=device
601
+ envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,
602
+ dtype=torch.uint8,
603
+ device=device,
604
)
605
606
if self._use_cudnn_prefill: