[ROCm][Compile] Fuse AITER static FP8 attention output (#58099)
Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <noreply@openai.com>
Andreas Karatzas committed
Sep 23, 2026 at 03:16 UTC
26879f3260425889401f2281f38b45303e0db70b
5 files changed
+251
-12
.buildkite/test-amd.yaml
+23
@@ -3452,6 +3452,29 @@ steps:
3452
commands:
3453
- pytest -v -s compile/h100/test_startup.py
3454
3455
+- label: ":amd: (MI355) Fusion and Compile FP8"
3456
+ timeout_in_minutes: 45
3457
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3458
+ dind: false
3459
+ agent_pool: mi355_2
3460
+ num_gpus: 2
3461
+ optional: true
3462
+ working_dir: "/vllm-workspace/"
3463
+ source_file_dependencies:
3464
+ - vllm/compilation/
3465
+ - vllm/model_executor/layers/
3466
+ - vllm/v1/attention/
3467
+ - vllm/distributed/
3468
+ - vllm/_aiter_ops.py
3469
+ - tests/compile/passes/
3470
+ - tests/compile/correctness_e2e/test_attn_quant.py
3471
+ - vllm/model_executor/kernels/linear/
3472
+ - vllm/model_executor/models/llama.py
3473
+ commands:
3474
+ - amd-smi
3475
+ - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py
3476
+ - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py
3477
+
3478
#-------------------------------------------------------- mi355 · distributed --------------------------------------------------------#
3479
3480
- label: ":amd: (MI355) Sharded RDT Weight Transfer"
.buildkite/test_areas/compile.yaml
+25
@@ -67,6 +67,31 @@ steps:
67
- pytest -v -s tests/compile/passes/test_silu_mul_quant_manual_fusion.py
68
# this runner has 2 GPUs available even though num_devices=2 is not set
69
- pytest -v -s tests/compile/passes/distributed/test_fusion_all_reduce.py
70
+ mirror:
71
+ amd:
72
+ label: ":amd: (MI355) Fusion and Compile FP8"
73
+ dind: false
74
+ device: mi355_2
75
+ num_devices: 2
76
+ working_dir: "/vllm-workspace/"
77
+ timeout_in_minutes: 45
78
+ depends_on:
79
+ - image-build-amd
80
+ source_file_dependencies:
81
+ - vllm/compilation/
82
+ - vllm/model_executor/layers/
83
+ - vllm/v1/attention/
84
+ - vllm/distributed/
85
+ - vllm/_aiter_ops.py
86
+ - tests/compile/passes/
87
+ - tests/compile/correctness_e2e/test_attn_quant.py
88
+ - vllm/model_executor/kernels/linear/
89
+ - vllm/model_executor/models/llama.py
90
+ commands:
91
+ - amd-smi
92
+ # MLA static FP8 and distributed AR/RMS are owned by the existing ROCm jobs.
93
+ - pytest -v -s tests/compile/passes/test_fusion_attn.py tests/compile/passes/test_silu_mul_quant_fusion.py tests/compile/passes/test_silu_mul_quant_manual_fusion.py
94
+ - pytest -v -s tests/compile/correctness_e2e/test_attn_quant.py
95
96
- label: ":nvidia: (H100) Fusion E2E Quick"
97
key: fusion-e2e-quick-h100
tests/compile/correctness_e2e/test_attn_quant.py
new
+88
@@ -0,0 +1,88 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import json
4
+import logging
5
+
6
+import pytest
7
+import regex as re
8
+
9
+from tests.utils import compare_two_settings
10
+from vllm._aiter_ops import is_aiter_found_and_supported
11
+from vllm.config import CompilationMode
12
+from vllm.platforms import current_platform
13
+from vllm.transformers_utils.repo_utils import hf_api
14
+
15
+
16
+@pytest.mark.skipif(
17
+ not current_platform.is_rocm() or not is_aiter_found_and_supported(),
18
+ reason="Requires ROCm AITER attention and quantization",
19
+)
20
+def test_rocm_aiter_static_attention_output_correctness(caplog_mp_spawn):
21
+ """Trained FP8 outputs agree when actual AITER output quantization is fused."""
22
+ model = hf_api().snapshot_download(
23
+ "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
24
+ revision="870a1177ed082c0ba5dbefd62573ab8f91803f6a",
25
+ allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt", "*.jinja"],
26
+ )
27
+ common_args = [
28
+ "--dtype",
29
+ "bfloat16",
30
+ "--load-format",
31
+ "auto",
32
+ "--max-model-len",
33
+ "512",
34
+ "--max-num-batched-tokens",
35
+ "512",
36
+ "--max-num-seqs",
37
+ "8",
38
+ "--kv-cache-memory-bytes",
39
+ str(256 * 1024**2),
40
+ "--gpu-memory-utilization",
41
+ "0.01",
42
+ "--attention-backend",
43
+ "ROCM_AITER_UNIFIED_ATTN",
44
+ ]
45
+
46
+ def settings(enabled):
47
+ return [
48
+ *common_args,
49
+ "--compilation-config",
50
+ json.dumps(
51
+ {
52
+ "mode": CompilationMode.VLLM_COMPILE,
53
+ "cudagraph_mode": "NONE",
54
+ "splitting_ops": [],
55
+ "inductor_compile_config": {"force_disable_caches": True},
56
+ "custom_ops": ["+quant_fp8", "-rms_norm"],
57
+ "pass_config": {
58
+ "fuse_attn_quant": enabled,
59
+ "fuse_norm_quant": False,
60
+ "fuse_act_quant": False,
61
+ "fuse_allreduce_rms": False,
62
+ },
63
+ }
64
+ ),
65
+ ]
66
+
67
+ env = {
68
+ "VLLM_ROCM_USE_AITER": "1",
69
+ "VLLM_ROCM_USE_AITER_LINEAR": "1",
70
+ "VLLM_DISABLE_COMPILE_CACHE": "1",
71
+ }
72
+ with caplog_mp_spawn(logging.DEBUG) as logs:
73
+ compare_two_settings(
74
+ model,
75
+ settings(False),
76
+ settings(True),
77
+ env1=env,
78
+ env2=env,
79
+ method="generate",
80
+ force_v1_runner=True,
81
+ )
82
+ matches = [
83
+ int(count) for count in re.findall(r"'attn_quant_fusion': (\d+)", logs.text)
84
+ ]
85
+ print(f"Attention output quantization rewrite counts: {matches}")
86
+ assert matches and all(count == 32 for count in matches), (
87
+ f"Expected output quantization fused in all 32 Llama layers, found {matches}"
88
+ )
tests/compile/passes/test_fusion_attn.py
+61
-5
@@ -12,6 +12,7 @@ from tests.v1.attention.utils import (
12
create_common_attn_metadata,
13
dense_kv_cache_views,
14
)
15
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
16
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
17
from vllm.compilation.passes.fusion.attn_quant_fusion import (
18
ATTN_OP,
@@ -189,6 +190,17 @@ class TestAttentionFp8StaticQuantPatternModel(AttentionQuantPatternModel):
190
return self.fp8_linear(attn_output)
191
192
193
+class TestAttentionAiterFp8StaticQuantPatternModel(
194
+ TestAttentionFp8StaticQuantPatternModel
195
+):
196
+ aiter_quant = True
197
+
198
+ def __init__(self, *args, **kwargs):
199
+ super().__init__(*args, **kwargs)
200
+ self.fp8_linear.kernel.quant_fp8.use_aiter = True
201
+ self.fp8_linear.input_scale.fill_(0.125)
202
+
203
+
204
class TestAttentionNvfp4QuantPatternModel(AttentionQuantPatternModel):
205
"""Test model for AttentionNvfp4QuantPattern fusion."""
206
@@ -455,11 +467,12 @@ def test_attention_quant_pattern(
467
)
468
469
# Check quantization ops in the graph before and after fusion
458
- quant_op = (
459
- torch.ops.aten.reciprocal
460
- if "-quant_fp8" in custom_ops_list
461
- else QUANT_OPS[quant_key]
462
- )
470
+ if "-quant_fp8" in custom_ops_list:
471
+ quant_op = torch.ops.aten.reciprocal
472
+ elif getattr(model_class, "aiter_quant", False):
473
+ quant_op = torch.ops.vllm.rocm_aiter_per_tensor_quant.default
474
+ else:
475
+ quant_op = QUANT_OPS[quant_key]
476
477
if expect_fusion:
478
# Note: for fp8, fully_replaced=False because query quant ops remain in
@@ -523,3 +536,46 @@ def test_attention_quant_pattern(
536
537
# Check that results are close
538
torch.testing.assert_close(result_unfused, result_fused, atol=1e-2, rtol=1e-2)
539
+ if getattr(model_class, "aiter_quant", False):
540
+ # Query quantization remains; exactly the output quant op is removed.
541
+ assert test_backend.op_count(quant_op, before=True) == (
542
+ test_backend.op_count(quant_op) + 1
543
+ )
544
+ for model in (model_unfused, model_fused):
545
+ torch.testing.assert_close(
546
+ model.fp8_linear.input_scale,
547
+ torch.full_like(model.fp8_linear.input_scale, 0.125),
548
+ atol=0,
549
+ rtol=0,
550
+ )
551
+
552
+
553
+@pytest.mark.skipif(
554
+ not current_platform.is_rocm() or not is_aiter_found_and_supported(),
555
+ reason="Requires ROCm AITER attention and quantization",
556
+)
557
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
558
+def test_rocm_aiter_static_quant_attention(
559
+ dtype, dist_init, disable_vllm_compile_cache, monkeypatch
560
+):
561
+ """Fuse the actual AITER op without changing its static scale or outputs."""
562
+ try:
563
+ with monkeypatch.context() as patch:
564
+ patch.setenv("VLLM_ROCM_USE_AITER", "1")
565
+ rocm_aiter_ops.refresh_env_variables()
566
+ test_attention_quant_pattern(
567
+ num_qo_heads=32,
568
+ num_kv_heads=8,
569
+ head_size=128,
570
+ batch_size=8,
571
+ dtype=dtype,
572
+ kv_cache_dtype="fp8",
573
+ custom_ops="+quant_fp8",
574
+ model_name="amd/Llama-3.1-8B-Instruct-FP8-KV",
575
+ model_class=TestAttentionAiterFp8StaticQuantPatternModel,
576
+ backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
577
+ dist_init=dist_init,
578
+ disable_vllm_compile_cache=disable_vllm_compile_cache,
579
+ )
580
+ finally:
581
+ rocm_aiter_ops.refresh_env_variables()
vllm/compilation/passes/fusion/attn_quant_fusion.py
+54
-7
@@ -7,6 +7,7 @@ from collections.abc import Callable
7
import torch
8
from torch._higher_order_ops.auto_functionalize import auto_functionalized
9
10
+from vllm._aiter_ops import rocm_aiter_ops
11
from vllm.config import VllmConfig, get_layers_from_vllm_config
12
from vllm.logger import init_logger
13
from vllm.model_executor.layers.attention import Attention
@@ -35,7 +36,9 @@ RESHAPE_OP = torch.ops.aten.reshape.default
36
_FP8_QUANT_KEY = QuantKey(dtype=FP8_DTYPE, scale=kStaticTensorScale, symmetric=True)
37
38
38
-class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
39
+class AttnFp8StaticQuantPattern(
40
+ VllmPatternReplacement[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]]
41
+):
42
"""Fusion for Attention+Fp8StaticQuant.
43
44
Only triggers when the attention implementation returns True in
@@ -51,8 +54,20 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
54
self._dtype = dtype
55
self._quant_matcher = MatcherQuantFP8(_FP8_QUANT_KEY)
56
57
+ def _quantize(
58
+ self, value: torch.Tensor, scale: torch.Tensor
59
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
60
+ return self._quant_matcher(value, scale)[0]
61
+
62
+ def _output(
63
+ self, value: torch.Tensor, scale: torch.Tensor
64
+ ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
65
+ return value
66
+
67
@property
55
- def pattern(self) -> Callable[..., torch.Tensor]:
68
+ def pattern(
69
+ self,
70
+ ) -> Callable[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]]:
71
# When _USE_LAYERNAME is enabled (torch >= 2.11), layer_name is
72
# passed as an explicit pattern input so the pattern matcher
73
# treats it as a wildcard matching hoisted LayerName placeholders.
@@ -78,7 +93,7 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
93
attn_out_view = RESHAPE_OP(
94
at1[1], [q.shape[0], self._num_heads * self._head_size]
95
)
81
- return self._quant_matcher(attn_out_view, scale)[0]
96
+ return self._quantize(attn_out_view, scale)
97
98
return _pattern_with_ln
99
@@ -97,12 +112,14 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
112
attn_out_view = RESHAPE_OP(
113
at1[1], [q.shape[0], self._num_heads * self._head_size]
114
)
100
- return self._quant_matcher(attn_out_view, scale)[0]
115
+ return self._quantize(attn_out_view, scale)
116
117
return _pattern
118
119
@property
105
- def replacement(self) -> Callable[..., torch.Tensor]:
120
+ def replacement(
121
+ self,
122
+ ) -> Callable[..., torch.Tensor | tuple[torch.Tensor, torch.Tensor]]:
123
_ln = _encode_layer_name(self._layer_name)
124
125
if _USE_LAYERNAME:
@@ -126,7 +143,9 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
143
output_block_scale=None,
144
kv_cache_dummy_dep=kv_cache_dummy_dep,
145
)
129
- return RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size])
146
+ return self._output(
147
+ RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size]), scale
148
+ )
149
150
return _replacement_with_ln
151
@@ -147,7 +166,9 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
166
output_block_scale=None,
167
kv_cache_dummy_dep=kv_cache_dummy_dep,
168
)
150
- return RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size])
169
+ return self._output(
170
+ RESHAPE_OP(at1[1], [-1, self._num_heads * self._head_size]), scale
171
+ )
172
173
return _replacement
174
@@ -168,6 +189,30 @@ class AttnFp8StaticQuantPattern(VllmPatternReplacement[..., torch.Tensor]):
189
return inputs
190
191
192
+class RocmAttnFp8StaticQuantPattern(AttnFp8StaticQuantPattern):
193
+ """Fuse AITER static quantization while preserving its scale output."""
194
+
195
+ def _quantize(
196
+ self, value: torch.Tensor, scale: torch.Tensor
197
+ ) -> tuple[torch.Tensor, torch.Tensor]:
198
+ out = torch.empty_like(value, dtype=FP8_DTYPE)
199
+ _, out, scale = auto_functionalized(
200
+ torch.ops.vllm.rocm_aiter_per_tensor_quant.default,
201
+ out=out,
202
+ x=value,
203
+ scale=scale,
204
+ is_dynamic=False,
205
+ )
206
+ # The shared AITER op declares scale mutable for its dynamic variant.
207
+ # Static quant preserves it; keep this graph output when replacing it.
208
+ return out, scale
209
+
210
+ def _output(
211
+ self, value: torch.Tensor, scale: torch.Tensor
212
+ ) -> tuple[torch.Tensor, torch.Tensor]:
213
+ return value, scale
214
+
215
+
216
class AttnNvfp4QuantPattern(
217
VllmPatternReplacement[..., tuple[torch.Tensor, torch.Tensor]]
218
):
@@ -387,6 +432,8 @@ class AttnQuantFusionPass(VllmFusionPatternMatcherPass):
432
for layer in layers:
433
if layer.impl.fused_output_quant_supported(_FP8_QUANT_KEY):
434
self.register(AttnFp8StaticQuantPattern(layer, dtype))
435
+ if current_platform.is_rocm() and rocm_aiter_ops.is_enabled():
436
+ self.register(RocmAttnFp8StaticQuantPattern(layer, dtype))
437
if _USE_LAYERNAME:
438
break
439