[Kernel] Fix awq error when n is not divisable by 128 (#13227)

jinzhen-lin · web-flow · commit 8c32b08a86bd · 2025-02-13T20:07:05.000-08:00
diff --git a/csrc/quantization/awq/gemm_kernels.cu b/csrc/quantization/awq/gemm_kernels.cu
@@ -334,7 +334,7 @@ __global__ void __launch_bounds__(64)
   }
 
   // TODO: Shang: Hoist loop invariance.
-  for (int ax1_0_1 = 0; ax1_0_1 < 4; ++ax1_0_1) {
+  for (int ax1_0_1 = 0; ax1_0_1 < (N / 32); ++ax1_0_1) {
     for (int local_id = 0; local_id < 8; ++local_id) {
       int row_offset = (((int)blockIdx_y) / j_factors1) * 16 +
                        ((int)threadIdx.x) / 4 + (local_id % 4) / 2 * 8;

Original file line number	Diff line number	Diff line change
`@@ -334,7 +334,7 @@ __global__ void __launch_bounds__(64)`
`334`	`334`	`}`
`335`	`335`
`336`	`336`	`// TODO: Shang: Hoist loop invariance.`
`337`		`- for (int ax1_0_1 = 0; ax1_0_1 < 4; ++ax1_0_1) {`
	`337`	`+ for (int ax1_0_1 = 0; ax1_0_1 < (N / 32); ++ax1_0_1) {`
`338`	`338`	`for (int local_id = 0; local_id < 8; ++local_id) {`
`339`	`339`	`int row_offset = (((int)blockIdx_y) / j_factors1) * 16 +`
`340`	`340`	`((int)threadIdx.x) / 4 + (local_id % 4) / 2 * 8;`