bugfix - get n

vllm-project · Jun 20, 2024 · 7d7b636 · 7d7b636
1 parent a5bc063
commit 7d7b636
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu b/csrc/quantization/cutlass_w8a8/scaled_mm_c2x.cu
@@ -348,7 +348,7 @@ void cutlass_gemm_sm80_dispatch(torch::Tensor& out, torch::Tensor const& a,
         out, a, b, std::forward<EpilogueArgs>(args)...);
   } else if (mp2 <= 128) {
     // M in (64, 128]
-    uint32_t const n = a.size(1);
+    uint32_t const n = out.size(1);
     bool const small_n = n < 8192;
     if (small_n) {
       return cutlass_gemm_caller<Cutlass2xGemmM128SmallN>(