From d5de40e79210cd05cfe3bd390b056f63ace58224 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 10:57:05 +0200 Subject: [PATCH 01/32] phase7-3b WP4: record pre-change exact-bits baseline and freeze design in findings --- req/plans/phase7-3b-plan.md | 8 +- req/plans/phase7-3b-reference-findings.md | 308 ++++++++++++++++++++++ 2 files changed, 313 insertions(+), 3 deletions(-) diff --git a/req/plans/phase7-3b-plan.md b/req/plans/phase7-3b-plan.md index e327f12..06983e1 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–WP3 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 in progress as of + 2026-10-20 — the pre-change exact-bits baseline is recorded and the core RoPE / decoder + seam design is frozen in `req/plans/phase7-3b-reference-findings.md` §6.9. 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, diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 7369d73..a788d2a 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -712,6 +712,314 @@ 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 static owner masks: three + float32 weight vectors of length `rotaryDims/2` (1.0 at the owned frequencies, 0.0 + elsewhere, derived from the sections at construction) give + `freqs_t = freqsT·wT + freqsH·wH + freqsW·wW` — a bit-identical copy for non-negative + finite frequencies (multiply by 1.0/0.0 and add in float32 are exact there) 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, + in `scope`, follows the frozen order: three float32 position arrays (built directly from + the int rows — the exact cast) × the `[1, rotaryDims/2]` `invFreq` array by broadcast + multiply → the per-axis `[B·T, rotaryDims/2]` freqs; owner-mask selection → `freqs_t`; + concatenate → `emb`; `MLXOps.cos`/`MLXOps.sin` → FLOAT32 arrays reshaped to + `[batch, 1, sequence, rotaryDims]` (broadcastable over the heads axis). +- `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` path keeps the old graph bit-for-bit, and the `rotation != null` + single-row path 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; `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 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): + ``` + zero = MLX.zeros(scope, hidden.shape(), hidden.dtype()) // [1, L, hidden] + source = concatenate([zero, features_i], 1) // [1, L + F_i, hidden] + mask = takeAxis(source, MLXArray(gather(L, destinations_i), [L]), 1) + updated = MLXOps.add(hidden, mask) + ``` + `gather` is the same index algorithm as the replacement (strictly increasing + destinations → row `L + f`), so one helper serves both; the add matches the reference's + masked clone + add + write-back (§6.6) and leaves every non-image row bit-identical + (an exact `+ 0.0` per element). + +**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. + +Adding the suite to `.github/workflows/ci.yml`'s required-execution list is deferred to +WP8 (plan §9); the `full-float32` tag selection into `:jmlx-core:float32GoldenTest` +happens in this WP's implementation units. + ## 7. Fast-vs-slow preprocessing divergence (measured once, both sides) - **Tool:** `tools/image-oracle/fast-comparison/` — an independent hash-locked CPython 3.12 venv From 2fe2926772a0dbec7bf0857312327026d46f14ee Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 13:30:00 +0200 Subject: [PATCH 02/32] phase7-3b WP4: add frozen MultiAxisRope core component with strict tests and pinned-vendor cross-check probe --- .../java/se/alipsa/jmlx/nn/MultiAxisRope.java | 209 +++++++ .../se/alipsa/jmlx/nn/MultiAxisRotation.java | 106 ++++ .../se/alipsa/jmlx/nn/MultiAxisRopeTest.java | 527 ++++++++++++++++++ .../nn/SingleRowDecodeEquivalenceTest.java | 101 ++++ tools/hf-reference/README.md | 48 ++ .../qwen3vl_rope_interleaved_check.jsh | 194 +++++++ .../qwen3vl_rope_interleaved_check.py | 273 +++++++++ 7 files changed, 1458 insertions(+) create mode 100644 jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java create mode 100644 jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java create mode 100644 jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java create mode 100644 jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java create mode 100644 tools/hf-reference/qwen3vl_rope_interleaved_check.jsh create mode 100644 tools/hf-reference/qwen3vl_rope_interleaved_check.py 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 0000000..d07b094 --- /dev/null +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java @@ -0,0 +1,209 @@ +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. + * + *

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 float[] ownerT; + private final float[] ownerH; + private final float[] ownerW; + + /** + * Validates the base and the sections and derives the per-frequency inverse and the static T/H/W + * owner masks. 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.ownerT = new float[half]; + this.ownerH = new float[half]; + this.ownerW = new float[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 masks realize that + // selection as a static 0/1 weight vector per axis. + if (f % 3 == 1 && f < 1L + 3L * mropeSection[1]) { + ownerH[f] = 1f; + } else if (f % 3 == 2 && f < 2L + 3L * mropeSection[2]) { + ownerW[f] = 1f; + } else { + ownerT[f] = 1f; + } + } + } + + /** 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) and no coordinate is negative. The frozen compute order then runs in + * {@code scope}: three float32 position arrays of shape {@code [B*T, 1]} (built directly from the + * int rows, an exact cast) broadcast-multiplied against the {@code [1, rotaryDims/2]} {@code + * invFreq} array give the per-axis frequencies; the static owner masks select {@code freqs_t = + * freqsT * wT + freqsH * wH + freqsW * wW} (multiplying by the static 0/1 masks and adding in + * float32 is a bit-identical copy for non-negative finite frequencies); {@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; + MLXArray invFreqArray = MLX.array(scope, invFreq, new int[] {1, half}); + MLXArray weightT = MLX.array(scope, ownerT, new int[] {1, half}); + MLXArray weightH = MLX.array(scope, ownerH, new int[] {1, half}); + MLXArray weightW = MLX.array(scope, ownerW, new int[] {1, half}); + MLXArray freqsT = axisFrequencies(scope, positions[0], flat, invFreqArray); + MLXArray freqsH = axisFrequencies(scope, positions[1], flat, invFreqArray); + MLXArray freqsW = axisFrequencies(scope, positions[2], flat, invFreqArray); + MLXArray selected = + MLXOps.add( + MLXOps.add(MLXOps.multiply(freqsT, weightT), MLXOps.multiply(freqsH, weightH)), + MLXOps.multiply(freqsW, weightW)); + MLXArray emb = MLXShape.concatenate(new MLXArray[] {selected, selected}, 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); + } + + private static MLXArray axisFrequencies(MLXScope scope, int[] row, int flat, MLXArray invFreq) { + float[] data = new float[flat]; + for (int i = 0; i < flat; i++) { + data[i] = (float) row[i]; + } + MLXArray positions = MLX.array(scope, data, new int[] {flat, 1}); + return MLXOps.multiply(positions, invFreq); + } + + /** + * 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 0000000..5fac8a0 --- /dev/null +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java @@ -0,0 +1,106 @@ +package se.alipsa.jmlx.nn; + +import java.util.Objects; +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; + + MultiAxisRotation(MLXArray cos, MLXArray sin, int rotaryDims) { + this.cos = cos; + this.sin = sin; + this.rotaryDims = rotaryDims; + } + + /** 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; + } + + /** + * 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) 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) { + Objects.requireNonNull(x, "x"); + int[] shape = x.shape(); + if (shape.length != 4) { + throw new IllegalArgumentException( + "rotate requires [batch, heads, sequence, headDim], got shape " + + java.util.Arrays.toString(shape)); + } + int[] cosShape = cos.shape(); + if (shape[0] != cosShape[0] || shape[2] != cosShape[2]) { + throw new IllegalArgumentException( + "x batch and sequence must match the prepared rotation (" + + shape[0] + + ", " + + shape[2] + + " against " + + cosShape[0] + + ", " + + cosShape[2] + + ")"); + } + if (rotaryDims > shape[3]) { + throw new IllegalArgumentException("rotaryDims exceeds head dimension"); + } + DType dtype = x.dtype(); + MLXArray cosInput = cos.dtype() == dtype ? cos : MLX.astype(cos, dtype); + MLXArray sinInput = sin.dtype() == dtype ? sin : MLX.astype(sin, dtype); + 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); + } +} 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 0000000..1a1e67a --- /dev/null +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java @@ -0,0 +1,527 @@ +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)})); + } + } + + @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(NullPointerException.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 0000000..71f62d1 --- /dev/null +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java @@ -0,0 +1,101 @@ +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 unpaddded 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)); + } + } + } + + 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/tools/hf-reference/README.md b/tools/hf-reference/README.md index 0ffeabb..866e2f5 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 0000000..cef8893 --- /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 0000000..3de00d3 --- /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() From c50b5cc32a81793c2eea078a489a8a2ca0afabce Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 15:32:17 +0200 Subject: [PATCH 03/32] phase7-3b WP4: add additive rotation/post-block decoder seam with old-path parity tests --- .../se/alipsa/jmlx/nn/DecoderAttention.java | 62 ++- .../java/se/alipsa/jmlx/nn/DecoderBlock.java | 28 + .../jmlx/nn/DecoderRotationSeamTest.java | 478 ++++++++++++++++++ .../se/alipsa/jmlx/models/DecoderModel.java | 218 +++++++- .../alipsa/jmlx/models/DecoderSeamTest.java | 329 ++++++++++++ 5 files changed, 1073 insertions(+), 42 deletions(-) create mode 100644 jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java create mode 100644 jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java 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 ee71134..77df6dd 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,31 @@ 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. + */ + public MLXArray forward( + MLXArray x, + KVCache cache, + MLXArray attentionMask, + int[] validLengths, + MultiAxisRotation rotation) { + Objects.requireNonNull(rotation, "rotation"); + 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 +224,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 +290,24 @@ 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. + 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 +320,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 128952e..c410f5d 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,32 @@ 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"); + Objects.requireNonNull(rotation, "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/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 0000000..0a38c84 --- /dev/null +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java @@ -0,0 +1,478 @@ +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) and a null rotation before any array is written, and the batch validation must + * reject a wrong number of lengths. + */ + @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( + NullPointerException.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)); + 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-models/src/main/java/se/alipsa/jmlx/models/DecoderModel.java b/jmlx-models/src/main/java/se/alipsa/jmlx/models/DecoderModel.java index 96ef392..e2cf2b5 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 @@ -19,6 +19,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 +169,22 @@ 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. + */ + @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 +199,30 @@ 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. 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); 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 = @@ -363,6 +401,28 @@ 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}). The batched mask is the same in both cases. 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 @@ -409,18 +469,39 @@ 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) { + x = postBlockHook.apply(i, x); + } } 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 +595,58 @@ 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. For a batch-1 unpadded row these are the same arithmetic as the + * old single-row path (the uniform append and the {@code [1,1,T,L]} batched mask match the + * single-row ones element for element, pinned by {@code SingleRowDecodeEquivalenceTest}), plus + * the hook firing after every block's residual output, before the final normalization. + */ + 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 +746,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 +781,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 +847,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 +934,23 @@ 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. 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); + MLXArray normalized = normalizedHiddenStates(ids, caches, rotation, postBlockHook); requirePostflight(caches); return tiedOutput ? embedding.project(normalized) : lmHead.forward(normalized); } 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 0000000..1217d19 --- /dev/null +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java @@ -0,0 +1,329 @@ +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 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 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); + } + } + + // --------------------------------------------------------------------- + // 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()); + } + } + } + + // --------------------------------------------------------------------- + // 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; + } +} From 4960bffe8acbf3a88541d4fab9e3f9d45688765c Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 16:22:15 +0200 Subject: [PATCH 04/32] phase7-3b WP4: visual injection helpers, exact-bits gate and acceptance docs --- AGENTS.md | 21 +- .../jmlx/models/VisualFeatureInjection.java | 140 +++++++ .../models/VisualFeatureInjectionTest.java | 363 ++++++++++++++++++ req/plans/phase7-3b-plan.md | 37 +- req/plans/phase7-3b-reference-findings.md | 49 +++ 5 files changed, 602 insertions(+), 8 deletions(-) create mode 100644 jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java create mode 100644 jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java diff --git a/AGENTS.md b/AGENTS.md index b967f06..aa4aa86 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 is deferred to WP8, and `Qwen3VlModel`, its loader +dispatch and the DeepStack consumption of `VisualFeatureInjection` are WP5/WP6. +No native pins, generated bindings, build files or CI changed. 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 0000000..01ae125 --- /dev/null +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java @@ -0,0 +1,140 @@ +package se.alipsa.jmlx.models; + +import java.util.Arrays; +import java.util.Objects; +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. They reuse {@link Idefics3VisionTower#mergeGatherIndices(int, + * java.util.List)} as the single shared gather-index algorithm — feature row {@code f} overwrites + * exactly the {@code f}-th destination and every other row reads itself — so one helper serves both + * compositions. All validation is checked before any MLX array is written, and every + * non-destination row comes back bit-identical (the masked add is an exact {@code + 0.0} per + * element there). 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 zero = MLX.zeros(scope, hidden.shape(), hidden.dtype()); + * source = concatenate([zero, 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 bit-identical through the exact {@code + 0.0} per + * element. + * + * @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) + * @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, + * 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); + int[] gather = gatherIndices(length, destinations); + MLXArray zero = MLX.zeros(scope, hidden.shape(), hidden.dtype()); + MLXArray source = MLXShape.concatenate(new MLXArray[] {zero, features}, 1); + MLXArray mask = + MLXShape.takeAxis(source, MLX.array(zero.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())); + } + + private static int validatePair( + String baseName, MLXArray base, String featuresName, MLXArray features, int[] destinations) { + Objects.requireNonNull(base, baseName); + Objects.requireNonNull(features, featuresName); + 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/VisualFeatureInjectionTest.java b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java new file mode 100644 index 0000000..264b8e3 --- /dev/null +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -0,0 +1,363 @@ +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 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 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 bit-for-bit, 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 + * 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} from the + * helper's own validation, before any MLX array is written. + */ +@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()); + } + } + + /** + * 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()); + } + } + + // --------------------------------------------------------------------- + // 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 shared gather-index algorithm, 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}. + */ + @Test + void rejectionsThrowBeforeAnyArrayIsWritten() { + try (MLXScope run = scope.newChild()) { + MLXArray embedded = array(run, rows(8, 0f)); + MLXArray features = array(run, rows(3, 10f)); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {3, 8})); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {-1})); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {4, 3})); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.replaceImageRows(embedded, features, null)); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {3, 4})); + assertThrows( + IllegalArgumentException.class, + () -> + VisualFeatureInjection.replaceImageRows( + embedded, array(run, rows(2, 10f)), new int[] {3, 4, 5})); + assertThrows( + IllegalArgumentException.class, + () -> + VisualFeatureInjection.replaceImageRows( + embedded, + array(run, flat(rows(3, 10f)), new int[] {3, HIDDEN}), + new int[] {3, 4, 5})); + assertThrows( + IllegalArgumentException.class, + () -> + VisualFeatureInjection.replaceImageRows( + embedded, + MLX.array(run, flat(rows(6, 10f)), new int[] {2, 3, HIDDEN}), + new int[] {3, 4, 5})); + assertThrows( + IllegalArgumentException.class, + () -> + VisualFeatureInjection.replaceImageRows( + embedded, + MLX.array(run, flat(rows(3, 10f)), new int[] {1, 3, HIDDEN + 1}), + new int[] {3, 4, 5})); + MLXArray half = MLX.astype(features, DType.FLOAT16); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.replaceImageRows(embedded, half, new int[] {3, 4, 5})); + + assertThrows( + IllegalArgumentException.class, + () -> + VisualFeatureInjection.addMaskedFeatures( + null, embedded, features, new int[] {3, 4, 5})); + assertThrows( + IllegalArgumentException.class, + () -> + VisualFeatureInjection.addMaskedFeatures( + run, + array(run, flat(rows(3, 10f)), new int[] {3, HIDDEN}), + features, + new int[] {3, 4, 5})); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.addMaskedFeatures(run, embedded, half, new int[] {3, 4, 5})); + assertThrows( + IllegalArgumentException.class, + () -> + VisualFeatureInjection.addMaskedFeatures( + run, embedded, array(run, rows(2, 10f)), new int[] {3, 4, 5})); + } + } + + // --------------------------------------------------------------------- + // 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 06983e1..f536b1e 100644 --- a/req/plans/phase7-3b-plan.md +++ b/req/plans/phase7-3b-plan.md @@ -1,10 +1,10 @@ # Phase 7.3b — Qwen3-VL image and text inference -- **Status:** WP0–WP3 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 in progress as of - 2026-10-20 — the pre-change exact-bits baseline is recorded and the core RoPE / decoder - seam design is frozen in `req/plans/phase7-3b-reference-findings.md` §6.9. WP5–WP8 - 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, @@ -399,6 +399,29 @@ 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, reusing +`Idefics3VisionTower`'s gather-index algorithm, 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 is WP8 per §9 — +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 @@ -605,8 +628,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 a788d2a..ce2c1c7 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -1020,6 +1020,55 @@ Adding the suite to `.github/workflows/ci.yml`'s required-execution list is defe WP8 (plan §9); 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: its component was last touched in Unit 2, where + its per-case max absolute error was ≤ 2.4e-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.** Adding the new native suites to + `.github/workflows/ci.yml`'s required-execution list is WP8 (plan §9); + `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 From 2bd62ce8b2d9e541762e91bde6e92c7ed45fac6b Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 16:58:00 +0200 Subject: [PATCH 05/32] phase7-3b WP4: fix two Javadoc typos from verification pass --- .../java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java | 2 +- .../src/main/java/se/alipsa/jmlx/models/DecoderModel.java | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) 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 index 71f62d1..320a42b 100644 --- a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java @@ -58,7 +58,7 @@ void uniformSingleRowAppendIsBitIdenticalToPlainAppend() { @Test void singleRowBatchedMaskMatchesTwoDimensionalMask() { try (MLXScope scope = new MLXScope()) { - // Prefill rows have T == S (no key padding), so the batched row is unpaddded and + // 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]; 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 e2cf2b5..52c9378 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 @@ -414,7 +414,7 @@ final MLXArray decoderStack( * 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. + * responsibilities of the 4-argument entry. */ final MLXArray decoderStack( MLXArray embeddings, From e01cfb3b948cd9144a8cc6a4470fab501479d428 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 17:24:46 +0200 Subject: [PATCH 06/32] review(WP4): correct VisualFeatureInjection -0.0 docs and record behaviour (PR #57) --- .../jmlx/models/VisualFeatureInjection.java | 17 +++--- .../models/VisualFeatureInjectionTest.java | 54 ++++++++++++++++--- req/plans/phase7-3b-reference-findings.md | 5 +- 3 files changed, 61 insertions(+), 15 deletions(-) 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 index 01ae125..d4c3f57 100644 --- a/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java @@ -23,11 +23,13 @@ * java.util.List)} as the single shared gather-index algorithm — feature row {@code f} overwrites * exactly the {@code f}-th destination and every other row reads itself — so one helper serves both * compositions. All validation is checked before any MLX array is written, and every - * non-destination row comes back bit-identical (the masked add is an exact {@code + 0.0} per - * element there). 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. + * 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 { @@ -60,8 +62,9 @@ static MLXArray replaceImageRows(MLXArray embedded, MLXArray features, int[] des * Masked residual add: the frozen {@code zero = MLX.zeros(scope, hidden.shape(), hidden.dtype()); * source = concatenate([zero, 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 bit-identical through the exact {@code + 0.0} per - * element. + * 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}). * * @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) 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 index 264b8e3..81fc2de 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -19,12 +19,13 @@ * 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 bit-for-bit, 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 - * 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} from the + * 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} from the * helper's own validation, before any MLX array is written. */ @EnabledIfNativeAvailable @@ -190,6 +191,47 @@ void zeroFeatureRowsAreAnExactPlusZeroIdentity() { } } + /** + * 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) // --------------------------------------------------------------------- diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index ce2c1c7..36ee6fd 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -964,8 +964,9 @@ joins) — no new public core API. ``` `gather` is the same index algorithm as the replacement (strictly increasing destinations → row `L + f`), so one helper serves both; the add matches the reference's - masked clone + add + write-back (§6.6) and leaves every non-image row bit-identical - (an exact `+ 0.0` per element). + 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).** From b050568ddf07b3bd62cec412ab3bd86720d16249 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 17:43:47 +0200 Subject: [PATCH 07/32] review(WP4): validate MultiAxisRotation shape before decoder allocations (PR #57) --- .../se/alipsa/jmlx/nn/DecoderAttention.java | 2 + .../se/alipsa/jmlx/nn/MultiAxisRotation.java | 45 ++++++++++++------- .../jmlx/nn/DecoderRotationSeamTest.java | 31 ++++++++++++- .../se/alipsa/jmlx/models/DecoderModel.java | 40 ++++++++++++----- .../alipsa/jmlx/models/DecoderSeamTest.java | 25 +++++++++++ 5 files changed, 113 insertions(+), 30 deletions(-) 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 77df6dd..27539fd 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 @@ -293,6 +293,8 @@ private MLXArray forwardInternal( 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 before any projection builds a graph node. + rotation.validate(batch, sequence, headDim); freqs = null; } else { freqs = staticFreqs != null ? staticFreqs : stepFrequencies; 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 index 5fac8a0..ff20401 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java @@ -48,6 +48,34 @@ 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 it 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"); + } + } + /** * 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 @@ -65,22 +93,7 @@ public MLXArray rotate(MLXArray x) { "rotate requires [batch, heads, sequence, headDim], got shape " + java.util.Arrays.toString(shape)); } - int[] cosShape = cos.shape(); - if (shape[0] != cosShape[0] || shape[2] != cosShape[2]) { - throw new IllegalArgumentException( - "x batch and sequence must match the prepared rotation (" - + shape[0] - + ", " - + shape[2] - + " against " - + cosShape[0] - + ", " - + cosShape[2] - + ")"); - } - if (rotaryDims > shape[3]) { - throw new IllegalArgumentException("rotaryDims exceeds head dimension"); - } + validate(shape[0], shape[2], shape[3]); DType dtype = x.dtype(); MLXArray cosInput = cos.dtype() == dtype ? cos : MLX.astype(cos, dtype); MLXArray sinInput = sin.dtype() == dtype ? sin : MLX.astype(sin, dtype); 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 index 0a38c84..ca88909 100644 --- a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java @@ -259,8 +259,9 @@ void blockRotationForwardMatchesManualComposition() { /** * The rotation overload must reject a missing {@code validLengths} (it always runs the batch - * bookkeeping) and a null rotation before any array is written, and the batch validation must - * reject a wrong number of lengths. + * 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. */ @Test void rotationForwardRejectsInvalidArguments() { @@ -283,6 +284,32 @@ void rotationForwardRejectsInvalidArguments() { 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")); assertEquals(0, cache.nextPosition()); } } 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 52c9378..8cbc347 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 @@ -208,8 +208,9 @@ final MLXArray stepLogits( * 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. The existing 4-argument overload delegates with {@code (null, null)} and keeps - * the identical graph. + * ones, unchanged. A non-null {@code rotation} is validated against the batch, the sequence width + * and the head dimension 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, @@ -219,6 +220,10 @@ final MLXArray stepLogits( 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.validate(tokenIds.shape()[0], tokenIds.shape()[1], descriptor.headDim()); + } try { MLXArray normalized = normalizedHiddenStatesBatch( @@ -409,12 +414,13 @@ final MLXArray decoderStack( * 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}). The batched mask is the same in both cases. 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. + * PostBlockHook}). The batched mask is the same in both cases. A non-null {@code rotation} is + * validated against the batch, the sequence width and the head dimension 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, @@ -449,6 +455,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.validate(batch, sequenceWidth, descriptor.headDim()); + } int[] queryStarts = new int[batch]; int[] keyStarts = new int[batch]; int keyWidth = 0; @@ -940,16 +950,22 @@ private MLXArray decode(MLXArray ids, List caches) { /** * 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. 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. + * seam-carrying stack and the same postflight and tied projection. A non-null {@code rotation} is + * validated against the batch, the sequence width and the head dimension 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"); } + if (rotation != null) { + // Reject a mismatched rotation before the embedding lookup builds any graph node. + rotation.validate(ids.shape()[0], ids.shape()[1], descriptor.headDim()); + } MLXArray normalized = normalizedHiddenStates(ids, caches, rotation, postBlockHook); requirePostflight(caches); return tiedOutput ? embedding.project(normalized) : lmHead.forward(normalized); 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 index 1217d19..7c264a5 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java @@ -293,6 +293,31 @@ void decodeRotationOverloadRejectsInvalidShapesBeforeCacheMutation() { } } + /** + * 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. + */ + @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)); + assertEquals(4, populated.getFirst().nextPosition()); + } + } + } + // --------------------------------------------------------------------- // helpers // --------------------------------------------------------------------- From 7699c348ac0da5f5a2fa57f05d52e402f6f28993 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 18:02:25 +0200 Subject: [PATCH 08/32] review(WP4): reject rotation width mismatching the configured rotary width (PR #57) --- .../se/alipsa/jmlx/nn/DecoderAttention.java | 10 ++++--- .../se/alipsa/jmlx/nn/MultiAxisRotation.java | 26 ++++++++++++++++--- .../jmlx/nn/DecoderRotationSeamTest.java | 15 ++++++++++- .../se/alipsa/jmlx/models/DecoderModel.java | 22 ++++++++++------ .../alipsa/jmlx/models/DecoderSeamTest.java | 26 ++++++++++++++++++- 5 files changed, 83 insertions(+), 16 deletions(-) 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 27539fd..4eb6b8f 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 @@ -203,7 +203,10 @@ public MLXArray forward( * 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. + * {@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, @@ -293,8 +296,9 @@ private MLXArray forwardInternal( 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 before any projection builds a graph node. - rotation.validate(batch, sequence, headDim); + // 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; 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 index ff20401..2e15c0e 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java @@ -53,9 +53,10 @@ public int rotaryDims() { * 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 it before any projection, mask or - * embedding builds a graph node, so a mismatched rotation is rejected before any array is - * written. + * 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(); @@ -76,6 +77,25 @@ public void validate(int batch, int sequence, int headDim) { } } + /** + * 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 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 index ca88909..0392a03 100644 --- a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java @@ -261,7 +261,9 @@ void blockRotationForwardMatchesManualComposition() { * 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. + * 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() { @@ -310,6 +312,17 @@ void rotationForwardRejectsInvalidArguments() { 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()); } } 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 8cbc347..9f1d5b1 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 @@ -208,9 +208,11 @@ final MLXArray stepLogits( * 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 - * and the head dimension before the embedding lookup builds any graph node. The existing - * 4-argument overload delegates with {@code (null, null)} and keeps the identical graph. + * 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, @@ -222,7 +224,8 @@ final MLXArray stepLogits( int[][] before = preflightBatch(tokenIds, caches, validLengths); if (rotation != null) { // Reject a mismatched rotation before the embedding lookup builds any graph node. - rotation.validate(tokenIds.shape()[0], tokenIds.shape()[1], descriptor.headDim()); + rotation.validateAgainst( + tokenIds.shape()[0], tokenIds.shape()[1], descriptor.headDim(), descriptor.rotaryDims()); } try { MLXArray normalized = @@ -415,7 +418,8 @@ final MLXArray decoderStack( * 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}). The batched mask is the same in both cases. A non-null {@code rotation} is - * validated against the batch, the sequence width and the head dimension before the mask or any + * 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 @@ -457,7 +461,7 @@ final MLXArray decoderStack( int batch = validLengths.length; if (rotation != null) { // Reject a mismatched rotation before the batched mask or any block builds a graph node. - rotation.validate(batch, sequenceWidth, descriptor.headDim()); + rotation.validateAgainst(batch, sequenceWidth, descriptor.headDim(), descriptor.rotaryDims()); } int[] queryStarts = new int[batch]; int[] keyStarts = new int[batch]; @@ -951,7 +955,8 @@ private MLXArray decode(MLXArray ids, List caches) { * 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 and the head dimension before the embedding + * 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 @@ -964,7 +969,8 @@ final MLXArray decode( } if (rotation != null) { // Reject a mismatched rotation before the embedding lookup builds any graph node. - rotation.validate(ids.shape()[0], ids.shape()[1], descriptor.headDim()); + rotation.validateAgainst( + ids.shape()[0], ids.shape()[1], descriptor.headDim(), descriptor.rotaryDims()); } MLXArray normalized = normalizedHiddenStates(ids, caches, rotation, postBlockHook); requirePostflight(caches); 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 index 7c264a5..f3740b8 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java @@ -3,6 +3,7 @@ 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; @@ -296,7 +297,11 @@ void decodeRotationOverloadRejectsInvalidShapesBeforeCacheMutation() { /** * 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. + * 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() { @@ -313,6 +318,25 @@ void seamEntriesRejectWrongWidthRotationBeforeArraysAreWritten() { () -> 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()); } } From e1ef27327a18904b3afd47fa28577ab1891049a3 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 18:25:46 +0200 Subject: [PATCH 09/32] review(WP4): pin hook-only single-row path against the old path (PR #57) --- .../se/alipsa/jmlx/models/DecoderModel.java | 15 +++- .../alipsa/jmlx/models/DecoderSeamTest.java | 77 ++++++++++++++++++- req/plans/phase7-3b-reference-findings.md | 12 ++- 3 files changed, 95 insertions(+), 9 deletions(-) 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 9f1d5b1..c4623c2 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 @@ -618,10 +618,17 @@ private MLXArray normalizedHiddenStates(MLXArray tokenIds, List caches) * 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. For a batch-1 unpadded row these are the same arithmetic as the - * old single-row path (the uniform append and the {@code [1,1,T,L]} batched mask match the - * single-row ones element for element, pinned by {@code SingleRowDecodeEquivalenceTest}), plus - * the hook firing after every block's residual output, before the final normalization. + * 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 (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, 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 index f3740b8..e8e99d0 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/DecoderSeamTest.java @@ -31,8 +31,10 @@ * 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 cases are tagged - * {@code full-float32} so they select into {@code :jmlx-models:float32GoldenTest} with {@code + * 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. @@ -256,6 +258,77 @@ void equalAxisRotationMatchesRopePathWithinKernelTolerance() { } } + /** + * 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) // --------------------------------------------------------------------- diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 36ee6fd..8a7f501 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -906,12 +906,18 @@ delegating to the new one with `null` rotation/hook so its graph is unchanged: (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` path keeps the old graph bit-for-bit, and the `rotation != null` - single-row path routes each block through the batch-aware overloads with + `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; `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 shared generation loop gets + 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); } From d684d9b0ac4ed00c76621718429c0d94774c4e78 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 18:40:08 +0200 Subject: [PATCH 10/32] review(WP4): align rotation-seam null-argument exceptions to IAE (PR #57) --- .../src/main/java/se/alipsa/jmlx/nn/DecoderAttention.java | 4 +++- jmlx-core/src/main/java/se/alipsa/jmlx/nn/DecoderBlock.java | 4 +++- .../src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java | 5 +++-- .../test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java | 2 +- .../src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java | 2 +- 5 files changed, 11 insertions(+), 6 deletions(-) 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 4eb6b8f..2b1ac8e 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 @@ -214,7 +214,9 @@ public MLXArray forward( MLXArray attentionMask, int[] validLengths, MultiAxisRotation rotation) { - Objects.requireNonNull(rotation, "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"); 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 c410f5d..f498925 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 @@ -102,7 +102,9 @@ public MLXArray forward( int[] validLengths, MultiAxisRotation rotation) { Objects.requireNonNull(x, "DecoderBlock.forward: x must not be null"); - Objects.requireNonNull(rotation, "DecoderBlock.forward: rotation 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"); 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 index 2e15c0e..ed57011 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java @@ -1,6 +1,5 @@ package se.alipsa.jmlx.nn; -import java.util.Objects; import se.alipsa.jmlx.core.DType; import se.alipsa.jmlx.core.MLX; import se.alipsa.jmlx.core.MLXArray; @@ -106,7 +105,9 @@ public void validateAgainst(int batch, int sequence, int headDim, int configured * allocates per the usual {@code scopeOf} rule (the innermost related scope of its operands). */ public MLXArray rotate(MLXArray x) { - Objects.requireNonNull(x, "x"); + if (x == null) { + throw new IllegalArgumentException("x must not be null"); + } int[] shape = x.shape(); if (shape.length != 4) { throw new IllegalArgumentException( 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 index 0392a03..f8405a7 100644 --- a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/DecoderRotationSeamTest.java @@ -281,7 +281,7 @@ void rotationForwardRejectsInvalidArguments() { () -> fixture.attention.forward(x, cache, null, null, rotation)); assertTrue(missingLengths.getMessage().contains("valid lengths")); assertThrows( - NullPointerException.class, + IllegalArgumentException.class, () -> fixture.attention.forward(x, cache, null, new int[] {3}, null)); assertThrows( IllegalArgumentException.class, 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 index 1a1e67a..7f17cd7 100644 --- a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java @@ -133,7 +133,7 @@ void rotateRejectsInvalidInputs() { () -> rotation.rotate( MLX.array(scope, values(1 * 1 * 3 * 16, 3), new int[] {1, 1, 3, 16}))); - assertThrows(NullPointerException.class, () -> rotation.rotate(null)); + assertThrows(IllegalArgumentException.class, () -> rotation.rotate(null)); } } From bd9e8c74d3285dbbcbc1f3cd47f6411492dd7cab Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 19:22:57 +0200 Subject: [PATCH 11/32] review(WP4): add the five native WP4 suites to the CI required-suite list (PR #57) --- .github/workflows/ci.yml | 11 ++++++++--- AGENTS.md | 6 +++--- req/plans/phase7-3b-plan.md | 4 ++-- req/plans/phase7-3b-reference-findings.md | 12 ++++++------ 4 files changed, 19 insertions(+), 14 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 1a15377..2005b11 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 aa4aa86..50f265d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -679,6 +679,6 @@ verified against the pre-change exact-bits baseline (`exactBitsVerifyTf32On`/`Of 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 is deferred to WP8, and `Qwen3VlModel`, its loader -dispatch and the DeepStack consumption of `VisualFeatureInjection` are WP5/WP6. -No native pins, generated bindings, build files or CI changed. +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/req/plans/phase7-3b-plan.md b/req/plans/phase7-3b-plan.md index f536b1e..6445a44 100644 --- a/req/plans/phase7-3b-plan.md +++ b/req/plans/phase7-3b-plan.md @@ -419,8 +419,8 @@ 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 is WP8 per §9 — -and no `Qwen3VlModel` or loader dispatch; those are WP5/WP6. +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) diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 8a7f501..0eaf01e 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -1023,8 +1023,8 @@ the rejection cases running in ordinary `:jmlx-core:test`): - head-dimension validation and the dtype-cast check that cos/sin are cast to the input dtype before the rotation. -Adding the suite to `.github/workflows/ci.yml`'s required-execution list is deferred to -WP8 (plan §9); the `full-float32` tag selection into `:jmlx-core:float32GoldenTest` +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 @@ -1071,10 +1071,10 @@ the acceptance commit (Unit 4 — `VisualFeatureInjection` and `VisualFeatureInj its per-case max absolute error was ≤ 2.4e-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.** Adding the new native suites to - `.github/workflows/ci.yml`'s required-execution list is WP8 (plan §9); - `Qwen3VlModel`, the `qwen3_vl` loader dispatch and the DeepStack consumption of - `VisualFeatureInjection` are WP5/WP6. +- **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) From 23d968b07797a8a07cf51031c75b5d1425873225 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 19:30:04 +0200 Subject: [PATCH 12/32] review(WP4): throw IAE for null arguments in VisualFeatureInjection (PR #57) --- .../jmlx/models/VisualFeatureInjection.java | 9 ++++++--- .../jmlx/models/VisualFeatureInjectionTest.java | 16 ++++++++++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) 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 index d4c3f57..f008172 100644 --- a/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java @@ -1,7 +1,6 @@ package se.alipsa.jmlx.models; import java.util.Arrays; -import java.util.Objects; import java.util.stream.Collectors; import java.util.stream.IntStream; import se.alipsa.jmlx.core.MLX; @@ -103,8 +102,12 @@ private static int[] gatherIndices(int length, int[] destinations) { private static int validatePair( String baseName, MLXArray base, String featuresName, MLXArray features, int[] destinations) { - Objects.requireNonNull(base, baseName); - Objects.requireNonNull(features, featuresName); + 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"); } 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 index 81fc2de..bd1c602 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -260,6 +260,14 @@ void rejectionsThrowBeforeAnyArrayIsWritten() { assertThrows( IllegalArgumentException.class, () -> VisualFeatureInjection.replaceImageRows(embedded, features, null)); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.replaceImageRows(null, features, new int[] {3, 4, 5}), + "embedded must not be null"); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.replaceImageRows(embedded, null, new int[] {3, 4, 5}), + "features must not be null"); assertThrows( IllegalArgumentException.class, () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {3, 4})); @@ -299,6 +307,14 @@ void rejectionsThrowBeforeAnyArrayIsWritten() { () -> VisualFeatureInjection.addMaskedFeatures( null, embedded, features, new int[] {3, 4, 5})); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.addMaskedFeatures(run, null, features, new int[] {3, 4, 5}), + "hidden must not be null"); + assertThrows( + IllegalArgumentException.class, + () -> VisualFeatureInjection.addMaskedFeatures(run, embedded, null, new int[] {3, 4, 5}), + "features must not be null"); assertThrows( IllegalArgumentException.class, () -> From af8d25f9c7229c54d79c5a7162790c05a792d12d Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 20:25:19 +0200 Subject: [PATCH 13/32] review(WP4): pin windowed evicted-cache single-row mask equivalence (PR #57) --- .../nn/SingleRowDecodeEquivalenceTest.java | 31 +++++ .../se/alipsa/jmlx/models/DecoderModel.java | 8 +- .../jmlx/models/DecoderSlidingWindowTest.java | 128 ++++++++++++++++++ req/plans/phase7-3b-reference-findings.md | 4 +- 4 files changed, 165 insertions(+), 6 deletions(-) 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 index 320a42b..64fd5b0 100644 --- a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/SingleRowDecodeEquivalenceTest.java @@ -84,6 +84,37 @@ void singleRowBatchedMaskMatchesTwoDimensionalMask() { 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)); + } } } 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 c4623c2..82f79c7 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 @@ -620,10 +620,10 @@ private MLXArray normalizedHiddenStates(MLXArray tokenIds, List caches) * #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 (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 + * 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 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 e96f99f..f06ecee 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/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 0eaf01e..24d5cdd 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -910,8 +910,8 @@ delegating to the new one with `null` rotation/hook so its graph is unchanged: 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; `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` + 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 From 557c200c635cff2449890b1119eabf577c3bb560 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 20:52:44 +0200 Subject: [PATCH 14/32] review(WP4): cache per-dtype rotation tables and simplify prepare uploads (PR #57) --- .../java/se/alipsa/jmlx/nn/MultiAxisRope.java | 58 ++++++++++--------- .../se/alipsa/jmlx/nn/MultiAxisRotation.java | 33 +++++++++-- req/plans/phase7-3b-reference-findings.md | 23 ++++---- 3 files changed, 73 insertions(+), 41 deletions(-) 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 index d07b094..ec9f515 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java @@ -114,16 +114,17 @@ public int rotaryDims() { * 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) and no coordinate is negative. The frozen compute order then runs in - * {@code scope}: three float32 position arrays of shape {@code [B*T, 1]} (built directly from the - * int rows, an exact cast) broadcast-multiplied against the {@code [1, rotaryDims/2]} {@code - * invFreq} array give the per-axis frequencies; the static owner masks select {@code freqs_t = - * freqsT * wT + freqsH * wH + freqsW * wW} (multiplying by the static 0/1 masks and adding in - * float32 is a bit-identical copy for non-negative finite frequencies); {@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. + * product cannot overflow) and no coordinate is negative. The frozen compute order then runs: the + * static owner masks select, 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) { @@ -158,29 +159,34 @@ public MultiAxisRotation prepare(MLXScope scope, int batch, int sequence, int[][ int flat = (int) tokens; int half = rotaryDims / 2; MLXArray invFreqArray = MLX.array(scope, invFreq, new int[] {1, half}); - MLXArray weightT = MLX.array(scope, ownerT, new int[] {1, half}); - MLXArray weightH = MLX.array(scope, ownerH, new int[] {1, half}); - MLXArray weightW = MLX.array(scope, ownerW, new int[] {1, half}); - MLXArray freqsT = axisFrequencies(scope, positions[0], flat, invFreqArray); - MLXArray freqsH = axisFrequencies(scope, positions[1], flat, invFreqArray); - MLXArray freqsW = axisFrequencies(scope, positions[2], flat, invFreqArray); - MLXArray selected = - MLXOps.add( - MLXOps.add(MLXOps.multiply(freqsT, weightT), MLXOps.multiply(freqsH, weightH)), - MLXOps.multiply(freqsW, weightW)); - MLXArray emb = MLXShape.concatenate(new MLXArray[] {selected, selected}, 1); + 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); } - private static MLXArray axisFrequencies(MLXScope scope, int[] row, int flat, MLXArray invFreq) { - float[] data = new float[flat]; + /** + * 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 masks 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++) { - data[i] = (float) row[i]; + for (int f = 0; f < half; f++) { + int coordinate = + ownerT[f] != 0f ? positions[0][i] : ownerH[f] != 0f ? positions[1][i] : positions[2][i]; + out[i * half + f] = (float) coordinate; + } } - MLXArray positions = MLX.array(scope, data, new int[] {flat, 1}); - return MLXOps.multiply(positions, invFreq); + return out; } /** 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 index ed57011..d4c5eed 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRotation.java @@ -1,5 +1,6 @@ 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; @@ -25,11 +26,15 @@ 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]}. */ @@ -100,9 +105,11 @@ public void validateAgainst(int batch, int sequence, int headDim, int configured * 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) 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). + * 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) { @@ -116,8 +123,8 @@ public MLXArray rotate(MLXArray x) { } validate(shape[0], shape[2], shape[3]); DType dtype = x.dtype(); - MLXArray cosInput = cos.dtype() == dtype ? cos : MLX.astype(cos, dtype); - MLXArray sinInput = sin.dtype() == dtype ? sin : MLX.astype(sin, 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); @@ -137,4 +144,20 @@ public MLXArray rotate(MLXArray x) { 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/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 24d5cdd..40fe225 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -787,11 +787,13 @@ mode with explicit `MLX_ENABLE_TF32=1` / `=0`) PASSES at the base commit with a `[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 static owner masks: three - float32 weight vectors of length `rotaryDims/2` (1.0 at the owned frequencies, 0.0 - elsewhere, derived from the sections at construction) give - `freqs_t = freqsT·wT + freqsH·wH + freqsW·wW` — a bit-identical copy for non-negative - finite frequencies (multiply by 1.0/0.0 and add in float32 are exact there) using only - the existing facades. + 0/1 weight vectors of length `rotaryDims/2` (1 at the owned frequencies, 0 elsewhere, + derived from the sections at construction) select, 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 @@ -830,11 +832,12 @@ source compatibility for exhaustive consumer switches (plan §6). 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, - in `scope`, follows the frozen order: three float32 position arrays (built directly from - the int rows — the exact cast) × the `[1, rotaryDims/2]` `invFreq` array by broadcast - multiply → the per-axis `[B·T, rotaryDims/2]` freqs; owner-mask selection → `freqs_t`; - concatenate → `emb`; `MLXOps.cos`/`MLXOps.sin` → FLOAT32 arrays reshaped to + 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). - `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 From 2c40b6d0ad9d7f1750175d21e0de06ca38bd40dc Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 21:34:14 +0200 Subject: [PATCH 15/32] review(WP4): guard the ownedPositions matrix size against int overflow (PR #57) --- .../java/se/alipsa/jmlx/nn/MultiAxisRope.java | 28 +++++++++++-------- .../se/alipsa/jmlx/nn/MultiAxisRopeTest.java | 21 ++++++++++++++ req/plans/phase7-3b-reference-findings.md | 6 +++- 3 files changed, 43 insertions(+), 12 deletions(-) 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 index ec9f515..cc385ce 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java @@ -114,17 +114,20 @@ public int rotaryDims() { * 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) and no coordinate is negative. The frozen compute order then runs: the - * static owner masks select, 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. + * 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 masks select, 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) { @@ -158,6 +161,9 @@ public MultiAxisRotation prepare(MLXScope scope, int batch, int sequence, int[][ } 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}); 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 index 7f17cd7..6f844b1 100644 --- a/jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java +++ b/jmlx-core/src/test/java/se/alipsa/jmlx/nn/MultiAxisRopeTest.java @@ -108,6 +108,27 @@ void prepareRejectsInvalidShapeAndPositions() { } } + /** + * 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()) { diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 40fe225..bc8fb4c 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -838,7 +838,11 @@ source compatibility for exhaustive consumer switches (plan §6). 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). + `[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 owner masks) — 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`: From 63f1cabf608837b2302f2801900da46bc0717ebd Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 21:47:02 +0200 Subject: [PATCH 16/32] review(WP4): record the round-3 probe rerun of the interleaved RoPE cross-check (PR #57) --- req/plans/phase7-3b-reference-findings.md | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index bc8fb4c..44525c6 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -1074,8 +1074,13 @@ the acceptance commit (Unit 4 — `VisualFeatureInjection` and `VisualFeatureInj 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: its component was last touched in Unit 2, where - its per-case max absolute error was ≤ 2.4e-7 against the 1e-5 threshold. + 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) and the row-length + overflow guard (review round 3), head `2c40b6d` — 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 From bf69877dc931682b25e48bcf8cf0f23282f91850 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 22:07:26 +0200 Subject: [PATCH 17/32] review(WP4): replace the three float owner masks with one byte owner table (PR #57) --- .../java/se/alipsa/jmlx/nn/MultiAxisRope.java | 32 ++++++++----------- req/plans/phase7-3b-reference-findings.md | 9 +++--- 2 files changed, 18 insertions(+), 23 deletions(-) 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 index cc385ce..801888a 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java @@ -35,13 +35,11 @@ public final class MultiAxisRope { private final int[] mropeSection; private final int rotaryDims; private final float[] invFreq; - private final float[] ownerT; - private final float[] ownerH; - private final float[] ownerW; + private final byte[] owner; /** * Validates the base and the sections and derives the per-frequency inverse and the static T/H/W - * owner masks. All validation is {@link IllegalArgumentException} and runs before any MLX array + * 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 @@ -75,21 +73,19 @@ public MultiAxisRope(float theta, int[] mropeSection) { this.mropeSection = mropeSection.clone(); this.rotaryDims = 2 * half; this.invFreq = deriveInvFrequencies(theta, half); - this.ownerT = new float[half]; - this.ownerH = new float[half]; - this.ownerW = new float[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 masks realize that - // selection as a static 0/1 weight vector per axis. + // 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]) { - ownerH[f] = 1f; + owner[f] = 1; } else if (f % 3 == 2 && f < 2L + 3L * mropeSection[2]) { - ownerW[f] = 1f; - } else { - ownerT[f] = 1f; + owner[f] = 2; } } } @@ -117,7 +113,7 @@ public int rotaryDims() { * 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 masks select, per (token, frequency) on + * 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 @@ -176,8 +172,8 @@ public MultiAxisRotation prepare(MLXScope scope, int batch, int sequence, int[][ /** * 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 masks before any array is written. This is the bit-identical copy the frozen mask + * 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 @@ -187,9 +183,7 @@ 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++) { - int coordinate = - ownerT[f] != 0f ? positions[0][i] : ownerH[f] != 0f ? positions[1][i] : positions[2][i]; - out[i * half + f] = (float) coordinate; + out[i * half + f] = (float) positions[owner[f]][i]; } } return out; diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 44525c6..0f2a4ef 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -786,10 +786,11 @@ mode with explicit `MLX_ENABLE_TF32=1` / `=0`) PASSES at the base commit with a 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 static owner masks: three - 0/1 weight vectors of length `rotaryDims/2` (1 at the owned frequencies, 0 elsewhere, - derived from the sections at construction) select, per (token, frequency) on the CPU - before any array is written, the owned row's position — the same copy the mask expression + 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 · From 3ac98ba81e65328965ca98a319ade6b6e4f3b17b Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 22:37:34 +0200 Subject: [PATCH 18/32] review(WP4): re-run the interleaved RoPE probe at the final head and update its record (PR #57) --- req/plans/phase7-3b-reference-findings.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 0f2a4ef..752c34c 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -1078,8 +1078,9 @@ the acceptance commit (Unit 4 — `VisualFeatureInjection` and `VisualFeatureInj 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) and the row-length - overflow guard (review round 3), head `2c40b6d` — Python exit 0, all 7 cases passed; + review rounds, against the restructured `prepare` (review round 2), the row-length + 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 From c13453876898afed2762c65164a508fc835b863a Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 22:41:36 +0200 Subject: [PATCH 19/32] review(WP4): attribute the pre-restructure 3*B*T host floats to the per-axis position arrays (PR #57) --- req/plans/phase7-3b-reference-findings.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 752c34c..863c8e2 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -841,7 +841,7 @@ source compatibility for exhaustive consumer switches (plan §6). 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 owner masks) — about 21× more for + 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` From ec4e7fda85569be1f57c34a18b7a32e8dd793524 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Tue, 20 Oct 2026 22:46:46 +0200 Subject: [PATCH 20/32] review(WP4): name the round-3 guard as the position-matrix size overflow guard (PR #57) --- req/plans/phase7-3b-reference-findings.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 863c8e2..c2a4325 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -1078,8 +1078,8 @@ the acceptance commit (Unit 4 — `VisualFeatureInjection` and `VisualFeatureInj 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 row-length - overflow guard (review round 3) and the byte owner-table refactor (review round 4), + 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. From e856b08dba99d4cf549ec82d6d2ee18255361cf4 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 10:15:53 +0200 Subject: [PATCH 21/32] review(WP4): document the partial-rotary tail as a jmlx extension and record the WP5 loader guard (PR #57) --- .../src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java | 9 +++++++++ req/plans/phase7-3b-plan.md | 7 +++++++ 2 files changed, 16 insertions(+) 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 index 801888a..a25575a 100644 --- a/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java +++ b/jmlx-core/src/main/java/se/alipsa/jmlx/nn/MultiAxisRope.java @@ -16,6 +16,15 @@ * 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 diff --git a/req/plans/phase7-3b-plan.md b/req/plans/phase7-3b-plan.md index 6445a44..780acfe 100644 --- a/req/plans/phase7-3b-plan.md +++ b/req/plans/phase7-3b-plan.md @@ -428,6 +428,13 @@ Add `Qwen3VlModel implements TextGenerationModel`, composing the decoder so the 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 From 4dfcdfe2c16060cfdf7eac894333b2c3a67d069c Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 10:28:35 +0200 Subject: [PATCH 22/32] review(WP4): pad the masked DeepStack add with a single zero row instead of a full-sequence zero tensor (PR #57) --- .../jmlx/models/VisualFeatureInjection.java | 103 ++++++++++++++---- .../models/VisualFeatureInjectionTest.java | 29 ++++- 2 files changed, 107 insertions(+), 25 deletions(-) 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 index f008172..c8f4a1f 100644 --- a/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java @@ -1,6 +1,7 @@ package se.alipsa.jmlx.models; import java.util.Arrays; +import java.util.Objects; import java.util.stream.Collectors; import java.util.stream.IntStream; import se.alipsa.jmlx.core.MLX; @@ -18,17 +19,21 @@ * *

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. They reuse {@link Idefics3VisionTower#mergeGatherIndices(int, - * java.util.List)} as the single shared gather-index algorithm — feature row {@code f} overwrites - * exactly the {@code f}-th destination and every other row reads itself — so one helper serves both - * compositions. 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. + * primitive is required. The row-gather replacement reuses {@link Idefics3VisionTower# + * mergeGatherIndices(int, java.util.List)} — feature row {@code f} overwrites exactly the {@code + * f}-th destination and every other row reads itself. The masked add instead uses a dedicated index + * builder whose non-destination rows all read the single zero row at the front of the source, 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 { @@ -58,12 +63,15 @@ static MLXArray replaceImageRows(MLXArray embedded, MLXArray features, int[] des } /** - * Masked residual add: the frozen {@code zero = MLX.zeros(scope, hidden.shape(), hidden.dtype()); - * source = concatenate([zero, 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}). + * 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 dedicated gather index of the source's {@code + * 1 + F} rows maps every non-destination row to the zero row (source index {@code 0}) and the + * {@code f}-th destination 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) @@ -81,11 +89,11 @@ static MLXArray addMaskedFeatures( throw new IllegalArgumentException("scope must not be null"); } int length = validatePair("hidden", hidden, "features", features, destinations); - int[] gather = gatherIndices(length, destinations); - MLXArray zero = MLX.zeros(scope, hidden.shape(), hidden.dtype()); - MLXArray source = MLXShape.concatenate(new MLXArray[] {zero, features}, 1); + int[] gather = maskedAddGatherIndices(length, features.shape()[1], destinations); + 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(zero.scope(), gather, new int[] {gather.length}), 1); + MLXShape.takeAxis(source, MLX.array(scope, gather, new int[] {gather.length}), 1); return MLXOps.add(hidden, mask); } @@ -100,6 +108,59 @@ private static int[] gatherIndices(int length, int[] destinations) { length, IntStream.of(destinations).boxed().collect(Collectors.toList())); } + /** + * The masked add's dedicated gather indices: the {@code length}-index array into a source of + * {@code 1 + featureCount} rows — one zero row followed by the feature rows — where every + * non-destination row reads the zero row (source index {@code 0}) and the {@code f}-th + * destination reads source row {@code 1 + f}. It validates the destinations exactly as {@link + * Idefics3VisionTower#mergeGatherIndices(int, java.util.List)} does (strictly increasing in + * {@code [0, length)}) and checks the {@code 1 + featureCount} source-row sum with checked + * arithmetic, where the shared helper checks the {@code length + featureCount} sum instead. + */ + private static int[] maskedAddGatherIndices(int length, int featureCount, int[] destinations) { + Objects.requireNonNull(destinations, "destinations"); + if (length < 0) { + throw new IllegalArgumentException("length must be non-negative"); + } + if (destinations.length != featureCount) { + throw new IllegalArgumentException( + "featureCount " + + featureCount + + " must equal the destination count " + + destinations.length); + } + int previous = -1; + for (int f = 0; f < featureCount; f++) { + int destination = destinations[f]; + if (destination < 0 || destination >= length) { + throw new IllegalArgumentException( + "feature destination " + destination + " is out of range [0, " + length + ")"); + } + if (destination <= previous) { + throw new IllegalArgumentException( + "feature destinations must be strictly increasing; index " + + f + + " has " + + destination + + " after " + + previous); + } + previous = destination; + } + Math.addExact(1, featureCount); + int[] indices = new int[length]; + int next = 0; + for (int i = 0; i < length; i++) { + if (next < featureCount && destinations[next] == i) { + indices[i] = 1 + next; + next++; + } else { + indices[i] = 0; + } + } + return indices; + } + private static int validatePair( String baseName, MLXArray base, String featuresName, MLXArray features, int[] destinations) { if (base == null) { 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 index bd1c602..0b5ecef 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -140,6 +140,27 @@ void singleLevelAddsFeaturesIntoExactlyTheImageRows() { } } + /** + * 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 @@ -238,10 +259,10 @@ void maskedAddNormalisesNegativeZeroInAnUntouchedRow() { /** * Every invalid input is rejected with {@link IllegalArgumentException} by the helper's own - * validation — destination semantics via the shared gather-index algorithm, 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}. + * 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}. */ @Test void rejectionsThrowBeforeAnyArrayIsWritten() { From 837728e382e4ba85719ef8e7e0a5d2658dab0ca1 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 10:38:25 +0200 Subject: [PATCH 23/32] review(WP4): validate addMaskedFeatures scope compatibility before allocating (PR #57) --- .../jmlx/models/VisualFeatureInjection.java | 35 +++++++++++++++++-- .../models/VisualFeatureInjectionTest.java | 24 +++++++++++++ 2 files changed, 57 insertions(+), 2 deletions(-) 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 index c8f4a1f..44e9783 100644 --- a/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java @@ -74,14 +74,17 @@ static MLXArray replaceImageRows(MLXArray embedded, MLXArray features, int[] des * {@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) + * 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, - * before any MLX array is written + * 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) { @@ -89,6 +92,7 @@ static MLXArray addMaskedFeatures( throw new IllegalArgumentException("scope must not be null"); } int length = validatePair("hidden", hidden, "features", features, destinations); + checkRelatedScopes("scope", scope, hidden, features); int[] gather = maskedAddGatherIndices(length, features.shape()[1], destinations); MLXArray zeroRow = MLX.zeros(scope, new int[] {1, 1, hidden.shape()[2]}, hidden.dtype()); MLXArray source = MLXShape.concatenate(new MLXArray[] {zeroRow, features}, 1); @@ -161,6 +165,33 @@ private static int[] maskedAddGatherIndices(int length, int featureCount, int[] return indices; } + /** + * 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)"); + } + } + private static int validatePair( String baseName, MLXArray base, String featuresName, MLXArray features, int[] destinations) { if (base == null) { 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 index 0b5ecef..58e756f 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -355,6 +355,30 @@ void rejectionsThrowBeforeAnyArrayIsWritten() { } } + /** + * 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() { + try (MLXScope child = scope.newChild(); + MLXScope sibling = scope.newChild()) { + MLXArray hidden = array(child, rows(8, 0f)); + MLXArray features = array(child, rows(3, 10f)); + IllegalArgumentException exception = + assertThrows( + IllegalArgumentException.class, + () -> + VisualFeatureInjection.addMaskedFeatures( + sibling, hidden, features, new int[] {3, 4, 5})); + assertEquals( + "scope must be related to hidden's and features' scopes (one an ancestor of the other)", + exception.getMessage()); + } + } + // --------------------------------------------------------------------- // data builders (Java-side expectations) // --------------------------------------------------------------------- From 14168923d01e7d7a526ec61578dfc4e77dfc1164 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 10:44:19 +0200 Subject: [PATCH 24/32] review(WP4): record the seam-path explicit-SDPA-mask decode benchmark as a pending WP6 gate item (PR #57) --- req/plans/phase7-3b-plan.md | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/req/plans/phase7-3b-plan.md b/req/plans/phase7-3b-plan.md index 780acfe..300589b 100644 --- a/req/plans/phase7-3b-plan.md +++ b/req/plans/phase7-3b-plan.md @@ -472,6 +472,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. @@ -612,7 +623,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 | From 1ebc2912b7750248fcf37759492a94a190097c5e Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 11:04:38 +0200 Subject: [PATCH 25/32] review(WP4): assert the rejection messages and the no-allocation promise in the pre-write rejection test (PR #57) --- .../models/VisualFeatureInjectionTest.java | 160 +++++++++--------- 1 file changed, 84 insertions(+), 76 deletions(-) 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 index 58e756f..5549c84 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -8,9 +8,11 @@ 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; @@ -25,8 +27,9 @@ * 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} from the - * helper's own validation, before any MLX array is written. + * 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 { @@ -262,99 +265,104 @@ void maskedAddNormalisesNegativeZeroInAnUntouchedRow() { * 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}. + * 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() { + void rejectionsThrowBeforeAnyArrayIsWritten() throws InterruptedException { try (MLXScope run = scope.newChild()) { - MLXArray embedded = array(run, rows(8, 0f)); - MLXArray features = array(run, rows(3, 10f)); - assertThrows( - IllegalArgumentException.class, - () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {3, 8})); - assertThrows( - IllegalArgumentException.class, - () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {-1})); - assertThrows( - IllegalArgumentException.class, - () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {4, 3})); - assertThrows( - IllegalArgumentException.class, - () -> VisualFeatureInjection.replaceImageRows(embedded, features, null)); - assertThrows( - IllegalArgumentException.class, + 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"); - assertThrows( - IllegalArgumentException.class, + assertRejection( () -> VisualFeatureInjection.replaceImageRows(embedded, null, new int[] {3, 4, 5}), "features must not be null"); - assertThrows( - IllegalArgumentException.class, - () -> VisualFeatureInjection.replaceImageRows(embedded, features, new int[] {3, 4})); - assertThrows( - IllegalArgumentException.class, - () -> - VisualFeatureInjection.replaceImageRows( - embedded, array(run, rows(2, 10f)), new int[] {3, 4, 5})); - assertThrows( - IllegalArgumentException.class, - () -> - VisualFeatureInjection.replaceImageRows( - embedded, - array(run, flat(rows(3, 10f)), new int[] {3, HIDDEN}), - new int[] {3, 4, 5})); - assertThrows( - IllegalArgumentException.class, + 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, - MLX.array(run, flat(rows(6, 10f)), new int[] {2, 3, HIDDEN}), - new int[] {3, 4, 5})); - assertThrows( - IllegalArgumentException.class, - () -> - VisualFeatureInjection.replaceImageRows( - embedded, - MLX.array(run, flat(rows(3, 10f)), new int[] {1, 3, HIDDEN + 1}), - new int[] {3, 4, 5})); - MLXArray half = MLX.astype(features, DType.FLOAT16); - assertThrows( - IllegalArgumentException.class, - () -> VisualFeatureInjection.replaceImageRows(embedded, half, new int[] {3, 4, 5})); - - assertThrows( - IllegalArgumentException.class, + 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})); - assertThrows( - IllegalArgumentException.class, + 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"); - assertThrows( - IllegalArgumentException.class, + assertRejection( () -> VisualFeatureInjection.addMaskedFeatures(run, embedded, null, new int[] {3, 4, 5}), "features must not be null"); - assertThrows( - IllegalArgumentException.class, + assertRejection( () -> - VisualFeatureInjection.addMaskedFeatures( - run, - array(run, flat(rows(3, 10f)), new int[] {3, HIDDEN}), - features, - new int[] {3, 4, 5})); - assertThrows( - IllegalArgumentException.class, - () -> VisualFeatureInjection.addMaskedFeatures(run, embedded, half, new int[] {3, 4, 5})); - assertThrows( - IllegalArgumentException.class, + 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, array(run, rows(2, 10f)), new int[] {3, 4, 5})); + 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 From 2ec446049aa541db5baa43ab34815258d525ccb1 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 11:15:39 +0200 Subject: [PATCH 26/32] review(WP4): validate the PostBlockHook result shape in decoderStack (PR #57) --- .../se/alipsa/jmlx/models/DecoderModel.java | 41 ++++++++-- .../jmlx/models/DecoderStackContractTest.java | 76 ++++++++++++++++++- 2 files changed, 109 insertions(+), 8 deletions(-) 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 82f79c7..9c2ace3 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 @@ -178,7 +178,11 @@ interface EmbeddingHook { * 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. + * Production ordinary decoders pass no-op or null hooks. Implementations must return a non-null + * array with exactly the same shape 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 { @@ -373,6 +377,24 @@ 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 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[] expected) { + if (hidden == null || !Arrays.equals(expected, hidden.shape())) { + throw new IllegalArgumentException( + "postBlockHook for block " + + blockIndex + + " must return the block's hidden state with the same shape " + + Arrays.toString(expected) + + ": got " + + (hidden == null ? "null" : Arrays.toString(hidden.shape()))); + } + } + /** * 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 @@ -417,12 +439,14 @@ final MLXArray decoderStack( * 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}). 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 + * PostBlockHook}). Its result is checked immediately: a null or differently shaped 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. */ @@ -496,7 +520,10 @@ final MLXArray decoderStack( ? 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(); x = postBlockHook.apply(i, x); + requirePostBlockResult(i, x, expected); } } return norm.forward(x); 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 2bedfb7..89d3add 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 @@ -20,7 +20,8 @@ * 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 + * advances, a post-block hook returning a null or differently shaped 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 @@ -238,6 +239,79 @@ 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 + + "]: got [" + + BATCH + + ", " + + WIDTH + + ", " + + (hidden + 1) + + "]", + failure.getMessage()); + } + } + + /** + * A post-block hook that returns null is rejected by the same check, naming the block index and + * the expected shape. + */ + @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 + + "]: got null", + failure.getMessage()); + } + } + @Test void shapePreflightRejectsPoisonedCache() { try (MLXScope run = scope.newChild()) { From 7ec882156587df727b19965398d0be0bbb8dcfd3 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 11:19:46 +0200 Subject: [PATCH 27/32] =?UTF-8?q?review(WP4):=20align=20the=20findings=20?= =?UTF-8?q?=C2=A76.9=20masked-add=20composition=20record=20with=20the=20si?= =?UTF-8?q?ngle-zero-row=20change=20(PR=20#57)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- req/plans/phase7-3b-plan.md | 5 +++-- req/plans/phase7-3b-reference-findings.md | 21 +++++++++++++-------- 2 files changed, 16 insertions(+), 10 deletions(-) diff --git a/req/plans/phase7-3b-plan.md b/req/plans/phase7-3b-plan.md index 300589b..fecd01d 100644 --- a/req/plans/phase7-3b-plan.md +++ b/req/plans/phase7-3b-plan.md @@ -410,8 +410,9 @@ post-block injection callback across `DecoderAttention`/`DecoderBlock`/`DecoderM 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, reusing -`Idefics3VisionTower`'s gather-index algorithm, no new primitive, no FFI from model +replacement and masked DeepStack residual add over existing core facades; the +replacement reuses `Idefics3VisionTower`'s gather-index algorithm, the masked add its +own dedicated builder 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 diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index c2a4325..2c07838 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -971,16 +971,21 @@ joins) — no new public core API. - **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): ``` - zero = MLX.zeros(scope, hidden.shape(), hidden.dtype()) // [1, L, hidden] - source = concatenate([zero, features_i], 1) // [1, L + F_i, hidden] - mask = takeAxis(source, MLXArray(gather(L, destinations_i), [L]), 1) + zeroRow = MLX.zeros(scope, {1, 1, hidden}, hidden.dtype()) // [1, 1, hidden] + source = concatenate([zeroRow, features_i], 1) // [1, 1 + F_i, hidden] + mask = takeAxis(source, MLXArray(maskedAddGatherIndices(L, F_i, destinations_i), [L]), 1) updated = MLXOps.add(hidden, mask) ``` - `gather` is the same index algorithm as the replacement (strictly increasing - destinations → row `L + f`), so one helper serves both; 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`). + The gather indices come from the masked add's own builder — every non-destination row + reads the zero row (source index `0`), the `f`-th destination (strictly increasing) + reads 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 row-gather replacement above + keeps the shared `mergeGatherIndices` helper. 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).** From 50559a3dd84f53029bd5c17a71c6477a8db83565 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 11:37:32 +0200 Subject: [PATCH 28/32] review(WP4): route the sibling-scope rejection through the no-allocation assertion helper (PR #57) --- .../jmlx/models/VisualFeatureInjectionTest.java | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) 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 index 5549c84..0cbb539 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -370,20 +370,16 @@ private static void assertRejection(Executable call, String message) throws Inte * allocated; the message names the parameter. */ @Test - void addMaskedFeaturesRejectsAUnrelatedSiblingScope() { + 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)); - IllegalArgumentException exception = - assertThrows( - IllegalArgumentException.class, - () -> - VisualFeatureInjection.addMaskedFeatures( - sibling, hidden, features, new int[] {3, 4, 5})); - assertEquals( - "scope must be related to hidden's and features' scopes (one an ancestor of the other)", - exception.getMessage()); + 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)"); } } From 7d8af664ea1fec1c825770ecf9ce259297110109 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 11:43:03 +0200 Subject: [PATCH 29/32] review(WP4): pin the operand-vs-operand scope rejection where scope is the common parent (PR #57) --- .../models/VisualFeatureInjectionTest.java | 21 +++++++++++++++++++ 1 file changed, 21 insertions(+) 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 index 0cbb539..442f431 100644 --- a/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java +++ b/jmlx-models/src/test/java/se/alipsa/jmlx/models/VisualFeatureInjectionTest.java @@ -383,6 +383,27 @@ void addMaskedFeaturesRejectsAUnrelatedSiblingScope() throws InterruptedExceptio } } + /** + * 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) // --------------------------------------------------------------------- From 7459f2018421cb2e58a9b0d15d0ea16f95a6ce93 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 11:44:58 +0200 Subject: [PATCH 30/32] review(WP4): preserve the underlying MLXScope error as the cause in checkRelatedScopes (PR #57) --- .../java/se/alipsa/jmlx/models/VisualFeatureInjection.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 index 44e9783..c1d2655 100644 --- a/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java @@ -188,7 +188,8 @@ private static void checkRelatedScopes( } catch (IllegalArgumentException e) { throw new IllegalArgumentException( paramName - + " must be related to hidden's and features' scopes (one an ancestor of the other)"); + + " must be related to hidden's and features' scopes (one an ancestor of the other)", + e); } } From 2ce3a33b43dc44d101ed264290d59f6c169cc510 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 11:59:46 +0200 Subject: [PATCH 31/32] review(WP4): remap the shared gather indices for the masked add instead of duplicating validation (PR #57) --- .../jmlx/models/VisualFeatureInjection.java | 103 +++++------------- req/plans/phase7-3b-plan.md | 6 +- req/plans/phase7-3b-reference-findings.md | 24 ++-- 3 files changed, 45 insertions(+), 88 deletions(-) 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 index c1d2655..673573d 100644 --- a/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java +++ b/jmlx-models/src/main/java/se/alipsa/jmlx/models/VisualFeatureInjection.java @@ -1,7 +1,6 @@ package se.alipsa.jmlx.models; import java.util.Arrays; -import java.util.Objects; import java.util.stream.Collectors; import java.util.stream.IntStream; import se.alipsa.jmlx.core.MLX; @@ -19,21 +18,24 @@ * *

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. The row-gather replacement reuses {@link Idefics3VisionTower# - * mergeGatherIndices(int, java.util.List)} — feature row {@code f} overwrites exactly the {@code - * f}-th destination and every other row reads itself. The masked add instead uses a dedicated index - * builder whose non-destination rows all read the single zero row at the front of the source, 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. + * 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 { @@ -68,10 +70,12 @@ static MLXArray replaceImageRows(MLXArray embedded, MLXArray features, int[] des * 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 dedicated gather index of the source's {@code - * 1 + F} rows maps every non-destination row to the zero row (source index {@code 0}) and the - * {@code f}-th destination 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. + * 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 @@ -93,7 +97,11 @@ static MLXArray addMaskedFeatures( } int length = validatePair("hidden", hidden, "features", features, destinations); checkRelatedScopes("scope", scope, hidden, features); - int[] gather = maskedAddGatherIndices(length, features.shape()[1], destinations); + 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 = @@ -112,59 +120,6 @@ private static int[] gatherIndices(int length, int[] destinations) { length, IntStream.of(destinations).boxed().collect(Collectors.toList())); } - /** - * The masked add's dedicated gather indices: the {@code length}-index array into a source of - * {@code 1 + featureCount} rows — one zero row followed by the feature rows — where every - * non-destination row reads the zero row (source index {@code 0}) and the {@code f}-th - * destination reads source row {@code 1 + f}. It validates the destinations exactly as {@link - * Idefics3VisionTower#mergeGatherIndices(int, java.util.List)} does (strictly increasing in - * {@code [0, length)}) and checks the {@code 1 + featureCount} source-row sum with checked - * arithmetic, where the shared helper checks the {@code length + featureCount} sum instead. - */ - private static int[] maskedAddGatherIndices(int length, int featureCount, int[] destinations) { - Objects.requireNonNull(destinations, "destinations"); - if (length < 0) { - throw new IllegalArgumentException("length must be non-negative"); - } - if (destinations.length != featureCount) { - throw new IllegalArgumentException( - "featureCount " - + featureCount - + " must equal the destination count " - + destinations.length); - } - int previous = -1; - for (int f = 0; f < featureCount; f++) { - int destination = destinations[f]; - if (destination < 0 || destination >= length) { - throw new IllegalArgumentException( - "feature destination " + destination + " is out of range [0, " + length + ")"); - } - if (destination <= previous) { - throw new IllegalArgumentException( - "feature destinations must be strictly increasing; index " - + f - + " has " - + destination - + " after " - + previous); - } - previous = destination; - } - Math.addExact(1, featureCount); - int[] indices = new int[length]; - int next = 0; - for (int i = 0; i < length; i++) { - if (next < featureCount && destinations[next] == i) { - indices[i] = 1 + next; - next++; - } else { - indices[i] = 0; - } - } - return indices; - } - /** * 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 diff --git a/req/plans/phase7-3b-plan.md b/req/plans/phase7-3b-plan.md index fecd01d..1e09e9a 100644 --- a/req/plans/phase7-3b-plan.md +++ b/req/plans/phase7-3b-plan.md @@ -411,9 +411,9 @@ existing overloads preserved and routed through the identical old graph, product 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 its -own dedicated builder 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 +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 diff --git a/req/plans/phase7-3b-reference-findings.md b/req/plans/phase7-3b-reference-findings.md index 2c07838..09f4f4b 100644 --- a/req/plans/phase7-3b-reference-findings.md +++ b/req/plans/phase7-3b-reference-findings.md @@ -973,19 +973,21 @@ joins) — no new public core API. ``` zeroRow = MLX.zeros(scope, {1, 1, hidden}, hidden.dtype()) // [1, 1, hidden] source = concatenate([zeroRow, features_i], 1) // [1, 1 + F_i, hidden] - mask = takeAxis(source, MLXArray(maskedAddGatherIndices(L, F_i, destinations_i), [L]), 1) + 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 come from the masked add's own builder — every non-destination row - reads the zero row (source index `0`), the `f`-th destination (strictly increasing) - reads 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 row-gather replacement above - keeps the shared `mergeGatherIndices` helper. 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`). + 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).** From df850308759a31f90ce5132e0d64a7783d18de54 Mon Sep 17 00:00:00 2001 From: Per Nyfelt Date: Wed, 21 Oct 2026 12:10:30 +0200 Subject: [PATCH 32/32] review(WP4): require the PostBlockHook result to keep the block's dtype as well as its shape (PR #57) --- .../se/alipsa/jmlx/models/DecoderModel.java | 39 ++++++++----- .../jmlx/models/DecoderStackContractTest.java | 55 ++++++++++++++++--- 2 files changed, 72 insertions(+), 22 deletions(-) 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 9c2ace3..29a5dda 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; @@ -179,10 +180,10 @@ interface EmbeddingHook { * 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 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. + * 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 { @@ -379,19 +380,26 @@ 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 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 + * 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[] expected) { - if (hidden == null || !Arrays.equals(expected, hidden.shape())) { + 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(expected) + + Arrays.toString(expectedShape) + + " and dtype " + + expectedDtype + ": got " - + (hidden == null ? "null" : Arrays.toString(hidden.shape()))); + + (hidden == null + ? "null" + : Arrays.toString(hidden.shape()) + " of dtype " + hidden.dtype())); } } @@ -439,10 +447,10 @@ final MLXArray decoderStack( * 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 or differently shaped 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 + * 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 @@ -522,8 +530,9 @@ final MLXArray decoderStack( 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); + requirePostBlockResult(i, x, expected, expectedDtype); } } return norm.forward(x); 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 89d3add..3469aff 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,9 +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 post-block hook returning a null or differently shaped 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. + * 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 { @@ -268,20 +269,20 @@ void postBlockHookReturningWrongShapeIsRejectedNamingTheBlock() { + WIDTH + ", " + hidden - + "]: got [" + + "] 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. + * the expected shape and dtype. */ @Test void postBlockHookReturningNullIsRejectedNamingTheBlock() { @@ -307,7 +308,47 @@ void postBlockHookReturningNullIsRejectedNamingTheBlock() { + WIDTH + ", " + hidden - + "]: got null", + + "] 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()); } }