Fix Gemma 4 singleton expert prefill and prepare 0.2.4

CHANGELOG.md

@@ -4,6 +4,13 @@ All notable changes to models are documented here.
 
 ## [Unreleased]
 
+## [0.2.4] - 2026-08-02
+
+### Fixed
+
+- Fixed Gemma 4 native batched prefill when MoE routing assigns a single token
+  to a Q4_K expert projection by falling back to the Java projection kernel.
+
 ## [0.2.3] - 2026-08-02
 
 ### Added

README.md

@@ -136,16 +136,16 @@ directly:
 
 ```kotlin
 dependencies {
-    implementation("com.integrallis:models:0.2.3")
-    implementation("com.integrallis:backend-java:0.2.3") // or backend-native
+    implementation("com.integrallis:models:0.2.4")
+    implementation("com.integrallis:backend-java:0.2.4") // or backend-native
 }
 ```
 
 Use Apple's on-device system model on a supported Apple Silicon Mac:
 
 ```kotlin
 dependencies {
-    implementation("com.integrallis:backend-apple:0.2.3")
+    implementation("com.integrallis:backend-apple:0.2.4")
 }
 ```
 

backend-java/src/main/java/com/integrallis/models/backend/purejava/gemma4/Gemma4ForwardPass.java

@@ -198,8 +198,10 @@ private record F32Scratch(int rows, int columns, float[][] inputs, float[][] out
     int maxAttentionOutputDim = maxLength(attentionOutputs);
     int maxProjectionInput =
         Math.max(
-            Math.max(dim, maxAttentionOutputDim),
-            Math.max(config.sharedHiddenDim(), config.expertHiddenDim()));
+            256,
+            Math.max(
+                Math.max(dim, maxAttentionOutputDim),
+                Math.max(config.sharedHiddenDim(), config.expertHiddenDim())));
 
     this.state = new float[dim];
     this.normalized = new float[dim];
@@ -1304,7 +1306,7 @@ private void projectBatched(
         matrix.data(), matrix.type(), matrix.rows(), matrix.columns(), input, batchSize, output);
   }
 
-  private void projectBatched(
+  void projectBatched(
       MemorySegment matrix,
       GgufTensorType type,
       int rows,
@@ -1316,6 +1318,10 @@ private void projectBatched(
       batchedMatrixKernel.multiply(output, input, matrix, type, batchSize, rows, columns);
       return;
     }
+    if (batchSize == 1) {
+      project(matrix, type, rows, columns, input, output);
+      return;
+    }
     if (type == GgufTensorType.F32) {
       F32Scratch scratch = findF32Scratch(rows, columns);
       if (scratch == null) {

backend-java/src/test/java/com/integrallis/models/backend/purejava/gemma4/Gemma4ForwardPassTest.java

@@ -383,6 +383,67 @@ public void multiplyRaggedIndependent(
     }
   }
 
+  @Test
+  void singletonQuantizedExpertProjectionFallsBackWhenNativeKernelRequiresBatching()
+      throws Exception {
+    ToyModel model = ToyModel.create();
+    Gemma4Weights weights = Gemma4Weights.fromGgufFile(model.file(), model.config());
+    Gemma4TensorLayout layout = weights.expertLayout();
+    GgufBatchedMatrixKernel batchOnlyKernel =
+        new GgufBatchedMatrixKernel() {
+          @Override
+          public boolean supports(GgufTensorType type) {
+            return type == GgufTensorType.Q4_K;
+          }
+
+          @Override
+          public boolean isEligible(GgufTensorType type, int batchSize, int rows, int cols) {
+            return batchSize > 1 && supports(type);
+          }
+
+          @Override
+          public void multiply(
+              float[] output,
+              float[] input,
+              MemorySegment matrix,
+              GgufTensorType type,
+              int batchSize,
+              int rows,
+              int cols) {
+            throw new AssertionError("singleton projection must use the Java fallback");
+          }
+        };
+
+    try (Gemma4ExpertCache experts =
+        new Gemma4ExpertCache(
+            new Gemma4ExpertLoader(reader(model.file().fileSegment())),
+            model.config().numLayers(),
+            model.config().numExperts(),
+            1,
+            (layer, expert) -> layout.layer(layer).expert(expert),
+            Gemma4ExpertCache.CachePolicy.LFU)) {
+      Gemma4ForwardPass forwardPass =
+          new Gemma4ForwardPass(
+              model.config(),
+              weights,
+              Gemma4KvCache.create(model.config(), 8, 2),
+              experts,
+              batchOnlyKernel);
+      float[] input = new float[256];
+      for (int index = 0; index < input.length; index++) {
+        input[index] = (index - 128) * 0.0078125f;
+      }
+      MemorySegment matrix = MemorySegment.ofArray(q4KBlock(0.125f, 0.0625f, 7));
+      float[] expected = new float[1];
+      float[] actual = new float[1];
+      TensorOps.ggufMatmul(expected, input, matrix, GgufTensorType.Q4_K, 1, 256);
+
+      forwardPass.projectBatched(matrix, GgufTensorType.Q4_K, 1, 256, input, 1, actual);
+
+      assertThat(actual).containsExactly(expected);
+    }
+  }
+
   private static void assertClose(float[] actual, float[] expected) {
     assertThat(actual).hasSameSizeAs(expected);
     for (int index = 0; index < actual.length; index++) {
@@ -790,6 +851,21 @@ private static float[] concat(float[] first, float[] second) {
     return result;
   }
 
+  private static byte[] q4KBlock(float scale, float minScale, int quant) {
+    byte[] block = new byte[144];
+    ByteBuffer buffer = ByteBuffer.wrap(block).order(ByteOrder.LITTLE_ENDIAN);
+    buffer.putShort(0, Float.floatToFloat16(scale));
+    buffer.putShort(2, Float.floatToFloat16(minScale));
+    for (int group = 0; group < 4; group++) {
+      block[4 + group] = 1;
+    }
+    for (int group = 4; group < 8; group++) {
+      block[4 + group + 4] = 1;
+    }
+    Arrays.fill(block, 16, block.length, (byte) (quant | (quant << 4)));
+    return block;
+  }
+
   private static final class FixtureBuilder {
     private final List<GgufTensorInfo> infos = new ArrayList<>();
     private final Map<String, float[]> values = new HashMap<>();

backend-native/src/main/rust/model-kernels/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "jmodels-kernels"
-version = "0.2.3"
+version = "0.2.4"
 edition = "2024"
 license = "Apache-2.0"
 publish = false

docs/content/antora.yml

@@ -1,7 +1,7 @@
 name: models
 title: Models
 version: 'current'
-display_version: '0.2.3'
+display_version: '0.2.4'
 prerelease: false
 start_page: ROOT:index.adoc
 nav:
@@ -10,7 +10,7 @@ asciidoc:
   attributes:
     source-language: java
     source-highlighter: highlight.js
-    models-version: '0.2.3'
+    models-version: '0.2.4'
     modeljars-version: '0.1.2'
     vectors-version: '0.1.4'
     url-models-github: https://github.com/integrallis/models

docs/landing/index.html

@@ -36,7 +36,7 @@
                 <a href="https://integrallis.com" class="brand-mark" aria-label="Integrallis"><img class="ilogo" src="assets/images/integrallis-logo.png" alt="Integrallis"></a>
                 <span class="sep">/</span>
                 <a href="#top" class="product">models</a>
-                <a href="https://github.com/integrallis/models/releases" class="version-pill" title="Release notes">v0.2.3</a>
+                <a href="https://github.com/integrallis/models/releases" class="version-pill" title="Release notes">v0.2.4</a>
             </div>
             <div class="nav-right">
                 <div class="nav-links">
@@ -182,7 +182,7 @@ <h2 class="section-title">Use Apple Intelligence from the JVM</h2>
                     </div>
                     <div class="apple-example reveal">
                         <div class="code-title">// Add the Apple backend</div>
-                        <pre><code>implementation("com.integrallis:backend-apple:0.2.3")</code></pre>
+                        <pre><code>implementation("com.integrallis:backend-apple:0.2.4")</code></pre>
                         <div class="code-title">// Check availability, then generate locally</div>
                         <pre><code class="language-java hljs" data-lang="java">try (var client = AppleFoundationModels.create()) {
     var status = client.availability();

docs/package.json

@@ -1,6 +1,6 @@
 {
   "name": "models-docs",
-  "version": "0.2.3",
+  "version": "0.2.4",
   "private": true,
   "description": "Documentation for Models",
   "dependencies": {

gradle.properties

@@ -1,2 +1,2 @@
-version = 0.2.3
+version = 0.2.4
 vectorsVersion = 0.1.5

models-backend-apple/README.md

@@ -5,7 +5,7 @@
 Java FFM bridge to Apple's on-device Foundation Models runtime.
 
 ```kotlin
-implementation("com.integrallis:backend-apple:0.2.3")
+implementation("com.integrallis:backend-apple:0.2.4")
 ```
 
 This module is intentionally separate from `backend-java`. The core

models-rag/README.md

@@ -8,7 +8,7 @@ generation, then enforces source attribution, abstention, and extractive fallbac
 on the generated answer.
 
 ```kotlin
-implementation("com.integrallis:models-rag:0.2.3")
+implementation("com.integrallis:models-rag:0.2.4")
 ```
 
 ```java

models-spring-boot-starter/README.md

@@ -4,7 +4,7 @@ Auto-configures a Models runtime as a Spring AI `ChatModel`.
 
 ```kotlin
 dependencies {
-    implementation("com.integrallis:models-spring-boot-starter:0.2.3")
+    implementation("com.integrallis:models-spring-boot-starter:0.2.4")
 }
 ```
 

notebooks/.env.example

@@ -1,4 +1,4 @@
 MODELS_NOTEBOOK_MODE=source
-MODELS_VERSION=0.2.3
+MODELS_VERSION=0.2.4
 MODELS_NOTEBOOK_REPOSITORY=
 MODELS_NOTEBOOK_PREPARE=true

notebooks/README.md

@@ -27,7 +27,7 @@ Prepare either classpath without Docker:
 ./gradlew prepareNotebookClasspath -PnotebookMode=source
 ./gradlew prepareNotebookClasspath \
   -PnotebookMode=release \
-  -PnotebookVersion=0.2.3
+  -PnotebookVersion=0.2.4
 ```
 
 Test a staged release before it reaches Maven Central:
@@ -36,7 +36,7 @@ Test a staged release before it reaches Maven Central:
 ./gradlew verifyStagedPublications
 ./gradlew prepareNotebookClasspath \
   -PnotebookMode=release \
-  -PnotebookVersion=0.2.3 \
+  -PnotebookVersion=0.2.4 \
   -PnotebookRepository=build/staging-deploy
 ```
 
@@ -55,7 +55,7 @@ and does not read the host Docker credential store. To use published artifacts:
 
 ```bash
 MODELS_NOTEBOOK_MODE=release \
-MODELS_VERSION=0.2.3 \
+MODELS_VERSION=0.2.4 \
 ./docker-compose.sh up --build
 ```
 
@@ -75,7 +75,7 @@ Release mode executes the same notebooks from the selected artifacts:
 
 ```bash
 MODELS_NOTEBOOK_MODE=release \
-MODELS_VERSION=0.2.3 \
+MODELS_VERSION=0.2.4 \
 ./docker-compose.sh run --rm --no-deps jupyter \
   bash /home/jovyan/work/models/notebooks/scripts/test-notebooks.sh
 ```

notebooks/docker-compose.yml

@@ -12,7 +12,7 @@ services:
       - models-notebook-gradle-cache:/home/jovyan/.gradle
     environment:
       - MODELS_NOTEBOOK_MODE=${MODELS_NOTEBOOK_MODE:-source}
-      - MODELS_VERSION=${MODELS_VERSION:-0.2.3}
+      - MODELS_VERSION=${MODELS_VERSION:-0.2.4}
       - MODELS_NOTEBOOK_REPOSITORY=${MODELS_NOTEBOOK_REPOSITORY:-}
       - MODELS_NOTEBOOK_PREPARE=${MODELS_NOTEBOOK_PREPARE:-true}
       - JUPYTER_ENABLE_LAB=yes

notebooks/jupyter/prepare-classpath.sh

@@ -6,7 +6,7 @@ fi
 
 repository=/home/jovyan/work/models
 mode=${MODELS_NOTEBOOK_MODE:-source}
-version=${MODELS_VERSION:-0.2.3}
+version=${MODELS_VERSION:-0.2.4}
 repository_url=${MODELS_NOTEBOOK_REPOSITORY:-}
 
 case "$mode" in