fix gemm_a8w8_bpreshuffle

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

Author: tjtanaavllm

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

@@ -49,9 +49,12 @@ 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
-        return gemm_a8w8_bpreshuffle(input, weight, scale_a, scale_b,
-                                        out_dtype)
+        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,