Skip to content

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

Open
giannisanni wants to merge 1 commit into
antirez:mainfrom
giannisanni:fix/cuda-iq2xxs-batch-lut
Open

cuda: fix uninitialized IQ2 dequant LUTs in batch MoE expert-tile kernels (n_embd > 4096)#513
giannisanni wants to merge 1 commit into
antirez:mainfrom
giannisanni:fix/cuda-iq2xxs-batch-lut

Conversation

@giannisanni

Copy link
Copy Markdown

Summary

The batch MoE expert-tile kernels (moe_gate_up_mid_expert_tile8_{row32,row2048,rowspan}) load the IQ2_XXS dequant tables (s_iq2_grid, s_iq2_signs) into shared memory inside the xq_blocks <= 16 activation-staging branch, but the dev_dot_iq2_xxs_q8_K_block8_deq_lut calls consume them unconditionally.

Models with n_embd <= 4096 (16 q8_K blocks, e.g. V4 Flash) always take the branch, so the bug is invisible there. Any model with n_embd > 4096 skips it and the entire batch gate/up dequant runs against uninitialized shared memory: multi-token batch encodes (prefill, MTP verify) produce fluent garbage at full speed, while single-token decode stays correct because the n_tokens == 1 path uses different kernels. That combination (fast, plausible-looking token stream, wrong content, decode fine) makes it easy to misattribute to model quality.

Fix

Hoist the two LUT loads out of the staging conditional in the three kernels and make the __syncthreads() barrier unconditional. The moe_gate_up_mid_decode_lut_qwarp32_kernel has the same textual pattern but its launcher already guards use_decode_lut_gate on xq_blocks <= 16, so it is left unchanged.

How it was found and tested

Found on a GLM-5.2 CUDA port (n_embd 7168 = 28 blocks) where long-prompt batch prefill output was corrupt while short prompts (token-major path) were fine. Bisected with per-stage kill switches down to the expert-tile gate/up kernels, then to the conditional LUT load.

Before/after on the same hardware (RTX 4060 Ti 16GB, SSD streaming, 600-token prompt, greedy):

prefill output
before 6.5 t/s fluent garbage
after 6.5 t/s correct (retrieval question about the prompt answered right)

Flash-class models are unaffected by the change (they take the same staging branch as before; the LUT loads just move a few lines up).

Context: this came out of the streaming experiments in #495; the same uninitialized LUTs also corrupted 2-token MTP verify logits there.

…kernels

The moe_gate_up_mid_expert_tile8_{row32,row2048,rowspan} kernels stage the
IQ2_XXS dequant tables (s_iq2_grid, s_iq2_signs) into shared memory inside
the xq_blocks <= 16 activation-staging branch, but the
dev_dot_iq2_xxs_q8_K_block8_deq_lut calls consume them unconditionally.

Models with n_embd <= 4096 (16 q8_K blocks, e.g. V4 Flash) always take the
branch, so the bug is invisible there. Any model with n_embd > 4096 skips
it and runs the entire batch gate/up dequant against uninitialized shared
memory: multi-token batch encodes (prefill, MTP verify) produce fluent
garbage at full speed while single-token decode stays correct, since the
n_tokens == 1 path uses different kernels.

Found on a GLM-5.2 port (n_embd 7168 = 28 blocks) where batch prefill
output was corrupt; with the fix, a 600-token prompt prefills correctly at
6.5 t/s vs 0.30 t/s token-major on the same hardware (RTX 4060 Ti, SSD
streaming). The fix hoists the two LUT loads out of the staging branch and
makes the barrier unconditional; the decode_lut kernel is left unchanged
because its launcher already guards xq_blocks <= 16.
giannisanni pushed a commit to giannisanni/neutronstar that referenced this pull request Jul 7, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant