Fix Qwen3 FP16 garbage on CUDA backend: cross-workgroup race in RMS-norm reduction

[mikepapadim]

Fix: Qwen3 FP16 emits garbage on the CUDA backend — cross-workgroup data race in the RMS-norm reduction

Symptom. Qwen3-1.7B FP16 single-token decode on the CUDA backend produces garbage (endless \n + stray tokens, argmax stuck) from the first sampled token; Qwen3-4B works. Deterministic, both CUDA-graph modes.

Root cause. reductionOneBlockWithLayer writes each workgroup's partial sum to output[groupId+1], then thread gid==0 (workgroup 0) sums output[1..N] in the same kernel with no inter-workgroup synchronization or memory fence — a data race. On the NON_NVIDIA scheduler this race doesn't exist (a separate reductionFinalNormalization task combines the partials); the NVIDIA path skips that task and relies on the racy in-kernel combine. Whether the race is lost is schedule/compilation dependent: Llama-1B and Qwen3-4B happen to win it, Qwen3-1.7B deterministically loses it — every layer's RMS-norm scale is computed from stale partials, corrupting the hidden state everywhere at once.

Diagnosis (bisection, all on RTX 4090, TornadoVM 5.0.1-jdk21-dev CUDA backend):

• Batched-decode Qwen3 path (PR Static batched decode (LLaMA + Qwen3): B independent sequences per step, up to 41x aggregate throughput #129 kernels) = coherent with the same weights → weights/tokenizer/logits fine.
• Batch-prefill mode produced the correct first token (<think>), then the decode graph degraded → single-token layer graph implicated.
• Swapped every Qwen3-specific fused kernel (QKV matmul, Q/K per-head norm, RoPE, split-KV attention) for plain @Parallel reference implementations → still garbage.
• Swapped the RMS reductions alone → correct output, exact logits match with the all-reference run (argmax='<think>', host-recomputed logits ≡ GPU logits).

Fix. New reductionOneBlockWithLayerSingleGroup: one workgroup strides over the input, reduces in local memory, lane 0 writes the final scale — race-free by construction, one kernel, no extra launch. The Qwen3 FP16 layer graphs use it on the NVIDIA scheduler path (grid global == local == localSize); the NON_NVIDIA path keeps the existing partials + finalize pair.

Verified (all coherent, greedy):

model	before	after
Qwen3-1.7B FP16	garbage (\n\n\n…9…avier)	<think> Okay, the user is asking for the capital of France… Paris (matches the validated batched-decode output)
Qwen3-1.7B + CUDA graphs	garbage	coherent
Qwen3-4B	coherent	coherent (unchanged)
Llama-3.2-1B / Mistral-7B	coherent	coherent (path untouched)

Perf unchanged: Qwen3-1.7B ~51 tok/s, Llama-1B ~90 tok/s before/after.

Note. The same racy combine is still used by the Llama layer graphs and the logits layer (TransformerComputeKernels.reductionOneBlockWithLayer) — they currently win the race on this hardware, but they are latent instances of the same bug class; happy to extend the fix there if you want it in this PR.

🤖 Generated with Claude Code

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.