bug fix ptpc invocation of aiter kernel and edge case of non aiter linear

Signed-off-by: tjtanaavllm <tunjian.tan@amd.com>

Author: tjtanaavllm

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

@@ -42,9 +42,13 @@ def rocm_aiter_gemm_a8w8_bpreshuffle_impl(
         # 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
-        return gemm_a8w8_bpreshuffle_CK(input, weight, scale_a, scale_b,
-                                        out_dtype)
+
+        m = input.shape[0]
+        n = weight.shape[0]
+        from aiter import gemm_a8w8_bpreshuffle_ck
+        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,

vllm/model_executor/layers/utils.py

@@ -101,6 +101,7 @@ def inner_function(layer: torch.nn.Module,
 
         if use_aiter:
             return aiter_ops.rocm_aiter_tuned_gemm(x, weight, bias)
+        return torch.nn.functional.linear(x, weight, bias)
 
     return inner_function