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

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]}.
+     *
+     * <p>The multi-workgroup variant ({@code reductionOneBlockWithLayer}) has workgroup 0
+     * combining the other workgroups' partial sums with <em>no inter-workgroup
+     * synchronization</em> — 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}.</p>
+     */
+    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;

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,15 +87,17 @@ 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);
             gridScheduler.addWorkerGrid("layer_" + i + ".attention", parallelAttentionWorker);
             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,