add triton fp8 gemm support

Author: k50112113

vllm/model_executor/layers/quantization/fp8.py

@@ -203,8 +203,7 @@ def __init__(self, quant_config: Fp8Config):
         # 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())
+                                           and envs.VLLM_ROCM_USE_AITER_LINEAR)
 
         self.block_quant = self.quant_config.weight_block_size is not None
         self.act_q_static = self.quant_config.activation_scheme == "static"

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

@@ -54,9 +54,12 @@ def rocm_aiter_gemm_w8a8_blockscale_impl(
     block_size: list[int],
     output_dtype: torch.dtype = torch.float16,
 ) -> torch.Tensor:
-    import aiter as rocm_aiter
+    # import aiter as rocm_aiter
 
-    return rocm_aiter.gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
+    # return rocm_aiter.gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
+    from aiter.ops.triton.gemm_a8w8_blockscale import gemm_a8w8_blockscale
+
+    return gemm_a8w8_blockscale(A, B, As, Bs, dtype=output_dtype)
 
 
 def rocm_aiter_gemm_w8a8_blockscale_fake(
@@ -185,7 +188,7 @@ def apply_w8a8_block_fp8_linear(
                                       block_size, input.dtype)
 
     else:
-        if use_aiter_and_is_supported:
+        if use_aiter_and_is_supported and current_platform.is_fp8_fnuz():
             q_input, x_scale = aiter_per1x128_quant(
                 input_2d.contiguous(), quant_dtype=rocm_aiter.dtypes.fp8)
         else: