Fix logits scaling in GraniteKernels: correct scaling order for hb and output writes.

Author: mikepapadim

src/main/java/org/beehive/gpullama3/inference/InferenceCore.java

@@ -663,7 +663,7 @@ public static FloatTensor forwardGranite(Model model, State state, int token, in
         weights.wcls.matmul(state.x, state.logits, config.vocabularySize(), dim);
 
         // Apply Granite logit scaling (divide by the scaling factor)
-        state.logits.mapInPlace(v -> v / logitScale);
+        state.logits.mapInPlace(v -> v * logitScale);
 
         return state.logits;
     }

src/main/java/org/beehive/gpullama3/tornadovm/kernels/GraniteKernels.java

@@ -77,7 +77,7 @@ public static void matrixVectorGenericWithGraniteScale(
 
         // Thread 0 in each workgroup writes the final result
         if (localId == 0) {
-            hb.set(rowId, sum);
+            hb.set(rowId,  sum * logitsScale);
         }
     }
 
@@ -156,7 +156,6 @@ public static void processHeadsFlashAttentionWithGraniteScale(KernelContext cont
                     score += q_shared[d] * k_tile[score_idx_in_tile * headSize + d];
                 }
                 score *= attentionScale;
-//                score /= TornadoMath.sqrt(headSize);
                 s_tile[score_idx_in_tile] = score;
             }
 
@@ -339,7 +338,7 @@ public static void matrixVectorGenericQ8ByteWithGraniteScale(KernelContext conte
 
         // Thread 0 writes the result
         if (localId == 0) {
-            output.set(rowId, logitsScale * sum);
+            output.set(rowId, sum * logitsScale);
         }
     }