@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f32cfd7d97
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, 00:00 UTC

[ROCm][FEAT] Support AITER RMSNorm quantization fusion pass (#26575)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>

vllmellm committed Dec 23, 2025 at 11:07 UTC f32cfd7d979796fbe1af4bdce3bc3b7bfe512607
5 files changed +663 -219
tests/compile/test_fusion.py
+189 -127
@@ -1,7 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -import itertools
4
5 import pytest
6 import torch
@@ -53,37 +52,61 @@ class TestModel(torch.nn.Module):
52 hidden_size: int,
53 eps: float,
54 group_shape: GroupShape,
56 - cuda_force_torch: bool,
55 + use_aiter: bool = False,
56 + cuda_force_torch: bool = False,
57 + use_aiter_quant_op: bool = True,
58 *args,
59 **kwargs,
60 ):
61 super().__init__(*args, **kwargs)
62 + self.use_aiter = use_aiter
63 + self.use_aiter_quant_op = use_aiter_quant_op
64 self.cuda_force_torch = cuda_force_torch
65 + self.group_shape = group_shape
66 + self.enable_quant_fp8_custom_op = None # Will be set later if applicable
67 +
68 self.norm = [RMSNorm(hidden_size, eps) for _ in range(4)]
63 - if group_shape.is_per_group():
64 - self.wscale = [
65 - torch.rand(
66 - (hidden_size // group_shape[1], hidden_size // group_shape[1]),
67 - dtype=torch.float32,
68 - )
69 - for _ in range(3)
70 - ]
71 - else:
72 - self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
73 - static = group_shape == GroupShape.PER_TENSOR
69 +
70 + # Setup quantization scale descriptor
71 + static = group_shape == GroupShape.PER_TENSOR and not use_aiter
72 quant_scale = ScaleDesc(torch.float32, static, group_shape)
73 self.quant_key = QuantKey(dtype=FP8_DTYPE, scale=quant_scale, symmetric=True)
74 +
75 + # Setup scales
76 if static:
77 self.scale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
78 else:
79 self.scale = [None for _ in range(3)]
80 +
81 + # Setup weights
82 self.w = [
83 torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE) for _ in range(3)
84 ]
83 - if not group_shape.is_per_group():
85 + if not group_shape.is_per_group() or use_aiter:
86 self.w = [self.w[0].t() for _ in range(3)]
87
88 + # Setup weight scales
89 if group_shape.is_per_group():
90 + scale_size = (
91 + (hidden_size + 128 - 1) // 128
92 + if use_aiter
93 + else hidden_size // group_shape[1]
94 + )
95 + wscale_shape: tuple[int, ...] = (scale_size, scale_size)
96 + else:
97 + wscale_shape = (1,)
98 + self.wscale = [torch.rand(wscale_shape, dtype=torch.float32) for _ in range(3)]
99 +
100 + # Setup FP8 linear operation
101 + is_per_group = group_shape.is_per_group()
102 + if is_per_group and use_aiter:
103 + self.fp8_linear = W8A8BlockFp8LinearOp(
104 + weight_group_shape=GroupShape(128, 128),
105 + act_quant_group_shape=group_shape,
106 + use_aiter_and_is_supported=use_aiter_quant_op,
107 + )
108 + # AITER blockwise doesn't use enable_quant_fp8_custom_op
109 + elif is_per_group:
110 self.fp8_linear = W8A8BlockFp8LinearOp(
111 weight_group_shape=GroupShape(group_shape[1], group_shape[1]),
112 act_quant_group_shape=group_shape,
@@ -91,6 +114,13 @@ class TestModel(torch.nn.Module):
114 use_aiter_and_is_supported=False,
115 )
116 self.enable_quant_fp8_custom_op = self.fp8_linear.input_quant_op.enabled()
117 + elif use_aiter:
118 + self.fp8_linear = Fp8LinearOp(
119 + act_quant_static=False,
120 + act_quant_group_shape=group_shape,
121 + )
122 + self.fp8_linear.quant_fp8.use_aiter = use_aiter_quant_op
123 + self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
124 else:
125 with override_cutlass_fp8_supported(not cuda_force_torch):
126 self.fp8_linear = Fp8LinearOp(
@@ -100,7 +130,6 @@ class TestModel(torch.nn.Module):
130 self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
131
132 self.enable_rms_norm_custom_op = self.norm[0].enabled()
103 - self.group_shape = group_shape
133
134 def forward(self, x):
135 # avoid having graph input be an arg to a pattern directly
@@ -126,19 +155,49 @@ class TestModel(torch.nn.Module):
155 y4, resid = self.norm[3](x4, resid) # use resid here
156 return y4
157
158 + def ops_in_model_before(self):
159 + if (
160 + self.use_aiter
161 + and self.group_shape.is_per_group()
162 + and current_platform.is_fp8_fnuz()
163 + ):
164 + return [rocm_aiter_ops.get_group_quant_op()]
165 + if self.use_aiter and self.group_shape.is_per_group():
166 + return [torch.ops.vllm.triton_per_token_group_quant_fp8.default]
167 + if self.use_aiter and self.use_aiter_quant_op:
168 + return [rocm_aiter_ops.get_per_token_quant_op()]
169 + if self.use_aiter:
170 + return [QUANT_OPS[self.quant_key]]
171 + if self.enable_quant_fp8_custom_op:
172 + return [QUANT_OPS[self.quant_key]]
173 + return [torch.ops.aten.reciprocal]
174 +
175 def ops_in_model_after(self):
176 + if self.use_aiter and self.group_shape.is_per_group():
177 + from vllm.compilation.rocm_aiter_fusion import (
178 + AiterFusedAddRMSFp8GroupQuantPattern,
179 + AiterRMSFp8GroupQuantPattern,
180 + )
181 +
182 + return [
183 + AiterFusedAddRMSFp8GroupQuantPattern.FUSED_OP,
184 + AiterRMSFp8GroupQuantPattern.FUSED_OP,
185 + ]
186 + if self.use_aiter:
187 + from vllm.compilation.rocm_aiter_fusion import (
188 + AiterFusedAddRMSNormDynamicQuantPattern,
189 + AiterRMSNormDynamicQuantPattern,
190 + )
191 +
192 + return [
193 + AiterFusedAddRMSNormDynamicQuantPattern.FUSED_OP,
194 + AiterRMSNormDynamicQuantPattern.FUSED_OP,
195 + ]
196 return [
197 FUSED_OPS[FusedRMSQuantKey(self.quant_key, True)],
198 FUSED_OPS[FusedRMSQuantKey(self.quant_key, False)],
199 ]
200
135 - def ops_in_model_before(self):
136 - return (
137 - [QUANT_OPS[self.quant_key]]
138 - if self.enable_quant_fp8_custom_op
139 - else [torch.ops.aten.reciprocal]
140 - )
141 -
201 def ops_in_model_before_partial(self):
202 return (
203 [RMS_OP, RMS_ADD_OP]
@@ -155,67 +214,45 @@ GROUP_SHAPES = [
214 ]
215
216
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 - ]
217 +def _run_fusion_test(
218 + model,
219 + fusion_pass,
220 + vllm_config,
221 + dtype,
222 + hidden_size,
223 + num_tokens,
224 +):
225 + """Helper function for common fusion test logic.
226
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 - ]
227 + Must be called within vllm_config context.
228 + """
229 + noop_pass = NoOpEliminationPass(vllm_config)
230 + cleanup_pass = PostCleanupPass(vllm_config)
231
178 - self.norm_weight = [torch.ones(hidden_size) for _ in range(4)]
179 - self.eps = eps
232 + backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)
233 + backend2 = TestBackend(noop_pass, cleanup_pass)
234
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)
235 + x = torch.rand(num_tokens, hidden_size)
236 + torch._dynamo.mark_dynamic(x, 0)
237
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 - )
238 + model_fused = torch.compile(model, backend=backend)
239 + result_fused = model_fused(x)
240
192 - x3 = self.w8a8_block_fp8_linear.apply(y2, self.w[1], self.wscale[1])
241 + model_unfused = torch.compile(model, backend=backend2)
242 + result_unfused = model_unfused(x)
243
194 - y3, resid = rocm_aiter_ops.rms_norm2d_with_add(
195 - x3, resid, self.norm_weight[2], self.eps
196 - )
244 + if dtype == torch.float16:
245 + ATOL, RTOL = (2e-3, 2e-3)
246 + else:
247 + ATOL, RTOL = (1e-2, 1e-2)
248
198 - x4 = self.w8a8_block_fp8_linear.apply(y3, self.w[2], self.wscale[2])
249 + torch.testing.assert_close(result_fused, result_unfused, atol=ATOL, rtol=RTOL)
250
200 - y4, resid = rocm_aiter_ops.rms_norm2d_with_add(
201 - x4, resid, self.norm_weight[3], self.eps
202 - )
203 - return y4
251 + assert fusion_pass.matched_count == 3
252 + backend.check_before_ops(model.ops_in_model_before())
253 + backend.check_after_ops(model.ops_in_model_after())
254
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 - ]
255 + return backend, backend2
256
257
258 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@@ -223,11 +260,8 @@ class TestRmsnormGroupFp8QuantModel(torch.nn.Module):
260 @pytest.mark.parametrize("num_tokens", [257])
261 @pytest.mark.parametrize("eps", [1e-5, 1e-6])
262 @pytest.mark.parametrize("group_shape", GROUP_SHAPES)
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 -)
263 +@pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
264 +@pytest.mark.parametrize("enable_quant_fp8_custom_op", [True, False])
265 # cuda_force_torch used to test torch code path on platforms that
266 # cutlass_fp8_supported() == True.
267 @pytest.mark.parametrize(
@@ -242,23 +276,13 @@ def test_fusion_rmsnorm_quant(
276 num_tokens,
277 eps,
278 group_shape,
245 - model_class,
279 enable_rms_norm_custom_op,
280 enable_quant_fp8_custom_op,
281 cuda_force_torch,
282 ):
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)
256 - maybe_create_device_identity() # needed for certain non-cutlass fp8 paths
257 -
283 if not enable_quant_fp8_custom_op and group_shape.is_per_group():
284 pytest.skip("Unsupported unwrapped quant fp8 op for blockwise quantization")
285
261 - # Skip test for 64-bit group shape when running with cutlass or deepgemm
286 if group_shape == GroupShape(1, 64) and (
287 cutlass_block_fp8_supported() or is_deep_gemm_supported()
288 ):
@@ -269,6 +293,7 @@ def test_fusion_rmsnorm_quant(
293 custom_ops.append("+rms_norm")
294 if enable_quant_fp8_custom_op:
295 custom_ops.append("+quant_fp8")
296 +
297 vllm_config = VllmConfig(
298 model_config=ModelConfig(dtype=dtype),
299 compilation_config=CompilationConfig(
@@ -279,60 +304,97 @@ def test_fusion_rmsnorm_quant(
304 ),
305 ),
306 )
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)
285 - if model_class is TestRmsnormGroupFp8QuantModel:
286 - from vllm.compilation.rocm_aiter_fusion import (
287 - RocmAiterRMSNormFp8GroupQuantFusionPass,
288 - )
307
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)
297 - model = model_class(
308 + with vllm.config.set_current_vllm_config(vllm_config):
309 + # Setup device before model creation
310 + torch.set_default_device("cuda")
311 + torch.set_default_dtype(dtype)
312 + torch.manual_seed(1)
313 + maybe_create_device_identity()
314 +
315 + fusion_pass = RMSNormQuantFusionPass(vllm_config)
316 + model = TestModel(
317 hidden_size=hidden_size,
318 eps=eps,
319 group_shape=group_shape,
320 + use_aiter=False,
321 cuda_force_torch=cuda_force_torch,
322 )
303 - # First dimension dynamic
304 - x = torch.rand(num_tokens, hidden_size)
305 - torch._dynamo.mark_dynamic(x, 0)
306 -
307 - model_fused = torch.compile(model, backend=backend)
308 - result_fused = model_fused(x)
309 -
310 - model_unfused = torch.compile(model, backend=backend2)
311 - result_unfused = model_unfused(x)
312 -
313 - if dtype == torch.float16:
314 - ATOL, RTOL = (2e-3, 2e-3)
315 - else:
316 - ATOL, RTOL = (1e-2, 1e-2)
317 -
318 - torch.testing.assert_close(result_fused, result_unfused, atol=ATOL, rtol=RTOL)
323
320 - assert fusion_pass.matched_count == 3
321 - backend.check_before_ops(model.ops_in_model_before())
324 + backend, _ = _run_fusion_test(
325 + model, fusion_pass, vllm_config, dtype, hidden_size, num_tokens
326 + )
327 backend.check_before_ops(
328 model.ops_in_model_before_partial(), fully_replaced=False
329 )
325 - backend.check_after_ops(model.ops_in_model_after())
330
331 # If RMSNorm custom op is disabled (native/torch impl used),
332 # there's a risk that the fused add doesn't get included in the
333 # replacement and only the rms part gets fused with quant.
334 # Hence, we check only 2 add nodes are left (final fused rmsnorm add).
331 - if (
332 - not enable_rms_norm_custom_op
333 - and model_class is not TestRmsnormGroupFp8QuantModel
334 - ):
335 + if not enable_rms_norm_custom_op:
336 n_add_nodes = lambda g: sum(1 for _ in find_op_nodes(torch.ops.aten.add, g))
337 # 7 = 1 (RMS) + 3x2 (3xRMS_ADD, 2 each)
338 assert n_add_nodes(backend.graph_pre_pass) == 7
339 assert n_add_nodes(backend.graph_post_pass) == 2
340 +
341 +
342 +GROUP_SHAPE_QUANT_OPS_MATCHS = [
343 + (GroupShape.PER_TOKEN, True),
344 + (GroupShape.PER_TOKEN, False),
345 + (GroupShape(1, 128), True),
346 +]
347 +
348 +
349 +@pytest.mark.parametrize("dtype", [torch.bfloat16])
350 +@pytest.mark.parametrize("hidden_size", [256])
351 +@pytest.mark.parametrize("num_tokens", [257])
352 +@pytest.mark.parametrize("eps", [1e-5, 1e-6])
353 +@pytest.mark.parametrize(
354 + "group_shape, use_aiter_quant_op", GROUP_SHAPE_QUANT_OPS_MATCHS
355 +)
356 +@pytest.mark.skipif(
357 + (not current_platform.is_rocm() or not IS_AITER_FOUND),
358 + reason="Only test on ROCm with aiter package installed",
359 +)
360 +def test_aiter_fusion_rmsnorm_quant(
361 + dtype: torch.dtype,
362 + hidden_size: int,
363 + num_tokens: int,
364 + eps: float,
365 + group_shape: GroupShape,
366 + use_aiter_quant_op: bool,
367 + monkeypatch: pytest.MonkeyPatch,
368 +):
369 + vllm_config = VllmConfig(
370 + model_config=ModelConfig(dtype=dtype),
371 + compilation_config=CompilationConfig(
372 + mode=CompilationMode.VLLM_COMPILE,
373 + custom_ops=["+rms_norm", "+quant_fp8"],
374 + pass_config=PassConfig(fuse_norm_quant=True, eliminate_noops=True),
375 + ),
376 + )
377 +
378 + with vllm.config.set_current_vllm_config(vllm_config), monkeypatch.context() as m:
379 + from vllm.compilation.rocm_aiter_fusion import RocmAiterRMSNormFusionPass
380 +
381 + m.setenv("VLLM_ROCM_USE_AITER", "1")
382 + rocm_aiter_ops.refresh_env_variables()
383 +
384 + torch.set_default_device("cuda")
385 + torch.set_default_dtype(dtype)
386 + torch.manual_seed(1)
387 + maybe_create_device_identity()
388 +
389 + fusion_pass = RocmAiterRMSNormFusionPass(vllm_config)
390 + model = TestModel(
391 + hidden_size=hidden_size,
392 + eps=eps,
393 + group_shape=group_shape,
394 + use_aiter=True,
395 + use_aiter_quant_op=use_aiter_quant_op,
396 + )
397 +
398 + _run_fusion_test(
399 + model, fusion_pass, vllm_config, dtype, hidden_size, num_tokens
400 + )
vllm/_aiter_ops.py
+144 -13
@@ -4,6 +4,7 @@ import functools
4 from collections.abc import Callable
5
6 import torch
7 +from torch._ops import OpOverload
8
9 import vllm.envs as envs
10 from vllm.platforms import current_platform
@@ -433,16 +434,16 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_impl(
434 from aiter import rmsnorm2d_fwd_with_add
435
436 residual_out = torch.empty_like(residual)
436 - output = torch.empty_like(x)
437 + out = torch.empty_like(x)
438 rmsnorm2d_fwd_with_add(
438 - output, # output
439 + out, # output
440 x, # input
441 residual, # residual input
442 residual_out, # residual output
443 weight,
444 variance_epsilon,
445 )
445 - return output, residual_out
446 + return out, residual_out
447
448
449 def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
@@ -451,7 +452,84 @@ def _rocm_aiter_rmsnorm2d_fwd_with_add_fake(
452 weight: torch.Tensor,
453 variance_epsilon: float,
454 ) -> tuple[torch.Tensor, torch.Tensor]:
454 - return torch.empty_like(x), torch.empty_like(residual)
455 + residual_out = torch.empty_like(residual)
456 + out = torch.empty_like(x)
457 + return out, residual_out
458 +
459 +
460 +def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl(
461 + x: torch.Tensor,
462 + residual: torch.Tensor,
463 + weight: torch.Tensor,
464 + epsilon: float,
465 + quant_dtype: torch.dtype,
466 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
467 + import aiter as rocm_aiter
468 +
469 + assert quant_dtype in [torch.int8, _FP8_DTYPE]
470 +
471 + y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
472 + out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
473 + residual_out = torch.empty_like(x)
474 +
475 + rocm_aiter.rmsnorm2d_fwd_with_add_dynamicquant(
476 + out,
477 + x,
478 + residual,
479 + residual_out,
480 + y_scale,
481 + weight,
482 + epsilon,
483 + use_model_sensitive_rmsnorm=0,
484 + )
485 +
486 + return out, residual_out, y_scale
487 +
488 +
489 +def _rocm_aiter_rmsnorm_fused_add_dynamic_quant_fake(
490 + x: torch.Tensor,
491 + residual: torch.Tensor,
492 + weight: torch.Tensor,
493 + epsilon: float,
494 + quant_dtype: torch.dtype,
495 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
496 + y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
497 + out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
498 + residual_out = torch.empty_like(x)
499 +
500 + return out, residual_out, y_scale
501 +
502 +
503 +def _rocm_aiter_rmsnorm_fused_dynamic_quant_impl(
504 + x: torch.Tensor,
505 + weight: torch.Tensor,
506 + epsilon: float,
507 + quant_dtype: torch.dtype,
508 +) -> tuple[torch.Tensor, torch.Tensor]:
509 + import aiter as rocm_aiter
510 +
511 + assert quant_dtype in [torch.int8, _FP8_DTYPE]
512 +
513 + y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
514 + out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
515 +
516 + rocm_aiter.rmsnorm2d_fwd_with_dynamicquant(
517 + out, x, y_scale, weight, epsilon, use_model_sensitive_rmsnorm=0
518 + )
519 +
520 + return out, y_scale
521 +
522 +
523 +def _rocm_aiter_rmsnorm_fused_dynamic_quant_fake(
524 + x: torch.Tensor,
525 + weight: torch.Tensor,
526 + epsilon: float,
527 + quant_dtype: torch.dtype,
528 +) -> tuple[torch.Tensor, torch.Tensor]:
529 + y_scale = torch.empty(x.shape[0], 1, dtype=torch.float32, device=x.device)
530 + out = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
531 +
532 + return out, y_scale
533
534
535 def _rocm_aiter_per_tensor_quant_impl(
@@ -527,7 +605,11 @@ def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl(
605 dtype_quant=AITER_FP8_DTYPE,
606 res1=residual,
607 )
530 - return (x_quant, x_quant_scales, res)
608 + return (
609 + x_quant,
610 + res,
611 + x_quant_scales,
612 + )
613
614
615 def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake(
@@ -541,8 +623,8 @@ def _rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake(
623 scale_shape = (M, (N + group_size - 1) // group_size)
624 return (
625 torch.empty_like(x, dtype=AITER_FP8_DTYPE, device=x.device),
544 - torch.empty(scale_shape, dtype=torch.float32, device=x.device),
626 torch.empty_like(residual, device=residual.device),
627 + torch.empty(scale_shape, dtype=torch.float32, device=x.device),
628 )
629
630
@@ -901,6 +983,20 @@ class rocm_aiter_ops:
983 dispatch_key=current_platform.dispatch_key,
984 )
985
986 + direct_register_custom_op(
987 + op_name="rocm_aiter_rmsnorm_fused_dynamic_quant",
988 + op_func=_rocm_aiter_rmsnorm_fused_dynamic_quant_impl,
989 + fake_impl=_rocm_aiter_rmsnorm_fused_dynamic_quant_fake,
990 + dispatch_key=current_platform.dispatch_key,
991 + )
992 +
993 + direct_register_custom_op(
994 + op_name="rocm_aiter_rmsnorm_fused_add_dynamic_quant",
995 + op_func=_rocm_aiter_rmsnorm_fused_add_dynamic_quant_impl,
996 + fake_impl=_rocm_aiter_rmsnorm_fused_add_dynamic_quant_fake,
997 + dispatch_key=current_platform.dispatch_key,
998 + )
999 +
1000 direct_register_custom_op(
1001 op_name="rocm_aiter_rmsnorm_fp8_group_quant",
1002 op_func=_rocm_aiter_rmsnorm_fp8_group_quant_impl,
@@ -936,13 +1032,54 @@ class rocm_aiter_ops:
1032 direct_register_custom_op(
1033 op_name="rocm_aiter_per_token_quant",
1034 op_func=_rocm_aiter_per_token_quant_impl,
939 - mutates_args=["scale"],
1035 fake_impl=_rocm_aiter_per_token_quant_fake,
1036 dispatch_key=current_platform.dispatch_key,
1037 )
1038
1039 _OPS_REGISTERED = True
1040
1041 + @staticmethod
1042 + def get_rmsnorm_fused_add_op() -> OpOverload:
1043 + return torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add.default
1044 +
1045 + @staticmethod
1046 + def get_rmsnorm_op() -> OpOverload:
1047 + return torch.ops.vllm.rocm_aiter_rms_norm.default
1048 +
1049 + @staticmethod
1050 + def get_rmsnorm_fused_add_dynamic_quant_op() -> OpOverload:
1051 + return torch.ops.vllm.rocm_aiter_rmsnorm_fused_add_dynamic_quant.default
1052 +
1053 + @staticmethod
1054 + def get_rmsnorm_fused_dynamic_quant_op() -> OpOverload:
1055 + return torch.ops.vllm.rocm_aiter_rmsnorm_fused_dynamic_quant.default
1056 +
1057 + @staticmethod
1058 + def get_rmsnorm_group_fused_quant_op() -> OpOverload:
1059 + return torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default
1060 +
1061 + @staticmethod
1062 + def get_rmsnorm_group_add_fused_quant_op() -> OpOverload:
1063 + return torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant.default
1064 +
1065 + @staticmethod
1066 + def get_per_token_quant_op() -> OpOverload:
1067 + return torch.ops.vllm.rocm_aiter_per_token_quant.default
1068 +
1069 + @staticmethod
1070 + def get_group_quant_op() -> OpOverload:
1071 + return torch.ops.vllm.rocm_aiter_group_fp8_quant.default
1072 +
1073 + @staticmethod
1074 + def get_act_mul_fused_fp8_group_quant_op() -> OpOverload:
1075 + return torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant.default
1076 +
1077 + @staticmethod
1078 + def rms_norm(
1079 + x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
1080 + ) -> torch.Tensor:
1081 + return torch.ops.vllm.rocm_aiter_rms_norm(x, weight, variance_epsilon)
1082 +
1083 @staticmethod
1084 def rms_norm2d_with_add(
1085 x: torch.Tensor,
@@ -954,12 +1091,6 @@ class rocm_aiter_ops:
1091 x, residual, weight, variance_epsilon
1092 )
1093
957 - @staticmethod
958 - def rms_norm(
959 - x: torch.Tensor, weight: torch.Tensor, variance_epsilon: float
960 - ) -> torch.Tensor:
961 - return torch.ops.vllm.rocm_aiter_rms_norm(x, weight, variance_epsilon)
962 -
1094 @staticmethod
1095 def gemm_a8w8(
1096 A: torch.Tensor,
vllm/compilation/matcher_utils.py
+103 -11
@@ -6,11 +6,13 @@ import torch
6 from torch._higher_order_ops import auto_functionalized
7 from torch._ops import OpOverload
8
9 +from vllm._aiter_ops import rocm_aiter_ops
10 from vllm.config import get_current_vllm_config
11 from vllm.model_executor.layers.activation import SiluAndMul
12 from vllm.model_executor.layers.layernorm import RMSNorm
13 from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
14 from vllm.model_executor.layers.quantization.utils.quant_utils import (
15 + GroupShape,
16 QuantKey,
17 _normalize_quant_group_shape,
18 kFp8Dynamic64Sym,
@@ -150,26 +152,50 @@ class MatcherRotaryEmbedding(MatcherCustomOp):
152
153
154 class MatcherRMSNorm(MatcherCustomOp):
153 - def __init__(self, epsilon: float, enabled: bool | None = None):
155 + def __init__(
156 + self,
157 + epsilon: float,
158 + enabled: bool | None = None,
159 + match_rocm_aiter: bool = False,
160 + ):
161 if enabled is None:
162 enabled = RMSNorm.enabled()
163
164 super().__init__(enabled)
165 self.epsilon = epsilon
166 + self._rmsnorm_op = RMS_OP
167 + self.match_rocm_aiter = match_rocm_aiter
168 +
169 + if match_rocm_aiter:
170 + self._rmsnorm_op = rocm_aiter_ops.get_rmsnorm_op()
171
172 def inputs(self):
173 input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)
174 weight = self.empty(16)
175 return [input, weight]
176
177 + def forward_rocm_aiter(
178 + self,
179 + input: torch.Tensor,
180 + weight: torch.Tensor,
181 + ) -> torch.Tensor:
182 + return self._rmsnorm_op(
183 + x=input,
184 + weight=weight,
185 + variance_epsilon=self.epsilon,
186 + )
187 +
188 def forward_custom(
189 self,
190 input: torch.Tensor,
191 weight: torch.Tensor,
192 ) -> torch.Tensor:
193 + if self.match_rocm_aiter:
194 + return self.forward_rocm_aiter(input, weight)
195 +
196 result = torch.empty_like(input)
197 _, result = auto_functionalized(
172 - RMS_OP,
198 + self._rmsnorm_op,
199 result=result,
200 input=input,
201 weight=weight,
@@ -189,12 +215,23 @@ class MatcherRMSNorm(MatcherCustomOp):
215
216
217 class MatcherFusedAddRMSNorm(MatcherCustomOp):
192 - def __init__(self, epsilon: float, enabled: bool | None = None):
218 + def __init__(
219 + self,
220 + epsilon: float,
221 + enabled: bool | None = None,
222 + match_rocm_aiter: bool = False,
223 + ):
224 if enabled is None:
225 enabled = RMSNorm.enabled()
226
227 super().__init__(enabled)
228 self.epsilon = epsilon
229 + self.match_rocm_aiter = match_rocm_aiter
230 +
231 + self._rmsnorm_op = RMS_ADD_OP
232 +
233 + if match_rocm_aiter:
234 + self._rmsnorm_op = rocm_aiter_ops.get_rmsnorm_fused_add_op()
235
236 def inputs(self):
237 input = self.empty(5, 16) if self.enabled else self.empty_f32(5, 16)
@@ -202,14 +239,27 @@ class MatcherFusedAddRMSNorm(MatcherCustomOp):
239 residual = self.empty(5, 16)
240 return [input, weight, residual]
241
242 + def forward_rocm_aiter(
243 + self,
244 + input: torch.Tensor,
245 + weight: torch.Tensor,
246 + residual: torch.Tensor,
247 + ) -> tuple[torch.Tensor, torch.Tensor]:
248 + return self._rmsnorm_op(
249 + x=input, residual=residual, weight=weight, variance_epsilon=self.epsilon
250 + )
251 +
252 def forward_custom(
253 self,
254 input: torch.Tensor,
255 weight: torch.Tensor,
256 residual: torch.Tensor,
257 ) -> tuple[torch.Tensor, torch.Tensor]:
258 + if self.match_rocm_aiter:
259 + return self.forward_rocm_aiter(input, weight, residual)
260 +
261 _, result, residual = auto_functionalized(
212 - RMS_ADD_OP,
262 + self._rmsnorm_op,
263 input=input,
264 residual=residual,
265 weight=weight,
@@ -236,22 +286,46 @@ class MatcherQuantFP8(MatcherCustomOp):
286 enabled: bool | None = None,
287 has_col_major_scales: bool = False,
288 is_e8m0: bool = False,
289 + match_rocm_aiter: bool = False,
290 ):
291 if enabled is None:
292 enabled = QuantFP8.enabled()
293
294 super().__init__(enabled)
295 self.quant_key = quant_key
245 - assert quant_key in QUANT_OPS, f"unsupported quantization scheme {quant_key}"
246 - self.QUANT_OP = QUANT_OPS[quant_key]
247 -
296 self.has_col_major_scales = has_col_major_scales
297 self.is_e8m0 = is_e8m0
298 + self.match_rocm_aiter = match_rocm_aiter
299 +
300 + if match_rocm_aiter:
301 + assert not quant_key.scale.group_shape.is_per_tensor(), (
302 + "ROCm aiter fusion pass does not support per tensor quantization"
303 + )
304 + if quant_key.scale.group_shape.is_per_token():
305 + self.QUANT_OP = rocm_aiter_ops.get_per_token_quant_op()
306 + else:
307 + assert quant_key.scale.group_shape.col == 128, (
308 + "ROCm aiter fusion pass currently supports "
309 + "quantization operation with group_size 128"
310 + )
311 + if current_platform.is_fp8_fnuz():
312 + self.QUANT_OP = rocm_aiter_ops.get_group_quant_op()
313 + else:
314 + self.QUANT_OP = (
315 + torch.ops.vllm.triton_per_token_group_quant_fp8.default
316 + )
317 +
318 + else:
319 + assert quant_key in QUANT_OPS, (
320 + f"unsupported quantization scheme {quant_key}"
321 + )
322 + self.QUANT_OP = QUANT_OPS[quant_key]
323 +
324 + assert quant_key.dtype == current_platform.fp8_dtype(), (
325 + "Only QuantFP8 supported by"
326 + )
327 + assert quant_key.scale2 is None
328
251 - assert quant_key.dtype == current_platform.fp8_dtype(), (
252 - "Only QuantFP8 supported by"
253 - )
254 - assert quant_key.scale2 is None
329 self.quant_fp8 = QuantFP8(
330 quant_key.scale.static,
331 quant_key.scale.group_shape,
@@ -259,11 +333,29 @@ class MatcherQuantFP8(MatcherCustomOp):
333 use_ue8m0=is_e8m0,
334 )
335
336 + def forward_rocm_aiter(
337 + self,
338 + input: torch.Tensor,
339 + scale: torch.Tensor | None = None,
340 + ) -> tuple[torch.Tensor, torch.Tensor]:
341 + quant_key_group_shape = self.quant_key.scale.group_shape
342 + if quant_key_group_shape == GroupShape.PER_TOKEN:
343 + return self.QUANT_OP(
344 + x=input,
345 + quant_dtype=self.quant_key.dtype,
346 + scale=scale,
347 + )
348 + else:
349 + return self.QUANT_OP(input, quant_key_group_shape.col)
350 +
351 def forward_custom(
352 self,
353 input: torch.Tensor,
354 scale: torch.Tensor | None = None,
355 ) -> tuple[torch.Tensor, torch.Tensor]:
356 + if self.match_rocm_aiter:
357 + return self.forward_rocm_aiter(input, scale)
358 +
359 result = torch.empty(
360 input.shape, device=input.device, dtype=self.quant_key.dtype
361 )
vllm/compilation/pass_manager.py
+4 -2
@@ -16,7 +16,7 @@ from .vllm_inductor_pass import VllmInductorPass
16
17 if rocm_aiter_ops.is_enabled():
18 from vllm.compilation.rocm_aiter_fusion import (
19 - RocmAiterRMSNormFp8GroupQuantFusionPass,
19 + RocmAiterRMSNormFusionPass,
20 RocmAiterSiluMulFp8GroupQuantFusionPass,
21 )
22
@@ -117,7 +117,9 @@ class PostGradPassManager(CustomGraphPass):
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)]
120 + self.passes += [
121 + RocmAiterRMSNormFusionPass(config),
122 + ]
123 if self.pass_config.fuse_act_quant:
124 self.passes += [ActivationQuantFusionPass(config)]
125 if rocm_aiter_ops.is_enabled():
vllm/compilation/rocm_aiter_fusion.py
+223 -66
@@ -9,60 +9,195 @@ 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._aiter_ops import rocm_aiter_ops
13 from vllm.compilation.activation_quant_fusion import ActivationQuantPattern
14 from vllm.config import VllmConfig
15 from vllm.logger import init_logger
16 +from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 + GroupShape,
18 + QuantKey,
19 + ScaleDesc,
20 +)
21 from vllm.platforms import current_platform
22
17 -from .fusion import empty_bf16
23 +from .fusion import (
24 + FusedRMSQuantKey,
25 +)
26 from .inductor_pass import enable_fake_mode
19 -from .matcher_utils import MatcherSiluAndMul
27 +from .matcher_utils import (
28 + MatcherFusedAddRMSNorm,
29 + MatcherQuantFP8,
30 + MatcherRMSNorm,
31 + MatcherSiluAndMul,
32 +)
33 from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass
34
35 logger = init_logger(__name__)
36 FP8_DTYPE = current_platform.fp8_dtype()
37
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 -)
38
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
39 +class AiterRMSNormQuantPattern:
40 + def __init__(
41 + self, epsilon: float, key: FusedRMSQuantKey, match_aiter_quant: bool = True
42 + ):
43 + self.epsilon = epsilon
44 + self.quant_dtype = key.quant.dtype
45 +
46 + self.rmsnorm_matcher = (
47 + MatcherRMSNorm(epsilon, match_rocm_aiter=True)
48 + if not key.fused_add
49 + else MatcherFusedAddRMSNorm(epsilon, match_rocm_aiter=True)
50 + )
51 + self.quant_matcher = MatcherQuantFP8(
52 + key.quant,
53 + match_rocm_aiter=match_aiter_quant,
54 + )
55 +
56
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
57 +class AiterRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
58 + """AITER RMSNorm + Dynamic Quantization pattern."""
59
36 -FUSED_SILU_MUL_QUANT_OP = torch.ops.vllm.rocm_aiter_act_mul_and_fp8_group_quant.default
60 + FUSED_OP = rocm_aiter_ops.get_rmsnorm_fused_dynamic_quant_op()
61
62 + def __init__(
63 + self,
64 + epsilon: float,
65 + quant_dtype: torch.dtype,
66 + match_aiter_quant: bool = True,
67 + group_shape: GroupShape = GroupShape.PER_TOKEN,
68 + symmetric=True,
69 + ):
70 + scale = ScaleDesc(torch.float32, False, group_shape)
71 + key = FusedRMSQuantKey(
72 + fused_add=False,
73 + quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
74 + )
75 +
76 + super().__init__(epsilon, key, match_aiter_quant)
77
39 -class AiterRMSFp8GroupQuantPattern:
78 + def register(self, pm_pass):
79 + def pattern(
80 + input: torch.Tensor,
81 + weight: torch.Tensor,
82 + ):
83 + result_rms = self.rmsnorm_matcher(input, weight)
84 + result, scale = self.quant_matcher(result_rms)
85 + return result, scale
86 +
87 + def replacement(
88 + input: torch.Tensor,
89 + weight: torch.Tensor,
90 + ):
91 + result = self.FUSED_OP(
92 + x=input,
93 + weight=weight,
94 + epsilon=self.epsilon,
95 + quant_dtype=self.quant_dtype,
96 + )
97 +
98 + return result[0], result[1]
99 +
100 + pm.register_replacement(
101 + pattern,
102 + replacement,
103 + self.rmsnorm_matcher.inputs(),
104 + pm.fwd_only,
105 + pm_pass,
106 + )
107 +
108 +
109 +class AiterFusedAddRMSNormDynamicQuantPattern(AiterRMSNormQuantPattern):
110 + """AITER RMSNorm Fused Add + Dynamic Quantization pattern."""
111 +
112 + FUSED_OP = rocm_aiter_ops.get_rmsnorm_fused_add_dynamic_quant_op()
113 +
114 + def __init__(
115 + self,
116 + epsilon: float,
117 + quant_dtype: torch.dtype,
118 + match_aiter_quant: bool = True,
119 + group_shape: GroupShape = GroupShape.PER_TOKEN,
120 + symmetric=True,
121 + ):
122 + scale = ScaleDesc(torch.float32, False, group_shape)
123 + key = FusedRMSQuantKey(
124 + fused_add=True,
125 + quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
126 + )
127 +
128 + super().__init__(epsilon, key, match_aiter_quant)
129 +
130 + def register(self, pm_pass):
131 + def pattern(
132 + input: torch.Tensor,
133 + weight: torch.Tensor,
134 + residual: torch.Tensor,
135 + ):
136 + result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
137 + result, scale = self.quant_matcher(result_rms)
138 +
139 + return result, residual_out, scale
140 +
141 + def replacement(
142 + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
143 + ):
144 + result = self.FUSED_OP(
145 + x=input,
146 + residual=residual,
147 + weight=weight,
148 + epsilon=self.epsilon,
149 + quant_dtype=self.quant_dtype,
150 + )
151 +
152 + return result[0], result[1], result[2]
153 +
154 + pm.register_replacement(
155 + pattern,
156 + replacement,
157 + self.rmsnorm_matcher.inputs(),
158 + pm.fwd_only,
159 + pm_pass,
160 + )
161 +
162 +
163 +class AiterRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
164 """
165 This pattern fuses aiter rms_norm & group fp8 quant custom
166 ops into an aiter rms_norm_group_fp8_quant op.
167 """
168
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
169 + FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_fused_quant_op()
170 +
171 + def __init__(
172 + self,
173 + epsilon: float,
174 + quant_dtype: torch.dtype,
175 + group_shape: GroupShape,
176 + match_aiter_quant: bool = True,
177 + symmetric=True,
178 + ):
179 + scale = ScaleDesc(torch.float32, False, group_shape)
180 + key = FusedRMSQuantKey(
181 + fused_add=False,
182 + quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
183 + )
184 +
185 + super().__init__(epsilon, key, match_aiter_quant)
186
187 def register(self, pm_pass: PatternMatcherPass):
188 def pattern(
189 input: torch.Tensor,
190 weight: torch.Tensor,
191 ):
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]
192 + result_rms = self.rmsnorm_matcher(input, weight)
193 + result, scale = self.quant_matcher(result_rms)
194 + return result, scale
195
196 def replacement(
197 input: torch.Tensor,
198 weight: torch.Tensor,
199 ):
65 - at = AITER_RMS_GROUP_QUANT_OP(
200 + at = self.FUSED_OP(
201 x=input,
202 weight=weight,
203 variance_epsilon=self.epsilon,
@@ -71,49 +206,52 @@ class AiterRMSFp8GroupQuantPattern:
206
207 return at[0], at[1]
208
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)
209 + pm.register_replacement(
210 + pattern, replacement, self.rmsnorm_matcher.inputs(), pm.fwd_only, pm_pass
211 + )
212
213
82 -class AiterFusedAddRMSFp8GroupQuantPattern:
214 +class AiterFusedAddRMSFp8GroupQuantPattern(AiterRMSNormQuantPattern):
215 """
216 This pattern fuses aiter rms_norm_with_add & group fp8 quant custom ops
217 into a aiter rms_norm_with_add_group_fp8_quant op.
218 """
219
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
220 + FUSED_OP = rocm_aiter_ops.get_rmsnorm_group_add_fused_quant_op()
221 +
222 + def __init__(
223 + self,
224 + epsilon: float,
225 + quant_dtype: torch.dtype,
226 + group_shape: GroupShape,
227 + match_aiter_quant: bool = True,
228 + symmetric=True,
229 + ):
230 + scale = ScaleDesc(torch.float32, False, group_shape)
231 + key = FusedRMSQuantKey(
232 + fused_add=True,
233 + quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
234 + )
235 +
236 + super().__init__(epsilon, key, match_aiter_quant)
237
238 def register(self, pm_pass: PatternMatcherPass):
239 def pattern(
240 input: torch.Tensor,
96 - residual: torch.Tensor,
241 weight: torch.Tensor,
242 + residual: torch.Tensor,
243 ):
99 - at1 = AITER_RMS_ADD_OP(
100 - x=input,
101 - residual=residual,
102 - weight=weight,
103 - variance_epsilon=self.epsilon,
104 - )
244 + result_rms, residual_out = self.rmsnorm_matcher(input, weight, residual)
245 + result, scale = self.quant_matcher(result_rms)
246
106 - at2 = self.quant_op(at1[0], 128)
107 -
108 - # result, scale, residual
109 - return at2[0], at2[1], at1[1]
247 + return result, residual_out, scale
248
249 def replacement(
250 input: torch.Tensor,
113 - residual: torch.Tensor,
251 weight: torch.Tensor,
252 + residual: torch.Tensor,
253 ):
116 - at = AITER_RMS_ADD_GROUP_QUANT_OP(
254 + at = self.FUSED_OP(
255 x=input,
256 residual=residual,
257 weight=weight,
@@ -124,18 +262,15 @@ class AiterFusedAddRMSFp8GroupQuantPattern:
262 # result, scale, residual
263 return at[0], at[1], at[2]
264
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)
265 + pm.register_replacement(
266 + pattern, replacement, self.rmsnorm_matcher.inputs(), pm.fwd_only, pm_pass
267 + )
268
269
136 -class RocmAiterRMSNormFp8GroupQuantFusionPass(VllmPatternMatcherPass):
270 +class RocmAiterRMSNormFusionPass(VllmPatternMatcherPass):
271 """
138 - This pass fuses rms_norm & quant custom ops into a fused rms_norm_quant op.
272 + This pass fuses aiter rms_norm & vllm/aiter quant custom ops
273 + into a fused rms_norm_quant op.
274 It also supports fused_add_rms_norm.
275 """
276
@@ -144,20 +279,33 @@ class RocmAiterRMSNormFp8GroupQuantFusionPass(VllmPatternMatcherPass):
279 super().__init__(config)
280
281 self.patterns: PatternMatcherPass = PatternMatcherPass(
147 - pass_name="rocm_aiter_rms_norm_fp8_group_quant_fusion_pass"
282 + pass_name="rocm_aiter_rms_norm_quant_fusion_pass"
283 )
284
285 # Make sure fused add patterns are before simple rms norm,
286 # as the latter is a subset of the former in torch ops
287 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
288 + # Fuse aiter rms_norm + aiter dynamic group fp8 quant
289 + AiterRMSFp8GroupQuantPattern(
290 + epsilon, FP8_DTYPE, GroupShape(1, 128)
291 + ).register(self.patterns)
292 +
293 + # Fuse aiter fused_add_rms_norm + aiter dynamic group fp8 quant
294 + AiterFusedAddRMSFp8GroupQuantPattern(
295 + epsilon, FP8_DTYPE, GroupShape(1, 128)
296 + ).register(self.patterns)
297 +
298 + for match_aiter_quant in [True, False]:
299 + # Fuse aiter rms_norm + (aiter / vllm built-in)
300 + # dynamic per-token fp8 quant
301 + AiterRMSNormDynamicQuantPattern(
302 + epsilon, FP8_DTYPE, match_aiter_quant=match_aiter_quant
303 + ).register(self.patterns)
304 +
305 + # Fuse aiter fused_add_rms_norm + (aiter / vllm built-in)
306 + # dynamic per-token fp8 quant
307 + AiterFusedAddRMSNormDynamicQuantPattern(
308 + epsilon, FP8_DTYPE, match_aiter_quant=match_aiter_quant
309 ).register(self.patterns)
310
311 self.dump_patterns(config, self.patterns)
@@ -169,6 +317,8 @@ class RocmAiterRMSNormFp8GroupQuantFusionPass(VllmPatternMatcherPass):
317
318 def uuid(self) -> Any:
319 fusion_patterns = [
320 + AiterRMSNormDynamicQuantPattern,
321 + AiterFusedAddRMSNormDynamicQuantPattern,
322 AiterRMSFp8GroupQuantPattern,
323 AiterFusedAddRMSFp8GroupQuantPattern,
324 ]
@@ -181,6 +331,8 @@ class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
331 ops into an aiter silu_and_mul_group_fp8_quant op.
332 """
333
334 + FUSED_SILU_MUL_QUANT_OP = rocm_aiter_ops.get_act_mul_fused_fp8_group_quant_op()
335 +
336 def __init__(self, quant_op: OpOverload):
337 self.silu_and_mul_matcher = MatcherSiluAndMul()
338 self.quant_op = quant_op
@@ -196,7 +348,7 @@ class AiterSiluMulFp8GroupQuantPattern(ActivationQuantPattern):
348 def replacement(
349 input: torch.Tensor,
350 ):
199 - at = FUSED_SILU_MUL_QUANT_OP(x=input, group_size=128)
351 + at = self.FUSED_SILU_MUL_QUANT_OP(x=input, group_size=128)
352 return at[0], at[1]
353
354 inputs = [
@@ -216,6 +368,11 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
368 https://github.com/pytorch/pytorch/pull/139321#issuecomment-2452354980
369 """
370
371 + AITER_GROUP_FP8_QUANT_OP = rocm_aiter_ops.get_group_quant_op()
372 + TRITON_GROUP_FP8_QUANT_OP = torch.ops.vllm.triton_per_token_group_quant_fp8.default
373 +
374 + QUANT_OPS = [AITER_GROUP_FP8_QUANT_OP, TRITON_GROUP_FP8_QUANT_OP]
375 +
376 @enable_fake_mode
377 def __init__(self, config: VllmConfig):
378 super().__init__(config)
@@ -224,7 +381,7 @@ class RocmAiterSiluMulFp8GroupQuantFusionPass(VllmPatternMatcherPass):
381 pass_name="rocm_aiter_silu_mul_fp8_group_quant_fusion_pass"
382 )
383
227 - for quant_op in [AITER_GROUP_FP8_QUANT_OP, TRITON_GROUP_FP8_QUANT_OP]:
384 + for quant_op in self.QUANT_OPS:
385 AiterSiluMulFp8GroupQuantPattern(quant_op).register(self.patterns)
386
387 self.dump_patterns(config, self.patterns)