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 14, 05:30 UTC
[Bugfix] Add fake mode around passes (#23349)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Aug 28, 2025 at 08:25 UTC
db74d604900d397e4ee524f93bcb256537679ce4
6 files changed
+66
-41
vllm/compilation/activation_quant_fusion.py
+2
@@ -10,6 +10,7 @@ from vllm.config import VllmConfig
10
from vllm.logger import init_logger
11
from vllm.platforms import current_platform
12
13
+from .inductor_pass import enable_fake_mode
14
from .vllm_inductor_pass import VllmInductorPass
15
16
logger = init_logger(__name__)
@@ -61,6 +62,7 @@ class ActivationQuantFusionPass(VllmInductorPass):
62
https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980
63
"""
64
65
+ @enable_fake_mode
66
def __init__(self, config: VllmConfig):
67
super().__init__(config)
68
vllm/compilation/collective_fusion.py
+6
@@ -19,6 +19,7 @@ from vllm.logger import init_logger
19
from vllm.platforms import current_platform
20
from vllm.utils import direct_register_custom_op
21
22
+from .inductor_pass import enable_fake_mode
23
from .vllm_inductor_pass import VllmInductorPass
24
25
FP8_DTYPE = current_platform.fp8_dtype()
@@ -349,6 +350,7 @@ class AllGatherCutlassScaledMMPattern(BasePattern):
350
351
class AsyncTPPass(VllmInductorPass):
352
353
+ @enable_fake_mode
354
def __init__(self, config: VllmConfig):
355
super().__init__(config)
356
@@ -1121,6 +1123,10 @@ class AllReduceFusionPass(VllmInductorPass):
1123
# in fallback path, when we don't use flashinfer
1124
fuse_rms_quant=config.compilation_config.pass_config.enable_fusion)
1125
1126
+ self.register_patterns()
1127
+
1128
+ @enable_fake_mode
1129
+ def register_patterns(self):
1130
for epsilon in [1e-5, 1e-6]:
1131
AllReduceFusedRMSNormStaticQuantFP8Pattern(
1132
epsilon,
vllm/compilation/fusion.py
+2
@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
from vllm.platforms import current_platform
18
19
from .fx_utils import find_getitem_maybe
20
+from .inductor_pass import enable_fake_mode
21
from .multi_output_match import MultiOutputMatch
22
from .vllm_inductor_pass import VllmInductorPass
23
@@ -528,6 +529,7 @@ class FusionPass(VllmInductorPass):
529
cls._instance.pass_config = config.compilation_config.pass_config
530
return cls._instance
531
532
+ @enable_fake_mode
533
def __init__(self, config: VllmConfig):
534
assert self.__class__._instance is None, \
535
"FusionPass singleton instance already exists"
vllm/compilation/fusion_attn.py
+34
-41
@@ -7,8 +7,6 @@ import torch
7
import torch._inductor.pattern_matcher as pm
8
from torch._higher_order_ops.auto_functionalize import auto_functionalized
9
from torch._inductor.pattern_matcher import PatternMatcherPass
10
-from torch._subclasses.fake_tensor import (FakeTensorMode,
11
- unset_fake_temporarily)
10
11
from vllm.attention import Attention
12
from vllm.config import VllmConfig, get_layers_from_vllm_config
@@ -19,6 +17,7 @@ from vllm.platforms import current_platform
17
from vllm.utils import round_up
18
19
from .fusion import QUANT_OPS, empty_bf16, empty_fp32, empty_i32
20
+from .inductor_pass import enable_fake_mode
21
from .vllm_inductor_pass import VllmInductorPass
22
23
logger = init_logger(__name__)
@@ -139,24 +138,21 @@ class AttentionFp8StaticQuantPattern(AttentionQuantPattern):
138
output_block_scale=None)
139
return RESHAPE_OP(at1[1], [-1, self.num_heads * self.head_size])
140
142
- # Need custom fake mode, otherwise tracing happens with real tensors.
143
- # That would not work for the unified_attention custom op.
144
- with unset_fake_temporarily(), FakeTensorMode():
145
- inputs = [
146
- empty_bf16(5, self.num_heads, self.head_size), # q
147
- empty_bf16(5, self.num_heads, self.head_size), # k
148
- empty_bf16(5, self.num_heads, self.head_size), # v
149
- empty_bf16(5, self.num_heads, self.head_size), # attn_output
150
- self.empty_quant(5, self.num_heads *
151
- self.head_size), # quant_output
152
- empty_fp32(1, 1) # scale
153
- ]
154
-
155
- pm.register_replacement(
156
- pattern, replacement, inputs,
157
- AttentionQuantPattern.wrap_trace_fn(
158
- AttentionQuantPattern.fx_view_to_reshape, pm.fwd_only),
159
- pm_pass)
141
+ inputs = [
142
+ empty_bf16(5, self.num_heads, self.head_size), # q
143
+ empty_bf16(5, self.num_heads, self.head_size), # k
144
+ empty_bf16(5, self.num_heads, self.head_size), # v
145
+ empty_bf16(5, self.num_heads, self.head_size), # attn_output
146
+ self.empty_quant(5,
147
+ self.num_heads * self.head_size), # quant_output
148
+ empty_fp32(1, 1) # scale
149
+ ]
150
+
151
+ pm.register_replacement(
152
+ pattern, replacement, inputs,
153
+ AttentionQuantPattern.wrap_trace_fn(
154
+ AttentionQuantPattern.fx_view_to_reshape, pm.fwd_only),
155
+ pm_pass)
156
157
158
class AttentionNvfp4QuantPattern(AttentionQuantPattern):
@@ -219,27 +215,23 @@ class AttentionNvfp4QuantPattern(AttentionQuantPattern):
215
[-1, self.num_heads * self.head_size // 2])
216
return output, at2[2]
217
222
- # Need custom fake mode, otherwise tracing happens with real tensors.
223
- # That would not work for the unified_attention custom op.
224
- with unset_fake_temporarily(), FakeTensorMode():
225
- inputs = [
226
- empty_bf16(5, self.num_heads, self.head_size), # q
227
- empty_bf16(5, self.num_heads, self.head_size), # k
228
- empty_bf16(5, self.num_heads, self.head_size), # v
229
- empty_bf16(5, self.num_heads, self.head_size), # output_attn
230
- self.empty_quant(5, self.num_heads * self.head_size //
231
- 2), # output_quant
232
- empty_i32(128,
233
- round_up(self.num_heads * self.head_size // 16,
234
- 4)), # output_scale
235
- empty_fp32(1, 1), # input_scale
236
- ]
237
-
238
- pm.register_replacement(
239
- pattern, replacement, inputs,
240
- AttentionQuantPattern.wrap_trace_fn(
241
- AttentionQuantPattern.fx_view_to_reshape, pm.fwd_only),
242
- pm_pass)
218
+ inputs = [
219
+ empty_bf16(5, self.num_heads, self.head_size), # q
220
+ empty_bf16(5, self.num_heads, self.head_size), # k
221
+ empty_bf16(5, self.num_heads, self.head_size), # v
222
+ empty_bf16(5, self.num_heads, self.head_size), # output_attn
223
+ self.empty_quant(5, self.num_heads * self.head_size //
224
+ 2), # output_quant
225
+ empty_i32(128, round_up(self.num_heads * self.head_size // 16,
226
+ 4)), # output_scale
227
+ empty_fp32(1, 1), # input_scale
228
+ ]
229
+
230
+ pm.register_replacement(
231
+ pattern, replacement, inputs,
232
+ AttentionQuantPattern.wrap_trace_fn(
233
+ AttentionQuantPattern.fx_view_to_reshape, pm.fwd_only),
234
+ pm_pass)
235
236
237
class AttnFusionPass(VllmInductorPass):
@@ -255,6 +247,7 @@ class AttnFusionPass(VllmInductorPass):
247
support are attention kernels, which need to support fusing output quant.
248
"""
249
250
+ @enable_fake_mode
251
def __init__(self, config: VllmConfig):
252
super().__init__(config)
253
vllm/compilation/inductor_pass.py
+20
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import functools
5
import hashlib
6
import inspect
7
import json
@@ -10,6 +11,8 @@ from typing import Any, Callable, Optional, Union
11
12
import torch
13
from torch import fx
14
+from torch._subclasses.fake_tensor import (FakeTensorMode,
15
+ unset_fake_temporarily)
16
17
from vllm.utils import is_torch_equal_or_newer
18
@@ -114,3 +117,20 @@ class CallableInductorPass(InductorPass):
117
118
def uuid(self) -> Any:
119
return self._uuid
120
+
121
+
122
+def enable_fake_mode(fn: Callable[..., Any]) -> Callable[..., Any]:
123
+ """
124
+ Applies a FakeTensorMode context. This is useful when you don't want to
125
+ create or run things with real tensors.
126
+ """
127
+
128
+ @functools.wraps(fn)
129
+ def fn_new(*args, **kwargs) -> Any:
130
+ with torch._guards.tracing(
131
+ None), unset_fake_temporarily(), FakeTensorMode():
132
+ result = fn(*args, **kwargs)
133
+
134
+ return result
135
+
136
+ return fn_new
vllm/compilation/sequence_parallelism.py
+2
@@ -14,6 +14,7 @@ from vllm.distributed.parallel_state import (
14
from vllm.logger import init_logger
15
from vllm.platforms import current_platform
16
17
+from .inductor_pass import enable_fake_mode
18
from .vllm_inductor_pass import VllmInductorPass
19
20
logger = init_logger(__name__)
@@ -436,6 +437,7 @@ class SequenceParallelismPass(VllmInductorPass):
437
performance.
438
"""
439
440
+ @enable_fake_mode
441
def __init__(self, config: VllmConfig):
442
super().__init__(config)
443