Skip to content

cuda: fix uninitialized IQ2 dequant LUTs in batch MoE expert-tile kernels (n_embd > 4096) #513

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
giannisanni wants to merge 1 commit into
base: main
Choose a base branch
from
Open

cuda: fix uninitialized IQ2 dequant LUTs in batch MoE expert-tile kernels (n_embd > 4096) #513

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
39 changes: 30 additions & 9 deletions ds4_cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10816,15 +10816,22 @@ __global__ static void moe_gate_up_mid_expert_tile8_row32_kernel(
slot[np] = pair[np] - tok[np] * n_expert;
xqb[np] = xq + (uint64_t)tok[np] * xq_blocks;
}
/* The IQ2 dequant LUTs are consumed unconditionally by the dot calls
* below, so they must load for every xq_blocks; only the xq staging is
* <=16-blocks specific. Models with n_embd > 4096 (more than 16 q8_K
* blocks) used to skip the LUT loads too, leaving s_iq2_grid and
* s_iq2_signs uninitialized for the whole gate/up dequant. */
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) {
uint32_t p = i / xq_blocks;
uint32_t b = i - p * xq_blocks;
sxq[p][b] = xqb[p][b];
}
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
__syncthreads();
}
__syncthreads();
if (xq_blocks <= 16u) {
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
if (row >= expert_mid_dim) return;
Expand Down Expand Up @@ -10906,15 +10913,22 @@ __global__ static void moe_gate_up_mid_expert_tile8_row2048_kernel(
slot[np] = pair[np] - tok[np] * n_expert;
xqb[np] = xq + (uint64_t)tok[np] * xq_blocks;
}
/* The IQ2 dequant LUTs are consumed unconditionally by the dot calls
* below, so they must load for every xq_blocks; only the xq staging is
* <=16-blocks specific. Models with n_embd > 4096 (more than 16 q8_K
* blocks) used to skip the LUT loads too, leaving s_iq2_grid and
* s_iq2_signs uninitialized for the whole gate/up dequant. */
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) {
uint32_t p = i / xq_blocks;
uint32_t b = i - p * xq_blocks;
sxq[p][b] = xqb[p][b];
}
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
__syncthreads();
}
__syncthreads();
if (xq_blocks <= 16u) {
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t rr = 0; rr < 64u; rr++) {
Expand Down Expand Up @@ -11000,15 +11014,22 @@ __global__ static void moe_gate_up_mid_expert_tile8_rowspan_kernel(
slot[np] = pair[np] - tok[np] * n_expert;
xqb[np] = xq + (uint64_t)tok[np] * xq_blocks;
}
/* The IQ2 dequant LUTs are consumed unconditionally by the dot calls
* below, so they must load for every xq_blocks; only the xq staging is
* <=16-blocks specific. Models with n_embd > 4096 (more than 16 q8_K
* blocks) used to skip the LUT loads too, leaving s_iq2_grid and
* s_iq2_signs uninitialized for the whole gate/up dequant. */
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
if (xq_blocks <= 16u) {
for (uint32_t i = threadIdx.x; i < np * xq_blocks; i += blockDim.x) {
uint32_t p = i / xq_blocks;
uint32_t b = i - p * xq_blocks;
sxq[p][b] = xqb[p][b];
}
for (uint32_t i = threadIdx.x; i < 256u; i += blockDim.x) s_iq2_grid[i] = cuda_iq2xxs_grid[i];
for (uint32_t i = threadIdx.x; i < 128u; i += blockDim.x) s_iq2_signs[i] = cuda_ksigns_iq2xs[i];
__syncthreads();
}
__syncthreads();
if (xq_blocks <= 16u) {
for (uint32_t p = 0; p < np; p++) xqb[p] = sxq[p];
}
for (uint32_t rr = 0; rr < ROW_SPAN / 32u; rr++) {
Expand Down