iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm] [fusion] support ROCm aiter fusion of per tensor quant + rmsnorm (#52831)

Browse Source 
Signed-off-by: Andy Friedrich <afriedri@amd.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
afriedri
2026-10-16 15:19:07 -05:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Claude Opus 5.5
parent 9eb7fb9157
commit cc2e4fecf2
4 changed files with 274 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/compile/passes/test_fusion.py
+16
View File
@@ -107,6 +107,8 @@ AITER_KERNEL_GROUPSHAPE_COMBINATIONS = [
# Per-token with ChannelWiseTorchFP8ScaledMMLinearKernel
(ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, True),
(ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, False),
# Static per-tensor: only the aiter quant op is matched by the pass
(ROCmFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR, True),
# Blockwise
(AiterFp8BlockScaledMMKernel, GroupShape(1, 128), True),
]
@@ -197,6 +199,9 @@ class TestModel(torch.nn.Module):
# Blockwise path
if self.use_aiter_fusion and self.use_aiter_quant_op:
return [rocm_aiter_ops.get_group_quant_op()]
elif self.group_shape.is_per_tensor():
if self.use_aiter_quant_op:
return [rocm_aiter_ops.get_per_tensor_quant_op()]
else:
if self.use_aiter_quant_op:
return [rocm_aiter_ops.get_per_token_quant_op()]
@@ -221,6 +226,17 @@ class TestModel(torch.nn.Module):
AiterFusedAddRMSFp8GroupQuantPattern.FUSED_OP,
AiterRMSFp8GroupQuantPattern.FUSED_OP,
]
elif self.group_shape.is_per_tensor():
# Static per-tensor aiter fusion
from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
AiterFusedAddRMSNormStaticQuantPattern,
AiterRMSNormStaticQuantPattern,
)
return [
AiterFusedAddRMSNormStaticQuantPattern.FUSED_OP,
AiterRMSNormStaticQuantPattern.FUSED_OP,
]
else:
# Per-token aiter fusion
from vllm.compilation.passes.fusion.rocm_aiter_fusion import (
vllm/_aiter_ops.py
+94
View File
@@ -1492,6 +1492,74 @@ def _rocm_aiter_rmsnorm_fp8_group_quant_fake(
)
def _rocm_aiter_fused_rms_fp8_per_tensor_static_quant_impl(
x: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
scale: torch.Tensor,
quant_dtype: torch.dtype,
) -> torch.Tensor:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_per_tensor_static_quant
x_quant, _, _, _ = fused_rms_fp8_per_tensor_static_quant(
x,
weight,
epsilon,
scale,
None,
None,
None,
dtype_quant=quant_dtype,
res1=None,
)
return x_quant
def _rocm_aiter_fused_rms_fp8_per_tensor_static_quant_fake(
x: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
scale: torch.Tensor,
quant_dtype: torch.dtype,
) -> torch.Tensor:
return torch.empty_like(x, dtype=quant_dtype)
def _rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant_impl(
x: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
scale: torch.Tensor,
quant_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_per_tensor_static_quant
x_quant, _, _, residual_out = fused_rms_fp8_per_tensor_static_quant(
x,
weight,
epsilon,
scale,
None,
None,
None,
dtype_quant=quant_dtype,
res1=residual,
)
return x_quant, residual_out
def _rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant_fake(
x: torch.Tensor,
residual: torch.Tensor,
weight: torch.Tensor,
epsilon: float,
scale: torch.Tensor,
quant_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
return torch.empty_like(x, dtype=quant_dtype), torch.empty_like(residual)
def _rocm_aiter_fused_rms_gated_fp8_group_quant_impl(
x: torch.Tensor,
weight: torch.Tensor,
@@ -3050,6 +3118,18 @@ class rocm_aiter_ops:
fake_impl=_rocm_aiter_rmsnorm_fp8_group_quant_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_fused_rms_fp8_per_tensor_static_quant",
op_func=_rocm_aiter_fused_rms_fp8_per_tensor_static_quant_impl,
fake_impl=_rocm_aiter_fused_rms_fp8_per_tensor_static_quant_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant",
op_func=_rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant_impl,
fake_impl=_rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant_fake,
)
direct_register_custom_op(
op_name="rocm_aiter_fused_rms_gated_fp8_group_quant",
op_func=_rocm_aiter_fused_rms_gated_fp8_group_quant_impl,
@@ -3179,6 +3259,16 @@ class rocm_aiter_ops:
def get_rmsnorm_group_fused_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default
@staticmethod
def get_fused_rms_fp8_per_tensor_static_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_fused_rms_fp8_per_tensor_static_quant.default
@staticmethod
def get_fused_add_rms_fp8_per_tensor_static_quant_op() -> OpOverload:
return (
torch.ops.vllm.rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant.default
)
@staticmethod
def get_fused_rms_gated_fp8_group_quant_op() -> OpOverload:
"""Return the fused gated-RMSNorm + FP8 group quant custom op."""
@@ -3192,6 +3282,10 @@ class rocm_aiter_ops:
def get_per_token_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_per_token_quant.default
@staticmethod
def get_per_tensor_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_per_tensor_quant.default
@staticmethod
def get_group_quant_op() -> OpOverload:
return torch.ops.vllm.rocm_aiter_group_fp8_quant.default
vllm/compilation/passes/fusion/matcher_utils.py
+19 -6
View File
@@ -394,10 +394,13 @@ class MatcherQuantFP8(MatcherCustomOp):
self.is_tma_aligned = is_tma_aligned
if match_rocm_aiter:
assert not quant_key.scale.group_shape.is_per_tensor(), (
"ROCm aiter fusion pass does not support per tensor quantization"
)
if quant_key.scale.group_shape.is_per_token():
if quant_key.scale.group_shape.is_per_tensor():
assert quant_key.scale.static, (
"ROCm aiter fusion pass only supports static per-tensor "
"quantization"
)
self.QUANT_OP = rocm_aiter_ops.get_per_tensor_quant_op()
elif quant_key.scale.group_shape.is_per_token():
self.QUANT_OP = rocm_aiter_ops.get_per_token_quant_op()
else:
assert quant_key.scale.group_shape.col == 128, (
@@ -432,14 +435,24 @@ class MatcherQuantFP8(MatcherCustomOp):
scale: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
quant_key_group_shape = self.quant_key.scale.group_shape
if quant_key_group_shape == GroupShape.PER_TENSOR:
# rocm_aiter_per_tensor_quant mutates `out` and `scale` and returns
# None, so it has to be matched through auto_functionalized.
assert scale is not None
out = torch.empty_like(input, dtype=self.quant_key.dtype)
result = auto_functionalized(
self.QUANT_OP, out=out, x=input, scale=scale, is_dynamic=False
)
# `scale` is mutable, so its functionalized output is a distinct
# value; return it to keep the quant node inside the match.
return result[1], result[2]
if quant_key_group_shape == GroupShape.PER_TOKEN:
return self.QUANT_OP( # type: ignore[no-any-return]
x=input,
quant_dtype=self.quant_key.dtype,
scale=scale,
)
else:
return self.QUANT_OP(input, quant_key_group_shape.col) # type: ignore[no-any-return]
return self.QUANT_OP(input, quant_key_group_shape.col) # type: ignore[no-any-return]
def forward_custom(
self,
vllm/compilation/passes/fusion/rocm_aiter_fusion.py
+145
View File
@@ -19,6 +19,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
ScaleDesc,
kFp8Dynamic128Sym,
kStaticTensorScale,
)
from vllm.platforms import current_platform
@@ -561,6 +562,140 @@ class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern):
)
class AiterRMSNormStaticQuantPattern(AiterRMSNormQuantPattern):
"""This pattern fuses rms_norm & aiter static per-tensor fp8 quant
into an aiter fused_rms_fp8_per_tensor_static_quant op.
"""
FUSED_OP = rocm_aiter_ops.get_fused_rms_fp8_per_tensor_static_quant_op()
def __init__(
self,
epsilon: float,
quant_dtype: torch.dtype,
symmetric: bool = True,
) -> None:
key = FusedRMSQuantKey(
fused_add=False,
quant=QuantKey(
dtype=quant_dtype, scale=kStaticTensorScale, symmetric=symmetric
),
)
super().__init__(epsilon, key)
def register(self, pm_pass: PatternMatcherPass) -> None:
def pattern(
input: torch.Tensor,
weight: torch.Tensor,
scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon)
result, scale = self.quant_matcher(result_rms, scale)
return result, scale
def replacement(
input: torch.Tensor,
weight: torch.Tensor,
scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
result = self.FUSED_OP(
x=input,
weight=weight,
epsilon=self.epsilon,
scale=scale,
quant_dtype=self.quant_dtype,
)
# The static scale is an input, so it passes through unchanged.
return result, scale
inputs = [
self.empty(5, 16), # input
self.empty(16), # weight
self.empty_f32(1, 1), # scale
]
pm.register_replacement(
pattern,
replacement,
inputs,
pm.fwd_only,
pm_pass,
)
class AiterFusedAddRMSNormStaticQuantPattern(AiterRMSNormQuantPattern):
"""This pattern fuses fused_add_rms_norm & aiter static per-tensor fp8
quant into an aiter fused_add_rms_fp8_per_tensor_static_quant op.
"""
FUSED_OP = rocm_aiter_ops.get_fused_add_rms_fp8_per_tensor_static_quant_op()
def __init__(
self,
epsilon: float,
quant_dtype: torch.dtype,
symmetric: bool = True,
) -> None:
key = FusedRMSQuantKey(
fused_add=True,
quant=QuantKey(
dtype=quant_dtype, scale=kStaticTensorScale, symmetric=symmetric
),
)
super().__init__(epsilon, key)
def register(self, pm_pass: PatternMatcherPass) -> None:
def pattern(
input: torch.Tensor,
weight: torch.Tensor,
residual: torch.Tensor,
scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm(
input, residual, weight, self.epsilon
)
result, scale = self.quant_matcher(result_rms, scale)
return result, residual_out, scale
def replacement(
input: torch.Tensor,
weight: torch.Tensor,
residual: torch.Tensor,
scale: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
result = self.FUSED_OP(
x=input,
residual=residual,
weight=weight,
epsilon=self.epsilon,
scale=scale,
quant_dtype=self.quant_dtype,
)
# The static scale is an input, so it passes through unchanged.
return result[0], result[1], scale
inputs = [
self.empty(5, 16), # input
self.empty(16), # weight
self.empty(5, 16), # residual
self.empty_f32(1, 1), # scale
]
pm.register_replacement(
pattern,
replacement,
inputs,
pm.fwd_only,
pm_pass,
)
class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
"""This pass fuses aiter rms_norm & vllm/aiter quant custom ops
into a fused rms_norm_quant op.
@@ -630,6 +765,14 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op
).register(self.patterns)
# Fuse aiter rms_norm + aiter static per-tensor fp8 quant
AiterRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(self.patterns)
# Fuse aiter fused_add_rms_norm + aiter static per-tensor fp8 quant
AiterFusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(
self.patterns
)
# When quant_fp8 custom ops are disabled, both AITER and native
# quant matchers trace through QuantFP8's native implementation.
# Registering both variants would create duplicate Inductor
@@ -690,6 +833,8 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass):
AiterFusedAddRMSNormDynamicQuantPattern,
AiterRMSFp8GroupQuantPattern,
AiterFusedAddRMSFp8GroupQuantPattern,
AiterRMSNormStaticQuantPattern,
AiterFusedAddRMSNormStaticQuantPattern,
DoubleAiterRMSFp8GroupQuantPattern,
DoubleAiterRMSFp8GroupQuantViewPattern,
AiterRMSNormGatedFp8GroupQuantPattern,