From e70da6153e523feb88f1a6da99779786c00e4c22 Mon Sep 17 00:00:00 2001 From: mikepapadim Date: Tue, 14 Jul 2026 10:33:09 +0100 Subject: [PATCH] Fix Qwen3 FP16 garbage output on CUDA: RMS-norm reduction had a cross-workgroup data race (workgroup 0 combined other groups' partials with no synchronization); NVIDIA path now uses a race-free single-workgroup reduction --- .../TransformerComputeKernelsLayered.java | 38 +++++++++++++++++++ .../layers/type/fp16/Qwen3FP16FFNLayers.java | 31 ++++++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java index 0b7482e3..1b400e61 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/kernels/TransformerComputeKernelsLayered.java @@ -2489,6 +2489,44 @@ public static float matrixVectorRowMajorOptimizedx(KernelContext context, int lo return localSum[0]; } + /** + * Race-free RMS-norm reduction in ONE workgroup: threads stride over the input, reduce in + * local memory, and lane 0 writes the final scale factor to {@code output[0]}. + * + *

The multi-workgroup variant ({@code reductionOneBlockWithLayer}) has workgroup 0 + * combining the other workgroups' partial sums with no inter-workgroup + * synchronization — a data race. On the CUDA backend the race outcome is + * schedule/compilation dependent: e.g. Qwen3-1.7B reads stale partials on every layer + * (wrong normalization scale → garbage output) while Llama-1B and Qwen3-4B happen to win + * the race. The NON_NVIDIA scheduler avoids the race with a separate + * {@code reductionFinalNormalization} task; this kernel is the NVIDIA-path equivalent. + * Launch with exactly one workgroup: {@code global == local == localMemSize}.

+ */ + public static void reductionOneBlockWithLayerSingleGroup(KernelContext context, FloatArray output, FloatArray x, int size, float ermsNorm, int localMemSize) { + int lid = context.localIdx; + int groupSize = context.localGroupSizeX; + float[] localX = context.allocateFloatLocalArray(localMemSize); + + float partial = 0.0f; + for (int j = lid; j < size; j += groupSize) { + float v = x.get(j); + partial += v * v; + } + localX[lid] = partial; + + for (int stride = groupSize / 2; stride > 0; stride /= 2) { + context.localBarrier(); + if (lid < stride) { + localX[lid] += localX[lid + stride]; + } + } + + if (lid == 0) { + float ss = localX[0] / size + ermsNorm; + output.set(0, 1.0f / TornadoMath.sqrt(ss)); + } + } + // Second kernel - Combines partial sums and computes final normalization public static void reductionFinalNormalization(KernelContext context, FloatArray output, int size, float ermsNorm) { int gid = context.globalIdx; diff --git a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java index ee2faccb..f0da8b54 100644 --- a/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java +++ b/src/main/java/org/beehive/gpullama3/tornadovm/layers/type/fp16/Qwen3FP16FFNLayers.java @@ -59,6 +59,8 @@ public Qwen3FP16FFNLayers(String taskGraphName, Qwen3State state, Qwen3TornadoWe @Override public GridScheduler updateGridScheduler(GridScheduler gridScheduler) { WorkerGrid rmsNormWorker = WorkerGridFactory.createRmsNormWorker(config.dim(), state.localSize); + // Single-workgroup grid for the race-free NVIDIA-path RMS reduction (global == local). + WorkerGrid rmsSingleGroupWorker = WorkerGridFactory.createRmsNormWorker(state.localSize, state.localSize); WorkerGrid ropeWorker = WorkerGridFactory.createRoPEWorker(config.numberOfHeads(), nEmbdHead); // Split-KV attention launches nHeads*nSplits workgroups (one per head-split) followed by a combine // pass over nHeads workgroups. @@ -85,7 +87,8 @@ public GridScheduler updateGridScheduler(GridScheduler gridScheduler) { // Map workers to tasks for each layer (in task execution order) for (int i = 0; i < config.numberOfLayers(); i++) { // === Attention Block === - gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_reduce", rmsNormWorker); + gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_reduce", + shouldUseFinalNormalization() ? rmsNormWorker : rmsSingleGroupWorker); gridScheduler.addWorkerGrid("layer_" + i + ".attn_rms_qkv_projection", fusedQKVWorker); gridScheduler.addWorkerGrid("layer_" + i + ".qk_rmsnorm", qkRmsNormWorker); gridScheduler.addWorkerGrid("layer_" + i + ".rope_and_kv_cache", ropeWorker); @@ -93,7 +96,8 @@ public GridScheduler updateGridScheduler(GridScheduler gridScheduler) { gridScheduler.addWorkerGrid("layer_" + i + ".attention_combine", attentionCombineWorker); gridScheduler.addWorkerGrid("layer_" + i + ".attn_output_proj", matmul1Worker); // === FFN Block === - gridScheduler.addWorkerGrid("layer_" + i + ".ffn_rms_reduce", rmsNormWorker); + gridScheduler.addWorkerGrid("layer_" + i + ".ffn_rms_reduce", + shouldUseFinalNormalization() ? rmsNormWorker : rmsSingleGroupWorker); if (shouldUseFinalNormalization()) { gridScheduler.addWorkerGrid("layer_" + i + ".ffn_rms_finalize", rmsNormWorker); } @@ -230,6 +234,18 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { // ═══════════════════════════════════════════════════════════════════════ // RMS Normalization - compute scale factor + if (!shouldUseFinalNormalization()) { + // NVIDIA path: single-workgroup reduction. The multi-workgroup kernel relies on a + // racy cross-workgroup combine (no finalize task on this path) — see kernel javadoc. + unifiedLayer.task("attn_rms_reduce", + TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, + context, + qwen3State.temp, + qwen3State.wrapX, + config.dim(), + config.rmsNormEps(), + qwen3State.localSize); + } else unifiedLayer.task("attn_rms_reduce", TransformerComputeKernelsLayered::reductionOneBlockWithLayer, context, @@ -361,6 +377,17 @@ protected TaskGraph createFFNLayerTaskGraph(int layerIndex) { // ═══════════════════════════════════════════════════════════════════════ // RMS Normalization - compute scale factor + if (!shouldUseFinalNormalization()) { + // NVIDIA path: single-workgroup reduction (race-free) — see attn_rms_reduce. + unifiedLayer.task("ffn_rms_reduce", + TransformerComputeKernelsLayered::reductionOneBlockWithLayerSingleGroup, + context, + qwen3State.tempFFN, + qwen3State.wrapX, + config.dim(), + config.rmsNormEps(), + qwen3State.localSize); + } else unifiedLayer.task("ffn_rms_reduce", TransformerComputeKernelsLayered::reductionOneBlockWithLayer, context,