Repository navigation
Enable bundle-aware Qwen speculation and dense mixed JANGH loading - #565
Conversation
Port of the vMLX Python DFlash2 speedups to the Swift 27B target. LaneQMM: Metal 4 tensor-ops matmul over raw affine codes with the dequant applied after the dot product (C += s*P + b*XS), widening 5/6/8-bit codes in threadgroup memory. Its cost is nearly flat from 1 to 16 rows, so wider verify blocks stop paying per row. Installed once per DFlash2 target on eligible language-model QuantizedLinear modules (untiled by default; fused readers use .weight directly). VMLX_LANE_QMM=0 disables it. DFlash2WidthChooser: picks the verify width per cycle by expected tokens per measured second. Acceptance is learned per position from every cycle; cost per width is an EMA shared per target. Candidate widths are (trained, 2x trained) when the target is lane-flat. VMLX_DFLASH2_BLOCK pins a width; VMLX_DFLASH2_WIDTHS=0 disables. Measured on max2, Qwen3.8-27B JANG_4D, greedy, 300 tokens (tok/s, prose/code/easy code/easy prose): DFlash2 before 24.5/52.8/68.1/67.1, after 35.8/75.0/147.1/145.8.
The fused GDN input projection used qwen_mixed_qmv, whose inner loops run to a runtime `valid` bound and cannot unroll. It was compute-bound even at one row (36 Flash-Next banks: 5.0 ms vs MLX qmv 3.6 ms) and scaled linearly with verify rows (36.9 ms at 8 rows). qwen_qmv_rows_full keeps every row's statement order (load_vector, qdot per output, += in k-block order, simd_sum) with constexpr bounds. Weight words, scales and biases are read once per k-block for all rows in the threadgroup. Each (kind, K, N, rows) is admitted only after a bitwise match against the caller's original kernel on the real weights; a mismatch keeps the original kernel. VMLX_BF16_AFFINE_ROWS=0 disables. The same kernel replaces the row-exact q8 projection. For 8 bits its per-row arithmetic matches MLX's mixed affine_qmv_fast_bf16_f16, and the previous batched M=1 addressing streamed the 636 MB lm_head once per verify row. This commit also enabled a batched GDN output tail. It broke row exactness (f32 promotion) and is reverted in a later commit. VMLX_FLASH_ROW_EXACT_OFF=<site,...> is a diagnostic to attribute cost per exact route. Staged exact verify, Flash-Next JANG_4S (ms at 1/2/4/8 rows): 26.3/32.7/53.3/95.8 before, 23.6/24.1/32.3/55.1 after. All shapes admitted bit-exact. Plain AR output is byte-identical to before and its speed is unchanged (55.5 vs 55.5 tok/s cool; the plain-AR path does not reach these kernels).
With F_NOCACHE and a cold cache, every decode token issued about 48 synchronous 16 KB reads to the model volume (iostat: 1,000-1,800 IOPS during plain AR on the external PCIe SSD), so decode depended on drive latency. vMLX Python (table_reader.py) and TensorFold read the table through the page cache and warm it once after load, so lookups become RAM hits. Python measured 3-6x faster cold prefill on 1.8k-6.7k token prompts. Pages are clean and reclaimable: under pressure macOS evicts them and lookups fall back to drive reads, never an error. With the warm cache, decode issues zero drive reads. Speed evidence is NOT a controlled A/B. The only cold-cache F_NOCACHE AR run measured 53.4 tok/s at 1617 MHz, against 55.5 page-cached. Once the table is cached, F_NOCACHE reads are also served from cache and the arms measure the same; a clean A/B needs a purged page cache. VMLX_QWEN4_PLE_NOCACHE=1 restores F_NOCACHE. VMLX_QWEN4_PLE_WARM=0 skips the warm.
Addressing verify rows as a batch over a [1, N, K] weight bank stops MLX from using the mixed bf16/f16 qmv. The output is promoted to float32, so the GDN block output in verify differed in dtype and value from the AR S1 tail. That moved every downstream layer, and greedy native-MTP no longer matched its single-token path. Found with the teacher-forced harness (RunBench BENCH_ROWEXACT=1). With the tail reverted, every verify row at W = 2..8 is bitwise identical to the AR S1 forward over 96 positions. The batched tail saved nothing measurable (4 rows: 53.6 vs 53.4 ms).
Port of vMLX Python native_mtp_copy.SuffixCopyProposer (TensorFold's SuffixLookupProposer). An incremental 4-gram index over prompt plus confirmed tokens proposes the continuation after the longest earlier match (at least 8 tokens). Width starts at 3, goes to 7 after a fully accepted window and back to 3 after a window under half; two first-token misses silence it for 16 proposals. With a proposal, the head still drafts one level to keep its aligned cache committed; the copies replace that draft. Copy cycles count toward AR safety but not the head depth controller, and their wall time is never billed to a head depth. Greedy staged verify only. VMLX_NATIVE_MTP_COPY=0 disables. Also adds NativeMTPDepthChooser (expected tokens per measured second), OPT-IN via VMLX_MTP_DEPTH_CHOOSER=1. It did not beat the existing governor in matched-clock A/B, so the default is unchanged. Flash-Next JANG_4S Adaptive, cool start at 1.6 GHz, tok/s (prose/code/easy code/easy prose): - copy on: 56.2/82.5/96.0/115.1 - copy off: 55.9/84.9/94.4/92.0 Adaptive with copy is identical to the iterator's reject-all (single-token) output on all 10 probe prompts.
BENCH_SPEED=1 is the Swift twin of the Python mtp_anatomy_probe: same prompts, greedy, 300 tokens, and decode rate from the engine's generateTime. Arms are ar, adaptive and dflash2; BENCH_SPEED_CATS selects categories and BENCH_SPEED_TEXT_OUT dumps full texts. BENCH_ROWCOST=1 times one target forward by row count, in the exact staged verify mode or plain. BENCH_ROWEXACT=1 is a teacher-forced comparison of every verify row's logits with the AR S1 step at the same position (bitwise, max diff, argmax), for each width. BENCH_ROWEXACT_LAYERS=1 bisects row 0 by layer and by layer-0 component through an opt-in diagnostic capture in Qwen4Exp.
Tiling policy:
- Default ("safe") tiles every lane module except Qwen35 GDN
linear_attn.in_proj* and out_proj. Fused S=1 kernels read those
weights' packed layout directly.
- VMLX_LANE_QMM_TILED=0 tiles nothing.
- VMLX_LANE_QMM_TILED=1 tiles everything (diagnostic only; it corrupts
S=1 output on models with direct readers).
Untiled lane modules now run MLX's quantized matmul below
LaneQMM.minRows rows (4; override with VMLX_LANE_QMM_MIN_ROWS), because
qmv is faster there.
Qwen3.8-27B JANG_4D target forward (ms at 1/8/16 rows):
- no lane: 48.9 / 83.3 / 119.2
- lane untiled: 61.4 / 74.0 / 77.2
- lane safe: 56.2 / 69.5 / 72.6
- lane all: 55.6 / 67.4 / 70.4
DFlash2 end to end, cool start, tok/s (prose/code/easy code/easy
prose): 35.5/80.9/157.4/155.4 untiled, 38.3/85.0/164.2/162.2 safe.
The row-exact hcup route required `combines`, so the final hyper-connection mixer (combines == false, same [10240, 320] up-projection) ran its verify rows through one multi-row matmul while AR ran it per row. Verify rows could then differ from AR in the last mixer, which changed greedy output. Found with BENCH_ROWEXACT_LAYERS_AT. On JANG_2L at position 52, every decoder layer was bitwise equal, the logits differed by 0.0625, and the first differing stage was the final mixer. The base commit fad8183 has the same defect. After the fix, verify rows equal the AR S1 rows bitwise at W = 2, 4, 5 and 8 on Flash-Next JANG_2L, 4S, 4M, 6S and Allosaurus JANGH2.
New `VMLXMTPServerMode.familyDefault` ("family_default") is the default
and the missing-key decode value. It resolves to Adaptive native MTP for
bundles whose family has a measured cold-start depth (Qwen3.8 Flash-Next,
`measuredFamilyAutoDepth`, keyed on model_type), and to Off for every
other family. Explicit Off and On (Adaptive) are unchanged and always win.
Contract version 5 migrates a stored `.off` written before version 5 to
`.familyDefault`, once. Every install since 2026-09-16 persisted the
then-default `.off` without a user choosing it.
Bundled DFlash 2: when no selected drafter fits and the mode is not Off,
`resolvedDFlash2Selection(configData:bundleDirectory:)` uses the bundle's
own `dflash2/` folder if it passes the same fit check. A Qwen 27B bundle
that ships its drafter therefore gets DFlash 2 without the user picking
a folder. `resolvedMTPDraftStrategy` takes `bundleDirectory:` too.
Tests: default and decode value, the one-time migration, family-scoped
launch, and bundled drafter discovery.
ProposalHeadStamp: the runtime derived a proposal-head verdict and wrote `vmlx_mtp_proposal_head.json` into the bundle. It wrote one into Qwen3.8-27B-JANG_4D, a bundle with no MTP head, which nearly shipped in a Hugging Face re-upload. The verdict is now derived in memory only (about 166 ms per load); only the converter writes stamps. An existing matching stamp is still honored. Qwen4 PLE: the full background warm runs only when the bundle is at most 60% of physical RAM. On a 128 GB Mac, Flash-Next JANG_6S (107 GB bundle, 30 GB table) with Adaptive MTP thrashed at 0.34 tok/s while the warm competed with 87 GB of Metal buffers. The same run with page-cached reads and no warm ran 45.0 prose / 90.3 easy prose tok/s. Rows are always read through the page cache, so only touched n-grams stay resident. VMLX_QWEN4_PLE_WARM=1/0 forces the warm on or off. This only skips a speedup; it never refuses or delays a load. RunBench: BENCH_MTPRESOLVE (osaurus-equivalent MTP resolution per bundle) and BENCH_ROWEXACT_DUMP.
"default Off prevents native MTP launch" now tests explicit Off; the shipped default is `.familyDefault`, which is on for Flash-Next. NativeMTPManualDepthTests used an untuned `qwen4_exp` fixture to exercise the tuning gate. Qwen3.8 Flash-Next has a measured family cold-start depth by design, so that fixture could never hit the gate; these tests already failed on fad8183. They now use a tuning-gated family (qwen3_5). Still failing, and failing identically on the base: - "an Auto-only block does not disable explicit diagnostics": the test says tuning `blocked` (auto-only) must still allow manual depth, but manualRecommendation documents that `blocked` vetoes manual depth too. Policy decision left open (handoff §4). - NativeMTPDepth3AutoLaunchTests: the published Qwen3.6 MXFP8-MTP tuning rows are measured on `deterministic_count_96_tokens`, which gate #343 (`isWorkloadGeneral`) rejects.
…rify
Qwen3.8-27B-JANGH2 is a dense qwen3_5 model whose MLP gate/up/down
projections are one-expert jangtq2 codebook banks (hadamard32 rotation,
2-bit). It could not load: the configuration partition only accepted
routed switch_mlp paths and K2 down banks, and the named mmap rejected
the bank names.
Loader
- JANGHConfigurationPartition accepts language_model.model.layers.L.mlp.
{gate,up,down}_proj for model_type qwen3_5.
- JANGHDenseModelPreparation gains a qwen3_5 initializer and
makeProjectionsByPath; Qwen35JANGHPreparation (MLXVLM) builds the model,
swaps the dense JANGH leaves in and excludes their tensors from the
ordinary safetensors load. Qwen35 MLP projections are UnaryLayer.
- LaneQMM skips models with JANGHDenseLinear leaves from flat-model
shortcuts.
- mlx submodule -> osaurus-ai/mlx claude/dense-jangh-gate-up-mmap
(named mmap admits dense gate/up/down with optional language_model.).
Speed (max2, cooldown gated, tok/s)
- JANGHDenseFastQMV: one-row codebook QMV, weights loaded once per
(row, block). Admitted once per shape only if bitwise equal to the
reference QMV. AR 12 -> 30.6 prose / 28.8 code, same text.
- Dense verify windows (>= 2 rows) use the sorted NAX QMM, and the NAX
kernel gains a 16-row tile (1x2 simdgroups of 16x32, 64 threads) for
<= 16 rows instead of paying MMA for 64-row tiles. It is admitted once
per shape only if bitwise equal to the 64-row tile. Target forward at
4-16 rows: ~100 ms -> 54-59 ms.
- 27B JANGH2 DFlash2: 24.6/50.3/116.0/122.1 -> 28.6/57.7/135.0/141.0
(prose/code/easy code/easy prose; Python 29.7/67.6/143.9/149.5).
Note: tile_matmad_nax issues no MMA when both TM and TN are odd; a 16x16
simdgroup tile silently produces garbage. Keep SN = 32.
Tests: JANGHDenseNamedMmapTests now pins dense gate/up/down acceptance
and the remaining rejections.
Media requests restored nothing and stored nothing, so every turn after an image re-prefilled the whole conversation (27B: 7-8 s TTFT per turn, measured in the live app). The M-RoPE delta is a pure function of the prompt's token ids and media grids (no pixels), so a restored media-salted entry is continued at positions recomputed from the FULL prompt: dflash2MediaPrefill(restoredPrefix:) resolves full-prompt positions once (setting the request's decode delta) and prefills only the text suffix at its slice of them, without the vision tower. A suffix that carries a media placeholder still resets to a full prefill. The media prefill also forces a chunk edge at the generation-suffix-stripped boundary and snapshots it, so a media turn stores the canonical cross-turn checkpoint like a text turn. Two-turn smoke (27B JANG_4D, image then text follow-up, disk cache): turn 2 restored 277 of 339 tokens and answered from the image correctly.
2690a74 to
a5a0689
Compare
Live app proof — osaurus dev build pinned to this engine headSetup.
AR references on the same machine: 27B 4D ~23, 27B JANGH2 ~29, Flash-Next ~50–55 tok/s. "Cold" = the first request App-side defect found and fixed in osaurus (
|
Dense JANGH verify windows (DFlash 2 / MTP, 2..16 rows) ran the 16-row NAX tile with N/64 threadgroups of 64 threads, each walking all of K serially (27B down_proj: 80 threadgroups x 272 barrier-separated steps). A 4..16-row verify cost ~2x a 1-row decode step. The new kernel slices K across grid z (each slice uses the same codebook tile loader and MMA sequence), writes FP32 partials and sums them. Applies only to single-bank, non-fused, non-rotated-output dense projections on NAX; other shapes keep the existing tile. VMLX_JANGH_SPLITK=0 restores the old tile. Summation order changes, so the result is admitted by agreement with the reference (dense 27B DFlash 2 is speed-first). New unit test covers an uneven K split, ragged rows, 2/4/8-bit and fp16/bf16. Qwen3.8-27B JANGH2, M5 Max, no phase syncs: verify rows 1/4/8/16: 32.8 / 62-65 / 63-65 / 64-65 ms -> 32.8 / 49.6 / 53.2 / 53.9 ms sampled DFlash 2 prose/code tok/s: 31.9/79.2 -> 44.0/97.0 (AR 27.4)
After any prefix-cache restore, init shifted the EMPTY drafter cache's offset to the restored count so drafter RoPE positions matched the target. The 27B drafter is all sliding-window (RotatingKVCache), and an offset ahead of the physical fill corrupts it below the window: updateInPlace returns keys[..<offset] including zero rows, and temporalOrder reads idx < offset as a wrapped ring. The drafter attended to zero rows between its context and the block, so every chat follow-up that hit the prefix cache drafted badly. The drafter cache now counts physical rows from 0 and the absolute base is kept in draftPositionBase (drafter RoPE is relative; only committed-row accounting uses the base). A context chunk longer than the drafter window now restarts the drafter cache at the first kept row instead of the in-attention clip advancing the offset past the fill. Qwen3.8-27B JANGH2, sampled, two-turn prose conversation with the disk prefix cache (RunBench BENCH_SPEED_PROSECONV, 28-token restore): turn 2: 22.0 tok/s, 0.22 accepted/cycle -> 41.6 tok/s, 1.73 control without the cache: 37.4 tok/s, 1.41 Prompts past the drafter window were not affected (offset >= window returns the full ring); greedy 3k-token prompts: 2.23/1.81/1.46 -> 2.25/2.22/1.78 accepted/cycle. RunBench: BENCH_SPEED_PROSECONV=1 two-turn prose conversation mode. Tests: re-anchor the DFlash 2 vs native MTP resolution-order scan after a reformat moved the call's first argument to its own line.
27B JANGH2: two fixes, live re-proof (engine
|
| 27B JANGH2 | turn-2 tok/s | accepted / cycle |
|---|---|---|
| before | 22.0 | 0.22 |
| control, no prefix cache | 37.4 | 1.41 |
| after | 41.6 | 1.73 |
27B JANG_4D after the fix: turn 2 1.78 accepted/cycle (turn 1 2.09).
Live app (osaurus dev build, isolated profile, Safe Auto, plain agent, effort None, bundle sampler)
One connected chat on Qwen3.8-27B JANGH2, real composer, decode tok/s from the app step log:
| turn | before (a5a0689c) |
after (e54e7388) |
|---|---|---|
| T1 code | 74.1 | 109.8 |
| T2 pasted image (red 7 / blue 4, both ✓) | 40.5 | 49.5 |
| T3 short follow-up (prefix restored) | 35.2 | 70.2 (0.70 s TTFT) |
| T4 400-word essay follow-up (prefix restored) | 23.7 | 37.2 |
| T5 second essay follow-up (prefix restored) | 18.7 | 36.1 |
Restored prefixes on T3–T5: 646 / 821 / 958 tokens. All answers coherent: no looping, no marker leakage, no empty
finalization.
Audit notes: do NOT
- Do not judge speculative speed from single fresh prompts. The follow-up turn with a restored prefix is the shape
users hit, and it hid this defect. Use the two-turn conversation mode with the disk cache. - Do not set a KV cache's
offsetto a logical position without filling those rows. Sliding-window caches treat
offset as fill below the window. - Do not "fix" a slow speculative lane by lowering widths or turning it off. Split acceptance per request first.
- Do not merge the split-K kernel into a contract that requires bitwise rows (Flash-Next native MTP). It is gated to
dense single-bank JANGH projections only.
Still open, not claimed
- 27B JANGH2 multi-row verify is still ~1.55x a decode step. The next lever is the verify tile itself.
- The first request after a large-bundle load is slow (page-in and kernel warm-up).
Variating pass (
|
Sampled requests verify through the staged batched verifier since 63de051 (VMLX_NATIVE_MTP_SAMPLED_STAGED=0 restores the sequential path). The recovery test still asserted the sequential path; it now checks whichever path is active and keeps the output, rejection, AR fallback and recovery assertions.
|
Ready for final audit at |
|
Independent audit at Validation executed locally with Metal:
Fresh sampled default Flash-Next4S medians: 58.38 prose /69.19 code /95.17 easy code /111.99 easy prose tok/s. Native MTP logs confirm staged sampled verify, acceptance, adaptive depth changes and AR recovery. Prose/easy-code/easy-prose are close to the provided target measurements; general code needs comparison beyond one run. 27B JANGH2 performance is not yet cleared against the fastest historical table: current full-matrix default35.36/93.29/149.91/114.55; a lower-clock explicit run22.34/47.16/119.17/79.03. Retained raw clocks/results rather than claiming parity. A historical binary control also runs slower today; no missing kernel/build optimization identified in source audit. The sustained restored-prose test preserves aligned drafter state, but its speed remains below the quoted best result. No fixed-depth, sampler, kernel or allocator changes to mask the difference. CI: six Linux jobs pass; formatting failure also occurs on main and is advisory; Mac/CUDA jobs remain queued without runners. Merge remains pending completion of the remaining independent gates. Receipts are in the private |
|
Independent audit found and reproduced a DFlash adaptive-cost lifetime bug. Tables were keyed only by ObjectIdentifier, so an allocator-reused model address could inherit a destroyed target's cost/warmup state. The original source reproduced this on the second target instance. Correction dc2c439 stores a weak target reference and reuses costs only for the same living object. It does not retain model weights or change width policy, sampling, verification, kernels or cache payloads. Verification:
No consistent throughput regression from this correction. These rows do not prove a speed improvement or eliminate the remaining historical-maximum and per-bundle proof gaps. The fresh-process speed variation is separate from the reproduced reload bug. App is being rebuilt on this engine revision before the remaining delegated/eval and per-quant checks. Merge remains held. Raw private receipts in pr565-pr3029-audit-20261007: width-cost-lifetime/, width-cost-integrated-tests.log, width-cost-release-build.log, baseline-binary-copy.json, width-cost-abba-receipt.json and width-cost-{A1,B1,B2,A2}.log/jsonl. |
|
Independent audit update at engine dc2c439 (not merge-ready yet):
Source evidence: MTPRuntime.swift inspector/completeness, JangLoader.swift declared tuple handling, VMLXServerRuntimeSettingsTests and DFlash2DrafterSelectionTests. |
|
Independent Allosaurus JANGH2 speed check completed on current production engine dc2c439 (84534fc adds tests only), M5 Max 128 GB. Four fresh processes, AR / Adaptive / Adaptive / AR, cooldown before each, no concurrent build or inference. Same ten campaign prompts, sampled T=1/top_p=.95/top_k=20, seed1234, thinking off, max300-token timing windows. Actual staged native-MTP verification confirmed.
These reproduce the historical bounded Allosaurus benchmark range. They do not prove completed-answer quality: capped outputs are retained as capped; AR ignored one easy-code instruction, and Adaptive omitted requested unit tests in one answer. Full output and binary/config hashes are preserved in the private audit artifacts. A subsequent real-app tool turn exposed an outstanding latency issue: native-MTP post-generation snapshot derivation replayed the prompt, taking8.46seconds on the first tool call and6.23seconds on the next. Actual SSD write/publication took about50–80ms. Stored prefixes were restored on the next tool steps (4434 then6239 tokens), but the synchronous re-derivation materially delays tool execution. This gate is still open; investigating reuse of the exact prefill-boundary capture already present in the AR/DFlash paths. No cache durability or speed-policy shortcut has been made. Evidence: private pr565-pr3029-audit-20261007/ALLOSAURUS-SPEED-CHECK.md, allosaurus-ar-adaptive-receipt.json, allosaurus-*.log and full text artifacts, app-dc2-live.log (gen/tail/derive and cache/fetch rows). Both PRs remain drafts pending the remaining live/cache gates. |
|
Independent audit update (not merge approval):
Evidence retained privately under |
|
Audit correction pushed at engine A live Allosaurus H2 tool chat exposed an older stable-prefix replay after the configured idle unload/reload. The response had already finished, but replaying the 4,343-token checkpoint added 10.68 seconds. The later 7,720-token SSD checkpoint restored successfully; acceptance on the code follow-up was high. This was a separate post-output delay, not evidence of low acceptance. The correction validates the exact saved native recurrent checkpoint through the existing disk deserializer and a scratch-cache restore before skipping replay. It preserves rejection/repair of malformed offsets, does not alter live target/drafter state, and does not add a RAM prefix tier. Stable-checkpoint eviction recency now uses the stored N-1 boundary, preserving the N=1 edge case. No sampler, kernel, verification arithmetic, or depth-policy changes. Independent real-model control: same saved full-precision SSD checkpoints, same 95-token prompt/answer task, fresh processes, three runs per binary, reversed order included, cooldown met for every arm. First-request post-output tails (seconds):
All twelve full answers matched and ended normally. Second-turn tails remained 0.030–0.034s. This is a cache-tail result, not a claim of higher sustained decode speed. The first baseline also had an unrelated first-use latency outlier; it is retained in the receipts and not attributed to this correction. Validation: 18 native/hybrid XCTest tests + 2 dense/hybrid quota tests passed on the final code. The earlier broader corruption/native-state run passed 18 Swift Testing tests; the explicit test-only Metal library setup now supports that runner. Release RunBench build passed. Parallel app correction: MiMo/N2 naming exceptions now also require the actual matching bundle architecture, so misleading aliases cannot bypass Qwen head inspection. 129 focused app tests passed before the mechanical engine repin, including alias, Off and actual-family controls. Exact repinned app build/live validation is in progress. SOURCE EVIDENCE: |
The inspector called a head complete from any MTP-named tensor plus a configured layer count. A single mtp.fc.weight, index entries whose shard is missing or does not contain them, and a header whose payload is absent all auto-launched native MTP (final audit sweep, 32 metadata decisions). - Tensor names now come from shard headers only, and only entries whose byte range lies inside the file. The index still selects which shards are inspected, and a head physically present but omitted from the index is still discovered. - MTPBundleStatus.incompleteTopologyReason: a head needs a fusion projection (fc, fc_hidden/fc_embedding, eh_proj, input_proj) and at least one compute block (attention, mixer, MLP or experts), matched on name components, never on a model name. Every installed MTP family satisfies it (Flash-Next affine/JANGH, Allosaurus, Ling, Nemotron Lightning, Ornith, Qwen3.8 sources). Real bundles resolve unchanged: Flash-Next 2L/4S/4M/6S/JANGH4 and Allosaurus -> native MTP D3, 1L -> AR, 27B JANG_4D/JANGH2 -> bundled DFlash 2. The audit's malformed fixtures now resolve to AR; DFlash 2 controls and explicit Off are unchanged. Tests: new physical/topology regression and family-layout cases; index fixtures in the MTP suites now write header-only shards; the native-MTP cancellation source scan follows the prefill helper introduced in ce2bfb2 (cancellation is still checked before prefill and between captured boundary chunks).
The automatic page-cache warm read the table's backing files into RAM whenever the bundle was under 60 % of physical memory (Flash-Next 4S/2L: 18-19 GB of n-gram rows; the gate only spared the largest bundles). The memory contract is model weights resident, n-gram rows read on demand, so the warm is now opt-in: VMLX_QWEN4_PLE_WARM=1 forces it, =auto keeps the old size gate, unset/0 leaves rows on SSD. Only cold prefill speed depends on it; decode reads the same rows either way.
LoadBundleFacts.totalSafetensorsBytes counts the whole bundle, including the Qwen4Exp per-layer n-gram embedding shards (Allosaurus 53.6 GB, Flash-Next JANGH4 53.7 GB, 4S/2L 18-19 GB). Those rows are read from SSD through pread and never become GPU weights, but three decisions sized against the total: - the load-time model-residency wired limit: Allosaurus (32.5 GB of weights, 88 GB bundle) and JANGH4 exceeded the safe ceiling and got NO residency, so their weights were unwired and evictable between turns (measured live: wired 43-47 GB during a turn, 5-8 GB right after); - Safe Auto's materialize-instead-of-mmap rule and its memory fraction; - the JANGH scheduling memory allowance. LoadBundleFacts now carries ssdResidentTableBytes (header-only scan of *.ple.ngram_embedding.* payloads) and gpuResidentWeightBytes; all three use the latter. Bundles without an n-gram table are unchanged. Tests: facts separate n-gram bytes from weights; no-table bundles keep every byte.
Audit follow-up — head
|
|
Final audit at f3963ef, retaining the full performance/follow-up-cache work from this branch. Two additional bundle-inspection corrections are committed: split fusion requires both fc_hidden and fc_embedding; n-gram residency accounting rejects negative, reversed and out-of-file spans before arithmetic. No generation kernels, sampler, adaptive controller, prefix storage, idle-unload policy or scratch-memory floor changed in this audit delta. Source evidence: Libraries/MLXLMCommon/SpecDec/MTPRuntime.swift; Libraries/MLXLMCommon/Cache/LoadConfiguration.swift. Focused tests exercise actual safetensors inspection, including missing either split projection, truncated payloads and Int.min/Int.max offsets. Live/test evidence:
Evidence: private merge-final-delta-20261007/{engine-build.log,engine-delta-tests.log,real-bundles-*.log,real-bundles-receipt.json}; prior audit claude-live-consolidation/README.md links frozen raw app artifacts. Previous per-quant speed tables remain historical measurements, not universal throughput guarantees. CI distinction: prior head's Linux builds passed; its lint job used swift-format 604 and failed repo-wide. Current-head CI may still be queued/running; those checks are not claimed passed. No release/tag/appcast. |
Bundle-aware speculative decoding for Qwen3.8, dense mixed JANGH loading, and the speed work that makes both pay off in the app.
What this PR does
dflash2/drafter use it. Explicit Off (AR) disables every speculative route. Every other family keeps the opt-in contract.VMLX_JANGH_SPLITK=0restores the old tile).e54e7388). Before this fix every restored chat turn drafted from a corrupted context and ran below AR.BENCH_SPEED_PROSECONV), and DFlash 2 / sampler arms for the variating pattern.Contracts
Evidence (M5 Max, osaurus dev build pinned to this branch, plus RunBench)
8e13b426has no new failures.main(5fed1819): local chat-template snapshots, the MLXPress source scan, and the Gemma4 audio bundle test. They depend on local bundles or are already stale; this PR does not touch them.Before merge
lintjob runs the latest swift-format (604) over the whole repository and flags ~800 files. It also failed on Add JSON Schema constrained decoding for local inference #564, which merged. Formatting is not part of this PR.osaurus-ai/mlxmain (5aba1efd), which includes the merged Tighten Nemotron required tool template #21.The draft flag stays for the final audit.