@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 59bcc5b6f2
Fork 0

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 9, 11:15 UTC

Use aiter triton fused_add_rmsnorm_pad for gpt-oss (#30976)

Signed-off-by: Rohan138 <rohanpotdar138@gmail.com>

Rohan Potdar committed Jan 28, 2026 at 14:47 UTC 59bcc5b6f2e68dab87426f4b77a3a72d72e6ec48
9 files changed +327 -11
tests/compile/test_fuse_act_padding.py new
+131
@@ -0,0 +1,131 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +
5 +import pytest
6 +import torch
7 +
8 +import vllm.config
9 +from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
10 +from vllm.compilation.noop_elimination import NoOpEliminationPass
11 +from vllm.compilation.post_cleanup import PostCleanupPass
12 +from vllm.config import (
13 + CompilationConfig,
14 + CompilationMode,
15 + ModelConfig,
16 + PassConfig,
17 + VllmConfig,
18 +)
19 +from vllm.model_executor.layers.layernorm import RMSNorm
20 +from vllm.model_executor.layers.utils import rocm_unquantized_gemm
21 +
22 +from .backend import TestBackend
23 +
24 +
25 +class TestModel(torch.nn.Module):
26 + def __init__(
27 + self,
28 + num_layers: int,
29 + hidden_size: int,
30 + num_local_experts: int,
31 + x_pad_to_multiple: int,
32 + ):
33 + super().__init__()
34 + self.num_layers = num_layers
35 + self.hidden_size = hidden_size
36 + self.x_pad_to_multiple = x_pad_to_multiple
37 + self.pad_dim = x_pad_to_multiple - (hidden_size % x_pad_to_multiple)
38 +
39 + self.norm = [RMSNorm(hidden_size, eps=1e-5) for _ in range(num_layers)]
40 + self.router = [
41 + torch.nn.Linear(hidden_size, num_local_experts) for _ in range(4)
42 + ]
43 +
44 + def forward(self, x):
45 + # avoid having graph input be an arg to a pattern directly
46 + x = resid = torch.relu(x)
47 + all_router_logits = []
48 + for layer in range(self.num_layers):
49 + x = x[:, : self.hidden_size]
50 + x, resid = self.norm[layer](x, resid)
51 + router_logits = rocm_unquantized_gemm(
52 + self, x, self.router[layer].weight, self.router[layer].bias
53 + )
54 + x = torch.nn.functional.pad(
55 + x, (0, self.pad_dim), mode="constant", value=0.0
56 + )
57 + all_router_logits.append(router_logits)
58 +
59 + return x, resid, *all_router_logits
60 +
61 + def ops_in_model_before(self):
62 + return [
63 + rocm_aiter_ops.get_rmsnorm_fused_add_op(),
64 + torch.ops.aten.constant_pad_nd,
65 + ]
66 +
67 + def ops_in_model_after(self):
68 + return [rocm_aiter_ops.get_triton_add_rmsnorm_pad_op()]
69 +
70 +
71 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
72 +@pytest.mark.parametrize("num_layers", [3])
73 +@pytest.mark.parametrize("hidden_size", [2880])
74 +@pytest.mark.parametrize("num_local_experts", [128])
75 +@pytest.mark.parametrize("x_pad_to_multiple", [256])
76 +@pytest.mark.skipif(
77 + not is_aiter_found_and_supported(),
78 + reason="Only test on ROCm with AITER installed and supported",
79 +)
80 +def test_fuse_act_padding(
81 + dtype: torch.dtype,
82 + num_layers: int,
83 + hidden_size: int,
84 + num_local_experts: int,
85 + x_pad_to_multiple: int,
86 + monkeypatch: pytest.MonkeyPatch,
87 +):
88 + vllm_config = VllmConfig(
89 + model_config=ModelConfig(dtype=dtype),
90 + compilation_config=CompilationConfig(
91 + mode=CompilationMode.VLLM_COMPILE,
92 + custom_ops=["+rms_norm"],
93 + pass_config=PassConfig(fuse_act_padding=True, eliminate_noops=True),
94 + ),
95 + )
96 +
97 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
98 + from vllm.compilation.rocm_aiter_fusion import (
99 + RocmAiterTritonAddRMSNormPadFusionPass,
100 + )
101 +
102 + torch.set_default_device("cuda")
103 + torch.set_default_dtype(dtype)
104 + torch.manual_seed(1)
105 +
106 + m.setenv("VLLM_ROCM_USE_AITER", "1")
107 + rocm_aiter_ops.refresh_env_variables()
108 +
109 + fusion_pass = RocmAiterTritonAddRMSNormPadFusionPass(vllm_config)
110 + passes = [
111 + NoOpEliminationPass(vllm_config),
112 + fusion_pass,
113 + PostCleanupPass(vllm_config),
114 + ]
115 + backend = TestBackend(*passes)
116 + model = TestModel(num_layers, hidden_size, num_local_experts, x_pad_to_multiple)
117 +
118 + x = torch.rand(1, hidden_size)
119 + torch._dynamo.mark_dynamic(x, 0)
120 +
121 + outputs_unfused = model(x)
122 +
123 + model_fused = torch.compile(model, backend=backend)
124 + outputs_fused = model_fused(x)
125 +
126 + torch.testing.assert_close(outputs_unfused, outputs_fused)
127 +
128 + assert fusion_pass.matched_count == num_layers
129 +
130 + backend.check_before_ops(model.ops_in_model_before())
131 + backend.check_after_ops(model.ops_in_model_after())
tests/compile/test_fusion.py
+2 -2
@@ -410,7 +410,7 @@ def test_aiter_fusion_rmsnorm_quant(
410 )
411
412 with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
413 - from vllm.compilation.rocm_aiter_fusion import RocmAiterRMSNormFusionPass
413 + from vllm.compilation.rocm_aiter_fusion import RocmAiterRMSNormQuantFusionPass
414
415 m.setenv("VLLM_ROCM_USE_AITER", "1")
416
@@ -420,7 +420,7 @@ def test_aiter_fusion_rmsnorm_quant(
420 torch.set_default_dtype(dtype)
421 torch.manual_seed(1)
422
423 - fusion_pass = RocmAiterRMSNormFusionPass(vllm_config)
423 + fusion_pass = RocmAiterRMSNormQuantFusionPass(vllm_config)
424
425 model = TestModel(
426 hidden_size=hidden_size,
vllm/_aiter_ops.py
+46
@@ -790,6 +790,41 @@ def _rocm_aiter_act_mul_and_fp8_group_quant_fake(
790 return x_fp8, out_bs
791
792
793 +def _rocm_aiter_triton_add_rmsnorm_pad_impl(
794 + x: torch.Tensor,
795 + weight: torch.Tensor,
796 + variance_epsilon: float,
797 + residual: torch.Tensor,
798 + x_pad_to_multiple: int,
799 +) -> tuple[torch.Tensor, torch.Tensor]:
800 + from aiter.ops.triton.fused_add_rmsnorm_pad import fused_add_rmsnorm_pad
801 +
802 + return fused_add_rmsnorm_pad(
803 + x,
804 + weight,
805 + variance_epsilon,
806 + residual,
807 + x_pad_to_multiple=x_pad_to_multiple,
808 + )
809 +
810 +
811 +def _rocm_aiter_triton_add_rmsnorm_pad_fake(
812 + x: torch.Tensor,
813 + weight: torch.Tensor,
814 + variance_epsilon: float,
815 + residual: torch.Tensor,
816 + x_pad_to_multiple: int,
817 +) -> tuple[torch.Tensor, torch.Tensor]:
818 + M, N = x.shape
819 + if x_pad_to_multiple > 0:
820 + N_out = (N + x_pad_to_multiple - 1) // x_pad_to_multiple * x_pad_to_multiple
821 + else:
822 + N_out = N
823 + out = torch.empty((M, N_out), dtype=x.dtype, device=x.device)
824 + residual_out = torch.empty_like(residual)
825 + return out, residual_out
826 +
827 +
828 # Global flag to ensure ops are registered only once
829 _OPS_REGISTERED = False
830
@@ -1108,6 +1143,13 @@ class rocm_aiter_ops:
1143 fake_impl=_rocm_aiter_act_mul_and_fp8_group_quant_fake,
1144 )
1145
1146 + direct_register_custom_op(
1147 + op_name="rocm_aiter_triton_add_rmsnorm_pad",
1148 + op_func=_rocm_aiter_triton_add_rmsnorm_pad_impl,
1149 + fake_impl=_rocm_aiter_triton_add_rmsnorm_pad_fake,
1150 + dispatch_key=current_platform.dispatch_key,
1151 + )
1152 +
1153 direct_register_custom_op(
1154 op_name="rocm_aiter_group_fp8_quant",
1155 op_func=_rocm_aiter_group_fp8_quant_impl,
@@ -1175,6 +1217,10 @@ class rocm_aiter_ops:
1217 def get_act_mul_fused_fp8_group_quant_op() -> OpOverload:
1218 return torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant.default
1219
1220 + @staticmethod
1221 + def get_triton_add_rmsnorm_pad_op() -> OpOverload:
1222 + return torch.ops.vllm.rocm_aiter_triton_add_rmsnorm_pad.default
1223 +
1224 @staticmethod
1225 def rms_norm(
1226 x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
vllm/compilation/pass_manager.py
+6 -2
@@ -18,8 +18,9 @@ from .vllm_inductor_pass import VllmInductorPass
18
19 if rocm_aiter_ops.is_enabled():
20 from vllm.compilation.rocm_aiter_fusion import (
21 - RocmAiterRMSNormFusionPass,
21 + RocmAiterRMSNormQuantFusionPass,
22 RocmAiterSiluMulFp8GroupQuantFusionPass,
23 + RocmAiterTritonAddRMSNormPadFusionPass,
24 )
25
26 if current_platform.is_cuda_alike():
@@ -123,13 +124,16 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
124 self.passes += [RMSNormQuantFusionPass(config)]
125 if rocm_aiter_ops.is_enabled():
126 self.passes += [
126 - RocmAiterRMSNormFusionPass(config),
127 + RocmAiterRMSNormQuantFusionPass(config),
128 ]
129 if self.pass_config.fuse_act_quant:
130 self.passes += [ActivationQuantFusionPass(config)]
131 if rocm_aiter_ops.is_enabled():
132 self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
133
134 + if self.pass_config.fuse_act_padding and rocm_aiter_ops.is_enabled():
135 + self.passes += [RocmAiterTritonAddRMSNormPadFusionPass(config)]
136 +
137 if self.pass_config.fuse_attn_quant:
138 self.passes += [AttnFusionPass(config)]
139
vllm/compilation/rocm_aiter_fusion.py
+104 -1
@@ -266,7 +266,7 @@ class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
266 )
267
268
269 -class RocmAiterRMSNormFusionPass(VllmPatternMatcherPass):
269 +class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
270 """
271 This pass fuses aiter rms_norm & vllm/aiter quant custom ops
272 into a fused rms_norm_quant op.
@@ -399,3 +399,106 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
399 AiterSiluMulFp8GroupQuantPattern,
400 ]
401 return VllmInductorPass.hash_source(self, *fusion_patterns)
402 +
403 +
404 +class AddAiterRMSNormPadPattern:
405 + """
406 + This pattern replaces an aiter_rmsnorm_with_add & a pad op
407 + with a custom triton_add_rmsnorm_pad op from AITER.
408 + """
409 +
410 + AITER_TRITON_ADD_RMSNORM_PAD_OP = rocm_aiter_ops.get_triton_add_rmsnorm_pad_op()
411 +
412 + def __init__(
413 + self,
414 + epsilon: float,
415 + hidden_size: int,
416 + x_pad_to_multiple: int,
417 + ):
418 + self.epsilon = epsilon
419 + self.hidden_size = hidden_size
420 + self.x_pad_to_multiple = x_pad_to_multiple
421 + self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon, match_rocm_aiter=True)
422 +
423 + def get_inputs(self) -> list[torch.Tensor]:
424 + input, weight, residual = self.rmsnorm_matcher.inputs()
425 + router_weight = torch.empty([8, 16], dtype=weight.dtype, device=weight.device)
426 + router_bias = torch.empty([8], dtype=weight.dtype, device=weight.device)
427 + return [input, weight, residual, router_weight, router_bias]
428 +
429 + def register(self, pm_pass: PatternMatcherPass) -> None:
430 + def pattern(
431 + input: torch.Tensor,
432 + weight: torch.Tensor,
433 + residual: torch.Tensor,
434 + router_weight: torch.Tensor,
435 + router_bias: torch.Tensor,
436 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
437 + pad_size = self.x_pad_to_multiple - (
438 + self.hidden_size % self.x_pad_to_multiple
439 + )
440 + result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
441 + router_logits = torch.ops.vllm.rocm_unquantized_gemm(
442 + result_rms, router_weight, router_bias
443 + )
444 + result = torch.nn.functional.pad(
445 + result_rms, (0, pad_size), mode="constant", value=0.0
446 + )
447 + return result, residual_out, router_logits
448 +
449 + def replacement(
450 + input: torch.Tensor,
451 + weight: torch.Tensor,
452 + residual: torch.Tensor,
453 + router_weight: torch.Tensor,
454 + router_bias: torch.Tensor,
455 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
456 + at = self.AITER_TRITON_ADD_RMSNORM_PAD_OP(
457 + x=input,
458 + weight=weight,
459 + variance_epsilon=self.epsilon,
460 + residual=residual,
461 + x_pad_to_multiple=self.x_pad_to_multiple,
462 + )
463 + result_padded = at[0]
464 + router_logits = torch.ops.vllm.rocm_unquantized_gemm(
465 + result_padded[:, : self.hidden_size], router_weight, router_bias
466 + )
467 + residual_out = at[1]
468 + return result_padded, residual_out, router_logits
469 +
470 + pm.register_replacement(
471 + pattern, replacement, self.get_inputs(), pm.fwd_only, pm_pass
472 + )
473 +
474 +
475 +class RocmAiterTritonAddRMSNormPadFusionPass(VllmPatternMatcherPass):
476 + """
477 + This pass replaces an AITER CK RMSNorm + residual add and a pad op
478 + with an triton_add_rmsnorm_pad op from AITER.
479 + """
480 +
481 + def __init__(self, config: VllmConfig):
482 + super().__init__(config)
483 + self.patterns: PatternMatcherPass = PatternMatcherPass(
484 + pass_name="rocm_aiter_triton_add_rmsnorm_pad_fusion_pass"
485 + )
486 +
487 + # gpt-oss has hidden size 2880
488 + # padded to a multiple of 128 on gfx942 and 256 on gfx950 respectively
489 + hidden_size = 2880
490 + for epsilon in [1e-5, 1e-6]:
491 + for x_pad_to_multiple in [128, 256]:
492 + AddAiterRMSNormPadPattern(
493 + epsilon, hidden_size, x_pad_to_multiple
494 + ).register(self.patterns)
495 +
496 + self.dump_patterns(config, self.patterns)
497 +
498 + @VllmInductorPass.time_and_log
499 + def __call__(self, graph: torch.fx.Graph) -> None:
500 + self.matched_count = self.patterns.apply(graph)
501 + logger.debug("Replaced %s patterns", self.matched_count)
502 +
503 + def uuid(self) -> str:
504 + return VllmInductorPass.hash_source(self, AddAiterRMSNormPadPattern)
vllm/config/compilation.py
+16
@@ -126,6 +126,10 @@ class PassConfig:
126 fuse_allreduce_rms: bool = Field(default=None)
127 """Enable flashinfer allreduce fusion."""
128
129 + # ROCm/AITER specific fusions
130 + fuse_act_padding: bool = Field(default=None)
131 + """Fuse the custom RMSNorm + padding ops."""
132 +
133 fi_allreduce_fusion_max_size_mb: float | None = None
134 """The threshold of the communicated tensor sizes under which
135 vllm should use flashinfer fused allreduce. Specified as a
@@ -194,6 +198,7 @@ class PassConfig:
198 "enable_sp",
199 "fuse_gemm_comms",
200 "fuse_allreduce_rms",
201 + "fuse_act_padding",
202 mode="wrap",
203 )
204 @classmethod
@@ -222,12 +227,23 @@ class PassConfig:
227 "Fusion enabled but reshape elimination disabled. "
228 "Allreduce + rms norm + quant (fp8) fusion might not work"
229 )
230 + if self.fuse_act_padding:
231 + logger.warning_once(
232 + "Fusion enabled but reshape elimination disabled. "
233 + "RMSNorm + padding fusion might not work"
234 + )
235 if self.enable_qk_norm_rope_fusion and not current_platform.is_cuda_alike():
236 logger.warning_once(
237 "QK Norm + RoPE fusion enabled but the current platform is not "
238 "CUDA or ROCm. The fusion will be disabled."
239 )
240 self.enable_qk_norm_rope_fusion = False
241 + if self.fuse_act_padding and not current_platform.is_rocm():
242 + logger.warning_once(
243 + "Padding fusion enabled but the current platform is not ROCm. "
244 + "The fusion will be disabled."
245 + )
246 + self.fuse_act_padding = False
247
248
249 class DynamicShapesType(str, enum.Enum):
vllm/config/vllm.py
+16
@@ -102,6 +102,18 @@ def enable_act_fusion(cfg: "VllmConfig") -> bool:
102 ) or cfg.compilation_config.is_custom_op_enabled("quant_fp8")
103
104
105 +def enable_norm_pad_fusion(cfg: "VllmConfig") -> bool:
106 + """Enable if using AITER RMSNorm and AITER Triton GEMMs
107 + and hidden size is 2880 i.e. gpt-oss; otherwise Inductor handles fusion."""
108 +
109 + return (
110 + envs.VLLM_ROCM_USE_AITER
111 + and envs.VLLM_ROCM_USE_AITER_RMSNORM
112 + and envs.VLLM_ROCM_USE_AITER_TRITON_GEMM
113 + and cfg.model_config.get_hidden_size() == 2880
114 + )
115 +
116 +
117 OPTIMIZATION_LEVEL_00 = {
118 "compilation_config": {
119 "pass_config": {
@@ -112,6 +124,7 @@ OPTIMIZATION_LEVEL_00 = {
124 "fuse_attn_quant": False,
125 "enable_sp": False,
126 "fuse_gemm_comms": False,
127 + "fuse_act_padding": False,
128 },
129 "cudagraph_mode": CUDAGraphMode.NONE,
130 "use_inductor_graph_partition": False,
@@ -127,6 +140,7 @@ OPTIMIZATION_LEVEL_01 = {
140 "fuse_attn_quant": False,
141 "enable_sp": False,
142 "fuse_gemm_comms": False,
143 + "fuse_act_padding": enable_norm_pad_fusion,
144 },
145 "cudagraph_mode": CUDAGraphMode.PIECEWISE,
146 "use_inductor_graph_partition": False,
@@ -142,6 +156,7 @@ OPTIMIZATION_LEVEL_02 = {
156 "fuse_attn_quant": IS_QUANTIZED,
157 "enable_sp": IS_DENSE,
158 "fuse_gemm_comms": IS_DENSE,
159 + "fuse_act_padding": enable_norm_pad_fusion,
160 },
161 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
162 "use_inductor_graph_partition": False,
@@ -157,6 +172,7 @@ OPTIMIZATION_LEVEL_03 = {
172 "fuse_attn_quant": IS_QUANTIZED,
173 "enable_sp": IS_DENSE,
174 "fuse_gemm_comms": IS_DENSE,
175 + "fuse_act_padding": enable_norm_pad_fusion,
176 },
177 "cudagraph_mode": CUDAGraphMode.FULL_AND_PIECEWISE,
178 "use_inductor_graph_partition": False,
vllm/model_executor/layers/utils.py
+5 -5
@@ -137,6 +137,11 @@ def rocm_unquantized_gemm_impl(
137
138 import math
139
140 + if use_aiter_triton_gemm(n, m, k, x.dtype):
141 + from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
142 +
143 + return gemm_a16w16(x, weight, bias)
144 +
145 use_skinny_reduce_counting = (
146 envs.VLLM_ROCM_USE_SKINNY_GEMM
147 and on_gfx950()
@@ -155,11 +160,6 @@ def rocm_unquantized_gemm_impl(
160 out = ops.wvSplitKrc(weight, x_view, cu_count, bias)
161 return out.reshape(*x.shape[:-1], weight.shape[0])
162
158 - if use_aiter_triton_gemm(n, m, k, x.dtype):
159 - from aiter.ops.triton.gemm_a16w16 import gemm_a16w16
160 -
161 - return gemm_a16w16(x, weight, bias)
162 -
163 use_skinny = (
164 envs.VLLM_ROCM_USE_SKINNY_GEMM
165 and on_gfx9()
vllm/model_executor/models/gpt_oss.py
+1 -1
@@ -187,7 +187,7 @@ class MLPBlock(torch.nn.Module):
187 )
188 else:
189 g = self.router(x)
190 - x = self.experts(hidden_states=x, router_logits=g)
190 + x = self.experts(hidden_states=x, router_logits=g)[:, : self.hidden_size]
191
192 if self.is_sequence_parallel:
193 x = tensor_model_parallel_all_gather(x.contiguous(), 0)