Phase 7.3b WP4: Core multi-axis RoPE and decoder integration seam

.github/workflows/ci.yml

@@ -161,7 +161,8 @@ jobs:
                             'GroupedQueryAttentionTest', 'MultiHeadAttentionTest',
                             'BatchAttentionMaskTest', 'BatchRopeProbeTest', 'StridedViewSafetyTest',
                             'Phase71OracleTest', 'Phase71DefaultModeTest', 'Phase71ContractsTest', 'EncoderAttentionTest',
-                            'MLXMemoryLeakTest'],
+                            'MLXMemoryLeakTest',
+                            'MultiAxisRopeTest', 'DecoderRotationSeamTest', 'SingleRowDecodeEquivalenceTest'],
               'jmlx-models': ['DecoderRefactorGoldenTest', 'MistralModelTest', 'GemmaModelTest',
                               'Phi3ModelTest', 'MixtralModelTest', 'Llama31ModelTest', 'Qwen3ModelTest',
                               'QuantizedCheckpointProbeTest', 'QuantizedDecoderTest', 'DecoderCacheLifecycleTest',
@@ -172,7 +173,8 @@ jobs:
                               'BatchSchedulerLifecycleTest', 'BatchSchedulerMemoryTest',
                               'BatchStepEquivalenceTest', 'BertGoldenTest', 'T5GoldenTest', 'Seq2SeqMemoryTest', 'Seq2SeqGenerationTest', 'EncoderDefaultPrecisionTest',
                               'ModalityRejectionTest', 'Idefics3VisionTowerSmokeTest',
-                              'SmolVlmGoldenTest', 'SmolVlmGenerationTest', 'SmolVlmDefaultPrecisionTest', 'VisionLanguageMemoryTest'],
+                              'SmolVlmGoldenTest', 'SmolVlmGenerationTest', 'SmolVlmDefaultPrecisionTest', 'VisionLanguageMemoryTest',
+                              'DecoderSeamTest', 'VisualFeatureInjectionTest'],
           }
           # Golden references run in their own full-float32 JVM. Batch equivalence must
           # execute in both precision modes. Probe (b) remains evidence only.
@@ -181,9 +183,12 @@ jobs:
                             'DecoderSlidingWindowTest', 'BertGoldenTest', 'T5GoldenTest', 'SmolVlmGoldenTest']
           model_results = {name: ['float32GoldenTest'] for name in float32_suites}
           model_results['BatchStepEquivalenceTest'] = ['test', 'float32GoldenTest']
+          model_results['DecoderSeamTest'] = ['test', 'float32GoldenTest']
+          model_results['VisualFeatureInjectionTest'] = ['test', 'float32GoldenTest']
           model_results['WorkerStreamProbeTest'] = ['threadProbeTest']
           core_results = {name: ['test', 'float32GoldenTest'] for name in
-                          ['DecoderAttentionTest', 'MultiHeadAttentionTest', 'MoeMlpTest', 'SwitchGluTest']}
+                          ['DecoderAttentionTest', 'MultiHeadAttentionTest', 'MoeMlpTest', 'SwitchGluTest',
+                           'MultiAxisRopeTest', 'DecoderRotationSeamTest']}
           core_results['Phase71OracleTest'] = ['float32GoldenTest']
           core_results['EncoderAttentionTest'] = ['float32GoldenTest']
           results_dirs = {'jmlx-models': model_results, 'jmlx-core': core_results}

AGENTS.md

@@ -336,7 +336,8 @@ jmlx-core        se.alipsa.jmlx.nn                 Module, Linear, QuantizedLine
                                                     MultiHeadAttention, GroupedQueryAttention,
                                                     DecoderAttention, DecoderBlock, AttentionMask,
                                                     RopeSpec, KVCache, KVCachePolicy, GatedMlp,
-                                                    SwiGLU, MoeMlp, SwitchGlu, ModuleGrad
+                                                    SwiGLU, MoeMlp, SwitchGlu, ModuleGrad,
+                                                    MultiAxisRope, MultiAxisRotation
                  se.alipsa.jmlx.core                MLX, MLXOps, MLXShape, MLXFast, MLXQuant,
                                                     MLXRandom, MLXGrad, MLXIO, MLXMemory, MLXArray,
                                                     DType, MLXException
