@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3c680f4a17
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 11, 17:30 UTC

[Rocm][torch.compile] Adding layernorm + fp8 block quant and silu + fp8 block quant for Aiter (#25693)

Signed-off-by: charlifu <charlifu@amd.com> Signed-off-by: Micah Williamson <micah.williamson@amd.com> Signed-off-by: Charlie Fu <Charlie.Fu@amd.com> Co-authored-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: wuhuikx <hattie.wu@amd.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com>

Charlie Fu committed Dec 9, 2025 at 16:39 UTC 3c680f4a17057d7994af8fbb1dc8c2d98307c890
6 files changed +610 -60
tests/compile/test_fusion.py
+93 -5
@@ -1,10 +1,13 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import itertools
5 +
6 import pytest
7 import torch
8
9 import vllm.plugins
10 +from vllm._aiter_ops import IS_AITER_FOUND, rocm_aiter_ops
11 from vllm.compilation.fusion import FUSED_OPS, FusedRMSQuantKey, RMSNormQuantFusionPass
12 from vllm.compilation.fx_utils import find_op_nodes
13 from vllm.compilation.matcher_utils import QUANT_OPS
@@ -152,13 +155,79 @@ GROUP_SHAPES = [
155 ]
156
157
158 +class TestRmsnormGroupFp8QuantModel(torch.nn.Module):
159 + def __init__(self, hidden_size: int, eps: float, **kwargs):
160 + super().__init__()
161 + self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
162 + weight_group_shape=GroupShape(128, 128),
163 + act_quant_group_shape=GroupShape(1, 128),
164 + cutlass_block_fp8_supported=False,
165 + use_aiter_and_is_supported=True,
166 + )
167 + self.w = [
168 + torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
169 + for _ in range(3)
170 + ]
171 +
172 + scale_hidden_size = (hidden_size + 128 - 1) // 128
173 + self.wscale = [
174 + torch.rand((scale_hidden_size, scale_hidden_size), dtype=torch.float32)
175 + for _ in range(3)
176 + ]
177 +
178 + self.norm_weight = [torch.ones(hidden_size) for _ in range(4)]
179 + self.eps = eps
180 +
181 + def forward(self, x):
182 + # avoid having graph input be an arg to a pattern directly
183 + x = resid = torch.relu(x)
184 + y = rocm_aiter_ops.rms_norm(x, self.norm_weight[0], self.eps)
185 +
186 + x2 = self.w8a8_block_fp8_linear.apply(y, self.w[0], self.wscale[0])
187 + # make sure resid is used for replacement to work
188 + y2, resid = rocm_aiter_ops.rms_norm2d_with_add(
189 + x2, resid, self.norm_weight[1], self.eps
190 + )
191 +
192 + x3 = self.w8a8_block_fp8_linear.apply(y2, self.w[1], self.wscale[1])
193 +
194 + y3, resid = rocm_aiter_ops.rms_norm2d_with_add(
195 + x3, resid, self.norm_weight[2], self.eps
196 + )
197 +
198 + x4 = self.w8a8_block_fp8_linear.apply(y3, self.w[2], self.wscale[2])
199 +
200 + y4, resid = rocm_aiter_ops.rms_norm2d_with_add(
201 + x4, resid, self.norm_weight[3], self.eps
202 + )
203 + return y4
204 +
205 + def ops_in_model_before(self):
206 + return [
207 + torch.ops.vllm.rocm_aiter_rms_norm,
208 + torch.ops.vllm.rocm_aiter_group_fp8_quant,
209 + ]
210 +
211 + def ops_in_model_before_partial(self):
212 + return []
213 +
214 + def ops_in_model_after(self):
215 + return [
216 + torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant,
217 + torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant,
218 + ]
219 +
220 +
221 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
222 @pytest.mark.parametrize("hidden_size", [256])
223 @pytest.mark.parametrize("num_tokens", [257])
224 @pytest.mark.parametrize("eps", [1e-5, 1e-6])
225 @pytest.mark.parametrize("group_shape", GROUP_SHAPES)
160 -@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
161 -@pytest.mark.parametrize("enable_quant_fp8_custom_op", [True, False])
226 +@pytest.mark.parametrize(
227 + "model_class, enable_rms_norm_custom_op, enable_quant_fp8_custom_op",
228 + list(itertools.product([TestModel], [True, False], [True, False]))
229 + + [(TestRmsnormGroupFp8QuantModel, False, False)],
230 +)
231 # cuda_force_torch used to test torch code path on platforms that
232 # cutlass_fp8_supported() == True.
233 @pytest.mark.parametrize(
@@ -173,10 +242,14 @@ def test_fusion_rmsnorm_quant(
242 num_tokens,
243 eps,
244 group_shape,
245 + model_class,
246 enable_rms_norm_custom_op,
247 enable_quant_fp8_custom_op,
248 cuda_force_torch,
249 ):
250 + if model_class is TestRmsnormGroupFp8QuantModel and not IS_AITER_FOUND:
251 + pytest.skip("AITER is not supported on this GPU.")
252 +
253 torch.set_default_device("cuda")
254 torch.set_default_dtype(dtype)
255 torch.manual_seed(1)
@@ -209,12 +282,24 @@ def test_fusion_rmsnorm_quant(
282 with vllm.config.set_current_vllm_config(vllm_config):
283 # Reshape pass is needed for the fusion pass to work
284 noop_pass = NoOpEliminationPass(vllm_config)
212 - fusion_pass = RMSNormQuantFusionPass(vllm_config)
285 + if model_class is TestRmsnormGroupFp8QuantModel:
286 + from vllm.compilation.rocm_aiter_fusion import (
287 + RocmAiterRMSNormFp8GroupQuantFusionPass,
288 + )
289 +
290 + fusion_pass = RocmAiterRMSNormFp8GroupQuantFusionPass(vllm_config)
291 + else:
292 + fusion_pass = RMSNormQuantFusionPass(vllm_config)
293 cleanup_pass = PostCleanupPass(vllm_config)
294
295 backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)
296 backend2 = TestBackend(noop_pass, cleanup_pass)
217 - model = TestModel(hidden_size, eps, group_shape, cuda_force_torch)
297 + model = model_class(
298 + hidden_size=hidden_size,
299 + eps=eps,
300 + group_shape=group_shape,
301 + cuda_force_torch=cuda_force_torch,
302 + )
303 # First dimension dynamic
304 x = torch.rand(num_tokens, hidden_size)
305 torch._dynamo.mark_dynamic(x, 0)
@@ -243,7 +328,10 @@ def test_fusion_rmsnorm_quant(
328 # there's a risk that the fused add doesn't get included in the
329 # replacement and only the rms part gets fused with quant.
330 # Hence, we check only 2 add nodes are left (final fused rmsnorm add).
246 - if not enable_rms_norm_custom_op:
331 + if (
332 + not enable_rms_norm_custom_op
333 + and model_class is not TestRmsnormGroupFp8QuantModel
334 + ):
335 n_add_nodes = lambda g: sum(1 for _ in find_op_nodes(torch.ops.aten.add, g))
336 # 7 = 1 (RMS) + 3x2 (3xRMS_ADD, 2 each)
337 assert n_add_nodes(backend.graph_pre_pass) == 7
tests/compile/test_silu_mul_quant_fusion.py
+57 -5
@@ -7,6 +7,7 @@ import torch
7
8 import vllm.envs as envs
9 from tests.kernels.quantization.nvfp4_utils import quant_nvfp4_tensor
10 +from vllm._aiter_ops import IS_AITER_FOUND
11 from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
12 from vllm.compilation.activation_quant_fusion import (
13 FUSED_OPS,
@@ -24,6 +25,7 @@ from vllm.config import (
25 set_current_vllm_config,
26 )
27 from vllm.model_executor.layers.activation import SiluAndMul
28 +from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8LinearOp
29 from vllm.model_executor.layers.quantization.utils.quant_utils import (
30 GroupShape,
31 kFp8StaticTensorSym,
@@ -126,6 +128,39 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
128 return [FUSED_OPS[kNvfp4Quant]]
129
130
131 +class TestSiluMulGroupFp8QuantModel(torch.nn.Module):
132 + def __init__(self, hidden_size: int, **kwargs):
133 + super().__init__()
134 + self.silu_and_mul = SiluAndMul()
135 + self.w8a8_block_fp8_linear = W8A8BlockFp8LinearOp(
136 + weight_group_shape=GroupShape(128, 128),
137 + act_quant_group_shape=GroupShape(1, 128),
138 + cutlass_block_fp8_supported=False,
139 + use_aiter_and_is_supported=True,
140 + )
141 + self.w = torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
142 +
143 + scale_hidden_size = (hidden_size + 128 - 1) // 128
144 + self.wscale = torch.rand(
145 + (scale_hidden_size, scale_hidden_size), dtype=torch.float32
146 + )
147 +
148 + self.enable_silu_mul_custom_op = self.silu_and_mul.enabled()
149 +
150 + def forward(self, x):
151 + y = self.silu_and_mul(x)
152 + x2 = self.w8a8_block_fp8_linear.apply(y, self.w, self.wscale)
153 + return x2
154 +
155 + def ops_in_model_before(self):
156 + return [
157 + SILU_MUL_OP if self.enable_silu_mul_custom_op else torch.ops.aten.mul,
158 + ]
159 +
160 + def ops_in_model_after(self):
161 + return [torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant]
162 +
163 +
164 @pytest.mark.parametrize("num_tokens", [32, 64])
165 @pytest.mark.parametrize("hidden_size", [128, 256])
166 @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
@@ -133,7 +168,10 @@ class TestSiluMulNvfp4QuantModel(torch.nn.Module):
168 @pytest.mark.parametrize(
169 "model_class, enable_quant_fp8_custom_op, cuda_force_torch",
170 list(itertools.product([TestSiluMulFp8QuantModel], [True, False], [True, False]))
136 - + [(TestSiluMulNvfp4QuantModel, False, False)],
171 + + [
172 + (TestSiluMulNvfp4QuantModel, False, False),
173 + (TestSiluMulGroupFp8QuantModel, False, False),
174 + ],
175 )
176 # cuda_force_torch used to test torch code path on platforms that
177 # cutlass_fp8_supported() == True.
@@ -144,13 +182,19 @@ def test_fusion_silu_and_mul_quant(
182 num_tokens: int,
183 hidden_size: int,
184 dtype: torch.dtype,
147 - model_class: type[TestSiluMulFp8QuantModel | TestSiluMulNvfp4QuantModel],
185 + model_class: type[
186 + TestSiluMulFp8QuantModel
187 + | TestSiluMulNvfp4QuantModel
188 + | TestSiluMulGroupFp8QuantModel
189 + ],
190 enable_silu_mul_custom_op: bool,
191 enable_quant_fp8_custom_op: bool,
192 cuda_force_torch: bool,
193 ):
194 if model_class is TestSiluMulNvfp4QuantModel and not is_nvfp4_supported():
195 pytest.skip("NVFP4 is not supported on this GPU.")
196 + if model_class is TestSiluMulGroupFp8QuantModel and not IS_AITER_FOUND:
197 + pytest.skip("AITER is not supported on this GPU.")
198
199 torch.set_default_device("cuda")
200 torch.set_default_dtype(dtype)
@@ -173,9 +217,15 @@ def test_fusion_silu_and_mul_quant(
217 )
218
219 with set_current_vllm_config(config):
176 - fusion_pass = ActivationQuantFusionPass(config)
220 + fusion_passes = [ActivationQuantFusionPass(config)]
221 + if IS_AITER_FOUND:
222 + from vllm.compilation.rocm_aiter_fusion import (
223 + RocmAiterSiluMulFp8GroupQuantFusionPass,
224 + )
225 +
226 + fusion_passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
227
178 - passes = [NoOpEliminationPass(config), fusion_pass, PostCleanupPass(config)]
228 + passes = [NoOpEliminationPass(config), *fusion_passes, PostCleanupPass(config)]
229 backend = TestBackend(*passes)
230 model = model_class(
231 hidden_size=hidden_size, cuda_force_torch=cuda_force_torch, x=x
@@ -194,12 +244,14 @@ def test_fusion_silu_and_mul_quant(
244 atol, rtol = 1e-3, 1e-3
245 elif model_class == TestSiluMulNvfp4QuantModel:
246 atol, rtol = 1e-1, 1e-1
247 + elif model_class == TestSiluMulGroupFp8QuantModel:
248 + atol, rtol = 5e-2, 5e-2
249
250 torch.testing.assert_close(
251 result[0].to(dtype=dtype), result2[0].to(dtype=dtype), atol=atol, rtol=rtol
252 )
253
202 - assert fusion_pass.matched_count == 1
254 + assert sum([p.matched_count for p in fusion_passes]) == 1
255
256 # In pre-nodes, quant op should be present and fused kernels should not
257 backend.check_before_ops(model.ops_in_model_before())
vllm/_aiter_ops.py
+170 -44
@@ -24,6 +24,15 @@ def is_aiter_found() -> bool:
24 # we keep this global outside to not cause torch compile breaks.
25 IS_AITER_FOUND = is_aiter_found()
26
27 +# Can't use dtypes.fp8 directly inside an op
28 +# because it returns wrong result on gfx942.
29 +# This is a workaround to get the correct FP8 dtype.
30 +# This might because that the get_gfx() is wrapped as a custom op.
31 +if IS_AITER_FOUND:
32 + from aiter import dtypes
33 +
34 + AITER_FP8_DTYPE = dtypes.fp8
35 +
36
37 def if_aiter_supported(func: Callable) -> Callable:
38 """Decorator that only executes the function if
@@ -45,36 +54,6 @@ def if_aiter_supported(func: Callable) -> Callable:
54 return wrapper
55
56
48 -def _rocm_aiter_group_fp8_quant_impl(
49 - x: torch.Tensor,
50 - group_size: int,
51 -) -> tuple[torch.Tensor, torch.Tensor]:
52 - assert x.shape[-1] % group_size == 0, "Input shape must be divisible by group size"
53 - from aiter import QuantType, dtypes, get_hip_quant
54 -
55 - aiter_per1x128_quant = get_hip_quant(QuantType.per_1x128)
56 - return aiter_per1x128_quant(x.contiguous(), quant_dtype=dtypes.fp8)
57 -
58 -
59 -def _rocm_aiter_group_fp8_quant_fake(
60 - x: torch.Tensor,
61 - group_size: int,
62 -) -> tuple[torch.Tensor, torch.Tensor]:
63 - from aiter import dtypes
64 -
65 - M, N = x.shape
66 - x_fp8 = torch.empty((M, N), dtype=dtypes.fp8, device=x.device)
67 - out_bs = torch.empty(
68 - (
69 - M,
70 - (N + group_size - 1) // group_size,
71 - ),
72 - dtype=torch.float32,
73 - device=x.device,
74 - )
75 - return x_fp8, out_bs
76 -
77 -
57 def _rocm_aiter_fused_moe_impl(
58 hidden_states: torch.Tensor,
59 w1: torch.Tensor,
@@ -522,6 +501,142 @@ def _rocm_aiter_per_token_quant_fake(
501 )
502
503
504 +def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl(
505 + x: torch.Tensor,
506 + residual: torch.Tensor,
507 + weight: torch.Tensor,
508 + variance_epsilon: float,
509 + group_size: int,
510 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
511 + from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
512 +
513 + (x_quant, x_quant_scales), _, _, res = fused_rms_fp8_group_quant(
514 + x,
515 + weight,
516 + variance_epsilon,
517 + None,
518 + None,
519 + None,
520 + group_size=group_size,
521 + dtype_quant=AITER_FP8_DTYPE,
522 + res1=residual,
523 + )
524 + return (x_quant, x_quant_scales, res)
525 +
526 +
527 +def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake(
528 + x: torch.Tensor,
529 + residual: torch.Tensor,
530 + weight: torch.Tensor,
531 + variance_epsilon: float,
532 + group_size: int,
533 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
534 + M, N = x.shape
535 + scale_shape = (M, (N + group_size - 1) // group_size)
536 + return (
537 + torch.empty_like(x, dtype=AITER_FP8_DTYPE, device=x.device),
538 + torch.empty(scale_shape, dtype=torch.float32, device=x.device),
539 + torch.empty_like(residual, device=residual.device),
540 + )
541 +
542 +
543 +def _rocm_aiter_rmsnorm_fp8_group_quant_impl(
544 + x: torch.Tensor,
545 + weight: torch.Tensor,
546 + variance_epsilon: float,
547 + group_size: int,
548 +) -> tuple[torch.Tensor, torch.Tensor]:
549 + from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
550 +
551 + (x_quant, x_quant_scales), _, _, res = fused_rms_fp8_group_quant(
552 + x,
553 + weight,
554 + variance_epsilon,
555 + None,
556 + None,
557 + None,
558 + group_size=group_size,
559 + dtype_quant=AITER_FP8_DTYPE,
560 + res1=None,
561 + )
562 + return (x_quant, x_quant_scales)
563 +
564 +
565 +def _rocm_aiter_rmsnorm_fp8_group_quant_fake(
566 + x: torch.Tensor,
567 + weight: torch.Tensor,
568 + variance_epsilon: float,
569 + group_size: int,
570 +) -> tuple[torch.Tensor, torch.Tensor]:
571 + M, N = x.shape
572 + scale_shape = (M, (N + group_size - 1) // group_size)
573 + return (
574 + torch.empty_like(x, dtype=AITER_FP8_DTYPE, device=x.device),
575 + torch.empty(scale_shape, dtype=torch.float32, device=x.device),
576 + )
577 +
578 +
579 +def _rocm_aiter_group_fp8_quant_impl(
580 + x: torch.Tensor,
581 + group_size: int,
582 +) -> tuple[torch.Tensor, torch.Tensor]:
583 + assert x.shape[-1] % group_size == 0, "Input shape must be divisible by group size"
584 + from aiter import QuantType, get_hip_quant
585 +
586 + aiter_per1x128_quant = get_hip_quant(QuantType.per_1x128)
587 + return aiter_per1x128_quant(x.contiguous(), quant_dtype=AITER_FP8_DTYPE)
588 +
589 +
590 +def _rocm_aiter_group_fp8_quant_fake(
591 + x: torch.Tensor,
592 + group_size: int,
593 +) -> tuple[torch.Tensor, torch.Tensor]:
594 + M, N = x.shape
595 + x_fp8 = torch.empty((M, N), dtype=AITER_FP8_DTYPE, device=x.device)
596 + out_bs = torch.empty(
597 + (
598 + M,
599 + (N + group_size - 1) // group_size,
600 + ),
601 + dtype=torch.float32,
602 + device=x.device,
603 + )
604 + return x_fp8, out_bs
605 +
606 +
607 +def _rocm_aiter_act_mul_and_fp8_group_quant_impl(
608 + x: torch.Tensor,
609 + group_size: int,
610 +) -> tuple[torch.Tensor, torch.Tensor]:
611 + from aiter.ops.triton.activation import act_mul_and_fp8_group_quant
612 +
613 + return act_mul_and_fp8_group_quant(
614 + x,
615 + activation="silu",
616 + group_size=group_size,
617 + dtype_quant=AITER_FP8_DTYPE,
618 + )
619 +
620 +
621 +def _rocm_aiter_act_mul_and_fp8_group_quant_fake(
622 + x: torch.Tensor,
623 + group_size: int,
624 +) -> tuple[torch.Tensor, torch.Tensor]:
625 + M, N = x.shape
626 + assert N % 2 == 0
627 + N_half = N // 2
628 + x_fp8 = torch.empty((M, N_half), dtype=AITER_FP8_DTYPE, device=x.device)
629 + out_bs = torch.empty(
630 + (
631 + M,
632 + (N_half + group_size - 1) // group_size,
633 + ),
634 + dtype=torch.float32,
635 + device=x.device,
636 + )
637 + return x_fp8, out_bs
638 +
639 +
640 # Global flag to ensure ops are registered only once
641 _OPS_REGISTERED = False
642
@@ -557,7 +672,7 @@ class rocm_aiter_ops:
672 @if_aiter_supported
673 def is_linear_fp8_enaled(cls) -> bool:
674 """ "Verifies device specs and availability of env variable."""
560 - return cls.is_linear_enabled() and current_platform.is_fp8_fnuz()
675 + return cls.is_linear_enabled()
676
677 @classmethod
678 @if_aiter_supported
@@ -632,14 +747,6 @@ class rocm_aiter_ops:
747 )
748
749 # register all the custom ops here
635 - direct_register_custom_op(
636 - op_name="rocm_aiter_group_fp8_quant",
637 - op_func=_rocm_aiter_group_fp8_quant_impl,
638 - mutates_args=[],
639 - fake_impl=_rocm_aiter_group_fp8_quant_fake,
640 - dispatch_key=current_platform.dispatch_key,
641 - )
642 -
750 direct_register_custom_op(
751 op_name="rocm_aiter_asm_moe_tkw1",
752 op_func=_rocm_aiter_asm_moe_tkw1_impl,
@@ -699,27 +806,46 @@ class rocm_aiter_ops:
806 direct_register_custom_op(
807 op_name="rocm_aiter_gemm_a8w8_blockscale",
808 op_func=_rocm_aiter_gemm_a8w8_blockscale_impl,
702 - mutates_args=[],
809 fake_impl=_rocm_aiter_gemm_a8w8_blockscale_fake,
704 - dispatch_key=current_platform.dispatch_key,
810 )
811
812 direct_register_custom_op(
813 op_name="rocm_aiter_rms_norm",
814 op_func=_rocm_aiter_rms_norm_impl,
710 - mutates_args=[],
815 fake_impl=_rocm_aiter_rms_norm_fake,
712 - dispatch_key=current_platform.dispatch_key,
816 )
817
818 direct_register_custom_op(
819 op_name="rocm_aiter_rmsnorm2d_fwd_with_add",
820 op_func=_rocm_aiter_rmsnorm2d_fwd_with_add_impl,
718 - mutates_args=[],
821 fake_impl=_rocm_aiter_rmsnorm2d_fwd_with_add_fake,
822 dispatch_key=current_platform.dispatch_key,
823 )
824
825 + direct_register_custom_op(
826 + op_name="rocm_aiter_rmsnorm_fp8_group_quant",
827 + op_func=_rocm_aiter_rmsnorm_fp8_group_quant_impl,
828 + fake_impl=_rocm_aiter_rmsnorm_fp8_group_quant_fake,
829 + )
830 +
831 + direct_register_custom_op(
832 + op_name="rocm_aiter_rmsnorm_with_add_fp8_group_quant",
833 + op_func=_rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl,
834 + fake_impl=_rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake,
835 + )
836 +
837 + direct_register_custom_op(
838 + op_name="rocm_aiter_act_mul_and_fp8_group_quant",
839 + op_func=_rocm_aiter_act_mul_and_fp8_group_quant_impl,
840 + fake_impl=_rocm_aiter_act_mul_and_fp8_group_quant_fake,
841 + )
842 +
843 + direct_register_custom_op(
844 + op_name="rocm_aiter_group_fp8_quant",
845 + op_func=_rocm_aiter_group_fp8_quant_impl,
846 + fake_impl=_rocm_aiter_group_fp8_quant_fake,
847 + )
848 +
849 direct_register_custom_op(
850 op_name="rocm_aiter_per_tensor_quant",
851 op_func=_rocm_aiter_per_tensor_quant_impl,
vllm/compilation/pass_manager.py
+11
@@ -5,6 +5,7 @@ import functools
5 from torch import fx as fx
6
7 from vllm import envs
8 +from vllm._aiter_ops import rocm_aiter_ops
9 from vllm.config import VllmConfig, set_current_vllm_config
10 from vllm.logger import init_logger
11 from vllm.platforms import current_platform
@@ -13,6 +14,12 @@ from vllm.utils.system_utils import set_env_var
14 from .post_cleanup import PostCleanupPass
15 from .vllm_inductor_pass import VllmInductorPass
16
17 +if rocm_aiter_ops.is_enabled():
18 + from vllm.compilation.rocm_aiter_fusion import (
19 + RocmAiterRMSNormFp8GroupQuantFusionPass,
20 + RocmAiterSiluMulFp8GroupQuantFusionPass,
21 + )
22 +
23 if current_platform.is_cuda_alike():
24 from .activation_quant_fusion import ActivationQuantFusionPass
25 from .fusion import RMSNormQuantFusionPass
@@ -109,8 +116,12 @@ class PostGradPassManager(CustomGraphPass):
116
117 if self.pass_config.fuse_norm_quant:
118 self.passes += [RMSNormQuantFusionPass(config)]
119 + if rocm_aiter_ops.is_enabled():
120 + self.passes += [RocmAiterRMSNormFp8GroupQuantFusionPass(config)]
121 if self.pass_config.fuse_act_quant:
122 self.passes += [ActivationQuantFusionPass(config)]
123 + if rocm_aiter_ops.is_enabled():
124 + self.passes += [RocmAiterSiluMulFp8GroupQuantFusionPass(config)]
125
126 if self.pass_config.fuse_attn_quant:
127 self.passes += [AttnFusionPass(config)]
vllm/compilation/rocm_aiter_fusion.py new
+242
@@ -0,0 +1,242 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +from typing import Any
4 +
5 +import torch
6 +import torch._inductor.pattern_matcher as pm
7 +from torch import fx
8 +from torch._inductor.pattern_matcher import PatternMatcherPass
9 +from torch._ops import OpOverload
10 +
11 +import vllm.model_executor.layers.quantization.utils.fp8_utils # noqa: F401
12 +from vllm.compilation.activation_quant_fusion import ActivationQuantPattern
13 +from vllm.config import VllmConfig
14 +from vllm.logger import init_logger
15 +from vllm.platforms import current_platform
16 +
17 +from .fusion import empty_bf16
18 +from .inductor_pass import enable_fake_mode
19 +from .matcher_utils import MatcherSiluAndMul
20 +from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
21 +
22 +logger = init_logger(__name__)
23 +FP8_DTYPE = current_platform.fp8_dtype()
24 +
25 +AITER_RMS_GROUP_QUANT_OP = torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default
26 +AITER_RMS_ADD_GROUP_QUANT_OP = (
27 + torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant.default
28 +)
29 +
30 +AITER_RMS_OP = torch.ops.vllm.rocm_aiter_rms_norm.default
31 +AITER_RMS_ADD_OP = torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add.default
32 +
33 +AITER_GROUP_FP8_QUANT_OP = torch.ops.vllm.rocm_aiter_group_fp8_quant.default
34 +TRITON_GROUP_FP8_QUANT_OP = torch.ops.vllm.triton_per_token_group_quant_fp8.default
35 +
36 +FUSED_SILU_MUL_QUANT_OP = torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant.default
37 +
38 +
39 +class AiterRMSFp8GroupQuantPattern:
40 + """
41 + This pattern fuses aiter rms_norm & group fp8 quant custom
42 + ops into an aiter rms_norm_group_fp8_quant op.
43 + """
44 +
45 + def __init__(self, epsilon: float, quant_dtype: torch.dtype, quant_op: OpOverload):
46 + self.epsilon = epsilon
47 + self.quant_dtype = quant_dtype
48 + self.quant_op = quant_op
49 +
50 + def register(self, pm_pass: PatternMatcherPass):
51 + def pattern(
52 + input: torch.Tensor,
53 + weight: torch.Tensor,
54 + ):
55 + at1 = AITER_RMS_OP(x=input, weight=weight, variance_epsilon=self.epsilon)
56 +
57 + at2 = self.quant_op(at1, 128)
58 +
59 + return at2[0], at2[1]
60 +
61 + def replacement(
62 + input: torch.Tensor,
63 + weight: torch.Tensor,
64 + ):
65 + at = AITER_RMS_GROUP_QUANT_OP(
66 + x=input,
67 + weight=weight,
68 + variance_epsilon=self.epsilon,
69 + group_size=128,
70 + )
71 +
72 + return at[0], at[1]
73 +
74 + inputs = [
75 + empty_bf16(5, 4), # input
76 + empty_bf16(1, 5), # weight
77 + ]
78 +
79 + pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
80 +
81 +
82 +class AiterFusedAddRMSFp8GroupQuantPattern:
83 + """
84 + This pattern fuses aiter rms_norm_with_add & group fp8 quant custom ops
85 + into a aiter rms_norm_with_add_group_fp8_quant op.
86 + """
87 +
88 + def __init__(self, epsilon: float, quant_dtype: torch.dtype, quant_op: OpOverload):
89 + self.epsilon = epsilon
90 + self.quant_dtype = quant_dtype
91 + self.quant_op = quant_op
92 +
93 + def register(self, pm_pass: PatternMatcherPass):
94 + def pattern(
95 + input: torch.Tensor,
96 + residual: torch.Tensor,
97 + weight: torch.Tensor,
98 + ):
99 + at1 = AITER_RMS_ADD_OP(
100 + x=input,
101 + residual=residual,
102 + weight=weight,
103 + variance_epsilon=self.epsilon,
104 + )
105 +
106 + at2 = self.quant_op(at1[0], 128)
107 +
108 + # result, scale, residual
109 + return at2[0], at2[1], at1[1]
110 +
111 + def replacement(
112 + input: torch.Tensor,
113 + residual: torch.Tensor,
114 + weight: torch.Tensor,
115 + ):
116 + at = AITER_RMS_ADD_GROUP_QUANT_OP(
117 + x=input,
118 + residual=residual,
119 + weight=weight,
120 + variance_epsilon=self.epsilon,
121 + group_size=128,
122 + )
123 +
124 + # result, scale, residual
125 + return at[0], at[1], at[2]
126 +
127 + inputs = [
128 + empty_bf16(5, 4), # input
129 + empty_bf16(5, 4), # residual
130 + empty_bf16(1, 5), # weight
131 + ]
132 +
133 + pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
134 +
135 +
136 +class RocmAiterRMSNormFp8GroupQuantFusionPass(VllmPatternMatcherPass):
137 + """
138 + This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op.
139 + It also supports fused_add_rms_norm.
140 + """
141 +
142 + @enable_fake_mode
143 + def __init__(self, config: VllmConfig):
144 + super().__init__(config)
145 +
146 + self.patterns: PatternMatcherPass = PatternMatcherPass(
147 + pass_name="rocm_aiter_rms_norm_fp8_group_quant_fusion_pass"
148 + )
149 +
150 + # Make sure fused add patterns are before simple rms norm,
151 + # as the latter is a subset of the former in torch ops
152 + for epsilon in [1e-5, 1e-6]:
153 + # Fuse rms_norm + dynamic group fp8 quant
154 + for quant_op in [AITER_GROUP_FP8_QUANT_OP, TRITON_GROUP_FP8_QUANT_OP]:
155 + AiterRMSFp8GroupQuantPattern(epsilon, FP8_DTYPE, quant_op).register(
156 + self.patterns
157 + )
158 +
159 + AiterFusedAddRMSFp8GroupQuantPattern(
160 + epsilon, FP8_DTYPE, quant_op
161 + ).register(self.patterns)
162 +
163 + self.dump_patterns(config, self.patterns)
164 +
165 + @VllmInductorPass.time_and_log
166 + def __call__(self, graph: fx.Graph):
167 + self.matched_count = self.patterns.apply(graph)
168 + logger.debug("Replaced %s patterns", self.matched_count)
169 +
170 + def uuid(self) -> Any:
171 + fusion_patterns = [
172 + AiterRMSFp8GroupQuantPattern,
173 + AiterFusedAddRMSFp8GroupQuantPattern,
174 + ]
175 + return self.hash_source(self, *fusion_patterns)
176 +
177 +
178 +class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
179 + """
180 + This pattern fuses aiter silu_and_mul & group fp8 quant custom
181 + ops into an aiter silu_and_mul_group_fp8_quant op.
182 + """
183 +
184 + def __init__(self, quant_op: OpOverload):
185 + self.silu_and_mul_matcher = MatcherSiluAndMul()
186 + self.quant_op = quant_op
187 +
188 + def register(self, pm_pass: PatternMatcherPass):
189 + def pattern(
190 + input: torch.Tensor,
191 + ):
192 + at1 = self.silu_and_mul_matcher(input)
193 + at2 = self.quant_op(at1, 128)
194 + return at2[0], at2[1]
195 +
196 + def replacement(
197 + input: torch.Tensor,
198 + ):
199 + at = FUSED_SILU_MUL_QUANT_OP(x=input, group_size=128)
200 + return at[0], at[1]
201 +
202 + inputs = [
203 + self.silu_and_mul_matcher.inputs()[0],
204 + ]
205 +
206 + pm.register_replacement(pattern, replacement, inputs, pm.fwd_only, pm_pass)
207 +
208 +
209 +class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
210 + """
211 + This pass fuses a pre-defined set of custom ops into fused ops.
212 + It uses the torch pattern matcher to find the patterns and replace them.
213 +
214 + Because patterns can only be registered once, the pass is a singleton.
215 + This will be addressed in a future version of PyTorch:
216 + https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980
217 + """
218 +
219 + @enable_fake_mode
220 + def __init__(self, config: VllmConfig):
221 + super().__init__(config)
222 +
223 + self.patterns: PatternMatcherPass = PatternMatcherPass(
224 + pass_name="rocm_aiter_silu_mul_fp8_group_quant_fusion_pass"
225 + )
226 +
227 + for quant_op in [AITER_GROUP_FP8_QUANT_OP, TRITON_GROUP_FP8_QUANT_OP]:
228 + AiterSiluMulFp8GroupQuantPattern(quant_op).register(self.patterns)
229 +
230 + self.dump_patterns(config, self.patterns)
231 +
232 + @VllmInductorPass.time_and_log
233 + def __call__(self, graph: torch.fx.Graph):
234 + self.matched_count = self.patterns.apply(graph)
235 + logger.debug("Replaced %s patterns", self.matched_count)
236 +
237 + def uuid(self):
238 + fusion_patterns = [
239 + ActivationQuantPattern,
240 + AiterSiluMulFp8GroupQuantPattern,
241 + ]
242 + return VllmInductorPass.hash_source(self, *fusion_patterns)
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+37 -6
@@ -196,6 +196,39 @@ direct_register_custom_op(
196 )
197
198
199 +def _triton_per_token_group_quant_fp8_impl(
200 + x: torch.Tensor,
201 + group_size: int,
202 +) -> tuple[torch.Tensor, torch.Tensor]:
203 + return per_token_group_quant_fp8(
204 + x, group_size, column_major_scales=False, use_ue8m0=False
205 + )
206 +
207 +
208 +def _triton_per_token_group_quant_fp8_fake(
209 + x: torch.Tensor,
210 + group_size: int,
211 +) -> tuple[torch.Tensor, torch.Tensor]:
212 + M, N = x.shape
213 + x_fp8 = torch.empty((M, N), dtype=current_platform.fp8_dtype(), device=x.device)
214 + out_bs = torch.empty(
215 + (
216 + M,
217 + (N + group_size - 1) // group_size,
218 + ),
219 + dtype=torch.float32,
220 + device=x.device,
221 + )
222 + return x_fp8, out_bs
223 +
224 +
225 +direct_register_custom_op(
226 + "triton_per_token_group_quant_fp8",
227 + _triton_per_token_group_quant_fp8_impl,
228 + fake_impl=_triton_per_token_group_quant_fp8_fake,
229 +)
230 +
231 +
232 # TODO fix ROCm->Triton custom path:
233 # https://github.com/vllm-project/vllm/issues/14397
234 class W8A8BlockFp8LinearOp:
@@ -341,17 +374,15 @@ class W8A8BlockFp8LinearOp:
374
375 if input_scale is not None:
376 q_input = input_2d
344 - # MI350 case uses triton kernel
377 elif use_triton:
346 - q_input, input_scale = per_token_group_quant_fp8(
378 + q_input, input_scale = torch.ops.vllm.triton_per_token_group_quant_fp8(
379 input_2d,
380 self.act_quant_group_shape.col,
349 - column_major_scales=False,
350 - use_ue8m0=False,
381 )
352 - # MI300 uses tuned AITER ASM/C++ kernel
382 else:
354 - q_input, input_scale = rocm_aiter_ops.group_fp8_quant(input_2d)
383 + q_input, input_scale = rocm_aiter_ops.group_fp8_quant(
384 + input_2d, self.act_quant_group_shape.col
385 + )
386
387 return gemm_a8w8_blockscale_op(
388 q_input,