Repository navigation
Geometry-general native Qwen3.5 dense loader #503
Description
Issue: geometry-general native Qwen3.5 dense loader (Metal side)
Backend-scale proposal per CONTRIBUTING ("a new backend ... open an issue first"): the
native Metal engine serves exactly one geometry today. Stacks on #490: that draft
already generalizes the shared family code — a geometry-level config.zig for "every
backend" plus a weights/ module tree — while explicitly keeping "Metal's admission of
the tied-head 2B MLX affine checkpoint". This issue is the Metal-side completion of that
split.
What is pinned, from source (main @ 1.0.2)
| Pin | Where |
|---|---|
| Geometry constants | zig/src/families/qwen3_5/config.zig (hidden = 2048, layers = 24, intermediate = 6144, ...); parse refuses any other integer (error.UnsupportedQwenGeometry) |
| Format | config.zig quantization check refuses anything but 4-bit / group 64 |
| Head | landed as configurable: the untied-head companion PR reads language_model.lm_head and cross-checks tiedness |
| Kernels | zig/kernels/metal/qwen3_5/kernels.zig lists one captured .metal source per projection shape (qmm_6144_2048, qmm_2048_2048, ...); Kernels.projection(n, k) is a comptime lookup, so a checkpoint with a shape outside the table cannot dispatch |
| Dispatch | forward.zig derives grids from the comptime constants (c.intermediate, c.conv_dim, head counts) |
The captured qmm kernels are MLX-derived, shape-baked Metal sources produced by the capture
tooling (tools/zig/capture_qwen35.py + the op-fixture harness). Recompiling MLX for a new
geometry is not a build step the repo has.
Why it matters
Qwen3.5-hybrid dense checkpoints now ship at several geometries (a 9B-class model is
24 layers -> 32, hidden 2048 -> 4096, untied, intermediate 12288). The Python line loads
them through mlx-lm with row kernels; the native engine — the faster path on every tested
Mac — refuses them at load. Untied heads are already unblocked by the companion PR; the
remaining pins are geometry and the per-shape kernel table.
Options, in ascending scope
Option A: qualified geometry set (smallest honest step)
Extend config.zig from one pinned geometry to an explicit table of admitted geometries;
generate and ship the captured kernel set for each (the capture tooling already exists).
Admission stays strict; kernels stay captured; nothing at runtime is new.
- Cost: a Mac session per geometry to capture and qualify (
tf-qwen35-check,
tf-qwen35-forward,tf-qwen35-exact). - Ships the 9B class if its shapes are admitted; every future geometry still needs a
capture session. - No change to forward.zig except reading geometry from
Configinstead of comptime
constants — a mechanical, reviewable diff.
Option B: runtime-shaped kernel compilation
Replace the per-shape captured qmm sources with one shape-parameterized Metal source per
kernel family, compiled at load with the checkpoint's dimensions
(mtl.Library.fromSource is already the load path; the release probes runtime-compiled vs
prebuilt libraries). The non-qmm kernels (norm, gdn_, attn_, mlp_act) already take their
dynamic extents from buffers/setValue — only the qmm projections bake shapes.
- Cost: a new Metal kernel family whose arithmetic must match the captured MLX kernels bit
for bit on the qualified path, or a documented deviation with fidelity receipts. - Risk: the captured kernels are MLX's, and matching them exactly is the load check's
contract (weights.zigreal-weight row checks,qualify.zigwindows).
Option C: a generic affine row decoder (what the Python line already does)
Port the packed-row affine decoder concept (simd_qmm, 2-8 bit, groups 32/64/128) into a
Metal kernel family that dispatches on runtime (n, k, bits, group). This also unlocks
mixed-width checkpoints natively, which the current 4/64-only admission refuses.
- Cost: the largest; a real kernel project with exactness gates per CONTRIBUTING.
- Payoff beyond geometry: mixed affine checkpoints (3/4/6-bit recipes) serve natively at
full speed instead of falling back.
Recommendation
A then C: land the geometry-table refactor (no arithmetic change, immediately testable)
and qualify the 9B shapes via capture; treat C as the follow-up that retires the capture
loop. B alone is strictly worse than C for this purpose.
Testing plan
- Config admission tests extend to the table (Linux-runnable today).
- Every new geometry: capture fixtures on the target GPU, then
tf-qwen35-check,
tf-qwen35-forward,tf-qwen35-exact, plustools/bench_concurrent.py,
tools/bench_openai.py,tools/prefill_cold.pyreceipts per CONTRIBUTING. - The untied-head PR's cross-check keeps admission honest per geometry.
Acceptance
A 9B-class untied checkpoint serves natively on an M-series Mac with drafted-vs-plain and
concurrent-vs-solo equality, and the recipe table grows a per-geometry qualification row.
Activity
[-]geometry-general native Qwen3.5 dense loader[/-][+]Geometry-general native Qwen3.5 dense loader[/+]ashhart commentedon Oct 10, 2026
Thank you, @akol1. Two parts of this are settled. #493 landed: the native 2B serves untied heads, and a
qwen3_5checkpoint opens its engine by geometry. The 27B stays on its own engine, with one path per model, so #520's second route for it was closed. Admitting another dense geometry, such as the 9B, still needs a qualification run on Metal for that checkpoint, so this issue stays open until one is scheduled.ashhart commentedon Oct 10, 2026
An update, @akol1: we're going to qualify the 9B now, starting with
mlx-community/Qwen3.5-9B-MLX-4bit(938d8919: hidden size 4096, 32 layers, an untied head, MLX affine 4-bit). On main aqwen3_5checkpoint that isn't the pinned 2B opens the 27B's engine, which admits only the 27B's shape, so the work is taking that shape check out in favour of the geometry from the config, then the same gate the 27B passed on Metal: drafted equal to"draft": false, resumed equal to fresh, concurrent equal to solo, and prefill and decode against mlx_lm on the same Mac. We'll post the results here.qwen: derive kept KV row bytes from the target config
qwen: bind target attention plans to parsed geometry
qwen: derive target attention preprocessing spans and grids
ashhart commentedon Oct 21, 2026
Thank you for the report. The 9B Metal path still needs pinned whole-model qualification, including reference agreement and draft, resume and concurrency checks. We are keeping this open rather than treating general geometry work as proof that this checkpoint is qualified.
- Synthema, Ash's AI Assistant.
ashhart commentedon Oct 27, 2026
Sorry for the silence here, @akol1. This is still open on main: a qwen3_5 checkpoint that isn't the 2B opens the 27B's Metal engine, which still admits only the 27B's shape (zig/src/families/qwen27/model.zig:162), so the 9B is refused. The next step is ours: take that shape from the checkpoint's config and add the 9B's projection shapes, then run the 9B's gate on an M5 (drafted equals plain, resumed equals fresh, concurrent equals solo, and prefill and decode against mlx_lm). We'll post the results here.
ashhart commentedon Oct 28, 2026
/atc claim
grantoverton commentedon Nov 1, 2026
A qualification run for the 9B is up as a PR family, in case it helps here:
Qwen3.5-9B support in Zig engine #1826 Qwen3.5-9B support in Zig engine: the 9B shape on the 27B engine at a KV group of 4 (on qwen27: serve Qwen3.5-MoE (Ornith-1.5-122B) on Metal: routed experts, its DFlash2 drafter and the publisher's 4-bit weights #1125's head-group commit), with exactness, fidelity against mlx_lm and M3/M5 receipts. Stacks on qwen27: chunk mixed-bit prompt passes at 128 rows #1822.
qwen27: chunk mixed-bit prompt passes at 128 rows #1822 qwen27: chunk mixed-bit prompt passes at 128 rows (needed by mixed 4/5-bit checkpoints on M5).
Qwen3.5-9B DFlash drafter in Zig engine #1851 Qwen3.5-9B DFlash drafter in Zig engine: the 9B's DFlash drafter per stream in the shared rounds, drafted equal to plain on both chips. Stacks on Qwen3.5-9B support in Zig engine #1826.
Qwen3.5-9B decisions in Zig engine #1871 Qwen3.5-9B decisions in Zig engine: a slot_lm decision checkpoint (
baselabs/elephant-1) answering/v1/decisionsand/v1/systemonethrough main's score queue, with block-bidirectional prompt attention, a time limit and client-gone cancellation, every answer checked bit for bit against a fresh pass. Stacks on Qwen3.5-9B DFlash drafter in Zig engine #1851.Qwen3.5-9B grow-as-used K/V in Zig engine #1896 Qwen3.5-9B grow-as-used K/V in Zig engine: opt-in
TF_Q27_KV_GROW; a slot's K/V grows with its stream and is given back when it ends, about 8.5–9.6 GB less held at--parallel 4 --context 65536, replies byte-identical. Stacks on Qwen3.5-9B DFlash drafter in Zig engine #1851.Qwen3.5-9B prompt rows in one product in Zig engine #1897 Qwen3.5-9B prompt rows in one product in Zig engine: mixed-width stacks and 5-bit matrices take every prompt row in one product, opt-in byte codes on M5; bit-identical. Stacks on Qwen3.5-9B grow-as-used K/V in Zig engine #1896.
Next in the family: more 9B kernels (M5 K slices, 5-bit byte codes, M3 rows). Happy to fold any of it into your branch or reshape it to fit what you have planned.