diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 934522a1b9..418d21f425 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -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 ( diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index fadb053915..6a6c99840a 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -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 diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 80be65a978..32cb0e35aa 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -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, diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index 849cd5d67d..16515f24bc 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py @@ -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,