diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1a153771..2005b110 100644 --- a/.github/workflows/ci.yml +++ b/.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} diff --git a/AGENTS.md b/AGENTS.md index b967f061..50f265de 100644 --- a/AGENTS.md +++ b/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. diff --git a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderAttention.java b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderAttention.java index ee711348..2b1ac8e8 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderAttention.java +++ b/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,7 +192,36 @@ 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( @@ -200,7 +229,8 @@ private MLXArray forwardInternal( 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 { diff --git a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderBlock.java b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderBlock.java index 128952ee..f498925b 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderBlock.java +++ b/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))); + } } diff --git a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java new file mode 100644 index 00000000..a25575a6 --- /dev/null +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java @@ -0,0 +1,224 @@ +package se.alipsa.jmlx.nn; + +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.core.MLXOps; +import se.alipsa.jmlx.core.MLXShape; +import se.alipsa.jmlx.memory.MLXScope; + +/** + * Immutable multi-axis (MRoPE) rotary-position specification for Hugging Face's interleaved layout + * ({@code mrope_interleaved: true}, {@code rope_type: default}), frozen for the Qwen3-VL text + * decoder in req/plans/phase7-3b-reference-findings.md §6.9. This is a dedicated component, not a + * new {@link RopeSpec} variant: the sealed {@code RopeSpec} hierarchy gains no permitted variant, + * and the ordinary decoders keep the unchanged {@code mlx_fast_rope} fast kernels via {@link + * RopeSpec#apply}. The fast kernels' offset-plus-arange contract cannot express unequal multimodal + * coordinates, which is why the Qwen3-VL path composes this component from existing core facades + * instead: no new FFI call. + * + *

The partial-rotary tail passthrough is a jmlx extension, not reference behaviour. The rotated + * width and the inverse-frequency denominator are both derived from {@code 2 * sum(mropeSection)} + * rather than the head dimension, so the frequency slots beyond that width pass through unrotated; + * in the pinned transformers Qwen3-VL reference, by contrast, both the denominator and the rotated + * width come from {@code head_dim}, and the frequency slots outside the H/W slices stay T-owned and + * are still rotated. This class agrees with the reference only when {@code 2 * sum(mropeSection) == + * head_dim} (the pinned config [24,20,20] with head_dim 128 satisfies this), which is why the plan + * §7 Qwen3-VL loader rejects a config violating that equality. + * + *

Deliberately not a {@link Module}: no parameters, no children, no owned scope. Every MLX array + * {@link #prepare} builds allocates into the caller's scope, and every array {@link + * MultiAxisRotation#rotate} builds allocates per the usual {@code scopeOf} rule (the innermost + * related scope of its operands). Positions and frequencies are recomputed on every {@code prepare} + * call: no position-derived state is cached. + * + *

{@code positions} is {@code int[3][batch * sequence]}: three rows of length {@code B*T} + * (row-major flattened {@code [B,T]}) where row 0 is the T (time) axis, row 1 the H (height) axis + * and row 2 the W (width) axis -- the reference's {@code [3, B*T]} {@code position_ids} layout as + * one validated object, in the {@code mrope_section} order. Coordinates are non-negative (decode + * steps pass {@code cachePosition + delta}, where the delta can be negative but the coordinate it + * produces must not be). + */ +public final class MultiAxisRope { + + private final float theta; + private final int[] mropeSection; + private final int rotaryDims; + private final float[] invFreq; + private final byte[] owner; + + /** + * Validates the base and the sections and derives the per-frequency inverse and the static T/H/W + * owner table. All validation is {@link IllegalArgumentException} and runs before any MLX array + * is written: {@code theta} must be finite and positive (the {@code RopeSpec} convention); {@code + * mropeSection} must be a non-null 3-entry array, every element non-negative, with a sum of at + * least 1 (so {@code rotaryDims = 2 * sum} is even by construction) and small enough to fit an + * {@code int} rotary dimension. + */ + public MultiAxisRope(float theta, int[] mropeSection) { + if (!Float.isFinite(theta) || theta <= 0) { + throw new IllegalArgumentException("theta must be positive and finite"); + } + if (mropeSection == null) { + throw new IllegalArgumentException("mropeSection must not be null"); + } + if (mropeSection.length != 3) { + throw new IllegalArgumentException("mropeSection must have exactly 3 entries (T, H, W)"); + } + long sum = 0; + for (int i = 0; i < 3; i++) { + if (mropeSection[i] < 0) { + throw new IllegalArgumentException("mropeSection[" + i + "] must be non-negative"); + } + sum += mropeSection[i]; + } + if (sum < 1) { + throw new IllegalArgumentException("mropeSection sum must be at least 1"); + } + if (sum > Integer.MAX_VALUE / 2) { + throw new IllegalArgumentException("mropeSection sum must fit an int rotary dimension"); + } + int half = (int) sum; + this.theta = theta; + this.mropeSection = mropeSection.clone(); + this.rotaryDims = 2 * half; + this.invFreq = deriveInvFrequencies(theta, half); + this.owner = new byte[half]; + for (int f = 0; f < half; f++) { + // The pinned reference (transformers 4.57.6 apply_interleaved_mrope) starts from the T row + // and overwrites in place: for dim 1 (H) idx = slice(1, mrope_section[1] * 3, 3) -- H owns + // 1 + 3j for j < mrope_section[1]; for dim 2 (W) idx = slice(2, mrope_section[2] * 3, 3) -- + // W owns 2 + 3j for j < mrope_section[2]; everything else stays T. The table records that + // selection as one static owner (0 = T, 1 = H, 2 = W) per frequency slot; the T slots stay + // at the byte zero, and the H and W conditions are disjoint, so every slot has exactly one + // owner. + if (f % 3 == 1 && f < 1L + 3L * mropeSection[1]) { + owner[f] = 1; + } else if (f % 3 == 2 && f < 2L + 3L * mropeSection[2]) { + owner[f] = 2; + } + } + } + + /** Returns the frequency base {@code theta}. */ + public float theta() { + return theta; + } + + /** Returns a defensive copy of the mrope sections (T, H, W). */ + public int[] mropeSection() { + return mropeSection.clone(); + } + + /** Returns {@code 2 * sum(mropeSection)}, the width of the rotated head prefix. */ + public int rotaryDims() { + return rotaryDims; + } + + /** + * Builds the per-step {@link MultiAxisRotation} for {@code batch} rows of {@code sequence} + * tokens. All validation is {@link IllegalArgumentException} and runs before any MLX array is + * written: {@code batch} and {@code sequence} are at least 1, {@code positions} has exactly 3 + * non-null rows, each row has length {@code batch * sequence} (computed in {@code long} so the + * product cannot overflow), no coordinate is negative and the position-matrix size {@code batch * + * sequence * rotaryDims/2} fits an {@code int} array (checked in {@code long} before the + * allocation, so an overflowing size is an IAE rather than a {@code NegativeArraySizeException}). + * The frozen compute order then runs: the static owner table selects, per (token, frequency) on + * the CPU before any array is written, the owned row's position, giving one float32 position + * matrix of shape {@code [B*T, rotaryDims/2]} (the exact int-to-float casts) that is + * broadcast-multiplied once against the {@code [1, rotaryDims/2]} {@code invFreq} array in {@code + * scope} -- a single float32 multiply per element with the same two operands the frozen mask + * expression {@code freqs_t = freqsT * wT + freqsH * wH + freqsW * wW} copies, which is a + * bit-identical copy for non-negative finite frequencies (findings §6.9, step 2); {@code emb = + * cat((freqs_t, freqs_t), -1)}; {@code MLXOps.cos}/{@code sin} in float32, with the attention + * scaling of 1.0 (the pinned {@code rope_type: default}) omitted as the identity on finite + * values; finally the two arrays are reshaped to {@code [batch, 1, sequence, rotaryDims]}, + * broadcastable over the heads axis. + */ + public MultiAxisRotation prepare(MLXScope scope, int batch, int sequence, int[][] positions) { + if (batch < 1) { + throw new IllegalArgumentException("batch must be at least 1"); + } + if (sequence < 1) { + throw new IllegalArgumentException("sequence must be at least 1"); + } + if (positions == null) { + throw new IllegalArgumentException("positions must not be null"); + } + if (positions.length != 3) { + throw new IllegalArgumentException("positions must have exactly 3 rows (T, H, W)"); + } + long tokens = (long) batch * sequence; + for (int axis = 0; axis < 3; axis++) { + if (positions[axis] == null) { + throw new IllegalArgumentException("positions[" + axis + "] must not be null"); + } + if ((long) positions[axis].length != tokens) { + throw new IllegalArgumentException( + "positions[" + axis + "] must have length batch * sequence = " + tokens); + } + } + for (int axis = 0; axis < 3; axis++) { + for (int coordinate : positions[axis]) { + if (coordinate < 0) { + throw new IllegalArgumentException("positions must be non-negative"); + } + } + } + int flat = (int) tokens; + int half = rotaryDims / 2; + if ((long) flat * half > Integer.MAX_VALUE) { + throw new IllegalArgumentException("batch * sequence * rotaryDims/2 must fit an int array"); + } + MLXArray invFreqArray = MLX.array(scope, invFreq, new int[] {1, half}); + MLXArray owned = + MLX.array(scope, ownedPositions(positions, flat, half), new int[] {flat, half}); + MLXArray freqs = MLXOps.multiply(owned, invFreqArray); + MLXArray emb = MLXShape.concatenate(new MLXArray[] {freqs, freqs}, 1); + MLXArray cos = MLXShape.reshape(MLXOps.cos(emb), new int[] {batch, 1, sequence, rotaryDims}); + MLXArray sin = MLXShape.reshape(MLXOps.sin(emb), new int[] {batch, 1, sequence, rotaryDims}); + return new MultiAxisRotation(cos, sin, rotaryDims); + } + + /** + * The per-token owned-position matrix in row-major {@code [flat, half]} order: element {@code (i, + * f)} is the exact float32 cast of {@code positions[owner[f]][i]}, selected on the CPU with the + * static owner table before any array is written. This is the bit-identical copy the frozen mask + * expression {@code freqsT * wT + freqsH * wH + freqsW * wW} makes for non-negative finite + * frequencies (multiplying by the 0/1 masks and adding in float32 is exact there), so multiplying + * this matrix once by {@code invFreq} performs the same per-element float32 multiply of the same + * two operands as the frozen per-axis broadcast-multiplies (findings §6.9, step 2). + */ + private float[] ownedPositions(int[][] positions, int flat, int half) { + float[] out = new float[flat * half]; + for (int i = 0; i < flat; i++) { + for (int f = 0; f < half; f++) { + out[i * half + f] = (float) positions[owner[f]][i]; + } + } + return out; + } + + /** + * Derives {@code invFreq[i] = 1 / theta^(2i/rotaryDims)} once, in the pinned reference's float32 + * arithmetic order: the exponent is the float32 division of the exact float32 numerator {@code + * 2i} by the float32 {@code rotaryDims}, the power is a float32 cast of the double power, and the + * reciprocal is a float32 division. The pinned transformers 4.57.6 {@code + * _compute_default_rope_parameters} builds exactly {@code 1.0 / (base ** (arange(0, head_dim, + * 2).float() / head_dim))} in float32, and this order reproduces it bit for bit on 63 of the 64 + * pinned 2B frequencies (theta 5,000,000, head_dim 128; the remaining one within 1e-9). The + * RopeSpec-style alternative -- double power, then cast -- deviates from the pinned reference by + * about 1e-5 in the rotate-half output already at position 89, inside the pinned 2B prefill + * envelope, and cannot meet the frozen 1e-5 cross-check threshold, so the float32 order stands; + * the pinned-venv cross-check (tools/hf-reference/qwen3vl_rope_interleaved_check) is the arbiter + * for this subtlety (findings §6.9). + */ + private static float[] deriveInvFrequencies(float theta, int half) { + int rotaryDims = 2 * half; + float[] invFreq = new float[half]; + for (int i = 0; i < half; i++) { + float exponent = (float) (2 * i) / (float) rotaryDims; + invFreq[i] = 1.0f / (float) Math.pow(theta, exponent); + } + return invFreq; + } +} diff --git a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java new file mode 100644 index 00000000..d4c5eedc --- /dev/null +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java @@ -0,0 +1,163 @@ +package se.alipsa.jmlx.nn; + +import java.util.EnumMap; +import se.alipsa.jmlx.core.DType; +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.core.MLXOps; +import se.alipsa.jmlx.core.MLXShape; + +/** + * Per-step rotation state one {@link MultiAxisRope#prepare} produces and every layer of that step + * shares: the float32 {@code cos}/{@code sin} tables of shape {@code [batch, 1, sequence, + * rotaryDims]} (broadcastable over the heads axis) plus the width of the rotated head prefix. + * Threading exactly like the existing per-step {@code stepFrequencies}: computed once per step, + * handed to each attention, never cached between steps. + * + *

{@link #rotate} applies the frozen composition of req/plans/phase7-3b-reference-findings.md + * §6.9 in the input array's dtype: the cos/sin tables are cast to {@code x.dtype()} (a no-op for + * the FLOAT32 model) and only the {@code rotaryDims} prefix of each head is rotated with the NeoX + * half-split {@code rotate_half = cat((−x[…, n:], x[…, :n]), −1)} -- the same layout the fast + * kernels use with {@code traditional=false} -- while the tail {@code [rotaryDims, headDim)} passes + * through bit-identical (partial rotary). + */ +public final class MultiAxisRotation { + + private final MLXArray cos; + private final MLXArray sin; + private final int rotaryDims; + private final EnumMap castCos; + private final EnumMap castSin; + + MultiAxisRotation(MLXArray cos, MLXArray sin, int rotaryDims) { + this.cos = cos; + this.sin = sin; + this.rotaryDims = rotaryDims; + this.castCos = new EnumMap<>(DType.class); + this.castSin = new EnumMap<>(DType.class); + } + + /** Returns the float32 cosine table, {@code [batch, 1, sequence, rotaryDims]}. */ + public MLXArray cos() { + return cos; + } + + /** Returns the float32 sine table, {@code [batch, 1, sequence, rotaryDims]}. */ + public MLXArray sin() { + return sin; + } + + /** Returns the width of the rotated head prefix. */ + public int rotaryDims() { + return rotaryDims; + } + + /** + * The rotation-shape preflight the decoder seam runs before any array is written: the batch and + * sequence axes of the heads to be rotated must match the prepared rotation, and the rotated + * prefix must fit within the head width. All validation is {@link IllegalArgumentException} and + * nothing is allocated. {@link #rotate} delegates its shape checks here, and the {@code + * DecoderAttention}/{@code DecoderModel} rotation entries call {@link #validateAgainst(int, int, + * int, int)} -- which additionally requires the rotation's own width to equal the decoder's + * configured rotary width -- before any projection, mask or embedding builds a graph node, so a + * mismatched rotation is rejected before any array is written. + */ + public void validate(int batch, int sequence, int headDim) { + int[] cosShape = cos.shape(); + if (batch != cosShape[0] || sequence != cosShape[2]) { + throw new IllegalArgumentException( + "rotation batch and sequence must match the prepared rotation (" + + batch + + ", " + + sequence + + " against " + + cosShape[0] + + ", " + + cosShape[2] + + ")"); + } + if (rotaryDims > headDim) { + throw new IllegalArgumentException("rotaryDims exceeds head dimension"); + } + } + + /** + * The width-aware rotation preflight the {@code DecoderAttention}/{@code DecoderModel} rotation + * entries run at their earliest point: {@link #validate(int, int, int)} plus the requirement that + * the rotation's own width equals the decoder's configured rotary width (the partial-rotary width + * from the rope spec or model descriptor) -- a rotation prepared from a different mrope-section + * sum would otherwise pass the prefix-fit check and silently rotate a different head prefix. All + * validation is {@link IllegalArgumentException} and nothing is allocated. + */ + public void validateAgainst(int batch, int sequence, int headDim, int configuredRotaryDims) { + validate(batch, sequence, headDim); + if (rotaryDims != configuredRotaryDims) { + throw new IllegalArgumentException( + "rotation rotaryDims " + + rotaryDims + + " does not match the configured rotary width " + + configuredRotaryDims); + } + } + + /** + * Rotates a {@code [batch, heads, sequence, headDim]} tensor. All validation is {@link + * IllegalArgumentException} and runs before any array is written: {@code x} is non-null with rank + * 4, its batch and sequence axes match the prepared rotation, and {@code rotaryDims} does not + * exceed the head dimension. Then the cos/sin tables are cast to {@code x.dtype()} (frozen step + * 5; a no-op for the FLOAT32 model, and otherwise cached per dtype in this per-step rotation so + * the step's repeated q/k rotates share one cast pair per dtype) and, in the input dtype, {@code + * out = x * cos + rotate_half(x) * sin} on the {@code rotaryDims} prefix with the untouched tail + * concatenated back (frozen step 6). Every built array allocates per the usual {@code scopeOf} + * rule (the innermost related scope of its operands). + */ + public MLXArray rotate(MLXArray x) { + if (x == null) { + throw new IllegalArgumentException("x must not be null"); + } + int[] shape = x.shape(); + if (shape.length != 4) { + throw new IllegalArgumentException( + "rotate requires [batch, heads, sequence, headDim], got shape " + + java.util.Arrays.toString(shape)); + } + validate(shape[0], shape[2], shape[3]); + DType dtype = x.dtype(); + MLXArray cosInput = cos.dtype() == dtype ? cos : castTable(cos, dtype, castCos); + MLXArray sinInput = sin.dtype() == dtype ? sin : castTable(sin, dtype, castSin); + int[] prefixStop = shape.clone(); + prefixStop[3] = rotaryDims; + MLXArray prefix = rotaryDims == shape[3] ? x : MLXShape.slice(x, new int[4], prefixStop); + int half = rotaryDims / 2; + int[] halfStop = shape.clone(); + halfStop[3] = half; + MLXArray front = MLXShape.slice(prefix, new int[4], halfStop); + MLXArray back = MLXShape.slice(prefix, new int[] {0, 0, 0, half}, prefixStop); + MLXArray rotated = MLXShape.concatenate(new MLXArray[] {MLXOps.negative(back), front}, 3); + MLXArray out = + MLXOps.add(MLXOps.multiply(prefix, cosInput), MLXOps.multiply(rotated, sinInput)); + if (rotaryDims == shape[3]) { + return out; + } + int[] tailStart = new int[4]; + tailStart[3] = rotaryDims; + return MLXShape.concatenate( + new MLXArray[] {out, MLXShape.slice(x, tailStart, shape.clone())}, 3); + } + + /** + * The {@code dtype} view of one of the float32 tables, cached in this per-step rotation: the + * first rotate in a non-table dtype builds the cast (allocated into the table's own scope, + * exactly where an uncached per-call cast would land) and every later rotate of the same dtype + * reuses it, so a step's 2 × layers q/k rotates pay for one cast pair per dtype instead of one + * per call. The rotation lives for exactly one step, so the cache never outlives its tables. + */ + private static MLXArray castTable(MLXArray table, DType dtype, EnumMap cache) { + MLXArray cast = cache.get(dtype); + if (cast == null) { + cast = MLX.astype(table, dtype); + cache.put(dtype, cast); + } + return cast; + } +} diff --git a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java new file mode 100644 index 00000000..f8405a7b --- /dev/null +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java @@ -0,0 +1,518 @@ +package se.alipsa.jmlx.nn; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.core.MLXFast; +import se.alipsa.jmlx.core.MLXOps; +import se.alipsa.jmlx.core.MLXShape; +import se.alipsa.jmlx.ffi.EnabledIfNativeAvailable; +import se.alipsa.jmlx.memory.MLXScope; + +/** + * Seam tests for the frozen WP4 decoder rotation integration (req/plans/phase7-3b-reference- + * findings.md §6.9): the batch-aware {@link DecoderAttention} forward that rotates the projected + * query and key heads from a request-supplied {@link MultiAxisRotation}, and the {@link + * DecoderBlock} passthrough. Every float32 case compares the seam against an operation-identical + * manual composition built from the same retained projection and normalization instances -- the + * same operations, the same order, the same shapes, the same scopes -- so those comparisons run at + * zero tolerance, and the cases are tagged {@code full-float32} so they select into {@code + * :jmlx-core:float32GoldenTest} with {@code MLX_ENABLE_TF32=0}. The single 1e-5 exception is the + * decode-versus-fresh-prefill comparison in {@link #rotationPrefillAndDecodeBitMatchFreshPrefill}, + * where the two runs use different kernel shapes. The rejection cases run in the ordinary {@code + * :jmlx-core:test} task. + */ +@EnabledIfNativeAvailable +class DecoderRotationSeamTest { + + private static final int HIDDEN = 8; + private static final int HEADS = 1; + private static final int KEY_VALUE_HEADS = 1; + private static final int HEAD_DIM = 8; + private static final int[] SECTIONS = {2, 1, 1}; // rotaryDims 8 == HEAD_DIM + private static final float SCALE = (float) (1.0 / Math.sqrt(HEAD_DIM)); + + // --------------------------------------------------------------------- + // float32-comparison cases (full-float32 tag, MLX_ENABLE_TF32=0) + // --------------------------------------------------------------------- + + /** + * A five-token prefill through the rotation seam must be bit-identical to the operation-identical + * reference composition on a second cache, the retained keys must equal the rotated projected + * keys (pinning rotate-before-append), and the one-token decode at position five must match the + * last row of a fresh six-token prefill at the cross-shape tolerance. The tolerance is needed + * because the decode (one query over a retained six-key cache) and the fresh prefill (six queries + * over six keys) drive different native attention and matmul kernel shapes: the last row's + * arithmetic is the same, but MLX does not guarantee bit-equal results across shapes. + */ + @Test + @Tag("full-float32") + void rotationPrefillAndDecodeBitMatchFreshPrefill() { + try (MLXScope scope = new MLXScope()) { + AttentionFixture fixture = new AttentionFixture(scope, null, null); + MultiAxisRope rope = new MultiAxisRope(10_000f, SECTIONS); + try (MLXScope run = scope.newChild()) { + KVCache cache = new KVCache(run); + KVCache referenceCache = new KVCache(run); + float[] prefill = values(1 * 5 * HIDDEN, 1); + MLXArray x = MLX.array(run, prefill, new int[] {1, 5, HIDDEN}); + MultiAxisRotation rotation = + rope.prepare( + run, 1, 5, new int[][] {{0, 1, 2, 3, 4}, {0, 2, 4, 2, 4}, {1, 1, 3, 3, 5}}); + MLXArray mask = + AttentionMask.batched(run, new int[] {0}, new int[] {0}, new int[] {5}, 5, 5, 0); + MLXArray actual = fixture.attention.forward(x, cache, mask, new int[] {5}, rotation); + MLXArray reference = + referenceAttention(fixture, x, referenceCache, mask, new int[] {5}, rotation); + assertArrayEquals(reference.toFloatArray(), actual.toFloatArray()); + // The retained keys are the rotated projected keys: rotation lands before the append. + MLXArray expectedKeys = + rotation.rotate( + AttentionHeads.toHeads( + fixture.keyProj.forward(x), 1, 5, KEY_VALUE_HEADS, HEAD_DIM)); + assertArrayEquals(expectedKeys.toFloatArray(), cache.keys().toFloatArray()); + + float[] decodeInput = values(1 * 1 * HIDDEN, 7); + MLXArray x1 = MLX.array(run, decodeInput, new int[] {1, 1, HIDDEN}); + MultiAxisRotation decodeRotation = rope.prepare(run, 1, 1, new int[][] {{5}, {6}, {4}}); + MLXArray decodeMask = + AttentionMask.batched(run, new int[] {5}, new int[] {0}, new int[] {1}, 1, 6, 0); + MLXArray decoded = + fixture.attention.forward(x1, cache, decodeMask, new int[] {1}, decodeRotation); + assertEquals(6, cache.nextPosition()); + MLXArray freshLast = freshSixTokenPrefillLastRow(fixture, rope, run, prefill, decodeInput); + assertArrayEquals(freshLast.toFloatArray(), decoded.toFloatArray(), 1e-5f); + } + } + } + + /** + * With non-unit QK normalization the seam must normalize before rotating: the reference applies + * the norms at the same position (after projection, before the rotation), so a different order + * would change both the retained keys and the output at zero tolerance. + */ + @Test + @Tag("full-float32") + void qkNormRunsBeforeRotation() { + try (MLXScope scope = new MLXScope()) { + float[] queryWeight = new float[HIDDEN]; + float[] keyWeight = new float[HIDDEN]; + for (int i = 0; i < HIDDEN; i++) { + queryWeight[i] = 2f; + keyWeight[i] = 0.5f; + } + AttentionFixture fixture = + new AttentionFixture( + scope, + new RMSNorm(scope, MLX.array(scope, queryWeight, new int[] {HIDDEN}), 1e-6f), + new RMSNorm(scope, MLX.array(scope, keyWeight, new int[] {HIDDEN}), 1e-6f)); + MultiAxisRope rope = new MultiAxisRope(10_000f, SECTIONS); + try (MLXScope run = scope.newChild()) { + KVCache cache = new KVCache(run); + KVCache referenceCache = new KVCache(run); + MLXArray x = MLX.array(run, values(1 * 3 * HIDDEN, 3), new int[] {1, 3, HIDDEN}); + MultiAxisRotation rotation = + rope.prepare(run, 1, 3, new int[][] {{0, 1, 2}, {1, 3, 5}, {2, 4, 6}}); + MLXArray mask = + AttentionMask.batched(run, new int[] {0}, new int[] {0}, new int[] {3}, 3, 3, 0); + MLXArray actual = fixture.attention.forward(x, cache, mask, new int[] {3}, rotation); + MLXArray reference = + referenceAttention(fixture, x, referenceCache, mask, new int[] {3}, rotation); + assertArrayEquals(reference.toFloatArray(), actual.toFloatArray()); + } + } + } + + /** + * A rotation prepared at all-zero coordinates is the bitwise identity: the angles are exactly + * zero, so the cos table is exactly 1 and the sin table exactly 0, and multiplying by 1 plus + * adding 0 is exact for every finite float. The test values never contain {@code -0.0} (the + * {@link #values} helper produces {@code 0/6.0f} for its zero, a positive zero), so the identity + * holds element for element, sign bits included. + */ + @Test + @Tag("full-float32") + void zeroPositionRotationIsBitwiseIdentity() { + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(10_000f, SECTIONS); + try (MLXScope run = scope.newChild()) { + float[] x = values(1 * 1 * 4 * HEAD_DIM, 0); + MultiAxisRotation rotation = + rope.prepare(run, 1, 4, new int[][] {zeros(4), zeros(4), zeros(4)}); + MLXArray array = MLX.array(run, x, new int[] {1, 1, 4, HEAD_DIM}); + assertArrayEquals(x, rotation.rotate(array).toFloatArray()); + } + } + } + + /** + * A ragged batch (valid lengths {2, 4} on a width-4 prefill, then {1, 1} on a width-1 decode) + * must bit-match the reference composition on a second cache and keep the per-row bookkeeping: + * the decode advances each row from its own position (row 0 from 2, row 1 from 4), and the + * rotation's per-row coordinates address the right valid suffix of every row. + */ + @Test + @Tag("full-float32") + void raggedBatchKeepsPerRowPositionsAndBitMatchesReference() { + try (MLXScope scope = new MLXScope()) { + AttentionFixture fixture = new AttentionFixture(scope, null, null); + MultiAxisRope rope = new MultiAxisRope(10_000f, SECTIONS); + try (MLXScope run = scope.newChild()) { + KVCache cache = new KVCache(run); + KVCache referenceCache = new KVCache(run); + // B=2 prefill, valid {2, 4}, width 4; row-major B*T positions, rows unequal on all axes. + MLXArray x = MLX.array(run, values(2 * 4 * HIDDEN, 2), new int[] {2, 4, HIDDEN}); + MultiAxisRotation rotation = + rope.prepare( + run, + 2, + 4, + new int[][] { + {0, 1, 0, 0, 2, 3, 4, 5}, {0, 2, 0, 0, 0, 1, 2, 3}, {1, 0, 0, 0, 1, 1, 2, 2} + }); + MLXArray mask = + AttentionMask.batched( + run, new int[] {0, 0}, new int[] {0, 0}, new int[] {2, 4}, 4, 4, 0); + MLXArray actual = fixture.attention.forward(x, cache, mask, new int[] {2, 4}, rotation); + MLXArray reference = + referenceAttention(fixture, x, referenceCache, mask, new int[] {2, 4}, rotation); + assertArrayEquals(reference.toFloatArray(), actual.toFloatArray()); + assertEquals(2, cache.nextPosition(0)); + assertEquals(4, cache.nextPosition(1)); + assertEquals(2, cache.rowLength(0)); + assertEquals(4, cache.rowLength(1)); + + // B=2 decode, valid {1, 1}, width 1; each row decodes from its own position. + MLXArray x1 = MLX.array(run, values(2 * 1 * HIDDEN, 5), new int[] {2, 1, HIDDEN}); + MultiAxisRotation decodeRotation = + rope.prepare(run, 2, 1, new int[][] {{2, 4}, {1, 2}, {3, 0}}); + MLXArray decodeMask = + AttentionMask.batched( + run, new int[] {2, 4}, new int[] {0, 0}, new int[] {1, 1}, 1, 5, 0); + MLXArray decoded = + fixture.attention.forward(x1, cache, decodeMask, new int[] {1, 1}, decodeRotation); + MLXArray referenceDecode = + referenceAttention( + fixture, x1, referenceCache, decodeMask, new int[] {1, 1}, decodeRotation); + assertArrayEquals(referenceDecode.toFloatArray(), decoded.toFloatArray()); + assertEquals(3, cache.nextPosition(0)); + assertEquals(5, cache.nextPosition(1)); + } + } + } + + /** + * The block passthrough must compose exactly input plus the rotated attention output, plus the + * MLP of the post-attention normalization -- bit-identical to the manual composition from the + * same retained instances. + */ + @Test + @Tag("full-float32") + void blockRotationForwardMatchesManualComposition() { + try (MLXScope scope = new MLXScope()) { + AttentionFixture fixture = new AttentionFixture(scope, null, null); + RMSNorm inputNorm = + new RMSNorm(scope, MLX.array(scope, values(HIDDEN, 10), new int[] {HIDDEN}), 1e-6f); + RMSNorm postNorm = + new RMSNorm(scope, MLX.array(scope, values(HIDDEN, 11), new int[] {HIDDEN}), 1e-6f); + SwiGLU mlp = + new SwiGLU( + scope, + MLX.array(scope, values(HIDDEN * 32, 12), new int[] {32, HIDDEN}), + null, + MLX.array(scope, values(HIDDEN * 32, 13), new int[] {32, HIDDEN}), + null, + MLX.array(scope, values(32 * HIDDEN, 14), new int[] {HIDDEN, 32}), + null); + DecoderBlock block = new DecoderBlock(scope, inputNorm, fixture.attention, postNorm, mlp); + MultiAxisRope rope = new MultiAxisRope(10_000f, SECTIONS); + try (MLXScope run = scope.newChild()) { + KVCache cache = new KVCache(run); + KVCache referenceCache = new KVCache(run); + MLXArray x = MLX.array(run, values(1 * 4 * HIDDEN, 8), new int[] {1, 4, HIDDEN}); + MultiAxisRotation rotation = + rope.prepare(run, 1, 4, new int[][] {{1, 2, 3, 4}, {0, 1, 2, 3}, {2, 0, 1, 2}}); + MLXArray mask = + AttentionMask.batched(run, new int[] {0}, new int[] {0}, new int[] {4}, 4, 4, 0); + MLXArray actual = block.forward(x, cache, mask, new int[] {4}, rotation); + MLXArray attention = + referenceAttention( + fixture, inputNorm.forward(x), referenceCache, mask, new int[] {4}, rotation); + MLXArray afterAttention = MLXOps.add(x, attention); + MLXArray reference = + MLXOps.add(afterAttention, mlp.forward(postNorm.forward(afterAttention))); + assertArrayEquals(reference.toFloatArray(), actual.toFloatArray()); + assertEquals(4, cache.nextPosition()); + } + } + } + + // --------------------------------------------------------------------- + // rejection cases (ordinary test task, TF32=1) + // --------------------------------------------------------------------- + + /** + * The rotation overload must reject a missing {@code validLengths} (it always runs the batch + * bookkeeping), a null rotation, a wrong number of lengths and a rotation prepared for a + * different sequence or batch or with a prefix wider than the head -- all {@link + * IllegalArgumentException} before any array is written, the cache left untouched -- and a + * rotation whose prefix fits the head but differs from the attention's configured rotary width + * (the case the prefix-fit check alone would let through). + */ + @Test + void rotationForwardRejectsInvalidArguments() { + try (MLXScope scope = new MLXScope()) { + AttentionFixture fixture = new AttentionFixture(scope, null, null); + MultiAxisRope rope = new MultiAxisRope(10_000f, SECTIONS); + try (MLXScope run = scope.newChild()) { + KVCache cache = new KVCache(run); + MLXArray x = MLX.array(run, values(1 * 3 * HIDDEN, 4), new int[] {1, 3, HIDDEN}); + MultiAxisRotation rotation = + rope.prepare(run, 1, 3, new int[][] {{0, 1, 2}, {1, 2, 3}, {2, 3, 4}}); + IllegalArgumentException missingLengths = + assertThrows( + IllegalArgumentException.class, + () -> fixture.attention.forward(x, cache, null, null, rotation)); + assertTrue(missingLengths.getMessage().contains("valid lengths")); + assertThrows( + IllegalArgumentException.class, + () -> fixture.attention.forward(x, cache, null, new int[] {3}, null)); + assertThrows( + IllegalArgumentException.class, + () -> fixture.attention.forward(x, cache, null, new int[] {3, 3}, rotation)); + // A rotation prepared for a different sequence or batch, or with a prefix wider than the + // head, is rejected before any projection builds a graph node. + MultiAxisRotation shortRotation = + rope.prepare(run, 1, 2, new int[][] {{0, 1}, {1, 2}, {2, 3}}); + IllegalArgumentException wrongSequence = + assertThrows( + IllegalArgumentException.class, + () -> fixture.attention.forward(x, cache, null, new int[] {3}, shortRotation)); + assertTrue(wrongSequence.getMessage().contains("must match the prepared rotation")); + MultiAxisRotation wideBatch = + rope.prepare( + run, + 2, + 3, + new int[][] {{0, 1, 0, 1, 2, 3}, {1, 2, 0, 1, 2, 3}, {2, 3, 0, 1, 2, 3}}); + assertThrows( + IllegalArgumentException.class, + () -> fixture.attention.forward(x, cache, null, new int[] {3}, wideBatch)); + MultiAxisRope wideRope = new MultiAxisRope(10_000f, new int[] {8, 1, 1}); // rotaryDims 10 + MultiAxisRotation wideRotation = + wideRope.prepare(run, 1, 3, new int[][] {{0, 1, 2}, {1, 2, 3}, {2, 3, 4}}); + IllegalArgumentException wideHead = + assertThrows( + IllegalArgumentException.class, + () -> fixture.attention.forward(x, cache, null, new int[] {3}, wideRotation)); + assertTrue(wideHead.getMessage().contains("rotaryDims exceeds head dimension")); + // A rotation prepared from a different mrope-section sum: its prefix (6) fits the head + // (8) but differs from the attention's configured rotary width (HEAD_DIM == 8), so only + // the width-mismatch check rejects it, and it does so before any projection is built. + MultiAxisRope narrowRope = new MultiAxisRope(10_000f, new int[] {1, 1, 1}); // rotaryDims 6 + MultiAxisRotation narrowRotation = + narrowRope.prepare(run, 1, 3, new int[][] {{0, 1, 2}, {1, 2, 3}, {2, 3, 4}}); + IllegalArgumentException wrongWidth = + assertThrows( + IllegalArgumentException.class, + () -> fixture.attention.forward(x, cache, null, new int[] {3}, narrowRotation)); + assertTrue(wrongWidth.getMessage().contains("does not match the configured rotary width")); + assertEquals(0, cache.nextPosition()); + } + } + } + + /** + * The block rotation overload must reject a block whose attention child is not a {@link + * DecoderAttention} (the only cached attention with the rotation overload) before any array is + * written, leaving the cache untouched. + */ + @Test + void blockRotationForwardRejectsNonDecoderAttentionChild() { + try (MLXScope scope = new MLXScope()) { + int embed = 16; + GroupedQueryAttention gqa = + new GroupedQueryAttention( + scope, + 2, + 1, + 10_000f, + MLX.array(scope, values(embed * embed, 0), new int[] {embed, embed}), + null, + MLX.array(scope, values(8 * embed, 1), new int[] {8, embed}), + null, + MLX.array(scope, values(8 * embed, 2), new int[] {8, embed}), + null, + MLX.array(scope, values(embed * embed, 3), new int[] {embed, embed}), + null); + SwiGLU mlp = + new SwiGLU( + scope, + MLX.array(scope, values(embed * 32, 4), new int[] {32, embed}), + null, + MLX.array(scope, values(embed * 32, 5), new int[] {32, embed}), + null, + MLX.array(scope, values(32 * embed, 6), new int[] {embed, 32}), + null); + DecoderBlock block = + new DecoderBlock( + scope, + new RMSNorm(scope, MLX.array(scope, values(embed, 7), new int[] {embed}), 1e-6f), + gqa, + new RMSNorm(scope, MLX.array(scope, values(embed, 8), new int[] {embed}), 1e-6f), + mlp); + MultiAxisRope rope = new MultiAxisRope(10_000f, SECTIONS); + try (MLXScope run = scope.newChild()) { + KVCache cache = new KVCache(run); + MLXArray x = MLX.array(run, values(1 * 3 * embed, 9), new int[] {1, 3, embed}); + MultiAxisRotation rotation = + rope.prepare(run, 1, 3, new int[][] {{0, 1, 2}, {1, 2, 3}, {2, 3, 4}}); + IllegalArgumentException failure = + assertThrows( + IllegalArgumentException.class, + () -> block.forward(x, cache, null, new int[] {3}, rotation)); + assertTrue( + failure + .getMessage() + .startsWith( + "DecoderBlock.forward with a MultiAxisRotation requires a DecoderAttention" + + " child")); + assertEquals(0, cache.nextPosition()); + } + } + } + + // --------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------- + + /** + * The operation-identical manual composition of {@link DecoderAttention#forward(MLXArray, + * KVCache, MLXArray, int[], MultiAxisRotation)} from the retained instances: projection to heads, + * the optional QK norms, the rotation, the cache append, the scaled dot-product attention over + * the retained keys, and the output projection -- the same order, shapes and scopes the seam + * itself runs. + */ + private static MLXArray referenceAttention( + AttentionFixture fixture, + MLXArray x, + KVCache cache, + MLXArray mask, + int[] validLengths, + MultiAxisRotation rotation) { + int batch = x.shape()[0]; + int sequence = x.shape()[1]; + MLXArray q = + AttentionHeads.toHeads(fixture.queryProj.forward(x), batch, sequence, HEADS, HEAD_DIM); + MLXArray k = + AttentionHeads.toHeads( + fixture.keyProj.forward(x), batch, sequence, KEY_VALUE_HEADS, HEAD_DIM); + if (fixture.queryNorm != null) { + q = fixture.queryNorm.forward(q); + } + if (fixture.keyNorm != null) { + k = fixture.keyNorm.forward(k); + } + q = rotation.rotate(q); + k = rotation.rotate(k); + MLXArray v = + AttentionHeads.toHeads( + fixture.valueProj.forward(x), batch, sequence, KEY_VALUE_HEADS, HEAD_DIM); + cache.append(k, v, validLengths); + MLXArray attended = + MLXFast.scaledDotProductAttention( + q, cache.keys(), cache.values(), SCALE, mask == null, mask, null); + MLXArray merged = MLXShape.flatten(MLXShape.transpose(attended, new int[] {0, 2, 1, 3}), 2, 3); + return fixture.outProj.forward(merged); + } + + /** + * The cross-shape reference for {@link #rotationPrefillAndDecodeBitMatchFreshPrefill}: a fresh + * six-token prefill (the five prefill tokens followed by the decode token) through the rotation + * seam on its own cache under {@code run}, returning the output's last row. + */ + private static MLXArray freshSixTokenPrefillLastRow( + AttentionFixture fixture, + MultiAxisRope rope, + MLXScope run, + float[] prefill, + float[] decodeInput) { + float[] fresh = new float[1 * 6 * HIDDEN]; + System.arraycopy(prefill, 0, fresh, 0, 5 * HIDDEN); + System.arraycopy(decodeInput, 0, fresh, 5 * HIDDEN, HIDDEN); + MLXArray x6 = MLX.array(run, fresh, new int[] {1, 6, HIDDEN}); + MultiAxisRotation freshRotation = + rope.prepare( + run, 1, 6, new int[][] {{0, 1, 2, 3, 4, 5}, {0, 2, 4, 2, 4, 6}, {1, 1, 3, 3, 5, 4}}); + MLXArray freshMask = + AttentionMask.batched(run, new int[] {0}, new int[] {0}, new int[] {6}, 6, 6, 0); + KVCache freshCache = new KVCache(run); + MLXArray freshPrefill = + fixture.attention.forward(x6, freshCache, freshMask, new int[] {6}, freshRotation); + return MLXShape.slice(freshPrefill, new int[] {0, 5, 0}, new int[] {1, 6, HIDDEN}); + } + + /** The shared attention fixture: identity projections plus the optional norms, retained. */ + private static final class AttentionFixture { + final DecoderAttention attention; + final Linear queryProj; + final Linear keyProj; + final Linear valueProj; + final Linear outProj; + final RMSNorm queryNorm; + final RMSNorm keyNorm; + + AttentionFixture(MLXScope scope, RMSNorm queryNorm, RMSNorm keyNorm) { + this.queryProj = new Linear(scope, identity(scope), null); + this.keyProj = new Linear(scope, identity(scope), null); + this.valueProj = new Linear(scope, identity(scope), null); + this.outProj = new Linear(scope, identity(scope), null); + this.queryNorm = queryNorm; + this.keyNorm = keyNorm; + this.attention = + new DecoderAttention( + scope, + HEADS, + KEY_VALUE_HEADS, + HEAD_DIM, + new RopeSpec.Base(10_000f), + HEAD_DIM, + null, + null, + queryProj, + keyProj, + valueProj, + outProj, + queryNorm, + keyNorm); + } + } + + /** The {@code [HIDDEN, HIDDEN]} identity projection weight. */ + private static MLXArray identity(MLXScope scope) { + float[] weight = new float[HIDDEN * HIDDEN]; + for (int i = 0; i < HIDDEN; i++) { + weight[i * HIDDEN + i] = 1f; + } + return MLX.array(scope, weight, new int[] {HIDDEN, HIDDEN}); + } + + /** Distinct deterministic values in [-1, 1] with a per-case seed offset. */ + private static float[] values(int n, int seed) { + float[] data = new float[n]; + for (int j = 0; j < n; j++) { + data[j] = (((j + seed) % 13) - 6) / 6.0f; + } + return data; + } + + private static int[] zeros(int n) { + return new int[n]; + } +} diff --git a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java new file mode 100644 index 00000000..6f844b17 --- /dev/null +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java @@ -0,0 +1,548 @@ +package se.alipsa.jmlx.nn; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import se.alipsa.jmlx.core.DType; +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.ffi.EnabledIfNativeAvailable; +import se.alipsa.jmlx.memory.MLXScope; + +/** + * Strict tests for the frozen WP4 multi-axis RoPE component (req/plans/phase7-3b-reference- + * findings.md §6.9). The float32-comparison cases are tagged {@code full-float32} so they select + * into {@code :jmlx-core:float32GoldenTest} with {@code MLX_ENABLE_TF32=0}; the rejection and + * structural cases run in the ordinary {@code :jmlx-core:test} task. + */ +@EnabledIfNativeAvailable +class MultiAxisRopeTest { + + private static final int[] PINNED_SECTIONS = {24, 20, 20}; + private static final float PINNED_THETA = 5_000_000f; + + // --------------------------------------------------------------------- + // structural / rejection cases (ordinary test task, TF32=1) + // --------------------------------------------------------------------- + + @Test + void gettersExposeThetaSectionsAndRotaryDims() { + MultiAxisRope rope = new MultiAxisRope(PINNED_THETA, PINNED_SECTIONS); + assertEquals(PINNED_THETA, rope.theta()); + assertArrayEquals(PINNED_SECTIONS, rope.mropeSection()); + assertEquals(128, rope.rotaryDims()); + int[] returned = rope.mropeSection(); + returned[0] = 0; + assertArrayEquals(PINNED_SECTIONS, rope.mropeSection()); + } + + @Test + void prepareReturnsFloat32TablesBroadcastableOverHeads() { + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(10_000f, new int[] {2, 1, 2}); + MultiAxisRotation rotation = rope.prepare(scope, 2, 3, rows3(6)); + assertArrayEquals(new int[] {2, 1, 3, 10}, rotation.cos().shape()); + assertArrayEquals(new int[] {2, 1, 3, 10}, rotation.sin().shape()); + assertEquals(DType.FLOAT32, rotation.cos().dtype()); + assertEquals(DType.FLOAT32, rotation.sin().dtype()); + assertEquals(10, rotation.rotaryDims()); + MLXArray x = MLX.array(scope, values(2 * 1 * 3 * 16, 9), new int[] {2, 1, 3, 16}); + assertArrayEquals(new int[] {2, 1, 3, 16}, rotation.rotate(x).shape()); + } + } + + @Test + void constructorRejectsInvalidThetaAndSections() { + assertThrows(IllegalArgumentException.class, () -> new MultiAxisRope(0f, PINNED_SECTIONS)); + assertThrows(IllegalArgumentException.class, () -> new MultiAxisRope(-1f, PINNED_SECTIONS)); + assertThrows( + IllegalArgumentException.class, () -> new MultiAxisRope(Float.NaN, PINNED_SECTIONS)); + assertThrows( + IllegalArgumentException.class, + () -> new MultiAxisRope(Float.POSITIVE_INFINITY, PINNED_SECTIONS)); + assertThrows( + IllegalArgumentException.class, + () -> new MultiAxisRope(Float.NEGATIVE_INFINITY, PINNED_SECTIONS)); + assertThrows(IllegalArgumentException.class, () -> new MultiAxisRope(1000f, null)); + assertThrows( + IllegalArgumentException.class, () -> new MultiAxisRope(1000f, new int[] {24, 20})); + assertThrows( + IllegalArgumentException.class, () -> new MultiAxisRope(1000f, new int[] {24, 20, 20, 4})); + assertThrows( + IllegalArgumentException.class, () -> new MultiAxisRope(1000f, new int[] {24, -1, 20})); + assertThrows( + IllegalArgumentException.class, () -> new MultiAxisRope(1000f, new int[] {0, 0, 0})); + assertThrows( + IllegalArgumentException.class, + () -> new MultiAxisRope(1000f, new int[] {Integer.MAX_VALUE, Integer.MAX_VALUE, 0})); + } + + @Test + void prepareRejectsInvalidShapeAndPositions() { + MultiAxisRope rope = new MultiAxisRope(1000f, new int[] {2, 1, 2}); + try (MLXScope scope = new MLXScope()) { + assertThrows(IllegalArgumentException.class, () -> rope.prepare(scope, 0, 4, rows3(4))); + assertThrows(IllegalArgumentException.class, () -> rope.prepare(scope, 1, 0, rows3(0))); + assertThrows(IllegalArgumentException.class, () -> rope.prepare(scope, 1, 4, null)); + assertThrows( + IllegalArgumentException.class, + () -> rope.prepare(scope, 1, 4, new int[][] {row(4), row(4)})); + assertThrows( + IllegalArgumentException.class, + () -> rope.prepare(scope, 1, 4, new int[][] {row(4), row(4), row(4), row(4)})); + assertThrows( + IllegalArgumentException.class, + () -> rope.prepare(scope, 1, 4, new int[][] {row(3), row(4), row(4)})); + assertThrows( + IllegalArgumentException.class, + () -> rope.prepare(scope, 1, 4, new int[][] {null, row(4), row(4)})); + int[] negative = row(4); + negative[1] = -1; + assertThrows( + IllegalArgumentException.class, + () -> rope.prepare(scope, 1, 4, new int[][] {row(4), negative, row(4)})); + } + } + + /** + * The int-overflow guard on the position-matrix size: sections {64000, 0, 0} give half = 64,000 + * and 33,555 tokens pass the row-length check (each row has exactly that length), so the long + * product is 33,555 * 64,000 = 2,147,520,000 > Integer.MAX_VALUE (2,147,483,647) while the plain + * int product wraps to -2,147,447,296 -- without the guard, the allocation itself would throw a + * NegativeArraySizeException instead of the documented IllegalArgumentException. The guard fires + * before any float[] allocation: only the three 33,555-int rows (about 402 KB total) are built. + */ + @Test + void prepareRejectsPositionMatrixSizeThatOverflowsInt() { + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(1000f, new int[] {64_000, 0, 0}); + int[] row = new int[33_555]; + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> rope.prepare(scope, 1, 33_555, new int[][] {row, row, row})); + assertTrue(thrown.getMessage().contains("fit an int array")); + } + } + + @Test + void rotateRejectsInvalidInputs() { + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(10_000f, new int[] {2, 1, 2}); // rotaryDims 10 + MultiAxisRotation rotation = rope.prepare(scope, 1, 2, rows3(2)); + IllegalArgumentException tooNarrow = + assertThrows( + IllegalArgumentException.class, + () -> + rotation.rotate( + MLX.array(scope, values(1 * 1 * 2 * 8, 0), new int[] {1, 1, 2, 8}))); + assertTrue(tooNarrow.getMessage().contains("rotaryDims exceeds head dimension")); + assertThrows( + IllegalArgumentException.class, + () -> rotation.rotate(MLX.array(scope, values(1 * 1 * 2 * 16, 1), new int[] {1, 1, 32}))); + assertThrows( + IllegalArgumentException.class, + () -> + rotation.rotate( + MLX.array(scope, values(2 * 1 * 2 * 16, 2), new int[] {2, 1, 2, 16}))); + assertThrows( + IllegalArgumentException.class, + () -> + rotation.rotate( + MLX.array(scope, values(1 * 1 * 3 * 16, 3), new int[] {1, 1, 3, 16}))); + assertThrows(IllegalArgumentException.class, () -> rotation.rotate(null)); + } + } + + // --------------------------------------------------------------------- + // float32-comparison cases (full-float32 tag, MLX_ENABLE_TF32=0) + // --------------------------------------------------------------------- + + @Test + @Tag("full-float32") + void equalAxesReduceToOneDimensionalRoPE() { + int[] axis = {2, 5, 0, 7}; + int[] sections = {2, 2, 2}; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(10_000f, sections); + float[] x = values(1 * 1 * 4 * 12, 0); + MultiAxisRotation rotation = rope.prepare(scope, 1, 4, new int[][] {axis, axis, axis}); + MLXArray arrayX = MLX.array(scope, x, new int[] {1, 1, 4, 12}); + assertArrayEquals( + expectedRotation(10_000, sections, 1, 1, 4, x, 12, axis, axis, axis), + rotation.rotate(arrayX).toFloatArray(), + 1e-5f); + float[][] tables = expectedTables(10_000, sections, axis, axis, axis); + assertArrayEquals(tables[0], rotation.cos().toFloatArray(), 1e-5f); + assertArrayEquals(tables[1], rotation.sin().toFloatArray(), 1e-5f); + } + } + + @Test + @Tag("full-float32") + void pinnedOwnerPatternWithIndexCodedAxes() { + int[] rowT = {0, 1, 2, 3}; + int[] rowH = {10, 11, 12, 13}; + int[] rowW = {20, 21, 22, 23}; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(PINNED_THETA, PINNED_SECTIONS); + float[] x = values(1 * 1 * 4 * 128, 0); + MultiAxisRotation rotation = rope.prepare(scope, 1, 4, new int[][] {rowT, rowH, rowW}); + MLXArray arrayX = MLX.array(scope, x, new int[] {1, 1, 4, 128}); + assertArrayEquals( + expectedRotation(5_000_000, PINNED_SECTIONS, 1, 1, 4, x, 128, rowT, rowH, rowW), + rotation.rotate(arrayX).toFloatArray(), + 1e-5f); + + // Per-frequency owner discrimination at token 0 (t=0, h=10, w=20). The candidates are + // derived from the component's own float32 invFreq order so this isolates the owner + // selection from the frequency derivation (the latter is pinned separately by the + // double-precision comparisons above and the pinned-venv probe). + double[] invFreq = invFrequenciesOptionB(PINNED_THETA, 64); + float[] sinTable = rotation.sin().toFloatArray(); + double minimumSeparation = Double.MAX_VALUE; + for (int f = 0; f < 64; f++) { + double[] candidates = { + Math.sin(0.0 * invFreq[f]), Math.sin(10.0 * invFreq[f]), Math.sin(20.0 * invFreq[f]) + }; + int owner = ownerAxis(f, PINNED_SECTIONS); + for (int axis = 0; axis < 3; axis++) { + if (axis != owner) { + minimumSeparation = + Math.min(minimumSeparation, Math.abs(candidates[axis] - candidates[owner])); + } + } + float actual = sinTable[f]; + assertTrue( + Math.abs(actual - candidates[owner]) <= 5e-7, + "frequency " + + f + + " should use the " + + "THW".charAt(owner) + + " axis, actual " + + actual + + " vs candidate " + + candidates[owner]); + for (int axis = 0; axis < 3; axis++) { + if (axis != owner) { + assertTrue( + Math.abs(actual - candidates[axis]) >= 1e-6, + "frequency " + f + " must not use the " + "THW".charAt(axis) + " axis"); + } + } + } + assertTrue(minimumSeparation >= 1e-6, "case is not discriminable: " + minimumSeparation); + } + } + + @Test + @Tag("full-float32") + void smallSectionsExerciseEveryOwnerAtConstantIndexCodedAxes() { + // [2,1,2]: owners T, H, W, T, T at constant index-coded coordinates. + int[] sections = {2, 1, 2}; + int[] rowT = {1, 1}; + int[] rowH = {2, 2}; + int[] rowW = {4, 4}; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(10_000f, sections); + float[] x = values(1 * 1 * 2 * 10, 0); + MultiAxisRotation rotation = rope.prepare(scope, 1, 2, new int[][] {rowT, rowH, rowW}); + MLXArray arrayX = MLX.array(scope, x, new int[] {1, 1, 2, 10}); + assertArrayEquals( + expectedRotation(10_000, sections, 1, 1, 2, x, 10, rowT, rowH, rowW), + rotation.rotate(arrayX).toFloatArray(), + 1e-5f); + } + // [1,1,1]: owners T, H, W on every frequency. + sections = new int[] {1, 1, 1}; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(1000f, sections); + float[] x = values(1 * 1 * 2 * 6, 5); + MultiAxisRotation rotation = rope.prepare(scope, 1, 2, new int[][] {rowT, rowH, rowW}); + MLXArray arrayX = MLX.array(scope, x, new int[] {1, 1, 2, 6}); + assertArrayEquals( + expectedRotation(1000, sections, 1, 1, 2, x, 6, rowT, rowH, rowW), + rotation.rotate(arrayX).toFloatArray(), + 1e-5f); + } + } + + @Test + @Tag("full-float32") + void decodePositionsWithNegativeDeltaAndBatchedRows() { + // B=1, T=1: cache position 82 plus delta -56 gives coordinate 26 on all three axes -- the + // pinned 2B decode step. + int[] row = {26}; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(PINNED_THETA, PINNED_SECTIONS); + float[] x = values(1 * 1 * 1 * 128, 0); + MultiAxisRotation rotation = rope.prepare(scope, 1, 1, new int[][] {row, row, row}); + MLXArray arrayX = MLX.array(scope, x, new int[] {1, 1, 1, 128}); + assertArrayEquals( + expectedRotation(5_000_000, PINNED_SECTIONS, 1, 1, 1, x, 128, row, row, row), + rotation.rotate(arrayX).toFloatArray(), + 1e-5f); + } + // B=2, T=2 with per-row positions (row 0 low, row 1 at the prefill-envelope edge). + int[] rowT = {0, 1, 30, 31}; + int[] rowH = {4, 5, 26, 27}; + int[] rowW = {8, 9, 22, 23}; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(PINNED_THETA, PINNED_SECTIONS); + float[] x = values(2 * 1 * 2 * 128, 0); + MultiAxisRotation rotation = rope.prepare(scope, 2, 2, new int[][] {rowT, rowH, rowW}); + MLXArray arrayX = MLX.array(scope, x, new int[] {2, 1, 2, 128}); + assertArrayEquals( + expectedRotation(5_000_000, PINNED_SECTIONS, 2, 1, 2, x, 128, rowT, rowH, rowW), + rotation.rotate(arrayX).toFloatArray(), + 1e-5f); + } + } + + @Test + @Tag("full-float32") + void partialRotaryTailPassesThroughBitIdentically() { + int[] sections = {2, 1, 2}; // rotaryDims 10 + int headDim = 16; + int[] rowT = {1, 1}; + int[] rowH = {2, 2}; + int[] rowW = {4, 4}; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(10_000f, sections); + float[] x = values(1 * 1 * 2 * headDim, 0); + for (int t = 0; t < 2; t++) { + for (int j = 0; j < headDim - 10; j++) { + x[t * headDim + 10 + j] = 90 + j; // distinctive, non-deterministic tail values + } + } + MultiAxisRotation rotation = rope.prepare(scope, 1, 2, new int[][] {rowT, rowH, rowW}); + MLXArray arrayX = MLX.array(scope, x, new int[] {1, 1, 2, headDim}); + float[] actual = rotation.rotate(arrayX).toFloatArray(); + assertArrayEquals( + expectedRotation(10_000, sections, 1, 1, 2, x, headDim, rowT, rowH, rowW), actual, 1e-5f); + for (int t = 0; t < 2; t++) { + for (int j = 0; j < headDim - 10; j++) { + assertEquals( + x[t * headDim + 10 + j], + actual[t * headDim + 10 + j], + "tail element (" + t + ", " + (10 + j) + ") must pass through bit-identically"); + } + } + } + } + + @Test + @Tag("full-float32") + void rotationCastsCosSinToInputDtype() { + int[] sections = {2, 1, 2}; + int[] rowT = {1, 1}; + int[] rowH = {2, 2}; + int[] rowW = {4, 4}; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(10_000f, sections); + MultiAxisRotation rotation = rope.prepare(scope, 1, 2, new int[][] {rowT, rowH, rowW}); + float[] x32 = values(1 * 1 * 2 * 10, 0); + MLXArray float32 = MLX.array(scope, x32, new int[] {1, 1, 2, 10}); + MLXArray out32 = rotation.rotate(float32); + assertEquals(DType.FLOAT32, out32.dtype()); + assertArrayEquals( + expectedRotation(10_000, sections, 1, 1, 2, x32, 10, rowT, rowH, rowW), + out32.toFloatArray(), + 1e-5f); + + MLXArray bfloat = MLX.astype(float32, DType.BFLOAT16); + MLXArray outBfloat = rotation.rotate(bfloat); + assertEquals(DType.BFLOAT16, outBfloat.dtype()); + assertArrayEquals( + roundedToDtype( + expectedRotation( + 10_000, sections, 1, 1, 2, bfloat.toFloatArray(), 10, rowT, rowH, rowW), + DType.BFLOAT16), + outBfloat.toFloatArray(), + 0.05f); + + MLXArray half = MLX.astype(float32, DType.FLOAT16); + MLXArray outHalf = rotation.rotate(half); + assertEquals(DType.FLOAT16, outHalf.dtype()); + assertArrayEquals( + roundedToDtype( + expectedRotation( + 10_000, sections, 1, 1, 2, half.toFloatArray(), 10, rowT, rowH, rowW), + DType.FLOAT16), + outHalf.toFloatArray(), + 0.02f); + } + } + + // --------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------- + + /** Distinct deterministic values in [-1, 1] with a per-case seed offset. */ + private static float[] values(int n, int seed) { + float[] data = new float[n]; + for (int j = 0; j < n; j++) { + data[j] = (((j + seed) % 13) - 6) / 6.0f; + } + return data; + } + + private static int[] row(int length) { + int[] data = new int[length]; + for (int i = 0; i < length; i++) { + data[i] = i; + } + return data; + } + + private static int[][] rows3(int length) { + return new int[][] {row(length), row(length), row(length)}; + } + + /** + * The reference's owner of frequency index {@code f} for the interleaved layout: H owns {@code f + * % 3 == 1} with {@code f < 1 + 3 * s1}, W owns {@code f % 3 == 2} with {@code f < 2 + 3 * s2}, + * everything else stays T. + */ + private static int ownerAxis(int f, int[] sections) { + if (f % 3 == 1 && f < 1 + 3 * sections[1]) { + return 1; + } + if (f % 3 == 2 && f < 2 + 3 * sections[2]) { + return 2; + } + return 0; + } + + /** + * The component's float32 invFreq derivation order (the documented option B, see {@code + * MultiAxisRope.deriveInvFrequencies}) recomputed here as a double-precision stand-in for the + * owner-candidate discrimination only. + */ + private static double[] invFrequenciesOptionB(float theta, int half) { + int rotaryDims = 2 * half; + double[] invFreq = new double[half]; + for (int i = 0; i < half; i++) { + float exponent = (float) (2 * i) / (float) rotaryDims; + invFreq[i] = 1.0 / (float) Math.pow(theta, exponent); + } + return invFreq; + } + + /** + * Java double-precision-derived float32 reference: float32 cos/sin tables of shape {@code + * [tokens, rotaryDims]} (the {@code cat((freqs, freqs), -1)} layout) for the given flattened axis + * position rows. + */ + private static float[][] expectedTables( + double theta, int[] sections, int[] rowT, int[] rowH, int[] rowW) { + int half = sections[0] + sections[1] + sections[2]; + int rotaryDims = 2 * half; + int tokens = rowT.length; + double[] invFreq = new double[half]; + for (int i = 0; i < half; i++) { + invFreq[i] = Math.pow(theta, -2.0 * i / (2.0 * half)); + } + float[] cos = new float[tokens * rotaryDims]; + float[] sin = new float[tokens * rotaryDims]; + for (int t = 0; t < tokens; t++) { + for (int f = 0; f < half; f++) { + int axis = ownerAxis(f, sections); + int position = axis == 0 ? rowT[t] : axis == 1 ? rowH[t] : rowW[t]; + double angle = position * invFreq[f]; + float c = (float) Math.cos(angle); + float s = (float) Math.sin(angle); + cos[t * rotaryDims + f] = c; + cos[t * rotaryDims + half + f] = c; + sin[t * rotaryDims + f] = s; + sin[t * rotaryDims + half + f] = s; + } + } + return new float[][] {cos, sin}; + } + + /** + * The frozen composition in double precision: for a flat {@code [batch, heads, sequence, + * headDim]} input, {@code out = x * cos + rotate_half(x) * sin} on the {@code rotaryDims} prefix + * ({@code rotate_half = cat((−x2, x1))}) with the tail passing through unchanged. + */ + private static float[] expectedRotation( + double theta, + int[] sections, + int batch, + int heads, + int sequence, + float[] x, + int headDim, + int[] rowT, + int[] rowH, + int[] rowW) { + float[][] tables = expectedTables(theta, sections, rowT, rowH, rowW); + int half = sections[0] + sections[1] + sections[2]; + int rotaryDims = 2 * half; + float[] out = new float[batch * heads * sequence * headDim]; + for (int b = 0; b < batch; b++) { + for (int h = 0; h < heads; h++) { + for (int t = 0; t < sequence; t++) { + int base = (b * heads + h) * sequence * headDim + t * headDim; + int table = (b * sequence + t) * rotaryDims; + for (int d = 0; d < headDim; d++) { + if (d >= rotaryDims) { + out[base + d] = x[base + d]; + continue; + } + int f = d % half; + float c = tables[0][table + f]; + float s = tables[1][table + f]; + if (d < half) { + out[base + d] = x[base + d] * c - x[base + d + half] * s; + } else { + out[base + d] = x[base + d] * c + x[base + d - half] * s; + } + } + } + } + } + return out; + } + + private static float[] roundedToDtype(float[] values, DType dtype) { + float[] out = new float[values.length]; + for (int i = 0; i < values.length; i++) { + out[i] = dtype == DType.BFLOAT16 ? toBf16(values[i]) : toF16(values[i]); + } + return out; + } + + /** + * Round-to-nearest-even from float32 to bfloat16 at the bit level: a bfloat16 value is the top 16 + * bits of the float32, so the low 16 bits are rounded with the half-ulp plus retained-LSB carry + * trick. + */ + private static float toBf16(float value) { + long bits = Float.floatToRawIntBits(value) & 0xFFFFFFFFL; + long lsb = (bits >>> 15) & 1; + long rounded = (bits + 0x7FFFL + lsb) & 0xFFFF0000L; + return Float.intBitsToFloat((int) rounded); + } + + /** + * Round-to-nearest-even from float32 to float16 via the exact double promotion: the low 48 double + * bits are rounded with the same half-ulp plus retained-LSB carry trick, sign handled separately. + */ + private static float toF16(float value) { + long bits = Double.doubleToRawLongBits((double) value); + boolean negative = bits < 0; + long magnitude = bits & 0x7FFFFFFFFFFFFFFFL; + long lsb = (magnitude >>> 48) & 1; + long rounded = (magnitude + 0x7FFFFFFFFFFFL + lsb) & 0x7FFF800000000000L; + if (negative) { + rounded |= 0x8000000000000000L; + } + return (float) Double.longBitsToDouble(rounded); + } +} diff --git a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java new file mode 100644 index 00000000..64fd5b0e --- /dev/null +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java @@ -0,0 +1,132 @@ +package se.alipsa.jmlx.nn; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; + +import org.junit.jupiter.api.Test; +import se.alipsa.jmlx.core.DType; +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.ffi.EnabledIfNativeAvailable; +import se.alipsa.jmlx.memory.MLXScope; + +/** + * Pins the single-row batched arithmetic that the WP4 rotation path relies on (plan §6, Unit 2): + * for a uniform batch row, {@code cache.append(k, v, {T})} is the same arithmetic as {@code + * cache.append(k, v)} -- bit-identical retained tensors -- and the B = 1 {@link + * AttentionMask#batched} mask is elementwise equal to the 2-D {@link AttentionMask#slidingWindow} + * mask. Both facts let the single-row rotation path route through the batch-aware overloads without + * changing the old graph's results. + */ +@EnabledIfNativeAvailable +class SingleRowDecodeEquivalenceTest { + + @Test + void uniformSingleRowAppendIsBitIdenticalToPlainAppend() { + try (MLXScope scope = new MLXScope()) { + KVCache plain = new KVCache(scope); + KVCache batched = new KVCache(scope); + MLXArray k1 = MLX.array(scope, values(1 * 2 * 4 * 8, 0), new int[] {1, 2, 4, 8}); + MLXArray v1 = MLX.array(scope, values(1 * 2 * 4 * 8, 5), new int[] {1, 2, 4, 8}); + plain.append(k1, v1); + batched.append(k1, v1, new int[] {4}); + assertArrayEquals(plain.keys().toFloatArray(), batched.keys().toFloatArray()); + assertArrayEquals(plain.values().toFloatArray(), batched.values().toFloatArray()); + assertEquals(plain.offset(), batched.offset()); + assertEquals(plain.nextPosition(0), batched.nextPosition(0)); + + MLXArray k2 = MLX.array(scope, values(1 * 2 * 3 * 8, 1), new int[] {1, 2, 3, 8}); + MLXArray v2 = MLX.array(scope, values(1 * 2 * 3 * 8, 6), new int[] {1, 2, 3, 8}); + plain.append(k2, v2); + batched.append(k2, v2, new int[] {3}); + assertArrayEquals(plain.keys().toFloatArray(), batched.keys().toFloatArray()); + assertArrayEquals(plain.values().toFloatArray(), batched.values().toFloatArray()); + assertEquals(plain.offset(), batched.offset()); + assertEquals(plain.nextPosition(0), batched.nextPosition(0)); + + MLXArray k3 = MLX.array(scope, values(1 * 2 * 1 * 8, 2), new int[] {1, 2, 1, 8}); + MLXArray v3 = MLX.array(scope, values(1 * 2 * 1 * 8, 7), new int[] {1, 2, 1, 8}); + plain.append(k3, v3); + batched.append(k3, v3, new int[] {1}); + assertArrayEquals(plain.keys().toFloatArray(), batched.keys().toFloatArray()); + assertArrayEquals(plain.values().toFloatArray(), batched.values().toFloatArray()); + assertEquals(plain.offset(), batched.offset()); + assertEquals(plain.nextPosition(0), batched.nextPosition(0)); + } + } + + @Test + void singleRowBatchedMaskMatchesTwoDimensionalMask() { + try (MLXScope scope = new MLXScope()) { + // Prefill rows have T == S (no key padding), so the batched row is unpadded and + // elementwise equal to the 2-D mask; decode rows are the queryStart = S - 1 case below. + for (int[] shape : new int[][] {{4, 4}, {90, 90}, {7, 7}}) { + int queryLength = shape[0]; + int keyLength = shape[1]; + MLXArray batched = + AttentionMask.batched( + scope, + new int[] {0}, + new int[] {0}, + new int[] {queryLength}, + queryLength, + keyLength, + 0); + MLXArray twoD = AttentionMask.slidingWindow(scope, 0, 0, queryLength, keyLength, keyLength); + assertArrayEquals(new int[] {1, 1, queryLength, keyLength}, batched.shape()); + assertArrayEquals(asFloat(batched), asFloat(twoD)); + } + for (int keyLength : new int[] {4, 90, 7}) { + MLXArray batched = + AttentionMask.batched( + scope, new int[] {keyLength - 1}, new int[] {0}, new int[] {1}, 1, keyLength, 0); + MLXArray twoD = + AttentionMask.slidingWindow(scope, keyLength - 1, 0, 1, keyLength, keyLength); + assertArrayEquals(asFloat(batched), asFloat(twoD)); + } + // Windowed batch-1 unpadded rows, parameterized as {queryStart, keyStart, queryLength, + // keyLength, window}: a prefill longer than the window, a one-token decode over an evicted + // cache (keyStart > 0) whose retained keys all sit within the window, and multi-token steps + // over evicted caches where the window actually masks keys. + for (int[] row : + new int[][] { + {0, 0, 9, 9, 4}, + {8, 5, 1, 4, 4}, + {8, 5, 2, 5, 4}, + {14, 11, 3, 6, 4} + }) { + int queryStart = row[0]; + int keyStart = row[1]; + int queryLength = row[2]; + int keyLength = row[3]; + int window = row[4]; + MLXArray batched = + AttentionMask.batched( + scope, + new int[] {queryStart}, + new int[] {keyStart}, + new int[] {queryLength}, + queryLength, + keyLength, + window); + MLXArray twoD = + AttentionMask.slidingWindow( + scope, queryStart, keyStart, queryLength, keyLength, window); + assertArrayEquals(new int[] {1, 1, queryLength, keyLength}, batched.shape()); + assertArrayEquals(asFloat(batched), asFloat(twoD)); + } + } + } + + private static float[] asFloat(MLXArray boolMask) { + return MLX.astype(boolMask, DType.FLOAT32).toFloatArray(); + } + + private static float[] values(int n, int seed) { + float[] data = new float[n]; + for (int j = 0; j < n; j++) { + data[j] = (((j + seed) % 13) - 6) / 6.0f; + } + return data; + } +} diff --git a/jmlx-models/src/main/java/se/alipsa/jmlx/models/DecoderModel.java b/jmlx-models/src/main/java/se/alipsa/jmlx/models/DecoderModel.java index 96ef3924..29a5dda2 100644 --- a/jmlx-models/src/main/java/se/alipsa/jmlx/models/DecoderModel.java +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/DecoderModel.java @@ -8,6 +8,7 @@ import java.util.Objects; import java.util.Set; import java.util.function.Consumer; +import se.alipsa.jmlx.core.DType; import se.alipsa.jmlx.core.MLX; import se.alipsa.jmlx.core.MLXArray; import se.alipsa.jmlx.core.MLXOps; @@ -19,6 +20,7 @@ import se.alipsa.jmlx.nn.KVCache; import se.alipsa.jmlx.nn.KVCachePolicy; import se.alipsa.jmlx.nn.Module; +import se.alipsa.jmlx.nn.MultiAxisRotation; import se.alipsa.jmlx.nn.RMSNorm; import se.alipsa.jmlx.nn.UnaryLayer; import se.alipsa.jmlx.tokenizer.HfTokenizer; @@ -168,6 +170,26 @@ interface EmbeddingHook { MLXArray apply(MLXArray embedded); } + /** + * Injects a transform after each decoder block's residual output and before the final + * normalization, invoked with the block's 0-based layer index. Qwen3-VL DeepStack (WP6/WP7) uses + * this seam to add DeepStack level {@code i} after text block {@code i}, where the level's + * features come from the vision tower's {@code deepstack_visual_indexes[i]} (findings §6.6/§6.9). + * Those are different counters that must never be mixed: the vision extraction indices address + * blocks of the 24-block vision stack (pinned 2B: {@code [5, 11, 17]}), while the injection index + * addresses blocks of the 28-block text stack. The hook is invoked after every block's + * residual output; selecting which layers actually receive an injection is the hook's own job. + * Production ordinary decoders pass no-op or null hooks. Implementations must return a non-null + * array with exactly the same shape and dtype as the hidden state they were given; {@code + * decoderStack} checks this immediately after each call and throws {@link + * IllegalArgumentException} naming the block index, so a bad result is reported at the hook + * rather than later inside the next block's projection. + */ + @FunctionalInterface + interface PostBlockHook { + MLXArray apply(int layerIndex, MLXArray hidden); + } + /** * Lazy batched step shared with the scheduler: runs one forward over a left-padded {@code [B,T]} * batch and returns the {@code [B,1,V]} logits of every row's last valid column, with no @@ -182,9 +204,38 @@ interface EmbeddingHook { */ final MLXArray stepLogits( MLXArray tokenIds, List caches, int[] validLengths, EmbeddingHook hook) { + return stepLogits(tokenIds, caches, validLengths, hook, null, null); + } + + /** + * The seam-carrying form of {@link #stepLogits(MLXArray, List, int[], EmbeddingHook)}: the + * embedding hook keeps its fourth, post-lookup/pre-scale position; a non-null {@code rotation} + * replaces the stack's derived step frequencies with the request's own multi-axis rotation, and + * {@code postBlockHook} fires after every block's residual output, before the final + * normalization. The preflight, postflight, evaluation and poisoning invariants are the batched + * ones, unchanged. A non-null {@code rotation} is validated against the batch, the sequence + * width, the head dimension and the configured rotary width ({@link + * MultiAxisRotation#validateAgainst(int, int, int, int)}) before the embedding lookup builds any + * graph node. The existing 4-argument overload delegates with {@code (null, null)} and keeps the + * identical graph. + */ + final MLXArray stepLogits( + MLXArray tokenIds, + List caches, + int[] validLengths, + EmbeddingHook hook, + MultiAxisRotation rotation, + PostBlockHook postBlockHook) { int[][] before = preflightBatch(tokenIds, caches, validLengths); + if (rotation != null) { + // Reject a mismatched rotation before the embedding lookup builds any graph node. + rotation.validateAgainst( + tokenIds.shape()[0], tokenIds.shape()[1], descriptor.headDim(), descriptor.rotaryDims()); + } try { - MLXArray normalized = normalizedHiddenStatesBatch(tokenIds, caches, validLengths, hook); + MLXArray normalized = + normalizedHiddenStatesBatch( + tokenIds, caches, validLengths, hook, rotation, postBlockHook); requirePostflightBatch(caches, validLengths.length); int[] shape = normalized.shape(); MLXArray last = @@ -327,6 +378,31 @@ private static void requirePostflightBatch(List caches, int batch) { } } + /** + * The post-block hook's contract: it must hand back the block's own hidden state, non-null and + * with the exact shape and dtype it was given, so the next block's projection sees the same + * tensor contract. Enforced immediately after the call, so a violation names the block instead of + * surfacing later inside the next block's projection. + */ + private static void requirePostBlockResult( + int blockIndex, MLXArray hidden, int[] expectedShape, DType expectedDtype) { + if (hidden == null + || !Arrays.equals(expectedShape, hidden.shape()) + || hidden.dtype() != expectedDtype) { + throw new IllegalArgumentException( + "postBlockHook for block " + + blockIndex + + " must return the block's hidden state with the same shape " + + Arrays.toString(expectedShape) + + " and dtype " + + expectedDtype + + ": got " + + (hidden == null + ? "null" + : Arrays.toString(hidden.shape()) + " of dtype " + hidden.dtype())); + } + } + /** * The embedding-input half of the decoder stack: the embedding lookup, the optional hook invoked * at the identical post-lookup/pre-scale point, and the optional sqrt-hidden scale. Both the @@ -363,6 +439,32 @@ private MLXArray decoderEmbeddings(MLXArray tokenIds, EmbeddingHook hook) { */ final MLXArray decoderStack( MLXArray embeddings, List caches, int[] validLengths, int sequenceWidth) { + return decoderStack(embeddings, caches, validLengths, sequenceWidth, null, null); + } + + /** + * The seam-carrying form of {@link #decoderStack(MLXArray, List, int[], int)}: a non-null {@code + * rotation} skips the step frequency derivation and rotates each block's query and key heads from + * the request's own multi-axis rotation instead; a non-null {@code postBlockHook} is invoked + * after every block's residual output and before the final {@code norm.forward} (see {@link + * PostBlockHook}). Its result is checked immediately: a null, differently shaped, or differently + * typed return is rejected as an {@link IllegalArgumentException} naming the block index. The + * batched mask is the same in both cases. A non-null {@code rotation} is validated against the + * batch, the sequence width, the head dimension and the configured rotary width ({@link + * MultiAxisRotation#validateAgainst(int, int, int, int)}) before the mask or any block builds a + * graph node. With {@code (null, null)} this runs the byte-identical stack the 4-argument + * overload always ran: one frequency derivation at the same graph position, the same block + * forwards and the same final normalization. This entry itself still neither evaluates nor + * poisons; embedding-start callers keep the {@link #preflightBatch(int, int, List, int[])}/{@link + * #poisonBatchIfMutated} and postflight/evaluation responsibilities of the 4-argument entry. + */ + final MLXArray decoderStack( + MLXArray embeddings, + List caches, + int[] validLengths, + int sequenceWidth, + MultiAxisRotation rotation, + PostBlockHook postBlockHook) { int[] shape = embeddings.shape(); if (shape.length != 3 || shape[0] != validLengths.length @@ -389,6 +491,10 @@ final MLXArray decoderStack( MLXArray x = embeddings; KVCache first = caches.getFirst(); int batch = validLengths.length; + if (rotation != null) { + // Reject a mismatched rotation before the batched mask or any block builds a graph node. + rotation.validateAgainst(batch, sequenceWidth, descriptor.headDim(), descriptor.rotaryDims()); + } int[] queryStarts = new int[batch]; int[] keyStarts = new int[batch]; int keyWidth = 0; @@ -409,18 +515,43 @@ final MLXArray decoderStack( sequenceWidth, keyWidth, window == null ? 0 : window); + // The null-rotation path derives the step frequencies exactly once, in the same scope and at + // the same graph position as before the seam existed; the rotation path skips the derivation + // and rotates from the request's own cos/sin tables. MLXArray frequencies = - descriptor.rope().stepFrequencies(x.scope(), descriptor.rotaryDims(), maxEnding); + rotation == null + ? descriptor.rope().stepFrequencies(x.scope(), descriptor.rotaryDims(), maxEnding) + : null; for (int i = 0; i < layers.size(); i++) { - x = layers.get(i).forward(x, caches.get(i), mask, frequencies, validLengths); + x = + rotation == null + ? layers.get(i).forward(x, caches.get(i), mask, frequencies, validLengths) + : layers.get(i).forward(x, caches.get(i), mask, validLengths, rotation); + if (postBlockHook != null) { + // MLXArray.shape() returns a fresh array on every call, so the hook cannot alias it. + int[] expected = x.shape(); + DType expectedDtype = x.dtype(); + x = postBlockHook.apply(i, x); + requirePostBlockResult(i, x, expected, expectedDtype); + } } return norm.forward(x); } private MLXArray normalizedHiddenStatesBatch( MLXArray tokenIds, List caches, int[] validLengths, EmbeddingHook hook) { + return normalizedHiddenStatesBatch(tokenIds, caches, validLengths, hook, null, null); + } + + private MLXArray normalizedHiddenStatesBatch( + MLXArray tokenIds, + List caches, + int[] validLengths, + EmbeddingHook hook, + MultiAxisRotation rotation, + PostBlockHook postBlockHook) { MLXArray x = decoderEmbeddings(tokenIds, hook); - return decoderStack(x, caches, validLengths, tokenIds.shape()[1]); + return decoderStack(x, caches, validLengths, tokenIds.shape()[1], rotation, postBlockHook); } private int[] preflight(MLXArray tokenIds, List caches) { @@ -514,33 +645,65 @@ private static MLXArray[] cacheArrays(List caches) { } private MLXArray normalizedHiddenStates(MLXArray tokenIds, List caches) { + return normalizedHiddenStates(tokenIds, caches, null, null); + } + + /** + * The single-row seam-carrying form of {@link #normalizedHiddenStates(MLXArray, List)}. With + * {@code (null, null)} it runs the old single-row graph byte for byte: the 2-D windowed mask (or + * the implicit causal one), the derived step frequencies and the non-batch block forwards. A + * non-null rotation or hook routes the row through the batch-aware overloads instead: the {@link + * #decoderStack(MLXArray, List, int[], int, MultiAxisRotation, PostBlockHook)} entry with one + * uniform valid length per row. With only a hook, that route is the same arithmetic as the old + * single-row path for a batch-1 unpadded row: the uniform append and the {@code [1,1,T,L]} + * batched mask match the single-row ones element for element, windowed and evicted included + * (pinned by {@code SingleRowDecodeEquivalenceTest}), the per-row RoPE form delegates to the + * offset form for a uniform row, and the explicit mask yields the same SDPA result as the + * implicit causal one -- pinned bit-exact at the logits, end to end, by {@code + * DecoderSeamTest.decodeWithNoOpPostBlockHookMatchesOldSingleRowPath} -- plus the hook firing + * after every block's residual output, before the final normalization. A non-null rotation + * deliberately changes the arithmetic instead: it rotates from the request's own cos/sin tables + * rather than the rope fast kernel, within the 1e-5 that {@code + * DecoderSeamTest.equalAxisRotationMatchesRopePathWithinKernelTolerance} pins. + */ + private MLXArray normalizedHiddenStates( + MLXArray tokenIds, + List caches, + MultiAxisRotation rotation, + PostBlockHook postBlockHook) { for (int i = 0; i < caches.size(); i++) { Objects.requireNonNull(caches.get(i), "cache " + i); } MLXArray x = decoderEmbeddings(tokenIds, null); - Integer window = descriptor.attention().slidingWindow(); - int keyLength = caches.get(0).length() + tokenIds.shape()[1]; - MLXArray mask = - window != null && (window < keyLength || caches.get(0).startPosition() > 0) - ? AttentionMask.slidingWindow( - x.scope(), - caches.get(0).nextPosition(), - caches.get(0).startPosition(), - tokenIds.shape()[1], - keyLength, - window) - : null; - MLXArray stepFrequencies = - descriptor - .rope() - .stepFrequencies( - x.scope(), - descriptor.rotaryDims(), - caches.get(0).nextPosition() + tokenIds.shape()[1]); - for (int i = 0; i < layers.size(); i++) { - x = layers.get(i).forward(x, caches.get(i), mask, stepFrequencies); + if (rotation == null && postBlockHook == null) { + Integer window = descriptor.attention().slidingWindow(); + int keyLength = caches.get(0).length() + tokenIds.shape()[1]; + MLXArray mask = + window != null && (window < keyLength || caches.get(0).startPosition() > 0) + ? AttentionMask.slidingWindow( + x.scope(), + caches.get(0).nextPosition(), + caches.get(0).startPosition(), + tokenIds.shape()[1], + keyLength, + window) + : null; + MLXArray stepFrequencies = + descriptor + .rope() + .stepFrequencies( + x.scope(), + descriptor.rotaryDims(), + caches.get(0).nextPosition() + tokenIds.shape()[1]); + for (int i = 0; i < layers.size(); i++) { + x = layers.get(i).forward(x, caches.get(i), mask, stepFrequencies); + } + return norm.forward(x); } - return norm.forward(x); + int sequence = tokenIds.shape()[1]; + int[] validLengths = new int[tokenIds.shape()[0]]; + Arrays.fill(validLengths, sequence); + return decoderStack(x, caches, validLengths, sequence, rotation, postBlockHook); } /** @@ -640,6 +803,18 @@ interface StepFailureGuard { void onStepFailure(List caches); } + /** + * The decode step {@link #runStepLoop} uses for every step at 1 or higher, returning the lazy + * {@code [1,1,V]} logits of the new token. The 8-argument entry's shared default is this model's + * own single-row {@code decode} path; a composing model (the Qwen3-VL work, WP6/WP7, this + * package) passes a step that calls the seam-carrying {@code decode} overload with its request's + * rotary state. The loop's own preflight, poisoning and evaluation ordering stay around the call. + */ + @FunctionalInterface + interface DecodeStep { + MLXArray apply(MLXArray ids, List caches); + } + /** * The shared generation step loop behind {@link #generate(GenerationRequest, Consumer)}: the * per-step activation scope, the per-request generation scope, sampler, caches, cancellation @@ -663,6 +838,37 @@ final GenerationResult runStepLoop( KVCachePolicy cachePolicy, Consumer listener, FirstStep firstStep) { + return runStepLoop( + request, + policy, + prompt, + frequencySeed, + promptPositions, + cachePolicy, + listener, + firstStep, + this::decode); + } + + /** + * The seam-carrying form of {@link #runStepLoop(GenerationRequest, GenerationConfig, int[], + * int[], int, KVCachePolicy, Consumer, FirstStep)}: steps at 1 or higher call {@code + * decodeStep.apply(ids, caches)} where the 8-argument entry calls the model's own single-row + * {@code decode}. Everything else -- the per-step activation scope, the first step, the sampler, + * cancellation, token and terminal events, the preflight, poisoning, evaluation ordering and the + * result construction -- is the identical loop body. The existing 8-argument entry delegates with + * {@code this::decode}, so its path is unchanged. + */ + final GenerationResult runStepLoop( + GenerationRequest request, + GenerationConfig policy, + int[] prompt, + int[] frequencySeed, + int promptPositions, + KVCachePolicy cachePolicy, + Consumer listener, + FirstStep firstStep, + DecodeStep decodeStep) { HfTokenizer tokenizer = request.tokenizer(); List generated = new ArrayList<>(); List logProbabilities = new ArrayList<>(); @@ -698,7 +904,7 @@ final GenerationResult runStepLoop( } else { MLXArray ids = MLX.array(activation, input, new int[] {1, input.length}); before = preflight(ids, caches); - lastLogits = decode(ids, caches); + lastLogits = decodeStep.apply(ids, caches); } selection = sampler.select( @@ -785,10 +991,31 @@ final GenerationResult runStepLoop( } private MLXArray decode(MLXArray ids, List caches) { + return decode(ids, caches, null, null); + } + + /** + * The single-row decode step with request-supplied rotary state and post-block injection: the + * same one-token and populated-cache validation as {@link #decode(MLXArray, List)}, then the + * seam-carrying stack and the same postflight and tied projection. A non-null {@code rotation} is + * validated against the batch, the sequence width, the head dimension and the configured rotary + * width ({@link MultiAxisRotation#validateAgainst(int, int, int, int)}) before the embedding + * lookup builds any graph node. Package-private (like the other seam entries) so a {@link + * DecodeStep} and the in-package seam tests can call it directly; production ordinary decoders + * reach it through the 8-argument {@link #runStepLoop} delegation with {@code (null, null)} and + * keep the old graph. + */ + final MLXArray decode( + MLXArray ids, List caches, MultiAxisRotation rotation, PostBlockHook postBlockHook) { if (ids.shape()[1] != 1 || caches.get(0).nextPosition() == 0) { throw new IllegalArgumentException("decode requires one token and a populated cache"); } - MLXArray normalized = normalizedHiddenStates(ids, caches); + if (rotation != null) { + // Reject a mismatched rotation before the embedding lookup builds any graph node. + rotation.validateAgainst( + ids.shape()[0], ids.shape()[1], descriptor.headDim(), descriptor.rotaryDims()); + } + MLXArray normalized = normalizedHiddenStates(ids, caches, rotation, postBlockHook); requirePostflight(caches); return tiedOutput ? embedding.project(normalized) : lmHead.forward(normalized); } diff --git a/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java new file mode 100644 index 00000000..673573d4 --- /dev/null +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java @@ -0,0 +1,194 @@ +package se.alipsa.jmlx.models; + +import java.util.Arrays; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.core.MLXOps; +import se.alipsa.jmlx.core.MLXShape; +import se.alipsa.jmlx.memory.MLXScope; + +/** + * Frozen WP4 visual-injection compositions for the Qwen3-VL text decoder (req/plans/ + * phase7-3b-reference-findings.md §6.9): the row-gather replacement that swaps the + * image-placeholder embeddings for the merged visual features, and the masked residual add that + * injects one DeepStack level into the hidden states at exactly the image-token rows (its add + * matches the pinned reference's masked clone + add + write-back, §6.6). + * + *

Both methods are pure compositions of existing core facades ({@code MLX.zeros}, {@code + * MLXShape.concatenate}, {@code MLXShape.takeAxis}, {@code MLXOps.add}) and never call FFI; no new + * primitive is required. Both compositions reuse {@link Idefics3VisionTower# + * mergeGatherIndices(int, java.util.List)} — the shared gather-index algorithm, where feature row + * {@code f} overwrites exactly the {@code f}-th destination and every other row reads itself. The + * row-gather replacement applies it directly; the masked add remaps its output over a source of a + * single zero row followed by the feature rows: every identity row (shared index {@code i < L}) + * remaps to the zero row (source index {@code 0}) and the {@code f}-th destination (shared index + * {@code L + f}) remaps to source row {@code 1 + f}, which is what lets its zero pad be one {@code + * [1, 1, hidden]} row instead of a full {@code [1, L, hidden]} zero tensor: a full-sequence zero + * allocation plus an {@code (L + F)}-wide concatenation per DeepStack level (three at prefill) + * would otherwise materialise for values that are exactly unchanged. All validation is checked + * before any MLX array is written, and every non-destination row comes back unchanged: the + * replacement leaves its untouched rows bit-identical, and the masked add is an exact {@code + 0.0} + * per element there, so its non-destination rows come back numerically identical (note IEEE {@code + * -0.0 + 0.0 = +0.0}, so a negative-zero element in a non-destination row is normalised to {@code + * +0.0}). The destinations follow the existing semantics: a flat, strictly increasing list in + * {@code [0, L)} whose runs are the images' contiguous placeholder ranges in prompt order. WP6's + * {@code Qwen3VlModel} is the consumer; the class is package-private on purpose, with no new public + * API. + */ +final class VisualFeatureInjection { + + private VisualFeatureInjection() {} + + /** + * Row-gather replacement: the frozen {@code source = concatenate([embedded, features], 1); gather + * = mergeGatherIndices(L, destinations); replaced = takeAxis(source, gather, 1)} composition that + * replaces each destination row with the corresponding feature row and leaves every other row + * untouched, the same composition {@code SmolVlmModel.merge} uses. + * + * @param embedded the original {@code [1, L, hidden]} token embeddings + * @param features the {@code [1, F, hidden]} merged visual feature rows, one block per image in + * prompt order + * @param destinations the image-token rows to replace, strictly increasing in {@code [0, L)}, + * whose runs are the images' contiguous placeholder ranges in prompt order + * @return the {@code [1, L, hidden]} embeddings with the destination rows replaced + * @throws IllegalArgumentException if the shapes, dtypes, row count or destinations are invalid, + * before any MLX array is written + */ + static MLXArray replaceImageRows(MLXArray embedded, MLXArray features, int[] destinations) { + int length = validatePair("embedded", embedded, "features", features, destinations); + int[] gather = gatherIndices(length, destinations); + MLXArray source = MLXShape.concatenate(new MLXArray[] {embedded, features}, 1); + return MLXShape.takeAxis( + source, MLX.array(embedded.scope(), gather, new int[] {gather.length}), 1); + } + + /** + * Masked residual add: the frozen {@code zeroRow = MLX.zeros(scope, {1, 1, hidden}, + * hidden.dtype()); source = concatenate([zeroRow, features], 1); mask = takeAxis(source, gather, + * 1); updated = MLXOps.add(hidden, mask)} composition that adds one DeepStack level into exactly + * the destination rows, leaving every other row numerically identical through the exact {@code + + * 0.0} per element (note IEEE {@code -0.0 + 0.0 = +0.0}, so a negative-zero element in an + * untouched row is normalised to {@code +0.0}). The gather indices are the shared {@code + * mergeGatherIndices} output remapped over the source's {@code 1 + F} rows: every identity row + * (shared index {@code i < L}) remaps to the zero row (source index {@code 0}) and the {@code + * f}-th destination (shared index {@code L + f}) remaps to source row {@code 1 + f}, which lets + * the zero pad be a single {@code [1, 1, hidden]} row instead of a full-sequence {@code [1, L, + * hidden]} zero tensor. + * + * @param scope the scope the zero pad and gather indices are allocated in; the results land per + * the standard scope rule (innermost scope among the operands); like the shapes, dtypes, row + * count and destinations, its relation to the operands' scopes (one an ancestor of the other) + * is validated before any MLX array is written + * @param hidden the current {@code [1, L, hidden]} hidden states + * @param features the {@code [1, F, hidden]} DeepStack level rows, in destination order + * @param destinations the image-token rows to add into, strictly increasing in {@code [0, L)}, + * whose runs are the images' contiguous placeholder ranges in prompt order + * @return the {@code [1, L, hidden]} hidden states with the level added into the destination rows + * @throws IllegalArgumentException if the shapes, dtypes, row count or destinations are invalid, + * or any of {@code scope} and the operands' scopes is unrelated (neither an ancestor of the + * other), in every case before any MLX array is written + */ + static MLXArray addMaskedFeatures( + MLXScope scope, MLXArray hidden, MLXArray features, int[] destinations) { + if (scope == null) { + throw new IllegalArgumentException("scope must not be null"); + } + int length = validatePair("hidden", hidden, "features", features, destinations); + checkRelatedScopes("scope", scope, hidden, features); + int[] gather = gatherIndices(length, destinations); + // Remap the shared indices: identity row i -> the zero row (0), destination L + f -> 1 + f. + for (int i = 0; i < gather.length; i++) { + gather[i] = gather[i] >= length ? gather[i] - length + 1 : 0; + } + MLXArray zeroRow = MLX.zeros(scope, new int[] {1, 1, hidden.shape()[2]}, hidden.dtype()); + MLXArray source = MLXShape.concatenate(new MLXArray[] {zeroRow, features}, 1); + MLXArray mask = + MLXShape.takeAxis(source, MLX.array(scope, gather, new int[] {gather.length}), 1); + return MLXOps.add(hidden, mask); + } + + /** + * The shared gather indices: the existing {@code Idefics3VisionTower} algorithm, which validates + * the destinations (strictly increasing in {@code [0, length)}), checks the {@code length + F} + * sum with checked arithmetic and returns the identity row mapping with each destination pointing + * at its feature row. + */ + private static int[] gatherIndices(int length, int[] destinations) { + return Idefics3VisionTower.mergeGatherIndices( + length, IntStream.of(destinations).boxed().collect(Collectors.toList())); + } + + /** + * The masked add's pre-write scope-compatibility check, applying {@link MLXScope#innermost} — the + * same relatedness test the ops' own {@code scopeOf} resolution applies — to the three pairs the + * method's chain requires to be related (one an ancestor of the other): {@code scope} against + * each operand (the final {@code add} resolves against {@code hidden}, {@code concatenate} + * against {@code features}), and the two operand scopes against each other. The operand-vs. + * operand pair is not implied by the other two: when {@code scope} is a common ancestor of both + * operand scopes, sibling operand scopes pass both operand-vs.-scope checks yet fail at the final + * {@code add}, whose mask lives in the deeper of {@code scope} and {@code features.scope()} — + * after the zero row was already allocated, breaking the pre-write validation promise. + * + * @throws IllegalArgumentException naming {@code paramName} if any of the three pairs is + * unrelated + */ + private static void checkRelatedScopes( + String paramName, MLXScope paramScope, MLXArray base, MLXArray features) { + try { + MLXScope.innermost(paramScope, base.scope()); + MLXScope.innermost(paramScope, features.scope()); + MLXScope.innermost(base.scope(), features.scope()); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException( + paramName + + " must be related to hidden's and features' scopes (one an ancestor of the other)", + e); + } + } + + private static int validatePair( + String baseName, MLXArray base, String featuresName, MLXArray features, int[] destinations) { + if (base == null) { + throw new IllegalArgumentException(baseName + " must not be null"); + } + if (features == null) { + throw new IllegalArgumentException(featuresName + " must not be null"); + } + if (destinations == null) { + throw new IllegalArgumentException("destinations must not be null"); + } + if (base.ndim() != 3 || base.shape()[0] != 1) { + throw new IllegalArgumentException( + baseName + " must be a [1, L, hidden] array, got " + Arrays.toString(base.shape())); + } + int length = base.shape()[1]; + if (features.ndim() != 3 + || features.shape()[0] != 1 + || features.shape()[2] != base.shape()[2] + || features.dtype() != base.dtype()) { + throw new IllegalArgumentException( + featuresName + + " must be a [1, F, " + + base.shape()[2] + + "] array of dtype " + + base.dtype() + + " to match " + + baseName + + ", got " + + Arrays.toString(features.shape()) + + " of dtype " + + features.dtype()); + } + if (features.shape()[1] != destinations.length) { + throw new IllegalArgumentException( + featuresName + + " row count " + + features.shape()[1] + + " must equal the destination count " + + destinations.length); + } + return length; + } +} diff --git a/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java new file mode 100644 index 00000000..e8e99d07 --- /dev/null +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java @@ -0,0 +1,451 @@ +package se.alipsa.jmlx.models; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.core.MLXOps; +import se.alipsa.jmlx.core.MLXShape; +import se.alipsa.jmlx.ffi.EnabledIfNativeAvailable; +import se.alipsa.jmlx.memory.MLXScope; +import se.alipsa.jmlx.nn.AttentionMask; +import se.alipsa.jmlx.nn.KVCache; +import se.alipsa.jmlx.nn.MultiAxisRope; +import se.alipsa.jmlx.nn.MultiAxisRotation; +import se.alipsa.jmlx.nn.RopeSpec; + +/** + * Seam tests for the frozen WP4 decoder integration (req/plans/phase7-3b-reference-findings.md + * §6.9) on the pinned llama checkpoint: the request-supplied {@link MultiAxisRotation} threaded + * through the model's single-row {@code decode} step and the batched {@code stepLogits} seam with a + * post-block hook, and the explicit {@code DecodeStep} form of {@code runStepLoop}. Every float32 + * case compares the seam against an independent manual composition over a forked cache set, built + * from a parallel {@link DecoderAssembler} component set with identical weight values and the same + * operation order and scopes, so those comparisons run at zero tolerance. The hook-only case + * ({@link #decodeWithNoOpPostBlockHookMatchesOldSingleRowPath}) instead compares the batched route + * against the old single-row path itself, bit-exact at the logits. The cases are tagged {@code + * full-float32} so they select into {@code :jmlx-models:float32GoldenTest} with {@code + * MLX_ENABLE_TF32=0}; the single exception is the fast-kernel-versus-explicit-rotation comparison + * in {@link #equalAxisRotationMatchesRopePathWithinKernelTolerance}, which uses the documented 1e-5 + * tolerance. The rejection cases run in the ordinary {@code :jmlx-models:test} task. + */ +@EnabledIfNativeAvailable +class DecoderSeamTest { + + private static final int[] SECTIONS = {4, 2, 2}; // rotaryDims 16 == llama headDim + + private static MLXScope scope; + private static DecoderModel model; + private static DecoderAssembler.Assembled assembled; + private static MultiAxisRope rope; + + @BeforeAll + static void loadCheckpoint() throws Exception { + scope = new MLXScope(); + model = (DecoderModel) TextGenerationModels.load(scope, SchedulerFixtures.checkpoint("llama")); + Path directory = SchedulerFixtures.checkpoint("llama"); + ArchitectureDescriptor descriptor = + ArchitectureMappings.parse(JsonFiles.read(directory.resolve("config.json"))); + assembled = + DecoderAssembler.assemble( + scope, + descriptor, + CheckpointLoader.load(scope, directory, ArchitectureMappings.tensorPlan(descriptor))); + rope = new MultiAxisRope(10_000f, SECTIONS); + assertEquals(descriptor.headDim(), rope.rotaryDims()); + assertEquals(descriptor.rotaryDims(), rope.rotaryDims()); + } + + @AfterAll + static void closeCheckpoint() { + scope.close(); + } + + // --------------------------------------------------------------------- + // float32-comparison cases (full-float32 tag, MLX_ENABLE_TF32=0) + // --------------------------------------------------------------------- + + /** + * A rotation prepared at cache position 82 with the pinned negative delta (82 - 56 = 26 on all + * three axes) must decode the next token to logits bit-identical to the manual composition over a + * forked cache set: the embedding, every block's rotation-threaded forward, the final + * normalization and the head, in the same scope. + */ + @Test + @Tag("full-float32") + void decodeWithNegativeDeltaRotationMatchesForkedReference() { + int[] prompt = tokens(82); + try (MLXScope run = scope.newChild()) { + List caches = caches(run); + model.forward(MLX.array(run, prompt, new int[] {1, prompt.length}), caches); + try (MLXScope gen = run.newChild()) { + List forks = forks(caches, gen); + MLXArray ids = MLX.array(gen, new int[] {token(82)}, new int[] {1, 1}); + MultiAxisRotation rotation = rope.prepare(gen, 1, 1, new int[][] {{26}, {26}, {26}}); + MLXArray mask = + AttentionMask.batched(gen, new int[] {82}, new int[] {0}, new int[] {1}, 1, 83, 0); + MLXArray x = assembled.embedding().forward(ids); + for (int i = 0; i < model.config().numHiddenLayers(); i++) { + x = assembled.layers().get(i).forward(x, forks.get(i), mask, new int[] {1}, rotation); + } + x = assembled.finalNorm().forward(x); + MLXArray expected = assembled.lmHead().forward(x); + MLXArray actual = model.decode(ids, caches, rotation, null); + assertArrayEquals(expected.toFloatArray(), actual.toFloatArray()); + assertEquals(83, caches.getFirst().nextPosition()); + assertEquals(83, forks.getFirst().nextPosition()); + } + } + } + + /** + * The seam overload with {@code (null, null)} must run the old single-row decode graph byte for + * byte: the manual reference is the same embedding, the same step-frequency derivation (null for + * the llama {@code Base} rope), the same four-argument block forwards with the implicit causal + * mask, the same final normalization and head, over a forked cache set. + */ + @Test + @Tag("full-float32") + void decodeWithNullSeamsMatchesOldSingleRowGraph() { + int[] prompt = tokens(40); + try (MLXScope run = scope.newChild()) { + List caches = caches(run); + model.forward(MLX.array(run, prompt, new int[] {1, prompt.length}), caches); + try (MLXScope gen = run.newChild()) { + List forks = forks(caches, gen); + MLXArray ids = MLX.array(gen, new int[] {token(40)}, new int[] {1, 1}); + MLXArray x = assembled.embedding().forward(ids); + MLXArray stepFrequencies = new RopeSpec.Base(10_000f).stepFrequencies(gen, 16, 41); + for (int i = 0; i < model.config().numHiddenLayers(); i++) { + x = assembled.layers().get(i).forward(x, forks.get(i), null, stepFrequencies); + } + x = assembled.finalNorm().forward(x); + MLXArray expected = assembled.lmHead().forward(x); + MLXArray actual = model.decode(ids, caches, null, null); + assertArrayEquals(expected.toFloatArray(), actual.toFloatArray()); + assertEquals(41, caches.getFirst().nextPosition()); + } + } + } + + /** + * The seam-carrying batched {@code stepLogits} must apply the rotation in every block and fire + * the post-block hook after every block's residual output and before the final normalization: the + * manual reference mirrors both, on a forked cache set, and the last-row logits match at zero + * tolerance. + */ + @Test + @Tag("full-float32") + void stepLogitsWithRotationAndPostBlockHookMatchesForkedReference() { + int[] prompt = tokens(6); + int layers = model.config().numHiddenLayers(); + int hidden = model.config().hiddenSize(); + float[][] pert = new float[layers][]; + for (int layer = 0; layer < layers; layer++) { + pert[layer] = new float[hidden]; + for (int j = 0; j < hidden; j++) { + pert[layer][j] = (j % 5 - 2) * 1e-3f * (layer + 1); + } + } + try (MLXScope run = scope.newChild()) { + List caches = caches(run); + try (MLXScope gen = run.newChild()) { + List forks = forks(caches, gen); + MLXArray ids = MLX.array(gen, prompt, new int[] {1, prompt.length}); + MultiAxisRotation rotation = + rope.prepare( + gen, + 1, + 6, + new int[][] {{0, 1, 2, 3, 4, 5}, {0, 2, 4, 2, 4, 6}, {1, 1, 3, 3, 5, 5}}); + MLXArray mask = + AttentionMask.batched(gen, new int[] {0}, new int[] {0}, new int[] {6}, 6, 6, 0); + MLXArray x = assembled.embedding().forward(ids); + for (int i = 0; i < layers; i++) { + x = assembled.layers().get(i).forward(x, forks.get(i), mask, new int[] {6}, rotation); + x = MLXOps.add(x, MLX.array(x.scope(), pert[i], new int[] {1, 1, hidden})); + } + x = assembled.finalNorm().forward(x); + MLXArray last = MLXShape.slice(x, new int[] {0, 5, 0}, new int[] {1, 6, hidden}); + MLXArray expected = assembled.lmHead().forward(last); + MLXArray actual = + model.stepLogits( + ids, + caches, + new int[] {6}, + null, + rotation, + (layerIndex, hiddenStates) -> + MLXOps.add( + hiddenStates, + MLX.array( + hiddenStates.scope(), pert[layerIndex], new int[] {1, 1, hidden}))); + assertArrayEquals(expected.toFloatArray(), actual.toFloatArray()); + assertEquals(6, caches.getFirst().nextPosition()); + } + } + } + + /** + * The explicit {@code DecodeStep} form of {@code runStepLoop} -- a first step running the plain + * batched prefill and every later step calling the seam's {@code decode} with {@code (null, + * null)} -- must generate exactly the tokens {@code generate} produces for the same request: the + * greedy selection is deterministic, and the loop body, sampler and cache policy are the + * identical loop. + */ + @Test + @Tag("full-float32") + void runStepLoopWithExplicitDecodeStepMatchesGenerate() { + int[] prompt = tokens(10); + GenerationRequest request = SchedulerFixtures.greedy(prompt, 8); + GenerationResult baseline = model.generate(request, ignored -> {}); + DecoderModel.FirstStep firstStep = + (activation, stepCaches) -> { + MLXArray ids = MLX.array(activation, prompt, new int[] {1, prompt.length}); + int[][] before = + model.preflightBatch(1, prompt.length, stepCaches, new int[] {prompt.length}); + MLXArray logits = model.stepLogits(ids, stepCaches, new int[] {prompt.length}, null); + return new DecoderModel.FirstStepResult( + logits, cs -> DecoderModel.poisonBatchIfMutated(cs, before, before[0].length)); + }; + GenerationResult explicit = + model.runStepLoop( + request, + request.config(), + prompt, + prompt, + prompt.length, + model.resolveCachePolicy(request.cachePolicy()), + ignored -> {}, + firstStep, + (ids, caches) -> model.decode(ids, caches, null, null)); + assertEquals(baseline.finishReason(), explicit.finishReason()); + assertEquals(baseline.generatedTokenIds(), explicit.generatedTokenIds()); + } + + /** + * With equal axes on all three rows the interleaved multi-axis rotation is mathematically the + * plain 1-D RoPE of the llama rope path, so the two {@code stepLogits} entries must agree within + * the documented 1e-5 tolerance. They are not expected to be bit-equal: the rope path runs the + * {@code mlx_fast_rope} kernel while the rotation path composes the explicit cos/sin tables, and + * the two do not share a bitwise-exact last-ulp contract (Unit 2 findings, §6.9). + */ + @Test + @Tag("full-float32") + void equalAxisRotationMatchesRopePathWithinKernelTolerance() { + int[] prompt = tokens(7); + int[] axis = {0, 1, 2, 3, 4, 5, 6}; + try (MLXScope run = scope.newChild()) { + List ropeCaches = caches(run); + List rotationCaches = caches(run); + MLXArray ids = MLX.array(run, prompt, new int[] {1, prompt.length}); + MLXArray ropeLogits = model.stepLogits(ids, ropeCaches, new int[] {7}, null); + MultiAxisRotation rotation = + rope.prepare(run, 1, 7, new int[][] {axis, axis.clone(), axis.clone()}); + MLXArray rotationLogits = + model.stepLogits(ids, rotationCaches, new int[] {7}, null, rotation, null); + assertArrayEquals(ropeLogits.toFloatArray(), rotationLogits.toFloatArray(), 1e-5f); + } + } + + /** + * The hook-only single-row decode -- rotation {@code null}, the hook a no-op lambda that returns + * its input unchanged -- must be bit-identical at the logits to the old single-row path for the + * same prompt, weights and cache contents. Both cache sets are prefilled by the same old-path + * forward (and asserted bit-equal before decoding), then every one-token decode step is compared + * element for element in raw IEEE bits and the greedy choices must stay in lockstep. This pins + * the {@code normalizedHiddenStates} claim that the hook-only batched route is "the same + * arithmetic as the old single-row path" end to end: for a uniform row the per-row RoPE form + * delegates to the offset form, and the explicit {@code [1,1,1,L]} mask selects the same SDPA + * result as the old path's implicit causal flag -- observed max absolute error 0.0, i.e. all 4 + * steps' 128-vocabulary logits bit-equal, full-float32. + */ + @Test + @Tag("full-float32") + void decodeWithNoOpPostBlockHookMatchesOldSingleRowPath() { + int[] prompt = tokens(32); + int steps = 4; + try (MLXScope run = scope.newChild()) { + List plain = caches(run); + List hooked = caches(run); + model.forward(MLX.array(run, prompt, new int[] {1, prompt.length}), plain); + model.forward(MLX.array(run, prompt, new int[] {1, prompt.length}), hooked); + for (int i = 0; i < plain.size(); i++) { + assertArrayEquals(plain.get(i).keys().toFloatArray(), hooked.get(i).keys().toFloatArray()); + assertArrayEquals( + plain.get(i).values().toFloatArray(), hooked.get(i).values().toFloatArray()); + } + try (MLXScope gen = run.newChild()) { + int input = token(32); + for (int step = 0; step < steps; step++) { + MLXArray ids = MLX.array(gen, new int[] {input}, new int[] {1, 1}); + MLXArray expected = model.decode(ids, plain, null, null); + MLXArray actual = model.decode(ids, hooked, null, (layer, hidden) -> hidden); + float[] oldLogits = expected.toFloatArray(); + float[] hookLogits = actual.toFloatArray(); + assertArrayEquals( + rawBits(oldLogits), + rawBits(hookLogits), + "step " + step + ": the no-op-hook decode must be bit-identical to the old path"); + int oldGreedy = argMax(oldLogits); + assertEquals( + oldGreedy, + argMax(hookLogits), + "step " + step + ": the greedy choice must stay in lockstep"); + input = oldGreedy; + } + assertEquals(prompt.length + steps, plain.getFirst().nextPosition()); + assertEquals(prompt.length + steps, hooked.getFirst().nextPosition()); + } + } + } + + /** The raw IEEE bits of an evaluated row, so {@code -0.0} cannot pass for {@code +0.0}. */ + private static int[] rawBits(float[] values) { + int[] bits = new int[values.length]; + for (int i = 0; i < values.length; i++) { + bits[i] = Float.floatToIntBits(values[i]); + } + return bits; + } + + private static int argMax(float[] values) { + int best = 0; + for (int i = 1; i < values.length; i++) { + if (values[i] > values[best]) { + best = i; + } + } + return best; + } + + // --------------------------------------------------------------------- + // rejection cases (ordinary test task, TF32=1) + // --------------------------------------------------------------------- + + /** + * The seam's {@code decode} must keep the original one-token and populated-cache validation: an + * empty cache is rejected before any array is written, and a two-token step is rejected with the + * populated caches left unadvanced. + */ + @Test + void decodeRotationOverloadRejectsInvalidShapesBeforeCacheMutation() { + try (MLXScope run = scope.newChild()) { + List empty = caches(run); + try (MLXScope gen = run.newChild()) { + MLXArray ids = MLX.array(gen, new int[] {5}, new int[] {1, 1}); + MultiAxisRotation rotation = rope.prepare(gen, 1, 1, new int[][] {{0}, {0}, {0}}); + IllegalArgumentException failure = + assertThrows( + IllegalArgumentException.class, () -> model.decode(ids, empty, rotation, null)); + assertEquals("decode requires one token and a populated cache", failure.getMessage()); + } + for (KVCache cache : empty) { + assertEquals(0, cache.nextPosition()); + } + + List populated = caches(run); + model.forward(MLX.array(run, tokens(4), new int[] {1, 4}), populated); + try (MLXScope gen = run.newChild()) { + MLXArray twoTokens = MLX.array(gen, new int[] {1, 2}, new int[] {1, 2}); + MultiAxisRotation rotation = rope.prepare(gen, 1, 1, new int[][] {{4}, {4}, {4}}); + assertThrows( + IllegalArgumentException.class, + () -> model.decode(twoTokens, populated, rotation, null)); + assertEquals(4, populated.getFirst().nextPosition()); + } + } + } + + /** + * A rotation prepared for a different sequence width must be rejected by the seam entries before + * any array is written: {@code stepLogits} rejects it before the embedding lookup, {@code decode} + * before its own embedding lookup, and the populated caches are left unadvanced. So must a + * rotation prepared from a different mrope-section sum whose prefix fits the head dimension but + * differs from the descriptor's configured rotary width: {@code stepLogits}, the embedding-start + * {@code decoderStack} and {@code decode} all reject it before the mask or any embedding lookup + * builds a graph node. + */ + @Test + void seamEntriesRejectWrongWidthRotationBeforeArraysAreWritten() { + int[] prompt = tokens(4); + try (MLXScope run = scope.newChild()) { + List populated = caches(run); + model.forward(MLX.array(run, prompt, new int[] {1, 4}), populated); + try (MLXScope gen = run.newChild()) { + MultiAxisRotation wrongWidth = + rope.prepare(gen, 1, 2, new int[][] {{0, 1}, {2, 3}, {4, 5}}); + MLXArray ids = MLX.array(gen, new int[] {token(4)}, new int[] {1, 1}); + assertThrows( + IllegalArgumentException.class, + () -> model.stepLogits(ids, populated, new int[] {1}, null, wrongWidth, null)); + assertThrows( + IllegalArgumentException.class, () -> model.decode(ids, populated, wrongWidth, null)); + // A rotation prepared from a different mrope-section sum: its prefix (14) fits the llama + // head dimension (16) but differs from the descriptor's configured rotary width (16), so + // only the width-mismatch check rejects it -- at each seam entry, before any graph node. + MultiAxisRope narrowRope = new MultiAxisRope(10_000f, new int[] {4, 2, 1}); // rotaryDims 14 + MultiAxisRotation wrongPrefix = narrowRope.prepare(gen, 1, 1, new int[][] {{4}, {4}, {4}}); + MLXArray embeddings = assembled.embedding().forward(ids); + IllegalArgumentException wrongPrefixStepLogits = + assertThrows( + IllegalArgumentException.class, + () -> model.stepLogits(ids, populated, new int[] {1}, null, wrongPrefix, null)); + assertTrue( + wrongPrefixStepLogits + .getMessage() + .contains("does not match the configured rotary width")); + assertThrows( + IllegalArgumentException.class, + () -> model.decoderStack(embeddings, populated, new int[] {1}, 1, wrongPrefix, null)); + assertThrows( + IllegalArgumentException.class, () -> model.decode(ids, populated, wrongPrefix, null)); + assertEquals(4, populated.getFirst().nextPosition()); + } + } + } + + // --------------------------------------------------------------------- + // helpers + // --------------------------------------------------------------------- + + private static List caches(MLXScope target) { + List caches = new ArrayList<>(); + for (int i = 0; i < model.config().numHiddenLayers(); i++) { + caches.add(new KVCache(target)); + } + return caches; + } + + /** One forked (evaluated) copy per cache, in the destination scope. */ + private static List forks(List caches, MLXScope destination) { + List forks = new ArrayList<>(); + for (KVCache cache : caches) { + forks.add(cache.fork(destination)); + } + return forks; + } + + private static int[] tokens(int n) { + int[] ids = new int[n]; + for (int i = 0; i < n; i++) { + ids[i] = token(i); + } + return ids; + } + + /** Deterministic in-vocabulary token at position {@code i}. */ + private static int token(int i) { + return (i * 7 + 3) % 128; + } +} diff --git a/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSlidingWindowTest.java b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSlidingWindowTest.java index e96f99fc..f06eceec 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSlidingWindowTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSlidingWindowTest.java @@ -14,8 +14,10 @@ import org.junit.jupiter.api.io.TempDir; import se.alipsa.jmlx.core.MLX; import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.core.MLXShape; import se.alipsa.jmlx.ffi.EnabledIfNativeAvailable; import se.alipsa.jmlx.memory.MLXScope; +import se.alipsa.jmlx.nn.AttentionMask; import se.alipsa.jmlx.nn.KVCache; import se.alipsa.jmlx.nn.KVCachePolicy; import tools.jackson.databind.JsonNode; @@ -100,6 +102,108 @@ void explicitSlidingRetentionMatchesFullMaskedAttention(@TempDir Path dir) throw } } + /** + * The hook-only single-row route -- rotation {@code null}, a no-op post-block hook -- is + * bit-identical at the logits to the old single-row path on a sliding-window model whose caches + * have already evicted keys (startPosition > 0): the old path's 2-D sliding-window mask and + * the batched route's {@code [1,1,T,L]} mask match element for element, so both routes select the + * same SDPA result. Four one-token {@code decode} steps pin the evicted-cache decode against the + * old path itself (every retained key within the window), and a final two-token step pins a step + * whose key length exceeds the window so the masks really mask keys: its old-path reference is + * composed manually over a forked cache set and sliced to the last row before the head, so both + * sides project the same shape (the {@code forward} and {@code stepLogits} entries differ in that + * projection shape, which shifts a few ulps in the head matmul alone). + */ + @Test + void noOpHookDecodeMatchesOldSingleRowPathOnEvictedSlidingCaches() throws Exception { + Path root = + Path.of(System.getProperty("jmlx.repository.root"), "tools", "hf-reference", "goldens"); + Path directory = root.resolve("checkpoints/mistral"); + try (MLXScope modelScope = new MLXScope()) { + DecoderModel model = MistralModel.load(modelScope, directory); + ArchitectureDescriptor descriptor = + ArchitectureMappings.parse(JsonFiles.read(directory.resolve("config.json"))); + DecoderAssembler.Assembled assembled = + DecoderAssembler.assemble( + modelScope, + descriptor, + CheckpointLoader.load( + modelScope, directory, ArchitectureMappings.tensorPlan(descriptor))); + int[] prompt = new int[8]; + for (int i = 0; i < prompt.length; i++) { + prompt[i] = token(i); + } + try (MLXScope inference = modelScope.newChild()) { + List plain = slidingCaches(inference, model); + List hooked = slidingCaches(inference, model); + model.forward(MLX.array(inference, prompt, new int[] {1, prompt.length}), plain); + model.forward(MLX.array(inference, prompt, new int[] {1, prompt.length}), hooked); + for (int i = 0; i < plain.size(); i++) { + assertArrayEquals( + plain.get(i).keys().toFloatArray(), hooked.get(i).keys().toFloatArray()); + assertArrayEquals( + plain.get(i).values().toFloatArray(), hooked.get(i).values().toFloatArray()); + } + // The 8-token prompt (window 4) evicted: three retained keys per layer, first at 5. + assertEquals(5, plain.getFirst().startPosition()); + assertEquals(8, plain.getFirst().nextPosition()); + try (MLXScope gen = inference.newChild()) { + int input = token(8); + for (int step = 0; step < 4; step++) { + MLXArray ids = MLX.array(gen, new int[] {input}, new int[] {1, 1}); + MLXArray expected = model.decode(ids, plain, null, null); + MLXArray actual = model.decode(ids, hooked, null, (layer, hidden) -> hidden); + float[] oldLogits = expected.toFloatArray(); + float[] hookLogits = actual.toFloatArray(); + assertArrayEquals( + rawBits(oldLogits), + rawBits(hookLogits), + "step " + step + ": the no-op-hook decode must be bit-identical to the old path"); + int oldGreedy = argMax(oldLogits); + assertEquals(oldGreedy, argMax(hookLogits), "step " + step + ": greedy lockstep"); + input = oldGreedy; + } + // A two-token step: keyLength 5 exceeds the window 4, so the masks really mask keys. + int[] twoTokens = {input, token(12)}; + MLXArray ids = MLX.array(gen, twoTokens, new int[] {1, twoTokens.length}); + List forks = new ArrayList<>(); + for (KVCache cache : plain) { + forks.add(cache.fork(gen)); + } + int next = plain.getFirst().nextPosition(); + int keyStart = plain.getFirst().startPosition(); + int keyLength = plain.getFirst().length() + twoTokens.length; + Integer window = descriptor.attention().slidingWindow(); + MLXArray x = assembled.embedding().forward(ids); + MLXArray mask = + AttentionMask.slidingWindow(gen, next, keyStart, twoTokens.length, keyLength, window); + MLXArray stepFrequencies = + descriptor + .rope() + .stepFrequencies(gen, descriptor.rotaryDims(), next + twoTokens.length); + int layers = model.config().numHiddenLayers(); + for (int i = 0; i < layers; i++) { + x = assembled.layers().get(i).forward(x, forks.get(i), mask, stepFrequencies); + } + x = assembled.finalNorm().forward(x); + MLXArray last = + MLXShape.slice(x, new int[] {0, 1, 0}, new int[] {1, 2, model.config().hiddenSize()}); + MLXArray expected = assembled.lmHead().forward(last); + MLXArray actual = + model.stepLogits( + ids, hooked, new int[] {twoTokens.length}, null, null, (layer, hidden) -> hidden); + assertArrayEquals( + rawBits(expected.toFloatArray()), + rawBits(actual.toFloatArray()), + "the two-token masked step must be bit-identical to the old graph"); + assertEquals(14, hooked.getFirst().nextPosition()); + assertEquals(11, hooked.getFirst().startPosition()); + assertEquals(3, hooked.getFirst().length()); + } + } + } + } + @Test void slidingGenerationCancellationStopsBetweenSteps(@TempDir Path dir) throws Exception { TinyCheckpoints.randomLlama(dir, 1, 2, false, false); @@ -139,6 +243,30 @@ private static List slidingCaches(MLXScope scope, DecoderModel model) { return caches; } + /** The raw IEEE bits of an evaluated row, so {@code -0.0} cannot pass for {@code +0.0}. */ + private static int[] rawBits(float[] values) { + int[] bits = new int[values.length]; + for (int i = 0; i < values.length; i++) { + bits[i] = Float.floatToIntBits(values[i]); + } + return bits; + } + + private static int argMax(float[] values) { + int best = 0; + for (int i = 1; i < values.length; i++) { + if (values[i] > values[best]) { + best = i; + } + } + return best; + } + + /** Deterministic in-vocabulary token at position {@code i}. */ + private static int token(int i) { + return (i * 7 + 3) % 128; + } + private static int[] ints(JsonNode node) { int[] result = new int[node.size()]; for (int i = 0; i < result.length; i++) { diff --git a/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderStackContractTest.java b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderStackContractTest.java index 2bedfb7a..3469aff0 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderStackContractTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderStackContractTest.java @@ -9,6 +9,7 @@ import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; +import se.alipsa.jmlx.core.DType; import se.alipsa.jmlx.core.MLX; import se.alipsa.jmlx.core.MLXArray; import se.alipsa.jmlx.ffi.EnabledIfNativeAvailable; @@ -20,8 +21,9 @@ * DecoderModel#decoderStack}) and for its shape-keyed preflight {@link * DecoderModel#preflightBatch(int, int, List, int[])}: the checks added with the entry point must * reject a wrong rank, batch, width or hidden size and a wrong cache count before any layer - * advances, a correctly shaped call must run the whole stack, and the preflight itself must keep - * the poisoned-cache and width-vs-valid-length guarantees a WP6 caller relies on. + * advances, a post-block hook returning a null, differently shaped, or differently typed state must + * be rejected naming the block, a correctly shaped call must run the whole stack, and the preflight + * itself must keep the poisoned-cache and width-vs-valid-length guarantees a WP6 caller relies on. */ @EnabledIfNativeAvailable class DecoderStackContractTest { @@ -238,6 +240,119 @@ void correctShapeRunsTheWholeStack() { } } + /** + * A post-block hook that returns a differently shaped state must be rejected immediately after + * the first block, naming the block index and the expected versus returned shape, instead of + * surfacing later inside the next block's projection. + */ + @Test + void postBlockHookReturningWrongShapeIsRejectedNamingTheBlock() { + try (MLXScope run = scope.newChild()) { + List caches = caches(run); + int hidden = model.config().hiddenSize(); + DecoderModel.PostBlockHook hook = (layer, state) -> embeddings(run, BATCH, WIDTH, hidden + 1); + IllegalArgumentException failure = + assertThrows( + IllegalArgumentException.class, + () -> + model.decoderStack( + embeddings(run, BATCH, WIDTH, hidden), + caches, + validLengths(), + WIDTH, + null, + hook)); + assertEquals( + "postBlockHook for block 0 must return the block's hidden state with the same shape [" + + BATCH + + ", " + + WIDTH + + ", " + + hidden + + "] and dtype FLOAT32: got [" + + BATCH + + ", " + + WIDTH + + ", " + + (hidden + 1) + + "] of dtype FLOAT32", + failure.getMessage()); + } + } + + /** + * A post-block hook that returns null is rejected by the same check, naming the block index and + * the expected shape and dtype. + */ + @Test + void postBlockHookReturningNullIsRejectedNamingTheBlock() { + try (MLXScope run = scope.newChild()) { + List caches = caches(run); + int hidden = model.config().hiddenSize(); + DecoderModel.PostBlockHook hook = (layer, state) -> null; + IllegalArgumentException failure = + assertThrows( + IllegalArgumentException.class, + () -> + model.decoderStack( + embeddings(run, BATCH, WIDTH, hidden), + caches, + validLengths(), + WIDTH, + null, + hook)); + assertEquals( + "postBlockHook for block 0 must return the block's hidden state with the same shape [" + + BATCH + + ", " + + WIDTH + + ", " + + hidden + + "] and dtype FLOAT32: got null", + failure.getMessage()); + } + } + + /** + * A post-block hook that returns a same-shape but differently typed state must be rejected + * immediately after the first block, naming the block index and the expected versus returned + * dtype, instead of silently retyping every later block's projection input. + */ + @Test + void postBlockHookReturningWrongDtypeIsRejectedNamingTheBlock() { + try (MLXScope run = scope.newChild()) { + List caches = caches(run); + int hidden = model.config().hiddenSize(); + DecoderModel.PostBlockHook hook = (layer, state) -> MLX.astype(state, DType.FLOAT16); + IllegalArgumentException failure = + assertThrows( + IllegalArgumentException.class, + () -> + model.decoderStack( + embeddings(run, BATCH, WIDTH, hidden), + caches, + validLengths(), + WIDTH, + null, + hook)); + assertEquals( + "postBlockHook for block 0 must return the block's hidden state with the same shape [" + + BATCH + + ", " + + WIDTH + + ", " + + hidden + + "] and dtype FLOAT32: got [" + + BATCH + + ", " + + WIDTH + + ", " + + hidden + + "] of dtype FLOAT16", + failure.getMessage()); + } + } + @Test void shapePreflightRejectsPoisonedCache() { try (MLXScope run = scope.newChild()) { diff --git a/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java new file mode 100644 index 00000000..442f4314 --- /dev/null +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -0,0 +1,491 @@ +package se.alipsa.jmlx.models; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.junit.jupiter.api.AfterAll; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.function.Executable; +import se.alipsa.jmlx.core.DType; +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.core.MLXMemoryTestSupport; +import se.alipsa.jmlx.ffi.EnabledIfNativeAvailable; +import se.alipsa.jmlx.memory.MLXScope; + +/** + * Strict tests for the frozen WP4 visual-injection compositions (req/plans/ + * phase7-3b-reference-findings.md §6.9): the row-gather replacement and the masked residual add. + * Every float32 case builds the expected result by hand in Java and compares it at zero tolerance + * after {@code toFloatArray()} forces evaluation: the destination rows must equal the feature rows + * bit-for-bit, every other row must equal the original row, and the masked add's exact {@code + + * 0.0} on the non-image rows stays observable (the zero-feature case uses exactly representable + * values such as {@code -1.5f}, where any non-exact add would change bits; the negative-zero case + * records that an untouched {@code -0.0} is normalised to {@code +0.0} while every non-zero value + * keeps its exact bits). The float32 cases are tagged {@code full-float32} so they select into + * {@code :jmlx-models:float32GoldenTest} with {@code MLX_ENABLE_TF32=0}; the rejection cases run in + * the ordinary {@code :jmlx-models:test} task and expect {@link IllegalArgumentException} with the + * exact message from the helper's own validation, before any MLX array is written — each case also + * pins the unchanged settled process-wide active-byte count across the rejected call. + */ +@EnabledIfNativeAvailable +class VisualFeatureInjectionTest { + + private static final int HIDDEN = 4; + + private static MLXScope scope; + + @BeforeAll + static void openScope() { + scope = new MLXScope(); + } + + @AfterAll + static void closeScope() { + scope.close(); + } + + // --------------------------------------------------------------------- + // row-gather replacement (full-float32) + // --------------------------------------------------------------------- + + /** A single three-row image in the middle of the sequence replaces exactly its rows. */ + @Test + @Tag("full-float32") + void singleImageMidSequenceReplacesExactlyItsRows() { + try (MLXScope run = scope.newChild()) { + int length = 8; + int[] destinations = {3, 4, 5}; + float[][] embedded = rows(length, 0f); + float[][] features = rows(3, 10f); + MLXArray replaced = + VisualFeatureInjection.replaceImageRows( + array(run, embedded), array(run, features), destinations); + assertArrayEquals(new int[] {1, length, HIDDEN}, replaced.shape()); + assertEquals(run, replaced.scope()); + assertArrayEquals( + flat(replacementRows(embedded, features, destinations)), replaced.toFloatArray()); + } + } + + /** An image starting at row 0 replaces the leading rows. */ + @Test + @Tag("full-float32") + void imageAtSequenceStartReplacesLeadingRows() { + try (MLXScope run = scope.newChild()) { + int length = 8; + int[] destinations = {0, 1, 2}; + float[][] embedded = rows(length, 0f); + float[][] features = rows(3, 10f); + MLXArray replaced = + VisualFeatureInjection.replaceImageRows( + array(run, embedded), array(run, features), destinations); + assertArrayEquals( + flat(replacementRows(embedded, features, destinations)), replaced.toFloatArray()); + } + } + + /** An image ending at the last row replaces the trailing rows. */ + @Test + @Tag("full-float32") + void imageAtSequenceEndReplacesTrailingRows() { + try (MLXScope run = scope.newChild()) { + int length = 8; + int[] destinations = {5, 6, 7}; + float[][] embedded = rows(length, 0f); + float[][] features = rows(3, 10f); + MLXArray replaced = + VisualFeatureInjection.replaceImageRows( + array(run, embedded), array(run, features), destinations); + assertArrayEquals( + flat(replacementRows(embedded, features, destinations)), replaced.toFloatArray()); + } + } + + /** Two differently sized images, placed in prompt order, both replace exactly their rows. */ + @Test + @Tag("full-float32") + void twoDifferentlySizedPlacementsReplaceBothImages() { + try (MLXScope run = scope.newChild()) { + int length = 10; + int[] destinations = {1, 2, 6, 7, 8}; + float[][] embedded = rows(length, 0f); + float[][] features = combine(rows(2, 10f), rows(3, 20f)); + MLXArray replaced = + VisualFeatureInjection.replaceImageRows( + array(run, embedded), array(run, features), destinations); + assertArrayEquals( + flat(replacementRows(embedded, features, destinations)), replaced.toFloatArray()); + } + } + + // --------------------------------------------------------------------- + // masked residual add (full-float32) + // --------------------------------------------------------------------- + + /** One DeepStack level adds into exactly the destination rows; every other row is untouched. */ + @Test + @Tag("full-float32") + void singleLevelAddsFeaturesIntoExactlyTheImageRows() { + try (MLXScope run = scope.newChild()) { + int length = 8; + int[] destinations = {2, 3, 4}; + float[][] hidden = rows(length, 0f); + float[][] features = rows(3, 10f); + MLXArray updated = + VisualFeatureInjection.addMaskedFeatures( + run, array(run, hidden), array(run, features), destinations); + assertArrayEquals(new int[] {1, length, HIDDEN}, updated.shape()); + assertArrayEquals(flat(addRows(hidden, features, destinations)), updated.toFloatArray()); + } + } + + /** + * Scattered (non-contiguous) destinations: single-row images at rows 1, 4 and 6 each receive + * exactly their own feature row, and every other row is the hidden row plus the exact {@code + + * 0.0}. + */ + @Test + @Tag("full-float32") + void scatteredDestinationsAddIntoExactlyTheirRows() { + try (MLXScope run = scope.newChild()) { + int length = 8; + int[] destinations = {1, 4, 6}; + float[][] hidden = rows(length, 0f); + float[][] features = rows(3, 10f); + MLXArray updated = + VisualFeatureInjection.addMaskedFeatures( + run, array(run, hidden), array(run, features), destinations); + assertArrayEquals(new int[] {1, length, HIDDEN}, updated.shape()); + assertArrayEquals(flat(addRows(hidden, features, destinations)), updated.toFloatArray()); + } + } + + /** + * Two levels at different destinations applied sequentially, mirroring the reference's per-level + * masked clone + add + write-back: level 0 into rows 1-2, then level 1 into rows 5-7, with the + * second add running on the first add's output. + */ + @Test + @Tag("full-float32") + void twoLevelsApplySequentiallyAtDifferentDestinations() { + try (MLXScope run = scope.newChild()) { + int length = 8; + float[][] hidden = rows(length, 0f); + float[][] level0 = rows(2, 10f); + float[][] level1 = rows(3, 20f); + MLXArray afterLevel0 = + VisualFeatureInjection.addMaskedFeatures( + run, array(run, hidden), array(run, level0), new int[] {1, 2}); + float[][] expectedAfterLevel0 = addRows(hidden, level0, new int[] {1, 2}); + assertArrayEquals(flat(expectedAfterLevel0), afterLevel0.toFloatArray()); + MLXArray afterLevel1 = + VisualFeatureInjection.addMaskedFeatures( + run, afterLevel0, array(run, level1), new int[] {5, 6, 7}); + assertArrayEquals( + flat(addRows(expectedAfterLevel0, level1, new int[] {5, 6, 7})), + afterLevel1.toFloatArray()); + } + } + + /** + * Zero-valued feature rows are an exact {@code + 0.0}: every row, image rows included, comes back + * bit-identical. The hidden values are exactly representable (multiples of 0.5 around {@code + * -1.5f}), so a non-exact addition would change bits and fail the comparison. + */ + @Test + @Tag("full-float32") + void zeroFeatureRowsAreAnExactPlusZeroIdentity() { + try (MLXScope run = scope.newChild()) { + int length = 8; + float[][] hidden = new float[length][HIDDEN]; + for (int r = 0; r < length; r++) { + for (int c = 0; c < HIDDEN; c++) { + hidden[r][c] = -1.5f + (r + c) * 0.5f; + } + } + MLXArray features = MLX.zeros(run, new int[] {1, 2, HIDDEN}, DType.FLOAT32); + MLXArray updated = + VisualFeatureInjection.addMaskedFeatures( + run, array(run, hidden), features, new int[] {2, 3}); + assertArrayEquals(flat(hidden), updated.toFloatArray()); + } + } + + /** + * The masked add's exact {@code + 0.0} on a non-image row is an identity numerically, not + * bit-for-bit: an IEEE {@code -0.0} element is normalised to {@code +0.0} (raw bits {@code + * 0x00000000}), while the row's non-zero values — negative non-zero included — keep their exact + * bits. + */ + @Test + @Tag("full-float32") + void maskedAddNormalisesNegativeZeroInAnUntouchedRow() { + try (MLXScope run = scope.newChild()) { + int length = 8; + float[][] hidden = rows(length, 0f); + hidden[5][0] = -0.0f; + hidden[5][1] = -1.5f; + hidden[5][2] = 2.25f; + hidden[5][3] = -0.125f; + assertEquals(0x80000000, Float.floatToIntBits(hidden[5][0])); + float[][] features = rows(2, 10f); + int[] destinations = {3, 4}; + MLXArray updated = + VisualFeatureInjection.addMaskedFeatures( + run, array(run, hidden), array(run, features), destinations); + float[] out = updated.toFloatArray(); + float[] expected = new float[length * HIDDEN]; + for (int r = 0; r < length; r++) { + boolean destination = destinations[0] == r || destinations[1] == r; + for (int c = 0; c < HIDDEN; c++) { + expected[r * HIDDEN + c] = + destination ? hidden[r][c] + features[r - destinations[0]][c] : hidden[r][c] + 0.0f; + } + } + assertArrayEquals(expected, out); + // The -0.0 element is normalised to +0.0: raw bits 0x00000000, not 0x80000000. + assertEquals(0x00000000, Float.floatToIntBits(out[5 * HIDDEN])); + // The row's non-zero values keep their exact bits. + assertEquals(Float.floatToIntBits(-1.5f), Float.floatToIntBits(out[5 * HIDDEN + 1])); + assertEquals(Float.floatToIntBits(2.25f), Float.floatToIntBits(out[5 * HIDDEN + 2])); + assertEquals(Float.floatToIntBits(-0.125f), Float.floatToIntBits(out[5 * HIDDEN + 3])); + } + } + + // --------------------------------------------------------------------- + // rejections (ordinary test task, MLX_ENABLE_TF32=1) + // --------------------------------------------------------------------- + + /** + * Every invalid input is rejected with {@link IllegalArgumentException} by the helper's own + * validation — destination semantics via the gather-index builders (the shared {@code + * mergeGatherIndices} and the masked add's mirror), count and shape checks via the pair + * validation — before any MLX array is written; a native failure from a mis-built index array + * would surface as an {@code MLXException}, not an {@code IllegalArgumentException}. Each case + * pins the exact rejection message and the "before any array is written" promise itself: the + * settled process-wide active-byte count ({@link MLXMemoryTestSupport#settledActiveBytes()}) is + * sampled before and after the rejected call and must be unchanged, so a stray allocation — the + * masked add's zero row in particular — would fail the test. + */ + @Test + void rejectionsThrowBeforeAnyArrayIsWritten() throws InterruptedException { + try (MLXScope run = scope.newChild()) { + final MLXArray embedded = array(run, rows(8, 0f)); + final MLXArray features = array(run, rows(3, 10f)); + final MLXArray twoRows = array(run, rows(2, 10f)); + final MLXArray twoDim = array(run, flat(rows(3, 10f)), new int[] {3, HIDDEN}); + final MLXArray badBatch = MLX.array(run, flat(rows(6, 10f)), new int[] {2, 3, HIDDEN}); + final MLXArray badHiddenDim = MLX.zeros(run, new int[] {1, 3, HIDDEN + 1}, DType.FLOAT32); + final MLXArray half = MLX.astype(features, DType.FLOAT16); + + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {3, 4, 8}), + "feature destination 8 is out of range [0, 8)"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {-1, 3, 4}), + "feature destination -1 is out of range [0, 8)"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {4, 3, 5}), + "feature destinations must be strictly increasing; index 1 has 3 after 4"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, features, null), + "destinations must not be null"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(null, features, new int[] {3, 4, 5}), + "embedded must not be null"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, null, new int[] {3, 4, 5}), + "features must not be null"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {3, 4}), + "features row count 3 must equal the destination count 2"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, twoRows, new int[] {3, 4, 5}), + "features row count 2 must equal the destination count 3"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, twoDim, new int[] {3, 4, 5}), + "features must be a [1, F, 4] array of dtype FLOAT32 to match embedded, " + + "got [3, 4] of dtype FLOAT32"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, badBatch, new int[] {3, 4, 5}), + "features must be a [1, F, 4] array of dtype FLOAT32 to match embedded, " + + "got [2, 3, 4] of dtype FLOAT32"); + assertRejection( + () -> + VisualFeatureInjection.replaceImageRows(embedded, badHiddenDim, new int[] {3, 4, 5}), + "features must be a [1, F, 4] array of dtype FLOAT32 to match embedded, " + + "got [1, 3, 5] of dtype FLOAT32"); + assertRejection( + () -> VisualFeatureInjection.replaceImageRows(embedded, half, new int[] {3, 4, 5}), + "features must be a [1, F, 4] array of dtype FLOAT32 to match embedded, " + + "got [1, 3, 4] of dtype FLOAT16"); + + assertRejection( + () -> + VisualFeatureInjection.addMaskedFeatures( + null, embedded, features, new int[] {3, 4, 5}), + "scope must not be null"); + assertRejection( + () -> VisualFeatureInjection.addMaskedFeatures(run, null, features, new int[] {3, 4, 5}), + "hidden must not be null"); + assertRejection( + () -> VisualFeatureInjection.addMaskedFeatures(run, embedded, null, new int[] {3, 4, 5}), + "features must not be null"); + assertRejection( + () -> + VisualFeatureInjection.addMaskedFeatures(run, twoDim, features, new int[] {3, 4, 5}), + "hidden must be a [1, L, hidden] array, got [3, 4]"); + assertRejection( + () -> VisualFeatureInjection.addMaskedFeatures(run, embedded, half, new int[] {3, 4, 5}), + "features must be a [1, F, 4] array of dtype FLOAT32 to match hidden, " + + "got [1, 3, 4] of dtype FLOAT16"); + assertRejection( + () -> + VisualFeatureInjection.addMaskedFeatures(run, embedded, twoRows, new int[] {3, 4, 5}), + "features row count 2 must equal the destination count 3"); + } + } + + /** + * Runs {@code call}, pinning both halves of the pre-write rejection contract: the exact {@link + * IllegalArgumentException} message and the unchanged settled process-wide active-byte count + * across the call (a plain read would race in-flight Metal completion-handler frees; the settled + * read is immune, and a stray allocation still moves it). + */ + private static void assertRejection(Executable call, String message) throws InterruptedException { + long before = MLXMemoryTestSupport.settledActiveBytes(); + IllegalArgumentException thrown = assertThrows(IllegalArgumentException.class, call); + assertEquals(message, thrown.getMessage()); + assertEquals(before, MLXMemoryTestSupport.settledActiveBytes(), "rejection must not allocate"); + } + + /** + * A sibling scope (the operands live in one child of the class scope, the parameter scope is the + * other child) is unrelated to both operand scopes — neither an ancestor of the other — and is + * rejected by the pre-write scope check, before the zero row or the gather index array is + * allocated; the message names the parameter. + */ + @Test + void addMaskedFeaturesRejectsAUnrelatedSiblingScope() throws InterruptedException { + try (MLXScope child = scope.newChild(); + MLXScope sibling = scope.newChild()) { + MLXArray hidden = array(child, rows(8, 0f)); + MLXArray features = array(child, rows(3, 10f)); + assertRejection( + () -> + VisualFeatureInjection.addMaskedFeatures( + sibling, hidden, features, new int[] {3, 4, 5}), + "scope must be related to hidden's and features' scopes (one an ancestor of the other)"); + } + } + + /** + * The scope is the common ancestor of both operands: each operand-vs.-scope pair is related (the + * operand's child scope is inside the parameter scope) and passes, so the rejection can only come + * from the operand-vs.-operand pair — the two sibling operand scopes — which the pre-write scope + * check still catches before the zero row or the gather index array is allocated. + */ + @Test + void addMaskedFeaturesRejectsSiblingOperandsUnderACommonParentScope() + throws InterruptedException { + try (MLXScope a = scope.newChild(); + MLXScope b = scope.newChild()) { + MLXArray hidden = array(a, rows(8, 0f)); + MLXArray features = array(b, rows(3, 10f)); + assertRejection( + () -> + VisualFeatureInjection.addMaskedFeatures( + scope, hidden, features, new int[] {3, 4, 5}), + "scope must be related to hidden's and features' scopes (one an ancestor of the other)"); + } + } + + // --------------------------------------------------------------------- + // data builders (Java-side expectations) + // --------------------------------------------------------------------- + + /** + * The replacement result by hand: destination rows take the feature rows, all others keep theirs. + */ + private static float[][] replacementRows( + float[][] embedded, float[][] features, int[] destinations) { + float[][] result = new float[embedded.length][]; + for (int row = 0; row < embedded.length; row++) { + float[] source = embedded[row]; + for (int f = 0; f < destinations.length; f++) { + if (destinations[f] == row) { + source = features[f]; + } + } + result[row] = source.clone(); + } + return result; + } + + /** + * The masked-add result by hand: destination rows are the exact float sums, all others keep + * theirs. + */ + private static float[][] addRows(float[][] hidden, float[][] features, int[] destinations) { + float[][] result = new float[hidden.length][]; + for (int row = 0; row < hidden.length; row++) { + float[] source = hidden[row]; + for (int f = 0; f < destinations.length; f++) { + if (destinations[f] == row) { + float[] sum = new float[HIDDEN]; + for (int c = 0; c < HIDDEN; c++) { + sum[c] = hidden[row][c] + features[f][c]; + } + source = sum; + } + } + result[row] = source.clone(); + } + return result; + } + + /** The {@code [1, rows, HIDDEN]} float32 array holding {@code values}. */ + private static MLXArray array(MLXScope scope, float[][] values) { + return MLX.array(scope, flat(values), new int[] {1, values.length, HIDDEN}); + } + + private static MLXArray array(MLXScope scope, float[] flat, int[] shape) { + return MLX.array(scope, flat, shape); + } + + /** + * Deterministic values at {@code (r, c)} of {@code offset + (r+1)/4 + (c+1)/16} — every value + * exactly representable in float32. + */ + private static float[][] rows(int rows, float offset) { + float[][] values = new float[rows][]; + for (int r = 0; r < rows; r++) { + values[r] = new float[HIDDEN]; + for (int c = 0; c < HIDDEN; c++) { + values[r][c] = offset + (r + 1) * 0.25f + (c + 1) * 0.0625f; + } + } + return values; + } + + private static float[][] combine(float[][] first, float[][] second) { + float[][] values = new float[first.length + second.length][]; + System.arraycopy(first, 0, values, 0, first.length); + System.arraycopy(second, 0, values, first.length, second.length); + return values; + } + + private static float[] flat(float[][] values) { + float[] flat = new float[values.length * HIDDEN]; + int i = 0; + for (float[] row : values) { + System.arraycopy(row, 0, flat, i, HIDDEN); + i += HIDDEN; + } + return flat; + } +} diff --git a/req/plans/phase7-3b-plan.md b/req/plans/phase7-3b-plan.md index e327f12a..1e09e9a6 100644 --- a/req/plans/phase7-3b-plan.md +++ b/req/plans/phase7-3b-plan.md @@ -1,8 +1,10 @@ # Phase 7.3b — Qwen3-VL image and text inference -- **Status:** WP0 and WP1 complete, 2026-10-15 — baseline, artifact pins, executable reference - probes and divergence evidence are recorded in `req/plans/phase7-3b-reference-findings.md`. - WP3–WP6 implementation and acceptance remain work below. +- **Status:** WP0–WP4 complete: WP0/WP1 2026-10-15 (reference findings), WP2 accepted + 2026-10-18 per PR #50, WP3 accepted 2026-10-19 per PR #55, WP4 completed 2026-10-20 + (local acceptance evidence in `req/plans/phase7-3b-reference-findings.md` §6.9's + acceptance block; acceptance rides this branch's PR). WP5–WP8 implementation and + acceptance remain work below. - **Parent:** `req/plans/phase7-plan.md` §7.3b, public decisions 1–4, 6–8 and Cross-cutting verification. - **Dependencies:** accepted 7.0 Qwen3 backbone and completed 7.3a multimodal contracts, @@ -397,12 +399,43 @@ update inventory overrides/call-site coverage and test scope/error behavior. Do add general indexing APIs. After the refactor, require exact old-path bit and greedy-ID equality under the same reproducibility tuple, not just tolerance-golden passes. +**2026-10-20 WP4 status:** WP4 is completed locally (the acceptance record is +`req/plans/phase7-3b-reference-findings.md` §6.9's acceptance block): the frozen core +components `MultiAxisRope` (validated `[3,B,T]` positions, the interleaved frequency / +rotate-half kernel with the pinned compute/cast order, equal and index-coded axes, +`cachePosition + delta` decode coordinates, partial-rotary tail passthrough, and +pre-write `IllegalArgumentException` rejections) and its `MultiAxisRotation` +family specification; the additive decoder seam — request-supplied rotary state and the +post-block injection callback across `DecoderAttention`/`DecoderBlock`/`DecoderModel`, +existing overloads preserved and routed through the identical old graph, production +decoders on no-op hooks — with the old-path parity suites; and the frozen +`VisualFeatureInjection` compositions in `se.alipsa.jmlx.models` (row-gather +replacement and masked DeepStack residual add over existing core facades; the +replacement reuses `Idefics3VisionTower`'s gather-index algorithm, the masked add +the same algorithm with its output remapped over a single zero row, no new primitive, +no FFI from model code, package-private with no new public API). Old-path evidence: the pre-change +baseline recorded at `128cd1e` in both TF32 modes then verified bit-exact +(`exactBitsVerifyTf32On`/`Off`, 14 variants × 9 captures, git-commit metadata the only +drift), `RopeSpec.java` and `MLXFast.java` empty diffs, root `check` and +`verifyMlxApiCallSites` green, and the pinned-vendor probe pair still matching (Python +exit 0 on all 7 cases; Java twin per-case max error ≤ 2.4e-7 against the 1e-5 +threshold from Unit 2, its component untouched since). Non-goals held: no public core +API, no build/CI change — the required-execution suite list addition landed in PR #57 +(review round 2) — and no `Qwen3VlModel` or loader dispatch; those are WP5/WP6. + ## 7. Vision tower and checkpoint assembly (WP5–WP6) Add `Qwen3VlModel implements TextGenerationModel`, composing the decoder so the scheduler's existing non-DecoderModel check rejects it. Dispatch `qwen3_vl` before generic decoder config parsing; explicitly reject `qwen3_vl_moe`. Metadata advertises TEXT and IMAGE, without exposing MLXArray results. Reuse shared generation mechanics and the 7.3a state/accounting contracts. +Derive the decoder's configured rotary width (`descriptor.rotaryDims()`) from `head_dim`, not +from the sum of `mrope_section`, and reject the config before model construction when +`2 * sum(mrope_section) != head_dim`: `MultiAxisRope` derives both its rotated width and its +inverse-frequency denominator from the section sum, and its partial-rotary tail passthrough — the +unrotated slots beyond `2 * sum(mrope_section)` — is a jmlx extension, not reference behaviour, so +a config violating the equality would load and silently diverge from the pinned reference (recorded +in the PR #57 review of the plan §6 component; plan §6, findings §6.9). Build family-specific config/descriptor/tensor plans for vision and nested text. Verify required, optional and forbidden names, shapes, biases, aliases and all unexpected tensors before model @@ -440,6 +473,17 @@ compatible artifact and record why. Synthetic quantized tests compare against th dequantized weights. Never compare quantized weights against unrelated HF float goldens or silently dequantize the entire 8B model to satisfy the manual run. +**2026-10-21 PR #57 review addition:** the WP4 seam is correct but its decode performance is +unmeasured: `normalizedHiddenStates(..., rotation, postBlockHook)` routes every single-row step +with a non-null rotation or hook through `decoderStack`, which passes an explicit `[1,1,1,L]` +BOOL batched mask to SDPA where the old path used `mask == null` (implicit causal). Qwen3-VL +decode always needs a rotation (`cachePosition + delta`), so once `Qwen3VlModel` exists every +generated token takes the masked path, and the bit-exact seam suites cover correctness but not +per-step cost. Before WP6 acceptance, run a one-off `:jmlx-examples:benchmarkDecode` comparison +of the seam path against the old implicit-causal path and record the result in +`req/phase6-4-benchmark.md` with the complete reproducibility tuple, as the other decode +benchmarks do. + ## 8. Request lifecycle and cache policies (WP7) Adopt parent decision 8's 2026-10-08 lifetime amendment, also adopted by 7.3a. @@ -580,7 +624,7 @@ SmolVLM fixtures are unchanged. | WP3 | Expansion, templates, three-axis position planner | WP0–WP2 | Pure tests and HF template/expanded-ID/position goldens | | WP4 | Core multi-axis RoPE and decoder injection seam | WP0, WP1, WP3 | Strict rotary tests; exact before/after decoder bits in both modes | | WP5 | Vision tower, main/DeepStack mergers | WP1, WP2, core primitives | HF intermediate-feature references | -| WP6 | Loader, composed model and affine scope | WP3–WP5 | Complete tensor validation; float/quantized synthetic loading tests | +| WP6 | Loader, composed model and affine scope | WP3–WP5 | Complete tensor validation; float/quantized synthetic loading tests; one-off benchmark of the seam's explicit SDPA mask against the implicit-causal decode path | | WP7 | Generation, ownership and cache accounting | WP6 | Cached/full equivalence, lifecycle/failure/memory tests | | WP8 | Tier-B, manual 8B, CI and documentation | WP7 | Parent numerical/artifact/performance/matrix gates | | WP9 | EXIF orientation policy decision (pre-WP2) and rotated-fixture parity (post-WP2 implementation) | WP1 | Policy recorded before WP2 acceptance; rotated-fixture parity with the slow processor once the Qwen path exists; SmolVLM fixtures unchanged | @@ -603,8 +647,8 @@ obtain supported macOS 26 execution evidence. Existing native skip gates remain - [x] Slow-PIL and effective-area deviations, processor-construction probe and guarded fast-comparison lock accepted. - [x] Locked `mlx-vlm` reference and separate HF/MLX provenance tensor plans accepted. - [ ] Qwen decode option, unchanged default conversion and exact preprocessing/prompt/position/template fixtures pass. -- [ ] Old one-dimensional decoder path remains bit-identical in both precision modes. -- [ ] Qwen kernel matrix, delta-shifted decode tests and unchanged `RopeSpec` variants verified. +- [x] Old one-dimensional decoder path remains bit-identical in both precision modes. +- [x] Qwen kernel matrix, delta-shifted decode tests and unchanged `RopeSpec` variants verified. - [ ] Vision/merger/DeepStack intermediate references and quantized synthetic tests pass. - [ ] All four Tier-A cases and forced generated-marker continuations pass strict/default cached/full-forward checks. - [ ] Resource validation, request reuse, cancellation/failure cleanup and memory bounds pass. diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 7369d731..09f4f4b8 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -712,6 +712,391 @@ no extra tokens. of this machine's terminal display stripping the 0x7c byte from displayed text (the true sentinels contain it). The hex-verified spellings above supersede those claims. +### 6.9 Core multi-axis RoPE and decoder integration seam (WP4 design freeze, 2026-10-20) + +This section freezes the WP4 design (plan §6) so the implementation units proceed without +re-deciding. Base commit at implementation start: `main` at +`128cd1ec2fb1004487015377183081827a29b956` ("Merge pull request #56 from +Alipsa/checkstyle-zero-warnings"); native pins `mlx-metal==0.31.2` / mlx-c +`fba4470b89073180056c9ea46c443051375f7399` (`native/install/lib/native-pin.properties`); +host Mac OS X 27.0.1, arm64, Apple M5 Max (MLX default device `gpu`). + +**Pre-change exact-bits baseline (recorded 2026-10-20, before any decoder or RoPE change).** +`./gradlew :jmlx-models:exactBitsRecord` (the 7.3a recorder, reused unchanged; one JVM per +mode with explicit `MLX_ENABLE_TF32=1` / `=0`) PASSES at the base commit with a clean tree: + +- spec hash `846396b1dbab5a770a4f2c36841fe5ed1570a586b908956ed6e224fe5248e687` (both modes; + identical to the 7.3a baseline — no spec constant changed), recorded git commit + `128cd1ec2fb1004487015377183081827a29b956`, `gitDirty: false`, + `decoderModelSha256 dee4004857eed4a08633f1c44fb3cca5cda7f90148c4a2f4ca0e47305e5809f9`; +- host tuple: device `gpu`, `Mac OS X` `27.0.1`, arm64, `mlxMetalVersion 0.31.2`, + `mlxcCommit fba4470b89073180056c9ea46c443051375f7399`; +- variants captured: 14 applicable per mode — llama, llama-q4, llama31, mistral, mistral-q4, + mixtral, phi3, phi3-q4, qwen2, qwen2-q4, qwen3, qwen3-q4, gemma, gemma-q4 (the 8 Tier-A + families plus the 6 derived q4/group-32 variants; mixtral is MoE and has no quantized + variant) with 9 captures each (126 per mode, 2,048 floats per variant): `direct-prefill` + (`model.forward` prefill logits), `decode-step-0/1` (single-row decode), `batch-prefill`, + `batch-decode-0/1`, `batch-prefill-hook` (batched `stepLogits` with the deterministic +1e-3 + row-1 `EmbeddingHook`, alongside the no-hook `batch-prefill`), `greedy` (exact generated + IDs) and `scheduler` (the gated end-to-end 2-request run, per-row IDs plus finish reason). +- Recordings: `jmlx-models/build/exact-bits/tf32-{1,0}/` (the module build directory the + tasks write to; never committed, wiped by `gradle clean`, device-specific PR evidence). The + WP4 units keep these trees in place: `:jmlx-models:exactBitsVerifyTf32{On,Off}` runs + against them on this machine. The stale 7.3a trees on disk (recorded at a dirty, earlier + commit) were deleted first, as the recorder's own re-record contract requires (a baseline + is never silently replaced). +- **No recorder extension is needed:** WP4 touches the decoder only additively (existing + overloads keep the identical graph by delegating with `null`), and the spec's captures + already exercise every path the change can reach — direct prefill bits, per-step decode + bits, batched `stepLogits` with and without the hook, exact greedy IDs and the full + scheduler run over all 8 families plus the quantized variants — with verify mode + (commit-only metadata drift allowed) as the bit-identity gate. + +**Per-path kernel matrix (frozen).** + +- Qwen3-VL — multimodal prefill, cached decode, text-only prefill and fresh full forward + alike — uses the dedicated explicit FLOAT32 frequency/cos/sin and half-rotation path below + (`MultiAxisRope`/`MultiAxisRotation`), a pure composition of existing core facades: no new + FFI call and no `mlx_fast_rope*` on this path. +- Ordinary decoders (llama, llama31, mistral, mixtral, phi3, qwen2, qwen3, gemma) and the + T5/encoder attention paths keep the unchanged `mlx_fast_rope`/`mlx_fast_rope_dynamic` fast + kernels via the existing `RopeSpec.apply` path (and the existing direct `MLXFast.rope` + call sites in `MultiHeadAttention`/`GroupedQueryAttention`). Absent multimodal inputs + route through the identical old graph: every new overload delegates with `null` + rotation/hook, and no existing overload is modified. +- No fast-kernel optimization of the Qwen3-VL path in this WP: the fast kernels' + offset-plus-arange contract cannot express unequal multimodal coordinates, and any later + optimization requires a dated amendment with shifted-offset probes, strict/default + cached/full equivalence and performance evidence (plan §6) — never a silent kernel mix + between the compared paths. + +**Compute/cast order (frozen; transcribed from §6.5).** + +1. Per-axis float32 frequencies: for d ∈ {0=T, 1=H, 2=W}, + `freqs[d][t, i] = float32(positions[d][t]) · invFreq[i]` — the position ids cast to + float32 (exact for the non-negative coordinates) and a single float32 multiply per + frequency, a length-1 reduction, exact. `invFreq[i] = 1 / theta^(2i/rotaryDims)` is + derived once at construction (pinned reference: `theta = 5,000,000`, + `rotaryDims = 2·Σ mrope_section = 128`, `Σ = 64 = head_dim/2`). +2. Interleaved overwrite starting from the T row (a per-frequency owner selection, bit- + identical to the reference's in-place overwrite): for dim 1 (H) `length = + mrope_section[1]·3` and `idx = slice(1, length, 3)` — H owns frequency indices + `1, 4, 7, …, 1+3(s1−1)` (count s1); for dim 2 (W) `length = mrope_section[2]·3` and + `idx = slice(2, length, 3)` — W owns `2, 5, 8, …, 2+3(s2−1)` (count s2); everything else + stays T (owner(f) = T when `f % 3 == 0` or `f >= 3·(s1+s2)`). Pinned reference + `[24,20,20]`: H owns `1, 4, …, 58` (20 frequencies), W owns `2, 5, …, 59` (20), T owns + the remaining 24 — a T-starting round-robin `T H W T H W …` with a terminal run of four + T frequencies. The implementation realizes the selection with a static owner table: a + single byte array of length `rotaryDims/2` holding one 0/1/2 owner (T/H/W) per frequency + slot, derived from the sections at construction (the H and W conditions are disjoint, so + every slot has exactly one owner) selects, per (token, frequency) on the CPU before any + array is written, the owned row's position — the same copy the mask expression + `freqsT·wT + freqsH·wH + freqsW·wW` is a bit-identical copy of for non-negative finite + frequencies (multiply by 1.0/0.0 and add in float32 are exact there) — so each `freqs_t` + element is still exactly one float32 multiply of the same two operands (owned position · + `invFreq[i]`) using only the existing facades. +3. `emb = cat((freqs_t, freqs_t), −1)` → per token `[rotaryDims]` (128 pinned, = head_dim). +4. `cos = cos(emb) · attention_scaling`, `sin = sin(emb) · attention_scaling`, computed in + float32; `attention_scaling = 1.0` for `rope_type: default` (the pinned artifact), so + the multiply is the identity on finite float32 values and is omitted. +5. Cast cos/sin to the input array's dtype (a no-op for the FLOAT32 model; kept for dtype + correctness of the composition). +6. In the input dtype, `out = x·cos + rotate_half(x)·sin`, where `rotate_half` is the NeoX + half-split form `(−x2, x1)`: `cat((−x[…, n:], x[…, :n]), −1)` — the same layout the fast + kernels use with `traditional=false`. Only the `rotaryDims` prefix of each head is + rotated; the tail `[rotaryDims, headDim)` passes through unchanged (partial rotary). + +**Frozen signatures.** + +New core components in `se.alipsa.jmlx.nn` — a dedicated component, **not** a new +`RopeSpec` variant: the sealed `RopeSpec` hierarchy gains no permitted variant, preserving +source compatibility for exhaustive consumer switches (plan §6). + +- `public final class MultiAxisRope` — immutable; holds theta, the cloned sections and the + derived `rotaryDims`/`invFreq`/owner masks. Deliberately not a `Module`: no parameters, + no children, no owned scope; every MLX array it builds allocates into the caller's scope. + ```java + public MultiAxisRope(float theta, int[] mropeSection) + public float theta() + public int[] mropeSection() // defensive copy + public int rotaryDims() // 2 · Σ mropeSection + public MultiAxisRotation prepare(MLXScope scope, int batch, int sequence, int[][] positions) + ``` + `positions` is `int[3][batch·sequence]` — three rows of length B·T (row-major flattened + `[B,T]`), row 0 = T (time), row 1 = H (height), row 2 = W (width). Chosen over three + separate `int[]` parameters: it is the reference's `[3, B·T]` `position_ids` layout as + one validated object, its row order is the `mrope_section` order and the §6.5 axis + indices 0/1/2, and `Qwen3VlPromptPlan`'s `temporalPositions()`/`heightPositions()`/ + `widthPositions()` map to the rows one-to-one in that order. Validation, all + `IllegalArgumentException`, all checked **before any MLX array is written**: constructor — + theta finite and positive (the `RopeSpec` convention), sections non-null (cloned), every + element ≥ 0, Σ ≥ 1 (rotaryDims is even by construction); `prepare` — batch ≥ 1, sequence ≥ + 1, exactly 3 rows, each row's length == batch·sequence (computed in `long` with overflow + checked before use), and no negative coordinate (rejected). Frequencies/cos/sin are + computed on every `prepare` call: no position-derived state is cached. The computation + follows the frozen order: one `[B·T, rotaryDims/2]` float32 position matrix (the + owner-mask-selected row per frequency, built directly from the int rows — the exact + cast) × the `[1, rotaryDims/2]` `invFreq` array by a single broadcast multiply → + `freqs_t` (bit-identical to the per-axis ×-then-owner-mask selection, step 2); + concatenate → `emb`; `MLXOps.cos`/`MLXOps.sin` in `scope` → FLOAT32 arrays reshaped to + `[batch, 1, sequence, rotaryDims]` (broadcastable over the heads axis). Host-memory trade-off + (accepted in review): the one position-matrix upload now holds `B·T·rotaryDims/2` float32 values + on the host where the pre-restructure path held `3·B·T` (three per-axis position arrays) — about 21× more for + the pinned 2B at a 32k prefill (8 MB vs 384 KB) — while decode (`B·T = 1`) is strictly cheaper + and the prefill cost is small next to the prefill compute. +- `public final class MultiAxisRotation` — the per-step rotation state one `prepare` + produces and every layer of that step shares, computed once per step and threaded exactly + like the existing per-step `stepFrequencies`: + ```java + public MLXArray cos() // FLOAT32 [batch, 1, sequence, rotaryDims] + public MLXArray sin() // FLOAT32 [batch, 1, sequence, rotaryDims] + public int rotaryDims() + public MLXArray rotate(MLXArray x) // x: [batch, heads, sequence, headDim] + ``` + `rotate` casts cos/sin to `x.dtype()` (frozen step 5), rotates the `rotaryDims` prefix + (frozen step 6) and concatenates the untouched tail back; it validates `x`'s rank/shape + against the stored cos/sin shape and `rotaryDims ≤ headDim` (the head-dimension + validation) before building anything. + +Decoder seam — all additive; **every existing signature stays**, each current overload +delegating to the new one with `null` rotation/hook so its graph is unchanged: + +- `DecoderAttention` (jmlx-core; the current forwards at L177-196 funnel into + `forwardInternal` L198-321): + ```java + public MLXArray forward( + MLXArray x, KVCache cache, MLXArray attentionMask, + int[] validLengths, MultiAxisRotation rotation) + ``` + funneled into `forwardInternal` with `stepFrequencies == null`; where `rope.apply` runs + today (after `AttentionHeads.toHeads` and the optional QK normalization), `rotation != + null` applies `q = rotation.rotate(q); k = rotation.rotate(k);` instead and skips the + frequency-array derivation. `validLengths` is required (non-null) on this path: the + rotation path always runs the batch-aware validation and per-row cache bookkeeping, which + reduce to the single-row arithmetic for B = 1. +- `DecoderBlock` (jmlx-core; the current forwards at L55-84): + ```java + public MLXArray forward( + MLXArray x, KVCache cache, MLXArray attentionMask, + int[] validLengths, MultiAxisRotation rotation) + ``` + the same residual structure as the existing batch overload, passing the rotation through + to `attention.forward(…, validLengths, rotation)`. +- `DecoderModel` (jmlx-models) gains a package-private functional interface beside + `EmbeddingHook`: + ```java + @FunctionalInterface + interface PostBlockHook { MLXArray apply(int layerIndex, MLXArray hidden); } + ``` + invoked after each selected early decoder block's residual output and before the final + `norm.forward` — where Qwen3-VL DeepStack adds level `i` after text block `i` (WP6/WP7). + Vision extraction indices and decoder injection indices are different counters and are + never mixed: the vision tower extracts its DeepStack features from vision blocks + `deepstack_visual_indexes = [5, 11, 17]` of the 24-block vision stack (§5), while the + text decoder injects level `i` (0-based) after text block `i` of the 28-block text stack + (§6.6). The additive overloads: + ```java + final MLXArray decoderStack( + MLXArray embeddings, List caches, int[] validLengths, int sequenceWidth, + MultiAxisRotation rotation, PostBlockHook postBlockHook) + final MLXArray stepLogits( + MLXArray tokenIds, List caches, int[] validLengths, + EmbeddingHook hook, MultiAxisRotation rotation, PostBlockHook postBlockHook) + ``` + (`decoderStack` keeps its package-private final embedding-start entry visibility; + `stepLogits` keeps the `EmbeddingHook` in its existing fourth position — the scheduler's + hook stays exactly where it is, post-lookup/pre-scale inside `decoderEmbeddings`.) The + existing `decoderStack(…, sequenceWidth)` (L364-418) and `stepLogits(…, hook)` (L183-198) + delegate with `(null, null)`. Inside `decoderStack`, `rotation != null` skips the + `stepFrequencies` derivation and calls the new block overload; the layer loop becomes + `x = blockᵢ.forward(…); if (postBlockHook != null) x = postBlockHook.apply(i, x);`, + ending in the unchanged `norm.forward(x)`. The single-row stack `normalizedHiddenStates` + (L516) and the single-row decode path `decode` (L787-794) get matching additive private + overloads `normalizedHiddenStates(tokenIds, caches, rotation, postBlockHook)` / + `decode(ids, caches, rotation, postBlockHook)` with the same delegation rule: the + `rotation == null && postBlockHook == null` path keeps the old graph bit-for-bit, and + the single-row path with a non-null `rotation` or `postBlockHook` routes each block + through the batch-aware overloads with + `validLengths = {sequenceWidth}` and a `[1,1,T,L]` batched mask (elementwise the same as + the 2-D mask for B = 1, windowed and evicted included; `cache.append(k, v, {T})` is the same + arithmetic as `cache.append(k, v)` for a uniform row — Unit 2 pins both against the `KVCache` + contract). The hook-only route (no rotation, a no-op hook) is pinned bit-exact + end-to-end against the old path at the logits by + `DecoderSeamTest.decodeWithNoOpPostBlockHookMatchesOldSingleRowPath` (full-float32; the + per-row RoPE form delegates to the offset form for a uniform row, so the two routes + build the same rope kernel call, and the explicit mask selects the same SDPA result as + the old path's implicit causal flag). The shared generation loop gets + ```java + @FunctionalInterface + interface DecodeStep { MLXArray apply(MLXArray ids, List caches); } + final GenerationResult runStepLoop( + GenerationRequest request, GenerationConfig policy, int[] prompt, + int[] frequencySeed, int promptPositions, KVCachePolicy cachePolicy, + Consumer listener, FirstStep firstStep, DecodeStep decodeStep) + ``` + (the existing 8-argument `runStepLoop` delegates with `this::decode`); steps ≥ 1 call + `decodeStep.apply(ids, caches)` where they call `decode(ids, caches)` today, and the + loop's own preflight, poisoning, evaluation ordering and result construction stay around + that call unchanged. Production ordinary decoders pass `null`/no-op for every new + parameter. Evaluation ordering, valid-length checks, cache preflight/postflight and + poisoning, and tied projection are all preserved and remain based on sequence/cache + positions — never on RoPE coordinates. +- Qwen3-VL wiring (WP6/WP7, stated here so the seam is closed): the request-supplied + rotary state is exactly what `Qwen3VlPromptPlan` already carries — the three prefill + position axes plus `positionDelta()`. Prefill uses + `mrope.prepare(scope, 1, L, plan.temporal/height/widthPositions())` (L = expanded + length); every decode step takes `p = plan.ropePositionAt(caches.get(0).nextPosition())` + — the logical cache position plus the delta, never a retained-KV length — as all three + axes (`new int[][] {{p}, {p}, {p}}`, B = 1, T = 1) through a `DecodeStep` that calls the + new `decode` overload with a `null` `PostBlockHook` (DeepStack injection is prefill- + only). + +**Visual replacement and masked residual add (frozen compositions; WP6/WP7 use them, WP4 +delivers the facades they build on).** Both are pure compositions of existing core +facades — `MLX.zeros`, `MLXShape.concatenate`, `MLXShape.takeAxis`, `MLXOps.add` (and +`slice` where a span is contiguous); model code never calls FFI, and no new primitive is +required. If a later unit ever needs one, the plan's rule applies: record the exact mlx-c +declaration first, probe the native semantics, update the inventory/call-site coverage and +test the scope/error behavior. They live as package-private static helpers in +`se.alipsa.jmlx.models` (beside `SmolVlmModel.merge`, the package WP6's `Qwen3VlModel` +joins) — no new public core API. + +- **Row-gather replacement** (merger features replace the image-placeholder embeddings) — + the same composition as `SmolVlmModel.merge` (L385-400): + ``` + source = concatenate([embedded, features], axis 1) // [1, L + F, hidden] + gather = mergeGatherIndices(L, destinations) // int[L], the existing algorithm + replaced = takeAxis(source, MLXArray(gather, [L]), 1) // [1, L, hidden] + ``` +- **Masked residual add** (DeepStack level i: add `features_i [1, F_i, hidden]` into the + `[1, L, hidden]` hidden states at exactly the image-token rows, zero elsewhere): + ``` + zeroRow = MLX.zeros(scope, {1, 1, hidden}, hidden.dtype()) // [1, 1, hidden] + source = concatenate([zeroRow, features_i], 1) // [1, 1 + F_i, hidden] + gather = gatherIndices(L, destinations_i) remapped: i -> 0, L + f -> 1 + f // int[L] + mask = takeAxis(source, MLXArray(gather, [L]), 1) + updated = MLXOps.add(hidden, mask) + ``` + The gather indices are the shared `mergeGatherIndices` helper's output (the same + one the row-gather replacement above applies directly) remapped over the source's + `1 + F_i` rows — every identity row (shared index `i < L`) remaps to the zero row + (source index `0`), the `f`-th destination (shared index `L + f`, strictly + increasing) remaps to source row `1 + f` — which is what lets the zero pad be a + single `[1, 1, hidden]` row instead of a full `[1, L, hidden]` zero tensor. The + caller `scope`'s relatedness to the operands' scopes (one an ancestor of the + other) is validated before any MLX array is written. The add matches the + reference's masked clone + add + write-back (§6.6) and leaves every non-image + row numerically identical (an exact `+ 0.0` per element; IEEE `-0.0 + 0.0 = +0.0`, + so a negative-zero element in a non-image row is normalised to `+0.0`). + +**Old-path bit-identity proof plan (the WP4 exit gate).** + +1. All existing suites green unchanged: `:jmlx-core:check` and `:jmlx-models:check`, + including `float32GoldenTest` (the full-float32 `DecoderAttentionTest` / + `MultiHeadAttentionTest` oracles and the family goldens). +2. `:jmlx-models:exactBitsVerifyTf32On` and `:jmlx-models:exactBitsVerifyTf32Off` pass + against the recordings above — 14 variants × 9 captures bit-exact in both TF32 modes on + this device (device-specific evidence; verify mode permits only the git-commit field to + drift). +3. `RopeSpec`'s sealed hierarchy and `MLXFast` are untouched: `git diff` over + `RopeSpec.java` and `MLXFast.java` is empty at WP4 completion (diff evidence). + +**Cross-check probe (planned; opt-in, same pattern as `qwen3vl_rope_multiframe_check.py`).** +`tools/hf-reference/qwen3vl_rope_interleaved_check.py` — pinned 4.57.6 venv +(`tools/hf-reference/.venv`: transformers 4.57.6, torch 2.6.0, CPU float32; never +Gradle-invoked, no golden output) — computes, for a fixed set of index-coded asymmetric +position cases (including the pinned 2B config theta 5,000,000 / sections `[24,20,20]` and +a small rotaryDims/sections case that exercises the owner pattern at every frequency), the +pinned `apply_interleaved_mrope` + rotate-half result in float32, prints one deterministic +case table at full float32 precision, and exits non-zero on any internal drift. The +matching Java probe (`tools/hf-reference/qwen3vl_rope_interleaved_check.jsh`, run manually +with the jmlx-core build classpath; also never Gradle-invoked) feeds the same cases to +`MultiAxisRope` and cross-checks the component element-wise with a strict max absolute +error of 1e-5 on the rotate-half output — an ulp-class tolerance for libm/Metal float32 +trigonometry, orders of magnitude below the O(1) difference any owner, frequency, +cast-order or rotation-layout error would produce; it prints the achieved max error per +case and exits non-zero past the threshold. + +**Test matrix the later units must deliver.** +`MultiAxisRopeTest` in `jmlx-core` (`@EnabledIfNativeAvailable`; the float32-comparison +cases tagged `full-float32` so they select into `:jmlx-core:float32GoldenTest` (MLX_ENABLE_TF32=0), +the rejection cases running in ordinary `:jmlx-core:test`): + +- identical axes (T = H = W) against a Java double-precision-derived float32 reference + with a tight bound (equal axes reduce mathematically to 1-D RoPE, but bit equality with + the fast kernels is **not** assumed — plan §6); +- unequal/index-coded axes: the pinned `[24,20,20]` owner pattern and a small sections + case, with small index-coded position integers so the output encodes the coordinates; +- nonzero decode positions at `cachePosition + delta` (B = 1, T = 1), including a negative + delta (e.g. cache position 82, delta −56 → coordinate 26), plus a batched case with + per-row positions; +- rejections: wrong row count (≠ 3), row-length mismatch (≠ B·T), negative coordinates, + negative or zero-sum sections, non-positive or non-finite theta, `rotaryDims > headDim` + and invalid `x` shapes — all `IllegalArgumentException` before any array is written; +- partial-rotary tail passthrough: the tail `[rotaryDims, headDim)` is bit-identical to + the input tail slice; +- head-dimension validation and the dtype-cast check that cos/sin are cast to the input + dtype before the rotation. + +The suite was added to `.github/workflows/ci.yml`'s required-execution list in PR #57 +(review round 2); the `full-float32` tag selection into `:jmlx-core:float32GoldenTest` +happens in this WP's implementation units. + +**WP4 acceptance (2026-10-20).** Delivered on `phase7-3b-wp4` in four commits over base +`128cd1e`: `d5de40e` (Unit 1 — the pre-change exact-bits baseline recorded in both TF32 +modes on this device, plus this §6.9 frozen design), `2fe2926` (Unit 2 — the +`MultiAxisRope`/`MultiAxisRotation` core components, `MultiAxisRopeTest`, and the +Python/jshell cross-check probe pair), `c50b5cc` (Unit 3 — the additive rotation/ +post-block seam across `DecoderAttention`/`DecoderBlock`/`DecoderModel`, with +`DecoderRotationSeamTest`, `SingleRowDecodeEquivalenceTest` and `DecoderSeamTest`), and +the acceptance commit (Unit 4 — `VisualFeatureInjection` and `VisualFeatureInjectionTest`). + +- **Validation matrix, all green.** Root `./gradlew check` BUILD SUCCESSFUL; + `./gradlew verifyMlxApiCallSites` BUILD SUCCESSFUL unchanged (the new code composes + existing facades only, so no inventory override or call-site mapping changed). Zero + failures in every suite; counts: `:jmlx-core:test` 49 suites / 414 tests with 1 + pre-existing skip (`MoeMlpBenchmarkTest.gatheredDecodeIsFasterThanDense`, the + untimed benchmark guard — not a regression), `:jmlx-core:float32GoldenTest` 8 / 30, + `:jmlx-models:test` 50 / 376, `:jmlx-models:float32GoldenTest` 13 / 43. The WP4 + suites (ordinary / full-float32): `MultiAxisRopeTest` 5 / 6, `DecoderRotationSeamTest` + 2 / 5, `SingleRowDecodeEquivalenceTest` 2 / 0 (no float32 cases), `DecoderSeamTest` + 1 / 5, `VisualFeatureInjectionTest` 1 / 7 (the one ordinary case is the pre-write + rejection test; the float32 cases compare `toFloatArray` bit-exactly). `./gradlew + -p buildSrc check` was not run: buildSrc is unchanged by this WP. +- **Old-path bit-identity gate, both TF32 modes, verbatim.** The Unit 1 baseline + recordings (spec hash + `846396b1dbab5a770a4f2c36841fe5ed1570a586b908956ed6e224fe5248e687`, recorded at + `128cd1ec2fb1004487015377183081827a29b956` before any WP4 change) were verified, not + re-recorded; each run re-derived its captures into a fresh + `build/exact-bits/tf32-{1,0}/derived-verify/` side effect. + `./gradlew :jmlx-models:exactBitsVerifyTf32On` → BUILD SUCCESSFUL with stdout + `verified 14 variants against .../tf32-1/manifest.json` and + `exact-bits verify tf32=1 (recorded at 128cd1ec2fb1004487015377183081827a29b956): + 14 applicable variants, 126 captures`; + `./gradlew :jmlx-models:exactBitsVerifyTf32Off` → the same two lines with `tf32-0` + and `tf32=0` (14 variants × 9 captures). Verify mode permits only the git-commit + metadata field to drift, and it is the only drift — the old one-dimensional decoder + path is bit-identical in both precision modes after the seam, the kernel and the + helpers landed. +- **Probe and diff evidence.** The Python probe was re-run once in the locked venv + after the final code state: all 7 cases (`pinned_prefill_indexed`, + `pinned_prefill_maxpos`, `pinned_decode_negative_delta`, `pinned_decode_equal_axes`, + `pinned_batched_rows`, `small_212_every_owner`, `small_111`) passed, exit 0. The + matching jshell twin was not re-run at acceptance (its component was last touched in + Unit 2, where its per-case max absolute error was ≤ 2.4e-7 against the 1e-5 + threshold); both sides were re-run 2026-10-20 in the locked environments after the + review rounds, against the restructured `prepare` (review round 2), the position-matrix + size overflow guard (review round 3) and the byte owner-table refactor (review round 4), + head `bf69877` — Python exit 0, all 7 cases passed; + the jshell twin exits 0 with every case OK and self-consistent, max absolute error + 2.38e-7 against the 1e-5 threshold. + `git diff main` over `RopeSpec.java` and `MLXFast.java` is empty: the sealed + hierarchy gained no variant and no fast kernel changed. +- **Deferrals, per plan.** The new native suites' addition to + `.github/workflows/ci.yml`'s required-execution list landed in PR #57 (review + round 2); `Qwen3VlModel`, the `qwen3_vl` loader dispatch and the DeepStack + consumption of `VisualFeatureInjection` are WP5/WP6. + ## 7. Fast-vs-slow preprocessing divergence (measured once, both sides) - **Tool:** `tools/image-oracle/fast-comparison/` — an independent hash-locked CPython 3.12 venv diff --git a/tools/hf-reference/README.md b/tools/hf-reference/README.md index 0ffeabbb..866e2f50 100644 --- a/tools/hf-reference/README.md +++ b/tools/hf-reference/README.md @@ -123,3 +123,51 @@ takes the upscale branch under either bound. Consumption contract for the Phase 7.3b WP5+ model-math tests: feed these bins directly as `pixel_values` (with the manifest's grids) without re-preprocessing, so a model-math error cannot be concealed by a preprocessing drift. + +### Qwen3-VL interleaved multi-axis RoPE cross-check (Phase 7.3b WP4) + +`qwen3vl_rope_interleaved_check.py` and its Java twin +`qwen3vl_rope_interleaved_check.jsh` are an opt-in pair, never invoked by +Gradle, that writes no golden: the Python probe runs the real pinned +`Qwen3VLTextRotaryEmbedding` plus `apply_rotary_pos_emb` (transformers 4.57.6, +CPU float32, single-threaded, deterministic) on fixed, index-coded asymmetric +cases -- prefill, prefill at the pinned maximum coordinate (89), decode at the +all-axes-equal shape, decode at the negative-delta coordinate (26), per-row +batched positions, and small sections that exercise every interleaved `f % 3` +owner -- and prints a deterministic full-precision float32 case table to +stdout. It exits non-zero if its internal gates fail: the pinned snapshot's +`text_config` must carry theta 5,000,000 / head_dim 128 / +`mrope_section [24, 20, 20]` / `mrope_interleaved true` / rope type default; +each module's `inv_freq` must be bit-exact against the documented float32 +derivation order; each case's explicit owner-gather recomputation must be +bit-exact against the module's own `apply_interleaved_mrope` output; the +equal-axes case's 2-D position-ids path must equal its 3-D path; the q and k +outputs must be bit-equal; and a double run must reproduce the table's +SHA-256. + +The Java twin reads that table through `JMLX_ROPE_REFERENCE` (jshell has no +argv), runs the frozen `se.alipsa.jmlx.nn.MultiAxisRope` component over every +case, and requires a max absolute error at or below 1e-5 against the +reference rotate-half output (plus bit-identical rotate self-consistency). +The pinned transformers output is authoritative: a divergence means the Java +component is fixed, not the probe. Run both from the repository root: + +```sh +./gradlew :jmlx-core:jar :jmlx-ffi:jar +tools/hf-reference/.venv/bin/python \ + tools/hf-reference/qwen3vl_rope_interleaved_check.py \ + > build/qwen3vl-rope-reference.txt +JMLX_LIBRARY_PATH=$PWD/native/install/lib \ +JMLX_ROPE_REFERENCE=$PWD/build/qwen3vl-rope-reference.txt \ + "$JDK25_HOME/bin/jshell" --execution local -s -J--enable-native-access=ALL-UNNAMED \ + --class-path jmlx-core/build/libs/jmlx-core-0.5.0-SNAPSHOT.jar:jmlx-ffi/build/libs/jmlx-ffi-0.5.0-SNAPSHOT.jar \ + tools/hf-reference/qwen3vl_rope_interleaved_check.jsh +``` + +where `$JDK25_HOME` is the Java 25 toolchain home the build uses. The Python +probe exits non-zero on any gate failure; the Java twin exits 0 when every +case is within 1e-5, 1 when any case exceeds the threshold or a rotate is not +self-consistent, 2 on a missing or unrecognized table, and 3 on any other +failure. Each per-case max error is printed. `--execution local` runs the +snippets inside the jshell process so the `System.exit` verdict is the probe's +exit code (the default child-engine mode swallows it). diff --git a/tools/hf-reference/qwen3vl_rope_interleaved_check.jsh b/tools/hf-reference/qwen3vl_rope_interleaved_check.jsh new file mode 100644 index 00000000..cef8893b --- /dev/null +++ b/tools/hf-reference/qwen3vl_rope_interleaved_check.jsh @@ -0,0 +1,194 @@ +// Phase 7.3b WP4 Java twin of qwen3vl_rope_interleaved_check.py: reads the +// pinned-reference float32 case table printed by that probe and checks +// jmlx's frozen se.alipsa.jmlx.nn.MultiAxisRope component against it with a +// strict 1e-5 max absolute error, plus rotate self-consistency (two rotates +// of the same input must be bit-identical). It is never invoked by Gradle +// and writes no golden; the pinned transformers output is authoritative, so +// a divergence means the Java component is fixed, not the probe. +// +// Invocation (from the repository root; the jars must be built first): +// +// ./gradlew :jmlx-core:jar :jmlx-ffi:jar +// tools/hf-reference/.venv/bin/python \ +// tools/hf-reference/qwen3vl_rope_interleaved_check.py \ +// > build/qwen3vl-rope-reference.txt +// JMLX_LIBRARY_PATH=$PWD/native/install/lib \ +// JMLX_ROPE_REFERENCE=$PWD/build/qwen3vl-rope-reference.txt \ +// "$JDK25_HOME/bin/jshell" --execution local -s -J--enable-native-access=ALL-UNNAMED \ +// --class-path jmlx-core/build/libs/jmlx-core-0.5.0-SNAPSHOT.jar:jmlx-ffi/build/libs/jmlx-ffi-0.5.0-SNAPSHOT.jar \ +// tools/hf-reference/qwen3vl_rope_interleaved_check.jsh +// +// $JDK25_HOME is the Java 25 toolchain home the build uses (on this machine: +// ~/.gradle/jdks/eclipse_adoptium-25-aarch64-os_x.2/jdk-25.0.4.1+1/Contents/Home). +// JMLX_LIBRARY_PATH is the staged native library directory that NativeLoader +// resolves after the jmlx.library.path property. jshell has no argv, so the +// table path comes from the JMLX_ROPE_REFERENCE environment variable. +// --execution local runs the snippets inside the jshell process, so the +// System.exit verdict is the probe's exit code (the default child-engine +// mode swallows it); -s suppresses jshell's prompts. +// +// Exit codes: 0 all cases within 1e-5; 1 any case over the threshold or a +// rotate self-consistency failure; 2 missing JMLX_ROPE_REFERENCE or an +// unrecognized table; 3 any other failure. + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import se.alipsa.jmlx.core.MLX; +import se.alipsa.jmlx.core.MLXArray; +import se.alipsa.jmlx.memory.MLXScope; +import se.alipsa.jmlx.nn.MultiAxisRope; +import se.alipsa.jmlx.nn.MultiAxisRotation; + +final class Cursor { + final List tokens; + int pos; + + Cursor(List tokens) { + this.tokens = tokens; + } + + String next() { + if (pos >= tokens.size()) { + throw new IllegalStateException("reference table ended unexpectedly"); + } + return tokens.get(pos++); + } + + int nextInt() { + return Integer.parseInt(next()); + } + + int[] nextInts(int count) { + int[] out = new int[count]; + for (int i = 0; i < count; i++) { + out[i] = nextInt(); + } + return out; + } + + float[] nextFloats(int count) { + float[] out = new float[count]; + for (int i = 0; i < count; i++) { + out[i] = (float) Double.parseDouble(next()); + } + return out; + } + + void expect(String key) { + String got = next(); + if (!got.equals(key)) { + throw new IllegalStateException("expected '" + key + "' but found '" + got + "'"); + } + } + + static double maxAbs(float[] actual, float[] expected) { + double max = 0.0; + for (int i = 0; i < actual.length; i++) { + max = Math.max(max, Math.abs((double) actual[i] - (double) expected[i])); + } + return max; + } +} + +String refPath = System.getenv("JMLX_ROPE_REFERENCE"); +if (refPath == null) { + System.err.println( + "set JMLX_ROPE_REFERENCE to the table printed by qwen3vl_rope_interleaved_check.py"); + System.exit(2); +} + +double threshold = 1e-5; +int caseCount = 0; +int failureCount = 0; +double worstError = 0.0; +String worstCase = ""; +try { + List tokens = new ArrayList<>(); + for (String line : Files.readAllLines(Path.of(refPath))) { + for (String part : line.split("\\s+")) { + if (!part.isEmpty()) { + tokens.add(part); + } + } + } + Cursor cursor = new Cursor(tokens); + if (!"#".equals(cursor.next()) + || !"qwen3vl-rope-interleaved-reference".equals(cursor.next()) + || !"v1".equals(cursor.next())) { + throw new IllegalStateException("unrecognized reference table schema"); + } + for (String key : new String[] {"transformers", "torch", "python", "platform"}) { + cursor.expect(key); + cursor.next(); + } + while (cursor.pos < cursor.tokens.size()) { + cursor.expect("case"); + String name = cursor.next(); + cursor.expect("theta"); + int theta = cursor.nextInt(); + cursor.expect("head_dim"); + int headDim = cursor.nextInt(); + cursor.expect("sections"); + int[] sections = cursor.nextInts(3); + cursor.expect("batch"); + int batch = cursor.nextInt(); + cursor.expect("heads"); + int heads = cursor.nextInt(); + cursor.expect("sequence"); + int sequence = cursor.nextInt(); + cursor.expect("positions_t"); + int[] t = cursor.nextInts(batch * sequence); + cursor.expect("positions_h"); + int[] h = cursor.nextInts(batch * sequence); + cursor.expect("positions_w"); + int[] w = cursor.nextInts(batch * sequence); + int count = batch * heads * sequence * headDim; + cursor.expect("x"); + float[] x = cursor.nextFloats(count); + cursor.expect("q"); + float[] q = cursor.nextFloats(count); + cursor.expect("k"); + float[] k = cursor.nextFloats(count); + + float[] y; + boolean selfConsistent; + try (MLXScope scope = new MLXScope()) { + MultiAxisRope rope = new MultiAxisRope(theta, sections); + MultiAxisRotation rotation = rope.prepare(scope, batch, sequence, new int[][] {t, h, w}); + MLXArray xArray = MLX.array(scope, x, new int[] {batch, heads, sequence, headDim}); + y = rotation.rotate(xArray).toFloatArray(); + float[] yRepeat = rotation.rotate(xArray).toFloatArray(); + selfConsistent = Arrays.equals(y, yRepeat); + } + double maxQ = Cursor.maxAbs(y, q); + double maxK = Cursor.maxAbs(y, k); + boolean pass = selfConsistent && maxQ <= threshold && maxK <= threshold; + caseCount++; + if (Math.max(maxQ, maxK) > worstError) { + worstError = Math.max(maxQ, maxK); + worstCase = name; + } + if (!pass) { + failureCount++; + } + System.out.printf( + "case %s: max_abs_error_q=%.3e max_abs_error_k=%.3e self_consistent=%b %s%n", + name, maxQ, maxK, selfConsistent, pass ? "OK" : "MISMATCH"); + } + if (caseCount == 0) { + throw new IllegalStateException("no cases found in the reference table"); + } + if (failureCount > 0) { + System.out.println( + failureCount + "/" + caseCount + " cases failed; worst " + worstError + " in " + worstCase); + System.exit(1); + } + System.out.println("all " + caseCount + " cases match MultiAxisRope within " + threshold); + System.exit(0); +} catch (Exception e) { + System.err.println("qwen3vl_rope_interleaved_check.jsh failed: " + e); + System.exit(3); +} diff --git a/tools/hf-reference/qwen3vl_rope_interleaved_check.py b/tools/hf-reference/qwen3vl_rope_interleaved_check.py new file mode 100644 index 00000000..3de00d3d --- /dev/null +++ b/tools/hf-reference/qwen3vl_rope_interleaved_check.py @@ -0,0 +1,273 @@ +#!/usr/bin/env python3 +"""Phase 7.3b WP4 opt-in cross-check: Qwen3-VL interleaved multi-axis text RoPE. + +Runs under tools/hf-reference/.venv (transformers 4.57.6, torch 2.6.0, +Pillow 12.3.0, NumPy 2.5.3, Python 3.12), CPU float32, single-threaded, +deterministic. It is never invoked by Gradle and writes no golden: it +evaluates the real pinned Qwen3VLTextRotaryEmbedding plus +apply_rotary_pos_emb on fixed, index-coded asymmetric cases and prints a +deterministic full-precision float32 case table to stdout. The Java twin +qwen3vl_rope_interleaved_check.jsh reads that table through +JMLX_ROPE_REFERENCE and compares jmlx's frozen se.alipsa.jmlx.nn.MultiAxisRope +component against it with a strict 1e-5 max absolute error; the pinned +transformers behavior is authoritative, so a divergence means the Java +component is fixed, not the probe. + +The pinned-2B cases carry the snapshot's theta 5,000,000 / head_dim 128 / +mrope_section [24, 20, 20] and cover prefill, prefill at the pinned maximum +coordinate (89), decode at the all-axes-equal shape (the only shape whose +2-D position-ids path is defined to agree with the 3-D one, so it gates +that path as well), decode at the negative-delta coordinate (26), and +per-row batched positions. The small cases (10,000 / 10 / [2, 1, 2] and +1,000 / 6 / [1, 1, 1]) exercise every interleaved f % 3 owner at constant +index-coded axes. Every case has head_dim = 2 * sum(mrope_section) because +the reference rotate_half rotates the whole head dimension. + +Internal gates -- any failure exits non-zero before the table is printed: + + * the pinned snapshot's text_config carries rope_theta 5,000,000, + head_dim 128, mrope_section [24, 20, 20], mrope_interleaved true and + rope_type "default"; + * every module's inv_freq is bit-exact against + 1.0 / (base ** (arange(0, head_dim, 2, int64).to(float32) / head_dim)), + the documented float32 derivation order (findings section 6.9), and + attention_scaling is 1.0; + * every case's owner-gather recomputation -- the same float32 matmul, an + explicit per-frequency owner selection, the cat((freqs, freqs), -1) + layout and the float32 cos/sin -- is bit-exact against the module's + own forward output, isolating apply_interleaved_mrope from the case + values; + * the equal-axes case's 2-D position-ids output is bit-exact against its + 3-D output (a 2-D input expands to all three axes, so the two paths + only agree when the axes are equal); + * the q output is bit-exact against the k output (the probe feeds + q = k = x), and a double run reproduces the printed table's SHA-256. +""" +from __future__ import annotations + +import argparse +import hashlib +import platform +import sys +from pathlib import Path + +import qwen3vl_probe as probe + +SCHEMA = ("#", "qwen3vl-rope-interleaved-reference", "v1") + +# (name, theta, head_dim, sections, batch, heads, sequence, t, h, w); the +# position rows are flattened in (batch, token) order. +CASES: list[tuple[str, int, int, list[int], int, int, int, list[int], list[int], list[int]]] = [ + ("pinned_prefill_indexed", 5_000_000, 128, [24, 20, 20], 1, 1, 4, + [0, 1, 2, 3], [10, 11, 12, 13], [20, 21, 22, 23]), + ("pinned_prefill_maxpos", 5_000_000, 128, [24, 20, 20], 1, 1, 2, + [88, 89], [80, 81], [70, 71]), + ("pinned_decode_negative_delta", 5_000_000, 128, [24, 20, 20], 1, 1, 1, + [26], [26], [26]), + ("pinned_decode_equal_axes", 5_000_000, 128, [24, 20, 20], 1, 1, 1, + [89], [89], [89]), + ("pinned_batched_rows", 5_000_000, 128, [24, 20, 20], 2, 1, 2, + [0, 1, 30, 31], [4, 5, 26, 27], [8, 9, 22, 23]), + ("small_212_every_owner", 10_000, 10, [2, 1, 2], 1, 1, 2, + [1, 1], [2, 2], [4, 4]), + ("small_111", 1_000, 6, [1, 1, 1], 1, 1, 2, + [0, 1], [2, 3], [4, 5]), +] + + +def deterministic_x(n: int) -> list[float]: + """The frozen probe input: distinct values in [-1, 1] over the flat index.""" + return [((i % 13) - 6) / 6.0 for i in range(n)] + + +def format_values(values: list[float]) -> str: + # 9 significant decimal digits round-trip a float32 exactly, so the Java + # twin reads back the reference's exact float32 values. + return " ".join(f"{v:.9e}" for v in values) + + +def owner_axis(f: int, sections: list[int]) -> int: + """The reference interleaved owner of frequency index f: H owns + f % 3 == 1 with f < 1 + 3 * s1, W owns f % 3 == 2 with f < 2 + 3 * s2, + everything else stays T (the slice(offset, length, 3) layout).""" + if f % 3 == 1 and f < 1 + 3 * sections[1]: + return 1 + if f % 3 == 2 and f < 2 + 3 * sections[2]: + return 2 + return 0 + + +def owner_gather_tables(module, head_dim: int, batch: int, sequence: int, + position_ids) -> tuple: + """Recompute the module's float32 cos/sin tables without + apply_interleaved_mrope: the same float32 matmul, an explicit + per-frequency owner selection, the cat((freqs, freqs), -1) layout.""" + import torch + + half = head_dim // 2 + freqs = ( + module.inv_freq[None, None, :, None].float().expand(3, batch, half, 1) + @ position_ids[:, :, None, :].float() + ).transpose(2, 3) + gathered = torch.zeros(batch, sequence, half) + for f in range(half): + gathered[:, :, f] = freqs[owner_axis(f, module.mrope_section), :, :, f] + emb = torch.cat((gathered, gathered), dim=-1) + return emb.cos() * module.attention_scaling, emb.sin() * module.attention_scaling + + +def gate_module(name: str, module, theta: int, head_dim: int, sections: list[int]) -> None: + import torch + + probe.check_facts(name, "attention_scaling", module.attention_scaling, 1.0) + probe.check_facts(name, "theta", float(module.config.rope_theta), float(theta)) + probe.check_facts(name, "head_dim", int(getattr(module.config, "head_dim")), head_dim) + probe.check_facts(name, "mrope_section", list(module.mrope_section), sections) + reference = 1.0 / ( + theta + ** (torch.arange(0, head_dim, 2, dtype=torch.int64).to(dtype=torch.float) / head_dim) + ) + if not torch.equal(module.inv_freq, reference): + sys.exit(f"{name}: module.inv_freq is not bit-exact against the documented " + "float32 derivation 1.0 / (base ** (arange(0, head_dim, 2).float() " + "/ head_dim))") + + +def compute_case(name: str, theta: int, head_dim: int, sections: list[int], batch: int, + heads: int, sequence: int, t: list[int], h: list[int], w: list[int], + module) -> dict: + import torch + from transformers.models.qwen3_vl.modeling_qwen3_vl import apply_rotary_pos_emb + + position_ids = torch.as_tensor([t, h, w], dtype=torch.long).reshape(3, batch, sequence) + x = torch.tensor(deterministic_x(batch * heads * sequence * head_dim), + dtype=torch.float32).reshape(batch, heads, sequence, head_dim) + cos, sin = module.forward(x, position_ids) + + cos_ref, sin_ref = owner_gather_tables(module, head_dim, batch, sequence, position_ids) + if not torch.equal(cos, cos_ref) or not torch.equal(sin, sin_ref): + sys.exit(f"{name}: apply_interleaved_mrope output is not bit-exact against the " + "explicit owner-gather recomputation") + if name == "pinned_decode_equal_axes": + cos_2d, sin_2d = module.forward(x, position_ids[0]) + if not torch.equal(cos, cos_2d) or not torch.equal(sin, sin_2d): + sys.exit(f"{name}: the 2-D position-ids path is not bit-exact against the " + "3-D path on the all-axes-equal shape") + q, k = apply_rotary_pos_emb(x, x, cos, sin) + if not torch.equal(q, k): + sys.exit(f"{name}: q and k outputs differ although the probe feeds q = k = x") + return { + "theta": theta, + "head_dim": head_dim, + "sections": sections, + "batch": batch, + "heads": heads, + "sequence": sequence, + "positions_t": t, + "positions_h": h, + "positions_w": w, + "x": x.flatten().tolist(), + "q": q.flatten().tolist(), + "k": k.flatten().tolist(), + } + + +def build_table(environment: dict, module_by_case: dict) -> str: + lines = [ + " ".join(SCHEMA), + f"transformers {environment['transformers']}", + f"torch {environment['torch']}", + f"python {environment['python']}", + f"platform {environment['platform']}", + ] + for (name, theta, head_dim, sections, batch, heads, sequence, t, h, w), module in ( + zip(CASES, module_by_case) + ): + case = compute_case(name, theta, head_dim, sections, batch, heads, sequence, + t, h, w, module) + lines.append(f"case {name}") + lines.append(f"theta {case['theta']}") + lines.append(f"head_dim {case['head_dim']}") + lines.append("sections " + " ".join(str(v) for v in case["sections"])) + lines.append(f"batch {case['batch']} heads {case['heads']} sequence {case['sequence']}") + lines.append("positions_t " + " ".join(str(v) for v in case["positions_t"])) + lines.append("positions_h " + " ".join(str(v) for v in case["positions_h"])) + lines.append("positions_w " + " ".join(str(v) for v in case["positions_w"])) + lines.append("x " + format_values(case["x"])) + lines.append("q " + format_values(case["q"])) + lines.append("k " + format_values(case["k"])) + return "\n".join(lines) + "\n" + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + here = Path(__file__).resolve().parent + parser.add_argument("--dir", type=Path, + default=here.parent.parent / ".tier-b" / "qwen3-vl-2b-instruct") + args = parser.parse_args() + directory: Path = args.dir + if not (directory / "config.json").is_file(): + sys.exit(f"no config.json under {directory}; pass --dir to the pinned snapshot") + + environment = probe.gate_versions() + environment = { + **environment, + "python": platform.python_version(), + "platform": f"{platform.system()}-{platform.machine()}", + } + import torch + from transformers.models.qwen3_vl.configuration_qwen3_vl import ( + Qwen3VLConfig, + Qwen3VLTextConfig, + ) + from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLTextRotaryEmbedding + + torch.set_num_threads(1) + torch.use_deterministic_algorithms(True) + + # The pinned snapshot's text config is the authority for the 2B cases. + rope_config = Qwen3VLConfig.from_pretrained(str(directory)) + text_config = rope_config.text_config + probe.check_facts("config", "rope_theta", float(text_config.rope_theta), 5_000_000.0) + probe.check_facts("config", "head_dim", int(text_config.head_dim), 128) + probe.check_facts("config", "mrope_section", + list(text_config.rope_scaling["mrope_section"]), [24, 20, 20]) + probe.check_facts("config", "mrope_interleaved", + bool(text_config.rope_scaling["mrope_interleaved"]), True) + probe.check_facts("config", "rope_type", + text_config.rope_scaling.get("rope_type", "default"), "default") + cpu = torch.device("cpu") + pinned_module = Qwen3VLTextRotaryEmbedding(text_config, device=cpu) + small_configs = { + (10_000, 10, (2, 1, 2)): Qwen3VLTextConfig( + head_dim=10, rope_theta=10_000.0, max_position_embeddings=32, + rope_scaling={"rope_type": "default", "mrope_section": [2, 1, 2], + "mrope_interleaved": True}), + (1_000, 6, (1, 1, 1)): Qwen3VLTextConfig( + head_dim=6, rope_theta=1_000.0, max_position_embeddings=32, + rope_scaling={"rope_type": "default", "mrope_section": [1, 1, 1], + "mrope_interleaved": True}), + } + modules = [] + for name, theta, head_dim, sections, batch, heads, sequence, _t, _h, _w in CASES: + probe.check_facts(name, "case_head_dim", head_dim, 2 * sum(sections), + "rotate_half rotates the whole head dimension") + if theta == 5_000_000: + module = pinned_module + else: + module = Qwen3VLTextRotaryEmbedding( + small_configs[(theta, head_dim, tuple(sections))], device=cpu) + gate_module(name, module, theta, head_dim, sections) + modules.append(module) + + # Double-run determinism: the printed table must be reproducible bit-exactly. + table = build_table(environment, modules) + table_again = build_table(environment, modules) + if hashlib.sha256(table.encode("utf-8")).hexdigest() != hashlib.sha256( + table_again.encode("utf-8")).hexdigest(): + sys.exit("the reference table is not reproducible across a double run") + sys.stdout.write(table) + + +if __name__ == "__main__": + main()