@@ -663,3 +664,21 @@ against the rotated fixtures (findings §4.8). No native pins or generated bindi
 no new Gradle tasks were added — the existing `verifyImageOracleFixtureDigests` and
 `verifyHfReferenceGoldens` guards cover the new fixtures, and `ImageOracleQwen3VlTest` (part of
 `:jmlx-vision:test`) is the Java-side byte verification.
+
+Phase 7.3b WP4 (2026-10-20) landed the frozen `MultiAxisRope`/`MultiAxisRotation`
+core components in `jmlx-core`'s `se.alipsa.jmlx.nn` (the Qwen3-VL interleaved
+three-axis RoPE kernel with its `cachePosition + delta` decode coordinates), the
+additive request-supplied rotary-state / post-block injection seam on the
+`DecoderModel` path (existing overloads untouched; production decoders use no-op
+hooks; `RopeSpec`'s sealed hierarchy and `MLXFast` are byte-identical — verified by
+empty diff), and the package-private `VisualFeatureInjection` compositions in
+`jmlx-models` (row-gather image-row replacement plus masked DeepStack residual add,
+composed from existing core facades over the Idefics3 gather-index algorithm — no new
+primitive, no FFI from model code). Old-path bit identity in both TF32 modes is
+verified against the pre-change exact-bits baseline (`exactBitsVerifyTf32On`/`Off`,
+14 variants × 9 captures), and the pinned-vendor probe pair
+(`tools/hf-reference/qwen3vl_rope_interleaved_check.{py,jsh}` — manually run in their
+locked environments, never Gradle-invoked) cross-checks the component. The CI
+required-suite list addition landed in PR #57 (review round 2), and `Qwen3VlModel`,
+its loader dispatch and the DeepStack consumption of `VisualFeatureInjection` are
+WP5/WP6. No native pins, generated bindings or build files changed.

jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderAttention.java

