Add PTPC Fused MoE and PTPC Gemm Support; AITER upgrade to 916bf3c

vllm/_custom_ops.py

@@ -1003,7 +1003,7 @@ def scaled_fp8_quant(
                                 device=input.device,
                                 dtype=torch.float32)
             torch.ops._C.dynamic_per_token_scaled_fp8_quant(
-                output, input, scale, scale_ub)
+                output, input.contiguous(), scale, scale_ub)
         else:
             scale = torch.zeros(1, device=input.device, dtype=torch.float32)
             torch.ops._C.dynamic_scaled_fp8_quant(output, input, scale)

vllm/model_executor/layers/fused_moe/rocm_aiter_fused_moe.py

@@ -248,7 +248,7 @@ def rocm_aiter_asm_moe_impl(
         fc2_smooth_scale=fc2_smooth_scale,
         a16=a16,
         per_tensor_quant_scale=per_tensor_quant_scale,
-        block_shape=tuple(block_shape),
+        block_shape=None if block_shape is None else tuple(block_shape),
         activation=activation,
         expert_mask=expert_mask,
     )

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

@@ -9,6 +9,7 @@
 from compressed_tensors.quantization import (ActivationOrdering,
                                              QuantizationStrategy)
 
+import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe  # noqa
 from vllm import _custom_ops as ops
 from vllm.logger import init_logger
@@ -106,6 +107,11 @@ def __init__(
                 "For FP8 Fused MoE layer, we require either per tensor or "
                 "channelwise, dynamic per token quantization.")
 
+        from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (
+            is_rocm_aiter_moe_enabled)
+
+        self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()
+
     def create_weights(self, layer: torch.nn.Module, num_experts: int,
                        hidden_size: int, intermediate_size_per_partition: int,
                        params_dtype: torch.dtype, **extra_weight_attrs):
@@ -249,6 +255,24 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
                     start += shard_size
             layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
                                                         requires_grad=False)
