Skip to content

nemotron hip: read oQ mixed-width MLX checkpoints (4, 5, 6 and 8 bits, a width per matrix) - #3068

Closed
gilby wants to merge 4 commits into
ashhart:mainfrom
gilby:strix/nemotron-hip-oq-widths
Closed

gilby wants to merge 4 commits into
ashhart:mainfrom
gilby:strix/nemotron-hip-oq-widths

Conversation

@gilby

@gilby gilby commented Nov 12, 2026 •

Copy link
Copy Markdown
Contributor

What this changes

Nemotron 3.5 Lightning on HIP read only mlx-q4g64. It now reads the oQ mixed-width MLX checkpoints the Metal backend of the family already serves: MLX affine at 4, 5, 6 or 8 bits in groups of 64, a width per matrix. No kernel source changes: the shared HIP tiles already decode these widths, and the single width lived in the Nemotron HIP admission, loader and memory plan.

Three commits, each a step that stands alone; say if you want them as three PRs and I'll split:

  1. 534825233 Each matrix at its own width. The loader's table returns the per-matrix width config.zig already parses (as qwen3_5 does on HIP); routed stacks are read at their own width and checked against it; the plan counts every matrix at its own width, the draft rows at lm_head's; admission admits 4/5/6/8 bits in groups of 64 and refuses the rest by name; the start-up line names the widths, e.g. Nemotron-H model mlx-q4g64+q5g64+q6g64+q8g64 18.14 GiB.
  2. b10870a78 The MTP head at its own file's width. The separate head file is read at the width its own metadata states (the stock head is 4-bit in groups of 64), beside a main model of any width, as the Metal recipe does.
  3. 963953d11 Kernel-contract records for the relu² stream tile at 5, 6 and 8 bits, in a new records file (hip_records.zig is at 600 lines), each wider width with a planted control.

Still refused by name, each with a host test: other bit widths and group sizes, unquantized (bf16) projections and embeddings, mtp.layers.* inside the main shards, a head file at an unsupported width.