@@ -182,7 +182,7 @@ public MLXArray forward(MLXArray x, KVCache cache, MLXArray attentionMask) {
   @Override
   public MLXArray forward(
       MLXArray x, KVCache cache, MLXArray attentionMask, MLXArray stepFrequencies) {
-    return forwardInternal(x, cache, attentionMask, stepFrequencies, null);
+    return forwardInternal(x, cache, attentionMask, stepFrequencies, null, null);
   }
 
   @Override
@@ -192,15 +192,45 @@ public MLXArray forward(
       MLXArray attentionMask,
       MLXArray stepFrequencies,
       int[] validLengths) {
-    return forwardInternal(x, cache, attentionMask, stepFrequencies, validLengths);
+    return forwardInternal(x, cache, attentionMask, stepFrequencies, validLengths, null);
+  }
+
+  /**
+   * Batch-aware forward that rotates the projected query and key heads from a request-supplied
+   * {@link MultiAxisRotation} instead of the {@link RopeSpec} fast kernels, skipping the
+   * frequency-array derivation entirely. The rotation lands where {@code rope.apply} runs today --
+   * after projection and the optional QK normalization, before the cache append -- so the retained
+   * keys are already rotated. {@code validLengths} is required: this path always runs the
+   * batch-aware validation and per-row cache bookkeeping, which reduce to the single-row arithmetic
+   * for a uniform B = 1 row (pinned by {@code SingleRowDecodeEquivalenceTest}). All validation is
+   * {@link IllegalArgumentException} before any array is written, including the rotation's own
+   * width against this attention's configured {@code rotaryDims} ({@link
+   * MultiAxisRotation#validateAgainst(int, int, int, int)}): a rotation prepared from a different
+   * mrope-section sum is rejected up front rather than silently rotating a different head prefix.
+   */
+  public MLXArray forward(
+      MLXArray x,
+      KVCache cache,
+      MLXArray attentionMask,
+      int[] validLengths,
+      MultiAxisRotation rotation) {
+    if (rotation == null) {
+      throw new IllegalArgumentException("rotation must not be null");
+    }
+    if (validLengths == null) {
+      throw new IllegalArgumentException(
+          "the rotation path requires batch valid lengths, one per row");
+    }
+    return forwardInternal(x, cache, attentionMask, null, validLengths, rotation);
   }
 
   private MLXArray forwardInternal(
       MLXArray x,
       KVCache cache,
       MLXArray attentionMask,
       MLXArray stepFrequencies,
-      int[] validLengths) {
+      int[] validLengths,
+      MultiAxisRotation rotation) {
     Objects.requireNonNull(x, "x");
     int[] shape = x.shape();
     if (shape.length != 3 || shape[0] <= 0 || shape[1] <= 0) {
@@ -265,18 +295,27 @@ private MLXArray forwardInternal(
       throw new IllegalArgumentException(
           "attentionMask must be BOOL " + Arrays.toString(expectedMask));
     }
-    MLXArray freqs = staticFreqs != null ? staticFreqs : stepFrequencies;
-    if (freqs == null) {
-      // Only dynamic specs reach here; computed once so q and k share one upload.
-      int ending = offset + sequence;
-      if (validLengths != null) {
-        ending = 0;
-        for (int row = 0; row < batch; row++) {
-          ending =
-              Math.max(ending, (cache == null ? 0 : cache.nextPosition(row)) + validLengths[row]);
+    MLXArray freqs;
+    if (rotation != null) {
+      // The rotation path rotates from its own cos/sin tables; no frequency array is derived.
+      // Validate the rotation's shape, and its width against this attention's configured rotary
+      // width, before any projection builds a graph node.
+      rotation.validateAgainst(batch, sequence, headDim, rotaryDims);
+      freqs = null;
+    } else {
+      freqs = staticFreqs != null ? staticFreqs : stepFrequencies;
+      if (freqs == null) {
+        // Only dynamic specs reach here; computed once so q and k share one upload.
+        int ending = offset + sequence;
+        if (validLengths != null) {
+          ending = 0;
+          for (int row = 0; row < batch; row++) {
+            ending =
+                Math.max(ending, (cache == null ? 0 : cache.nextPosition(row)) + validLengths[row]);
+          }
         }
+        freqs = rope.stepFrequencies(x.scope(), rotaryDims, ending);
       }
-      freqs = rope.stepFrequencies(x.scope(), rotaryDims, ending);
     }
     MLXArray q = AttentionHeads.toHeads(queryProj.forward(x), batch, sequence, numHeads, headDim);
     MLXArray k =
@@ -289,7 +328,10 @@ private MLXArray forwardInternal(
     if (keyNorm != null) {
       k = keyNorm.forward(k);
     }
-    if (validLengths == null) {
+    if (rotation != null) {
+      q = rotation.rotate(q);
+      k = rotation.rotate(k);
+    } else if (validLengths == null) {
       q = rope.apply(q, rotaryDims, offset, freqs);
       k = rope.apply(k, rotaryDims, offset, freqs);
     } else {

jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderBlock.java

@@ -86,4 +86,34 @@ public MLXArray forward(
                 inputNorm.forward(x), cache, attentionMask, stepFrequencies, validLengths));
     return MLXOps.add(afterAttention, mlp.forward(postAttentionNorm.forward(afterAttention)));
   }
+
+  /**
+   * Applies this block to a left-padded batch, rotating the attention heads from a request-supplied
+   * {@link MultiAxisRotation} instead of the block's rotary spec. Same residual structure as the
+   * batch overload: input plus the rotated attention output, plus the MLP of the post-attention
+   * normalization. The attention child must be a {@link DecoderAttention}, the only cached
+   * attention with the rotation overload; blocks built from another {@link CachedAttention}
+   * subclass are rejected before any array is written.
+   */
+  public MLXArray forward(
+      MLXArray x,
+      KVCache cache,
+      MLXArray attentionMask,
+      int[] validLengths,
+      MultiAxisRotation rotation) {
+    Objects.requireNonNull(x, "DecoderBlock.forward: x must not be null");
+    if (rotation == null) {
+      throw new IllegalArgumentException("DecoderBlock.forward: rotation must not be null");
+    }
+    if (!(attention instanceof DecoderAttention decoderAttention)) {
+      throw new IllegalArgumentException(
+          "DecoderBlock.forward with a MultiAxisRotation requires a DecoderAttention child");
+    }
+    MLXArray afterAttention =
+        MLXOps.add(
+            x,
+            decoderAttention.forward(
+                inputNorm.forward(x), cache, attentionMask, validLengths, rotation));
+    return MLXOps.add(afterAttention, mlp.forward(postAttentionNorm.forward(afterAttention)));
+  }
 }