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 8, 12:30 UTC
[torch.compile] Enable AR+rms fusion by default available for `-O2` (#34299)
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
Luka Govedič committed
Feb 11, 2026 at 03:30 UTC
addac0e65343e4c24a109975d54c4673bbfb029c
2 files changed
+16
-8
vllm/config/compilation.py
+1
-2
@@ -115,7 +115,7 @@ class PassConfig:
115
"""Fuse the custom SiluMul + quant ops."""
116
fuse_attn_quant: bool = Field(default=None)
117
"""Fuse the custom attention + quant ops."""
118
- eliminate_noops: bool = Field(default=None)
118
+ eliminate_noops: bool = Field(default=True)
119
"""Eliminate no-op ops."""
120
enable_sp: bool = Field(default=None)
121
"""Enable sequence parallelism."""
@@ -194,7 +194,6 @@ class PassConfig:
194
"fuse_norm_quant",
195
"fuse_act_quant",
196
"fuse_attn_quant",
197
- "eliminate_noops",
197
"enable_sp",
198
"fuse_gemm_comms",
199
"fuse_allreduce_rms",
vllm/config/vllm.py
+15
-6
@@ -102,6 +102,19 @@ def enable_act_fusion(cfg: "VllmConfig") -> bool:
102
) or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
103
104
105
+def enable_allreduce_rms_fusion(cfg: "VllmConfig") -> bool:
106
+ """Enable if TP > 1 and Hopper+ and flashinfer installed."""
107
+ from vllm.platforms import current_platform
108
+ from vllm.utils.flashinfer import has_flashinfer
109
+
110
+ return (
111
+ cfg.parallel_config.tensor_parallel_size > 1
112
+ and current_platform.is_cuda()
113
+ and current_platform.has_device_capability(90)
114
+ and has_flashinfer()
115
+ )
116
+
117
+
118
def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
119
"""Enable if using AITER RMSNorm and AITER Triton GEMMs
120
and hidden size is 2880 i.e. gpt-oss; otherwise Inductor handles fusion."""
@@ -118,7 +131,6 @@ def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
131
OPTIMIZATION_LEVEL_00 = {
132
"compilation_config": {
133
"pass_config": {
121
- "eliminate_noops": False,
134
"fuse_norm_quant": False,
135
"fuse_act_quant": False,
136
"fuse_allreduce_rms": False,
@@ -137,7 +149,6 @@ OPTIMIZATION_LEVEL_00 = {
149
OPTIMIZATION_LEVEL_01 = {
150
"compilation_config": {
151
"pass_config": {
140
- "eliminate_noops": True,
152
"fuse_norm_quant": enable_norm_fusion,
153
"fuse_act_quant": enable_act_fusion,
154
"fuse_allreduce_rms": False,
@@ -156,10 +167,9 @@ OPTIMIZATION_LEVEL_01 = {
167
OPTIMIZATION_LEVEL_02 = {
168
"compilation_config": {
169
"pass_config": {
159
- "eliminate_noops": True,
170
"fuse_norm_quant": enable_norm_fusion,
171
"fuse_act_quant": enable_act_fusion,
162
- "fuse_allreduce_rms": False,
172
+ "fuse_allreduce_rms": enable_allreduce_rms_fusion,
173
"fuse_attn_quant": IS_QUANTIZED,
174
"enable_sp": IS_DENSE,
175
"fuse_gemm_comms": IS_DENSE,
@@ -175,10 +185,9 @@ OPTIMIZATION_LEVEL_02 = {
185
OPTIMIZATION_LEVEL_03 = {
186
"compilation_config": {
187
"pass_config": {
178
- "eliminate_noops": True,
188
"fuse_norm_quant": enable_norm_fusion,
189
"fuse_act_quant": enable_act_fusion,
181
- "fuse_allreduce_rms": False,
190
+ "fuse_allreduce_rms": enable_allreduce_rms_fusion,
191
"fuse_attn_quant": IS_QUANTIZED,
192
"enable_sp": IS_DENSE,
193
"fuse_gemm_comms": IS_DENSE,