Fixes non-coalesced HBM access in marlin_int4_fp8_preprocess_kernel_awq (#47268)
Signed-off-by: xjx <493337577@qq.com> Signed-off-by: flutist-alibaba <30485581+flutist@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Co-authored-by: Shengqi Chen <harry-chen@outlook.com>
flutist-alibaba committed
Jul 21, 2026 at 09:30 UTC
0d9210a50259105d4ce889a78e8142433c273e28
1 file changed
+11
-7
csrc/libtorch_stable/quantization/marlin/marlin_int4_fp8_preprocess.cu
+11
-7
@@ -39,11 +39,15 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
39
// AWQ zeros: (size_k // group_size, size_n // 8)
40
const int32_t* __restrict__ qzeros, int32_t size_n, int32_t size_k,
41
int32_t group_size) {
42
- int32_t val =
43
- qweight[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y];
44
- int32_t zero =
45
- qzeros[(blockIdx.x * 32 + threadIdx.x) / group_size * size_n / 8 +
46
- blockIdx.y];
42
+ // Thread mapping: threadIdx.x -> column dim (coalesced read within a row),
43
+ // blockIdx.x -> row dim. Adjacent threads read consecutive int32 in the
44
+ // same row (stride 1) instead of striding across rows (stride size_n/8).
45
+ int col = blockIdx.y * 32 + threadIdx.x;
46
+ if (col >= size_n / 8) return;
47
+ (void)size_k;
48
+
49
+ int32_t val = qweight[blockIdx.x * (size_n / 8) + col];
50
+ int32_t zero = qzeros[blockIdx.x / group_size * (size_n / 8) + col];
51
int32_t new_val = 0;
52
53
#pragma unroll
@@ -58,7 +62,7 @@ __global__ void marlin_int4_fp8_preprocess_kernel_awq(
62
zero >>= 4;
63
}
64
61
- output[(blockIdx.x * 32 + threadIdx.x) * size_n / 8 + blockIdx.y] = new_val;
65
+ output[blockIdx.x * (size_n / 8) + col] = new_val;
66
}
67
68
torch::stable::Tensor marlin_int4_fp8_preprocess(
@@ -102,7 +106,7 @@ torch::stable::Tensor marlin_int4_fp8_preprocess(
106
"qweight.size(0) % qzeros.size(0) != 0");
107
STD_TORCH_CHECK(group_size % 8 == 0, "group_size % 8 != 0");
108
105
- dim3 blocks(size_k / 32, size_n / 8);
109
+ dim3 blocks(size_k, (size_n / 8 + 31) / 32);
110
marlin_int4_fp8_preprocess_kernel_awq<<<blocks, 32, 0, stream>>>(
111
reinterpret_cast<const int32_t*>(qweight.const_data_ptr()),
112
reinterpret_cast<int32_t*>(output.mutable_data_ptr()),