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 4, 10:00 UTC
[DSV4] Add knob to enable pre-attn gemm (#41443)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
May 1, 2026 at 15:23 UTC
bcf5cac9fb956788f649d1f5297b74c886a9d6d3
3 files changed
+24
-3
vllm/envs.py
+12
@@ -245,6 +245,7 @@ if TYPE_CHECKING:
245
VLLM_DEBUG_WORKSPACE: bool = False
246
VLLM_DISABLE_SHARED_EXPERTS_STREAM: bool = False
247
VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
248
+ VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD: int = 4096
249
VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
250
VLLM_USE_V2_MODEL_RUNNER: bool = False
251
VLLM_LOG_MODEL_INSPECTION: bool = False
@@ -1678,6 +1679,17 @@ environment_variables: dict[str, Callable[[], Any]] = {
1679
"VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD": lambda: int(
1680
int(os.getenv("VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD", 256))
1681
),
1682
+ # Token-count cutoff for multi-stream overlap of the attention input
1683
+ # GEMM with auxiliary GEMMs (e.g. fused_wqa_wkv overlapped with indexer
1684
+ # weights / kv-score projections in DeepSeek-V4). At or below this many
1685
+ # tokens the FP8 main GEMM has idle SMs to share with the bf16 aux GEMMs
1686
+ # and overlap is a 5-45% win; above it the FP8 GEMM saturates the device
1687
+ # and the cross-stream sync becomes pure overhead. Set to 0 to disable
1688
+ # the multi-stream path entirely. Empirical crossover on B300 (148 SMs)
1689
+ # is ~4096; B200 (132 SMs) is expected ~3072.
1690
+ "VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD": lambda: int(
1691
+ os.getenv("VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD", "4096")
1692
+ ),
1693
# Format for saving torch.compile cache artifacts
1694
# - "binary": saves as binary file
1695
# Safe for multiple vllm serve processes accessing the same torch compile cache.
vllm/model_executor/layers/deepseek_v4_attention.py
+3
@@ -13,6 +13,7 @@ import torch.nn as nn
13
import torch.nn.functional as F
14
from transformers import DeepseekV2Config, DeepseekV3Config
15
16
+import vllm.envs as envs
17
from vllm.model_executor.layers.linear import (
18
ReplicatedLinear,
19
)
@@ -385,6 +386,8 @@ class DeepseekV4MultiHeadLatentAttentionWrapper(PluggableLayer):
386
self.ln_events[0],
387
self.ln_events[1:4],
388
self.aux_stream_list[:3],
389
+ enable=hidden_states.shape[0]
390
+ <= envs.VLLM_MULTI_STREAM_GEMM_TOKEN_THRESHOLD,
391
)
392
393
return qr_kv, kv_score, indexer_kv_score, indexer_weights
vllm/utils/multi_stream_utils.py
+9
-3
@@ -64,6 +64,7 @@ def execute_in_parallel(
64
start_event: torch.cuda.Event,
65
done_events: list[torch.cuda.Event],
66
aux_streams: list[torch.cuda.Stream] | None = None,
67
+ enable: bool = False,
68
) -> tuple[Any, list[Any]]:
69
"""Run default_fn on the current stream and aux_fns concurrently on
70
aux_streams.
@@ -74,8 +75,9 @@ def execute_in_parallel(
75
76
start_event fans out from the current stream to every launched aux stream;
77
done_events[i] is recorded after aux_fns[i] so the current stream joins
77
- before returning. When aux_streams is None, all aux_fns run sequentially
78
- on the current stream.
78
+ before returning. Falls back to sequential execution on the current stream
79
+ when aux_streams is None or enable is False; in that case default_fn runs
80
+ first, then aux_fns in order.
81
82
Args:
83
default_fn: Callable for the default (current) stream.
@@ -86,13 +88,17 @@ def execute_in_parallel(
88
corresponding aux_fn. Length must match aux_fns.
89
aux_streams: Per-aux CUDA streams. Length must match aux_fns.
90
Multi-stream is disabled when None.
91
+ enable: Opt-in switch for the multi-stream path. Defaults to False,
92
+ so callers that pass aux_streams must also pass enable=True
93
+ (typically gated by an env var) to actually overlap. When False,
94
+ execution falls back to sequential on the current stream.
95
96
Returns:
97
Tuple of (default_result, aux_results) where aux_results[i] is the
98
result of aux_fns[i] (or None when skipped).
99
"""
100
aux_results: list[Any]
95
- if aux_streams is None:
101
+ if aux_streams is None or not enable:
102
default_result = default_fn()
103
aux_results = [fn() if fn is not None else None for fn in aux_fns]
104
return default_result, aux_results