+        self.rocm_aiter_use_asm = (self.rocm_aiter_moe_enabled
+                                   and envs.VLLM_ROCM_USE_AITER_ASMMOE)
+        # Property to determine if AITER is used
+        if self.rocm_aiter_moe_enabled:
+            from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (  # noqa E501
+                rocm_aiter_fused_experts, shuffle_weights)
+
+            # reshaping weights is required for aiter moe kernel.
+            shuffled_w13, shuffled_w2 = shuffle_weights(
+                layer.w13_weight.data, layer.w2_weight.data)
+            layer.w13_weight = torch.nn.Parameter(shuffled_w13,
+                                                  requires_grad=False)
+            layer.w2_weight = torch.nn.Parameter(shuffled_w2,
+                                                 requires_grad=False)
+            self.rocm_aiter_fused_experts_func = rocm_aiter_fused_experts
+        else:
+            from vllm.model_executor.layers.fused_moe import fused_experts
+            self.fused_experts_func = fused_experts
 
     def apply(
         self,
@@ -282,6 +306,25 @@ def apply(
             scoring_func=scoring_func,
             e_score_correction_bias=e_score_correction_bias)
 
+        if self.rocm_aiter_moe_enabled:
+            return self.rocm_aiter_fused_experts_func(
+                hidden_states=x,
+                w1=layer.w13_weight,
+                w2=layer.w2_weight,
+                topk_weights=topk_weights,
+                topk_ids=topk_ids,
+                activation=activation,
+                apply_router_weight_on_input=apply_router_weight_on_input,
+                use_fp8_w8a8=True,
+                per_channel_quant=self.weight_quant.strategy ==
+                QuantizationStrategy.CHANNEL,
+                w1_scale=layer.w13_weight_scale,
+                w2_scale=layer.w2_weight_scale,
+                a1_scale=layer.w13_input_scale,
+                a2_scale=layer.w2_input_scale,
+                expert_map=expert_map,
+                use_asm=self.rocm_aiter_use_asm)
+
         return fused_experts(
             x,
             layer.w13_weight,

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

@@ -6,6 +6,7 @@
 from compressed_tensors.quantization import QuantizationStrategy
 from torch.nn import Parameter
 
+from vllm import envs
 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
     CompressedTensorsScheme)
 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
@@ -26,6 +27,12 @@ def __init__(self, strategy: str, is_static_input_scheme: bool):
         self.out_dtype = torch.get_default_dtype()
         self.is_static_input_scheme = is_static_input_scheme
         self.fp8_linear = Fp8LinearOp(use_per_token_if_dynamic=True)
+        # AITER is only supported on ROCm and only for FP8_FNUZ
+        # and at the moment are MI300 series
+        self.use_aiter_and_is_supported = (current_platform.is_rocm()
+                                           and envs.VLLM_ROCM_USE_AITER
+                                           and envs.VLLM_ROCM_USE_AITER_LINEAR
+                                           and current_platform.is_fp8_fnuz())
 
     @classmethod
     def get_min_capability(cls) -> int:
@@ -75,7 +82,17 @@ def process_weights_after_loading(self, layer) -> None:
             else:
                 weight_scale = layer.weight_scale.data
 
-            layer.weight = Parameter(weight.t(), requires_grad=False)
+            if self.use_aiter_and_is_supported:
+                from aiter.ops.shuffle import shuffle_weight
+
+                # keep the weight as (N, K)
+                layer.weight = Parameter(shuffle_weight(weight,
+                                                        layout=(16, 16)),
+                                         requires_grad=False)
+            else:
+                # keep the weight as (K, N)
+                layer.weight = Parameter(weight.t(), requires_grad=False)
+
             # required by torch.compile to be torch.nn.Parameter
             layer.weight_scale = Parameter(weight_scale, requires_grad=False)
 

vllm/model_executor/layers/quantization/utils/fp8_utils.py

@@ -53,7 +53,7 @@ def rocm_aiter_gemm_w8a8_blockscale_impl(
 ) -> torch.Tensor:
     import aiter as rocm_aiter
 
-    return rocm_aiter.gemm_a8w8_blockscale_CK(A, B, As, Bs, dtype=output_dtype)
+    return rocm_aiter.gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
 
 
 def rocm_aiter_gemm_w8a8_blockscale_fake(

vllm/model_executor/layers/quantization/utils/w8a8_utils.py

@@ -9,6 +9,7 @@
 from vllm._aiter_ops import aiter_ops
 from vllm.config import CompilationLevel, get_current_vllm_config
 from vllm.platforms import current_platform
+from vllm.utils import direct_register_custom_op
 
 # Input scaling factors are no longer optional in _scaled_mm starting
 # from pytorch 2.5. Allocating a dummy tensor to pass as input_scale
@@ -33,6 +34,68 @@ def use_skinny_gemm() -> bool:
     return envs.VLLM_ROCM_USE_SKINNY_GEMM
 
 
+if current_platform.is_rocm():
+
+    def rocm_aiter_gemm_a8w8_bpreshuffle_impl(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            out_dtype: Optional[torch.dtype] = None,
+            scale_a: Optional[torch.Tensor] = None,
+            scale_b: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        # This AITER function can be used for
+        # - per-token activations + per-channel weights
+        #   e.g. vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+        # accept the weight as # keep the weight as (N, K)
+        # NOTE: The weight has to be shuffled in the
+        # process_weights_after_loading of the CompressedTensorsW8A8Fp8 class
+        from aiter import gemm_a8w8_bpreshuffle_ck
+        m = input.shape[0]
+        n = weight.shape[0]
+        Y = torch.empty(m, n, dtype=out_dtype, device=input.device)
+        gemm_a8w8_bpreshuffle_ck(input, weight, scale_a, scale_b, Y)
+        return Y
+
+    def rocm_aiter_gemm_a8w8_bpreshuffle_fake(
+            input: torch.Tensor,
+            weight: torch.Tensor,
+            out_dtype: Optional[torch.dtype] = None,
+            scale_a: Optional[torch.Tensor] = None,
+            scale_b: Optional[torch.Tensor] = None) -> torch.Tensor:
+
+        m = input.shape[0]
+        n = weight.shape[0]
+        if out_dtype is None:
+            out_dtype = input.dtype
+        return torch.empty((m, n), dtype=out_dtype, device=input.device)
+
+    if current_platform.is_rocm():
+        direct_register_custom_op(
+            op_name="rocm_aiter_gemm_a8w8_bpreshuffle",
+            op_func=rocm_aiter_gemm_a8w8_bpreshuffle_impl,
+            mutates_args=[],
+            fake_impl=rocm_aiter_gemm_a8w8_bpreshuffle_fake,
+            dispatch_key=current_platform.dispatch_key,
+        )
+
+
+def rocm_aiter_per_token_w8a8_scaled_mm(qinput: torch.Tensor,
+                                        weight: torch.Tensor,
+                                        out_dtype: torch.dtype,
+                                        scale_a: torch.Tensor,
+                                        scale_b: torch.Tensor,
+                                        bias: torch.Tensor,
+                                        input_2d: torch.Tensor,
+                                        output_shape: list) -> torch.Tensor:
+    output_shape = [*qinput.shape[:-1], weight.shape[0]]
+    output = torch.ops.vllm.rocm_aiter_gemm_a8w8_bpreshuffle(
+        qinput, weight, out_dtype=out_dtype, scale_a=scale_a, scale_b=scale_b)
+    if bias is not None:
+        output = output + bias
+
+    return torch.narrow(output, 0, 0, input_2d.shape[0]).view(*output_shape)
+
+
 def rocm_aiter_per_tensor_w8a8_scaled_mm(qinput: torch.Tensor,
                                          weight: torch.Tensor,
                                          out_dtype: torch.dtype,
@@ -311,6 +374,11 @@ def apply(
             elif (use_per_token_if_dynamic and not per_tensor_weights
                   and not per_tensor_activations
                   and USE_ROWWISE_TORCH_SCALED_MM):
+                if self.use_aiter_and_is_supported:
+                    return rocm_aiter_per_token_w8a8_scaled_mm(
+                        qinput, weight, out_dtype, x_scale, weight_scale, bias,
+                        input_2d, output_shape)
+
                 # For now validated on ROCm platform
                 # fp8 rowwise scaling in torch._scaled_mm is introduced in
                 # https://github.com/pytorch/pytorch/pull/144432 using hipBLASLt

vllm/platforms/rocm.py

@@ -135,6 +135,8 @@ class RocmPlatform(Platform):
         "fbgemm_fp8", "gguf", "quark", "ptpc_fp8"
     ]
 
+    _fp8_dtype: torch.dtype = None
+
     @classmethod
     def get_attn_backend_cls(cls, selected_backend, head_size, dtype,
                              kv_cache_dtype, block_size, use_v1,
@@ -328,10 +330,12 @@ def is_fp8_fnuz(cls) -> bool:
 
     @classmethod
     def fp8_dtype(cls) -> torch.dtype:
-        if cls.is_fp8_fnuz():
-            return torch.float8_e4m3fnuz
-        else:
-            return torch.float8_e4m3fn
+        if cls._fp8_dtype is None:
+            if cls.is_fp8_fnuz():
+                cls._fp8_dtype = torch.float8_e4m3fnuz
+            else:
+                cls._fp8_dtype = torch.float8_e4m3fn
+        return cls._fp8_dtype
 
     @classmethod
     def supports_v1(cls, model_config: ModelConfig) -> bool: