Geometry-general native Qwen3.5 dense loader

Open
  • Is something wrong?
#503 2 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 1 day

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
25/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python, zig
Domain
ai, backend

Research direction

Start with zig/src/families/qwen3_5/config.zig, zig/kernels/metal/qwen3_5/kernels.zig, and forward.zig; then review the capture tooling at tools/zig/capture_qwen35.py and the qualification commands listed in the issue. The issue recommends first adding a strict geometry table and capturing/qualifying the 9B shapes, while runtime-shaped or generic kernels are larger follow-ups. Done means a 9B-class untied checkpoint serves natively on an M-series Mac with the stated equality checks and a per-geometry qualification row.

Written by the indexing model from the issue text.

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 Config instead 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.zig real-weight row checks, qualify.zig windows).
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, plus tools/bench_concurrent.py,
    tools/bench_openai.py, tools/prefill_cold.py receipts 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.

Dominant language
Python
Stars
995
Forks
169
Avg merge
1d 10h
Merged PRs (30d)
449

Getting set up

This project ships no dev container, Dockerfile or contributing guide, so setting up is up to you: start from its README, and see our first-contribution guide for the general steps.

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from ashhart/TensorFold

All issues in ashhart/TensorFold

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.