Checks on this head

  • Head commit: 5474a338e: 963953d11 (the head of the receipt) plus the changelog line.
  • Lean checker: 345 problems on the base and 345 here, 0 new. python3 tools/pr_check.py --base origin/main --pr 3068: 0 errors, 0 warnings.
  • Tests: zig build test -Dhip=true --summary all: 13,676 of 13,790 tests passed, 113 skipped, 1 failed. The failure is dashboard_test "dashboard page polls the public stats snapshot without private details", which fails the same way on main without this change.
  • New host tests, each red with its line reverted: the loader's table reads each matrix at config.json's width for it (red when width ignores its key); every refusal by name (red when every width is accepted); the plan's bytes equal the real oQ6 and oQ8e header bytes (red with a 4-bit-only count); the head is planned at its own file's width (red when planned at the main model's); the stock 4-bit test still pins the checkpoint at exactly 17,775,339,392 bytes, so the uniform path is unchanged.
  • On the GPU (gfx1151): zig build hip-kernel-contract: SUMMARY gfx1151 162 kernels: 162 pass gates 1-4, 0 fail, 2 known findings; 44 controls: 44 red as planted, 0 missed; 173 limit problems (main: 159 kernels, 41 controls, 173 limit problems). zig build hip-kernel-test: 159 of 165 passed, 6 skipped, the same as main.
  • Prompt state against chunk size: teacher-forced rows on the oQ4e checkpoint over 1,000 ids in chunks of 2,048, 256 and 300 are bit-identical to the one-chunk run.
  • CI: I'll post the results for this head when the run finishes.
  • Not run: other HIP chips (no card; the new records' usage for gfx1030, gfx1100 and gfx1201 comes from the contract tool's compile, and on gfx1100 the 1×8 tile at 5, 6 and 8 bits reports 256 VGPRs with a scratch spill, recorded as measured); TP > 1 on the GPU with an oQ checkpoint (host tests cover the split); oQ5e and oQ6e checkpoints (none is published yet; 5-bit matrices are covered by the oQ4e checkpoint, which holds at least 25 of them); Metal and CUDA (untouched).

Receipt

One Strix Halo iGPU (gfx1151), fresh servers with the page cache dropped before each, --context 131072 --parallel 4 --prompt-cache-gib 0. Raw tool JSON for every arm: https://gist.github.com/gilby/00b2bb9f0acf63a422cd7168b0c322b8

A. The 4-bit checkpoint, main against this head (TensorFold/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-MLX-4bit at d9d758fb with mtp-4bit.safetensors; arms main, this, this, main; main is 6c5cef894). Aggregate tok/s, both arms of each side:

cell main this head change
bench_openai code, T=0 / T=1.0 122.9 / 124.6, 117.9 / 117.2 122.7 / 122.7, 117.1 / 117.6 −0.8% / −0.2%
bench_openai chat, T=0 / T=1.0 99.0 / 95.7, 91.6 / 92.0 95.4 / 99.1, 92.0 / 90.4 −0.1% / −0.7%
1 stream, T=0 / T=1.0 132.5 / 132.5, 125.7 / 124.5 131.5 / 131.6, 124.5 / 125.7 −0.7% / 0.0%
2 streams, T=0 / T=1.0 144.6 / 144.6, 135.8 / 136.9 144.8 / 146.5, 135.9 / 137.7 +0.7% / +0.3%
4 streams, T=0 / T=1.0 177.5 / 175.6, 168.8 / 168.6 176.1 / 175.9, 168.1 / 167.9 −0.3% / −0.4%
--no-drafts code / chat, T=0 87.9 / 87.8 87.9 / 87.9 level
cold first token, 2,048 / 7,099 / 28,333 tokens 0.729 / 0.732 s, 2.660 / 2.668 s, 13.06 / 13.16 s 0.728 / 0.730 s, 2.666 / 2.673 s, 13.15 / 13.17 s −0.2% / +0.2% / +0.4%
  • Every change is inside main's own spread between its two arms (up to 1.35% on code, 3.4% on chat at T=0, where both builds read either about 95.5 or about 99). I read all cells as unchanged, which is what a loader change intends.
  • Tokens. The 12 cold replies' token_sha256 equal main's in each arm; the server log's reply hashes for all 89 requests of an arm are the same set as main's (six positions differ in completion order between concurrent streams, and differ the same way between main's own two arms); 17 of 17 equal with --no-drafts. In every cell of every arm concurrent equals alone (42 equal, 0 unequal, 0 failed) and alone equals "draft": false (24 equal, 0 unequal, 0 failed). Tokens a round 2.31–2.32 in all four arms.
  • Memory on the device is 24.42 GiB in both builds; the start-up line reads Nemotron-H model 17.34 GiB on main and Nemotron-H model mlx-q4g64 17.34 GiB here.

B. oQ checkpoints on this head (they do not load on main: an oQ4e fails at load with backbone.layers.0.mixer.in_proj: packed words and group tables have incompatible or unrepresentable geometry, an 8-bit one is refused by name). Each with the stock mtp-4bit.safetensors beside it:

checkpoint widths, planned GiB on device plain code / chat, T=0 drafted code / chat, T=0 drafted code / chat, T=1.0 1 / 2 / 4 streams, T=0 1 / 2 / 4 streams, T=1.0 cold first token 2K / 7K / 28K
user3892r4 mlx-oQ4e (3948ddc5) 4+5+6+8, 18.14 25.19 GiB 67.1 / 67.0 112.3 / 90.9 96.2 / 81.6 117.4 / 141.8 / 170.0 118.7 / 130.1 / 161.2 0.732 / 2.681 / 13.24 s
TensorFold oQ6 (8cf893c6) 6+8, 25.03 32.03 GiB 61.3 / 61.3 104.2 / 73.2 89.0 / 72.8 101.4 / 122.9 / 145.6 103.5 / 102.7 / 130.8 0.736 / 2.712 / 13.29 s
user3892r4 mlx-oQ8e (5b942665) 8, 32.10 39.17 GiB 53.5 / 53.5 90.6 / 68.0 76.9 / 63.7 89.3 / 99.3 / 117.2 89.6 / 93.6 / 115.6 0.772 / 2.821 / 13.71 s
  • All three: concurrent equals alone (42 equal, 0 unequal, 0 failed) and alone equals "draft": false (24 equal, 0 unequal), greedy and sampled; every cold row is cold. Memory on the device less the planned bytes is a constant 7.0–7.1 GiB at every width, so the plan follows the real bytes.

  • A checkpoint with a bf16 embedding and bf16 attention stays refused by name, as Nemotron on HIP: read oQ mixed-width MLX checkpoints (4, 5, 6 and 8 bits, a width per tensor) #3061 asks about; so does a head stored inside the main shards.

  • New switch? None. No default changes; the 4-bit checkpoint's tokens, speed and memory are main's.

  • Boxes it reaches: HIP devices serving Nemotron-H. Shared code touched: none outside the family except the kernel-contract records and their test (core/contract/hip_records_relu2_widths.zig, tests/hip/contract_relu2*.zig, contract_main.zig); no launch or kernel changes for any family.

  • Environment: Zig 0.17.0, zig build native-hip -Dhip=true -Dhip-arch=gfx1151 -Doptimize=ReleaseFast, system hipcc (HIP 7.1, clang 21), Linux 7.0.

  • Licence of new weights: n/a (no weights added; the checkpoints named are public uploads read as published).

The loader, the memory plan and admission read config.json's per-module
widths: MLX affine 4, 5, 6 and 8 bits in groups of 64, one width per
tensor. Routed stacks are checked against their stated width. Other
widths, unquantized projections and an MTP head inside the model's
shards are refused by name. The uniform 4-bit checkpoint reads as before.
mtp-4bit.safetensors is read at the quantization its metadata states
(4-bit in groups of 64 when it states none, as the stock head), beside a
model of any width. Admission plans its bytes at that width and refuses a
head width the launches do not read, by name.
Gates 1-4 run the relu^2 stream tile at every width Nemotron's oQ
checkpoints store it (4, 5, 6 and 8 bits in groups of 64), each wider
width with a planted control that reads its codes as 4-bit. Gate 5 records
the 24 launches at 5, 6 and 8 bits on the four RDNA chips.
@ashhart

ashhart commented Nov 12, 2026

Copy link
Copy Markdown
Owner

Thanks @gilby. This is a thorough first receipt, and the loader change reads cleanly.

What I checked on 5474a338e:

  • Tables.width returns each matrix's own width from config.json (zig/src/families/nemotron/hip/load.zig:33). A routed stack's geometry is checked against that width before upload (load.zig:197).
  • affine() in hip/plan.zig:35 takes a width per matrix, and the draft rows are planned at lm_head's width (plan.zig:108). The oQ6 and oQ8 byte counts are pinned to the real headers.
  • The 4-bit path is unchanged. Table A's tokens equal main's in every arm, speed sits inside main's own spread, and the stock checkpoint stays pinned at 17,775,339,392 bytes.
  • In table B, device memory minus planned bytes is a constant 7.0 to 7.1 GiB at every width. That covers the plan-against-device ask on Nemotron on HIP: read oQ mixed-width MLX checkpoints (4, 5, 6 and 8 bits, a width per tensor) #3061.
  • CI is green on this head.

What's left. Items 1 and 3 repeat our reply on #3061, and 2, 4 and 5 are new:

  1. Split this into three PRs, records first: 963953d11 (relu² records), then 534825233 (loader, admission, plan), then b10870a78 (the head file). Give each its own CHANGELOG line and receipt. This PR can become one of them.
  2. On its own, the loader commit reads the head file at the model's width. loadMtp builds its table from l.t.w (load.zig:254 at 534825233), and admission plans the head at c.bits (admit.zig:208). So an oQ6 or oQ8 model next to the stock 4-bit head would pass admission, then fail at load. Until the head PR lands, the loader PR should refuse a head beside a model that isn't 4-bit in groups of 64, by name, with a host test.
  3. Fidelity for each new width class. oQ4e, oQ6 and oQ8e have no main to compare against. Run teacher-forced histories on a fixed public prompt set against an independent reference of the same checkpoint (mlx_lm is fine). List every top-token difference with its margin (CONTRIBUTING.md:78). This goes in the loader PR.
  4. In the records PR, list gfx1100's 1×8 tile at 5, 6 and 8 bits (256 VGPRs, a scratch spill) as a known finding, not a pass, and say its speed there is unmeasured.
  5. Name the seed you used for the T=1.0 cells in table A.

One overlap to know about. #3001's padded row stride, for gfx1151, will also change affine() in hip/plan.zig. Whichever lands first, the other rebases.

We have no AMD hardware, so your gfx1151 runs are the evidence here. Once the split PRs carry these, we'll review each and land them in order.

  • Synthena (Ash's Agent)

@gilby

gilby commented Nov 12, 2026

Copy link
Copy Markdown
Contributor Author

oQ5e and oQ6e are now tested too. The description listed them as not run because no such checkpoint was published. I built all four levels (oQ4e, oQ5e, oQ6e, oQ8e) from nvidia/NVIDIA-Nemotron-3.5-Lightning-30B-A3B-BF16 at a9904d24 with oMLX's oQ quantizer at commit e81d525 (stock settings: imatrix-calibrated, groups of 64, router left bf16, no MTP tensors in the main shards) and ran them on this head with the stock mtp-4bit.safetensors beside each. They are local builds, not published; the oQ4e and oQ8e rows agree with the public uploads in the description. Same box, flags and tools as the receipt above; raw JSON: https://gist.github.com/gilby/af27af1ea836ba5f89a1b2b237ee3edd

checkpoint widths (experts / Mamba / attention / lm_head) planned GiB, on device plain code, T=0 drafted code / chat, T=0 drafted code / chat, T=1.0 1 / 2 / 4 streams, T=0 1 / 2 / 4 streams, T=1.0 cold first token 2K / 7K / 28K
oQ4e 4 / 6+8 / 5+6+8 / 6 18.11, 25.14 67.3 115.9 / 87.5 95.4 / 83.4 125.6 / 143.5 / 168.4 97.5 / 123.7 / 157.9 0.733 / 2.680 / 13.19 s
oQ5e 5 / 6+8 / 6+8 / 8 21.81, 28.89 60.3 106.7 / 84.5 89.7 / 72.8 111.2 / 125.9 / 157.5 107.4 / 124.5 / 148.8 0.759 / 2.773 / 13.53 s
oQ6e 6 / 8 / 8 / 8 25.26, 32.34 56.8 94.1 / 77.6 84.5 / 72.4 100.3 / 116.0 / 143.4 99.1 / 107.2 / 134.0 0.758 / 2.786 / 13.57 s
oQ8e 8 / 8 / 8 / 8 32.10, 39.17 53.4 82.4 / 66.4 74.2 / 65.2 88.1 / 101.4 / 121.2 89.8 / 93.1 / 115.0 0.767 / 2.816 / 13.66 s
  • Exactness, all four: concurrent equals alone (42 equal, 0 unequal, 0 failed) and alone equals "draft": false (24 equal, 0 unequal, 0 failed), greedy and sampled; every cold row cold. On oQ5e, teacher-forced rows in prompt chunks of 256 and 300 are bit-identical to the one-chunk run.
  • Shared experts and the embedding are 8-bit at every level. Memory on the device less the planned bytes is 7.0–7.1 GiB at every level, as for the other checkpoints.
  • One sampled cell varies more than the rest: oQ4e at one stream, T=1.0, reads 97.5 here against 118.7 on the public oQ4e upload, while its greedy cell is higher (125.6 against 117.4). Draft acceptance at T=1.0 scatters between runs on this model (the 4-bit checkpoint's own sampled reps span 107 to 131 on the code prompt); I did not rerun it.
  • No refusals and no new warnings at any level.

@gilby

gilby commented Nov 12, 2026

Copy link
Copy Markdown
Contributor Author

CI is green on 5474a338e: 8 checks passed (fmt, zig, zig-windows, zig-macos-dev-tools, pytest, release-paths, both benchmark-clients) and the two nightly host jobs were skipped.

@ashhart

ashhart commented Nov 13, 2026

Copy link
Copy Markdown
Owner

Thanks @gilby for building oQ5e and oQ6e and running them, and for the green CI on 5474a338e. Our review above crossed with your comment by two minutes, so this reads the two together.

Done, at every width.

  • The 4-bit path is main's. Table A shows the same tokens, speed and memory.
  • Device memory less the planned bytes is 7.0 to 7.1 GiB at all five oQ checkpoints, so the plan follows the real bytes.
  • Concurrent equals alone, and drafted equals "draft": false, greedy and sampled, at every width. Prompt chunks match the one-chunk run on oQ4e and oQ5e.
  • hip-kernel-contract on gfx1151 passes 162 kernels, and all 44 controls go red, including one new control each for 5, 6 and 8 bits.
  • Every refusal is by name, with a host test.

Still open. This isn't ready to land yet.

  1. The split. The branch still holds all three commits at 5474a338e. Please open three PRs, in this order: 963953d11 (relu² records), 534825233 (loader, admission, plan), then b10870a78 (head file). Each gets its own CHANGELOG line. GitHub now reports a conflict with main, and it is only in CHANGELOG.md, so the split clears that too.
  2. Fidelity against an independent reference. No width has it yet. The new oQ5e and oQ6e rows compare the engine with itself. They show it is consistent, not that its tokens are the model's. Please run teacher-forced histories on a fixed public prompt set against an independent reference of the same checkpoint (mlx_lm is fine), and list every top-token difference with its margin (CONTRIBUTING.md:78). The public oQ4e upload holds 4-, 5-, 6- and 8-bit matrices, so oQ4e, oQ6 and oQ8e together cover every width. Your local oQ5e and oQ6e builds are useful extra data, but nobody else can fetch them, so they can't stand in for public uploads. This goes in the loader PR.
  3. In the loader PR on its own, loadMtp reads the head file at the model's width (load.zig:254 at 534825233), and admission plans it at c.bits (admit.zig:208). Until the head PR lands, the loader PR should refuse a head beside a model that isn't 4-bit in groups of 64, by name, with a host test.
  4. In the records PR, list gfx1100's 1×8 tile at 5, 6 and 8 bits as a known finding, not a pass. Its records show 256 VGPRs and scratch (hip_records_relu2_widths.zig:43, :115, :187). Say its speed there is unmeasured.
  5. Name the seed you used for the T=1.0 cells in table A.

One small fix for the records PR. The comment at zig/tests/hip/contract_relu2.zig:68 says gate 2 must turn the control red, but line 87 plants it on gate 3 (.shapes). Gate 3 is the right one, so only the comment needs to change.

The oQ4e one-stream T=1.0 cell you flagged is not a gate item. oQ checkpoints have no main to compare against, so it needs no rerun.

We have no AMD card, so your gfx1151 runs are the device evidence. Once the three PRs are up with items 2 to 5, we'll review each and land them in order.

  • Synthena (Ash's Agent)

@gilby

gilby commented Nov 13, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for both reviews. This PR is replaced by three stacked PRs, landing in this order. Parts 1 and 2 are open; I'll close this one when part 3 is open:

Your points, one by one, numbered as in your latest comment:

  1. The split. Done, in your order: 963953d11 became part 1, 534825233 part 2, b10870a78 part 3, each on main 223da7659 with its own Unreleased line in the current form and its own receipt in its body. Parts 1 and 2 have since taken main by merge commits to clear CHANGELOG.md conflicts; their receipts were measured before those merges, against 223da7659, and each body says so. Part 1 changes no served path, so its receipt is the kernel-contract SUMMARY line and the host tests; parts 2 and 3 carry served cells measured on their own heads.
  2. Fidelity against an independent reference. In Nemotron on HIP: read each matrix at its own oQ width (part 2 of 3) #3242's body, with every difference listed in its gist. Teacher-forced on the 16 prompts of tools/bench_prompts.json, each followed by the reference's own 256-token greedy reply, against mlx_lm 0.32.2 loading the same public files: the top token differs on 75 / 48 rows (prompt path / decode path) of about 4,600 / 4,100 on the oQ4e, 53 / 47 on the oQ6 and 64 / 44 on the oQ8e. The uniform 4-bit checkpoint that main serves reads 67 / 55 on the same harness, and a planted control (the 4-bit engine against the oQ4e reference) reads 343. The largest margin among the differences is 4.5 logits (oQ6, one row). So the oQ widths sit where main's 4-bit checkpoint sits against this reference, and I found nothing specific to a width.
  3. The head beside a model that isn't 4-bit in groups of 64, in the loader PR on its own. Part 2 refuses mtp-4bit.safetensors beside a model whose base width is not 4-bit in groups of 64, by name and before any weights load: mtp-4bit.safetensors sits beside a model stored 6-bit in groups of 64: Nemotron on HIP reads the MTP head beside a model stored 4-bit in groups of 64 only; move mtp-4bit.safetensors out of the folder to serve it without drafts. Its host test refuses the head beside a 6-bit and an 8-bit model and admits it beside a mixed model whose base is 4/64; it goes red with the refusal removed. On part 2, you serve such a checkpoint without drafts by moving the head file out of the folder. --mtp-drafts 0 does not help, because admission reads the folder before the policy decides whether to load the head. Part 3 removes the refusal, and its test becomes part 3's acceptance test: the head is admitted beside 6-bit, 8-bit and mixed models and planned at 4/64. That test goes red with the refusal put back.
  4. gfx1100's 1×8 tile at 5, 6 and 8 bits. Part 1 lists it as a known finding, not a pass, and says its speed there is unmeasured, since I have no gfx1100 card. Main's known entries cover gate 1–4 failures only, and gate 5's limits do not count scratch, so a known entry would not reach this case. I followed the rest of the mechanism instead. The finding is in the Known findings section of docs/kernel-contract.md, and spills_on_gfx1100 in hip_records_relu2_widths.zig names the three records. A host test fails when a record spills that the list does not name, or when a named one stops spilling, so a fix has to remove the entry, as with a known entry. If you'd rather gate 5 carried known entries of its own, I'll do that in a separate PR.
  5. The seed of the T=1.0 cells. Table A used the tools' default seeds. tools/bench_concurrent.py sends --seed, default 1234 (tools/bench_concurrent.py:210). In a T=1.0 cell, stream i gets 1234 + i (:269 with --prompts, which the receipt passes; :273 without it), and greedy cells send 1234 on every stream. tools/bench_openai.py sends 1234 + rep for its three reps and 1234 for the warm-up, at both temperatures (tools/bench_openai.py:113, :115). Each new body names these seeds.

The comment at zig/tests/hip/contract_relu2.zig:68 now says gate 3, which is where kernel() plants the control (.shapes). Only the comment changed, in part 1.

#3001: noted. Part 2 changes affine() in hip/plan.zig to take a width per matrix. Whichever lands first, the other rebases.

@gilby

gilby commented Nov 13, 2026

Copy link
Copy Markdown
Contributor Author

Part 3 of 3 is open as #3251 (the MTP head file read at its own width), stacked on #3242 and #3163. The repeats I mentioned are in its body: the three cells that read slower on one arm were repeated as interleaved pairs, and the one that still reads slower in both pairs is bench_openai chat at T=0 on the stock 4-bit checkpoint (−1.81% and −0.33%), a cell that flips between about 95 and 99 tok/s on main and part 2 as well. Closing this PR in favour of the three.

@gilby gilby closed this Nov 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants