Use fusion pass to select AITER group quant RMSNorm and w8a8 gemm (#707)

* initial aiter rms group quant fusion work

* aiter rms norm group quant fusion

* remove debug prints

Signed-off-by: charlifu <charlifu@amd.com>

Author: micah-wil

vllm/compilation/fusion.py

@@ -9,6 +9,7 @@
 from torch._inductor.pattern_matcher import PatternMatcherPass
 from torch._ops import OpOverload
 
+import vllm.envs as envs
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -76,6 +77,22 @@ def __str__(self):
     torch.ops._C.rms_norm_dynamic_per_token_quant.default,  # noqa: E501
 }
 
+if current_platform.is_rocm() and envs.VLLM_ROCM_USE_AITER:
+    AITER_RMS_GROUP_QUANT_OP = \
+        torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default
+    AITER_RMS_ADD_GROUP_QUANT_OP = \
+        torch.ops.vllm.rocm_aiter_rmsnorm_with_add_fp8_group_quant.default
+    
+    BLOCK_LINEAR_OP = torch.ops.vllm.apply_w8a8_block_fp8_linear.default
+    AITER_BLOCK_LINEAR_OP = \
+        torch.ops.vllm.rocm_aiter_gemm_w8a8_blockscale.default
+    
+    AITER_RMS_OP = torch.ops.vllm.rocm_aiter_rms_norm.default
+    AITER_RMS_ADD_OP = torch.ops.vllm.rocm_aiter_rmsnorm2d_fwd_with_add.default
+    
+    import aiter as rocm_aiter
+    rocm_aiter_fp8_dtype = rocm_aiter.dtypes.fp8
+    rocm_aiter_fp8_quant_group_size = 128
 
 class RMSNormQuantPattern:
 
@@ -336,6 +353,129 @@ def replacement(result: torch.Tensor, input: torch.Tensor,
             pm.fwd_only,
             pm_pass,
         )
+        
+        
+class AiterRMSGroupQuantFP8Pattern():
+
+    def __init__(self,
+                 epsilon: float,
+                 quant_dtype: torch.dtype):
+        self.epsilon = epsilon
+        self.quant_dtype = quant_dtype
+
+    def register(self, pm_pass: PatternMatcherPass):
+
+        def pattern(input: torch.Tensor, weight: torch.Tensor, #result_rms: torch.Tensor,
+                    linear_weight: torch.Tensor,
+                    linear_weight_scale: torch.Tensor):
+            at1 = AITER_RMS_OP(x=input,
+                               weight=weight,
+                               variance_epsilon=self.epsilon)
+            
+            at2 = BLOCK_LINEAR_OP(input=at1,
+                                    weight=linear_weight,
+                                    block_size=[128, 128],
+                                    weight_scale=linear_weight_scale,
+                                    input_scale=None,
+                                    bias=None,
+                                    cutlass_block_fp8_supported=False,
+                                    use_aiter_and_is_supported=True)
+
+            return at2
+
+        def replacement(input: torch.Tensor, weight: torch.Tensor,
+                        linear_weight: torch.Tensor,
+                        linear_weight_scale: torch.Tensor):
+            at1 = AITER_RMS_GROUP_QUANT_OP(x=input,
+                                           residual=None,
+                                           weight=weight,
+                                           variance_epsilon=self.epsilon)
+            
+            at2 = AITER_BLOCK_LINEAR_OP(A=at1[0],
+                                        B=linear_weight,
+                                        As=at1[1],
+                                        Bs=linear_weight_scale,
+                                        block_size=[128, 128],
+                                        output_dtype=input.dtype)
+
+            return at2
+
+        inputs = [
+            empty_bf16(5, 4),  # input
+            empty_bf16(1, 5),  # weight
+            torch.empty((2, 5), device="cuda", dtype=FP8_DTYPE), # linear_weight
+            empty_fp32(1, 1),  # linear_weight_scale
+        ]
+
+        pm.register_replacement(
+            pattern,
+            replacement,
+            inputs,
+            pm.fwd_only,
+            pm_pass)
+
+
+class AiterFusedAddRMSGroupQuantPattern():
+
+    def __init__(self,
+                 epsilon: float,
+                 quant_dtype: torch.dtype):
+        self.epsilon = epsilon
+        self.quant_dtype = quant_dtype
+
+    def register(self, pm_pass: PatternMatcherPass):
+
+        def pattern(input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+                    linear_weight: torch.Tensor,
+                    linear_weight_scale: torch.Tensor):
+            at1 = AITER_RMS_ADD_OP(x=input,
+                                   residual=residual,
+                                   weight=weight,
+                                   variance_epsilon=self.epsilon)
+            
+            at2 = BLOCK_LINEAR_OP(input=at1[0],
+                                  weight=linear_weight,
+                                  block_size=[128, 128],
+                                  weight_scale=linear_weight_scale,
+                                  input_scale=None,
+                                  bias=None,
+                                  cutlass_block_fp8_supported=False,
+                                  use_aiter_and_is_supported=True)
+            # result, residual
+            return at2, at1[1]
+
+        def replacement(input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+                        linear_weight: torch.Tensor,
+                        linear_weight_scale: torch.Tensor):
+
+            at1 = AITER_RMS_ADD_GROUP_QUANT_OP(x=input,
+                                               residual=residual,
+                                               weight=weight,
+                                               variance_epsilon=self.epsilon)
+            
+            at2 = AITER_BLOCK_LINEAR_OP(A=at1[0],
+                                        B=linear_weight,
+                                        As=at1[1],
+                                        Bs=linear_weight_scale,
+                                        block_size=[128, 128],
+                                        output_dtype=input.dtype)
+            # result, residual
+            return at2, at1[2]
+
+        inputs = [
+            empty_bf16(5, 4),  # input
+            empty_bf16(5, 4),  # residual
+            empty_bf16(1, 5),  # weight
+            torch.empty((2, 5), device="cuda", dtype=FP8_DTYPE), # linear_weight
+            empty_fp32(1, 1), # linear_weight_scale
+        ]
+
+        pm.register_replacement(
+            pattern,
+            replacement,
+            inputs,
+            pm.fwd_only,
+            pm_pass)
 
 
 class RMSNormQuantFusionPass(VllmPatternMatcherPass):
@@ -367,6 +507,14 @@ def __init__(self, config: VllmConfig):
             # Fuse fused_add_rms_norm + dynamic per-token fp8 quant
             FusedAddRMSNormDynamicQuantPattern(epsilon, FP8_DTYPE).register(
                 self.patterns)
+            
+            if envs.VLLM_ROCM_USE_AITER:
+                # Fuse rms_norm + dynamic group fp8 quant
+                AiterRMSGroupQuantFP8Pattern(epsilon, FP8_DTYPE).register(
+                    self.patterns)
+                
+                AiterFusedAddRMSGroupQuantPattern(epsilon, FP8_DTYPE).register(
+                    self.patterns)
 
         self.dump_patterns(config, self.patterns)
 
@@ -380,4 +528,6 @@ def uuid(self) -> Any:
                                 RMSNormStaticQuantPattern,
                                 RMSNormDynamicQuantPattern,
                                 FusedAddRMSNormStaticQuantPattern,
-                                FusedAddRMSNormDynamicQuantPattern)
+                                FusedAddRMSNormDynamicQuantPattern,
+                                AiterRMSGroupQuantFP8Pattern,
+                                AiterFusedAddRMSGroupQuantPattern)

vllm/model_executor/layers/layernorm.py

@@ -16,6 +16,11 @@ def is_rocm_aiter_rmsnorm_enabled() -> bool:
     return envs.VLLM_ROCM_USE_AITER_RMSNORM \
         and envs.VLLM_ROCM_USE_AITER
 
+if current_platform.is_rocm() and is_rocm_aiter_rmsnorm_enabled():
+    import aiter as rocm_aiter
+    from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
+    rocm_aiter_fp8_dtype = rocm_aiter.dtypes.fp8
+    rocm_aiter_fp8_quant_group_size = 128
 
 def rms_norm(x: torch.Tensor, weight: torch.Tensor,
              variance_epsilon: float) -> torch.Tensor:
@@ -88,6 +93,44 @@ def rocm_aiter_rmsnorm2d_fwd_with_add_impl(
     return output, residual_out
 
 
+def rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl(
+        x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+        variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    import aiter as rocm_aiter
+    (x_quant, x_quant_scales), _, _, res = \
+    fused_rms_fp8_group_quant(
+        x,
+        weight,
+        variance_epsilon,
+        None,
+        None,
+        None,
+        group_size=rocm_aiter_fp8_quant_group_size,
+        dtype_quant=rocm_aiter_fp8_dtype,
+        res1=residual,
+    )
+    return (x_quant, x_quant_scales, res)
+
+
+def rocm_aiter_rmsnorm_fp8_group_quant_impl(
+        x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+        variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
+    import aiter as rocm_aiter
+    (x_quant, x_quant_scales), _, _, res = \
+    fused_rms_fp8_group_quant(
+        x,
+        weight,
+        variance_epsilon,
+        None,
+        None,
+        None,
+        group_size=rocm_aiter_fp8_quant_group_size,
+        dtype_quant=rocm_aiter_fp8_dtype,
+        res1=residual,
+    )
+    return (x_quant, x_quant_scales)
+
+
 def rocm_aiter_rms_norm_fake(x: torch.Tensor, weight: torch.Tensor,
                              variance_epsilon: float) -> torch.Tensor:
     return torch.empty_like(x)
@@ -99,6 +142,24 @@ def rocm_aiter_rmsnorm2d_fwd_with_add_fake(
     return torch.empty_like(x), torch.empty_like(residual)
 
 
+def rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake(
+        x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+        variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    M, N = x.shape
+    scale_shape = (M, (N + rocm_aiter_fp8_quant_group_size - 1) // rocm_aiter_fp8_quant_group_size)
+    return (torch.empty_like(x, dtype=rocm_aiter_fp8_dtype, device=x.device), 
+            torch.empty(scale_shape, dtype=torch.float32, device=x.device),
+            torch.empty_like(residual, device=residual.device))
+
+
+def rocm_aiter_rmsnorm_fp8_group_quant_fake(
+        x: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor,
+        variance_epsilon: float) -> tuple[torch.Tensor, torch.Tensor]:
+    M, N = x.shape
+    scale_shape = (M, (N + rocm_aiter_fp8_quant_group_size - 1) // rocm_aiter_fp8_quant_group_size)
+    return (torch.empty_like(x, dtype=rocm_aiter_fp8_dtype, device=x.device), torch.empty(scale_shape, dtype=torch.float32, device=x.device))
+
+
 if current_platform.is_rocm():
     direct_register_custom_op(
         op_name="rocm_aiter_rms_norm",
@@ -115,6 +176,22 @@ def rocm_aiter_rmsnorm2d_fwd_with_add_fake(
         fake_impl=rocm_aiter_rmsnorm2d_fwd_with_add_fake,
         dispatch_key=current_platform.dispatch_key,
     )
+    
+    direct_register_custom_op(
+        op_name="rocm_aiter_rmsnorm_fp8_group_quant",
+        op_func=rocm_aiter_rmsnorm_fp8_group_quant_impl,
+        mutates_args=[],
+        fake_impl=rocm_aiter_rmsnorm_fp8_group_quant_fake,
+        dispatch_key=current_platform.dispatch_key,
+    )
+    
+    direct_register_custom_op(
+        op_name="rocm_aiter_rmsnorm_with_add_fp8_group_quant",
+        op_func=rocm_aiter_rmsnorm_with_add_fp8_group_quant_impl,
+        mutates_args=[],
+        fake_impl=rocm_aiter_rmsnorm_with_add_fp8_group_quant_fake,
+        dispatch_key=current_platform.dispatch_key,
+    )
 
 
 def dispatch_rocm_rmsnorm_func(with_fused_add: bool, dtype: torch.dtype):