修改cutlass 单测

26513bb5 · gaoqiong · 1a9775b8 · 26513bb5 · 26513bb5
Commit 26513bb5 authored Dec 03, 2024 by gaoqiong
Expand all Hide whitespace changes
Inline Side-by-side

Showing with 242 additions and 230 deletions

tests/kernels/test_cutlass.py tests/kernels/test_cutlass.py +240 -229

vllm/_custom_ops.py vllm/_custom_ops.py +2 -1

No files found.
--- a/tests/kernels/test_cutlass.py
+++ b/tests/kernels/test_cutlass.py
--- a/vllm/_custom_ops.py
+++ b/vllm/_custom_ops.py
@@ -706,7 +706,8 @@ def cutlass_scaled_mm(a: torch.Tensor,
    # torch.ops._C.cutlass_scaled_mm(out, a, b, scale_a, scale_b, bias)
    # return out
-    return quant_ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
+    #return quant_ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
+    return quant_ops.rocblas_scaled_mm_nn(a, b, scale_a, scale_b, out_dtype, bias)
 def rocblas_scaled_mm(a: torch.Tensor,
                      b: torch.Tensor,