cuda: fix uninitialized IQ2 dequant LUTs in batch MoE expert-tile kernels (n_embd > 4096)#513
Open
giannisanni wants to merge 1 commit into
Open
cuda: fix uninitialized IQ2 dequant LUTs in batch MoE expert-tile kernels (n_embd > 4096)#513giannisanni wants to merge 1 commit into
giannisanni wants to merge 1 commit into
Conversation
…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
This was referenced Jul 7, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 thexq_blocks <= 16activation-staging branch, but thedev_dot_iq2_xxs_q8_K_block8_deq_lutcalls 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 withn_embd > 4096skips 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 then_tokens == 1path 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. Themoe_gate_up_mid_decode_lut_qwarp32_kernelhas the same textual pattern but its launcher already guardsuse_decode_lut_gateonxq_blocks <= 16, so it is left unchanged.How it was found and tested
Found on a GLM-5.2 CUDA port (
n_embd7168 = 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):
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.