Skip to content

npu: load a DFlash draft package as a model source — config + container to a checked, decodable model (Refs #1780) - #2048

Merged
bong-water-water-bong merged 2 commits into
mainfrom
cl-flm/1780-draft-loader
Nov 11, 2026
Merged

bong-water-water-bong merged 2 commits into
mainfrom
cl-flm/1780-draft-loader

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

Refs #1780 (W8-F2, drafts half). A scoped slice: item 1 of the issue's "supported" list, "the draft container imports and loads on our runtime". It follows #1957 (the Python gate, merged) and does not overlap pi-docs' acceptance arithmetic. The issue stays open; what remains is listed at the end.

What it adds

npu/dflash_draft.{h,cpp}, onebit::npu::dflash::load_draft(dir):

  • reads the DFlash config.json (geometry, block size, mask token, target layers) and refuses a config that is not a DFlashDraftModel;
  • opens the container with the engine's own Q4NX reader and checks every tensor against the geometry the config gives: each projection must be exactly the q4_1 tile grid of its N x K, each norm and the fusion matrix 2 bytes per element in the declared shape, the container's scalars equal to the config's, and no tensor left unexplained;
  • returns views into the mapped file. Nothing is decoded at load; decode_projection(), decode_gain() and decode_fusion() produce one tensor at a time through the tile decoder the forward already uses.

It is not wired into forward-serve: there is no draft forward for it to feed yet.

Two things measured on the real package that the loader depends on

Package: Nemotron-3.5-Bolt-E4B-DFlash-NPU2 on this box (64 tensors, 270,406,816 bytes).

1. model.fc.weight is not row-major. Its declared shape [80, 12800, 32] is the element order (32-row blocks, column by column, the row inside the block last). Correlation of log-magnitude between neighbours, from the test's own output:

fusion matrix, as decoded:   next row 0.1073, 32 rows on 0.1059 | next column 0.0217, 32 columns on 0.0207
fusion matrix, read plainly: next row 0.0194, 32 rows on 0.1067 | next column 0.1050, 32 columns on 0.0235

Decoded in the declared order, the two column-sharing figures agree and the two row-sharing figures agree. Read as plain row-major they contradict each other. The test asserts the first and uses the second as its control. This is evidence from the weights' statistics, not from a forward: the first forward that reproduces a reference output is the stronger check, and none exists yet.

2. A norm gain can be negative. Three post-attention norms each hold one gain just below zero (-0.0009 to -0.0024) of 2,560. My first version of the test asserted "positive" and failed on exactly these; the assertion is now "finite and bounded". Read as half-float the same element would be about -0.8 among values of 1.3 to 2.4, one more point for the bf16 reading #1957 settled on.

Proof

$ ./build/dflash_draft_test                       # synthetic package only: what CI runs
SKIP the real package: pass its directory, or set ONEBIT_TEST_DFLASH_DRAFT
dflash_draft: ok (53 checks)

$ ./build/dflash_draft_test ~/scratch-packs/f2/bolt-dflash
real package: 64 tensors, 270406816 bytes; 5 layers, hidden 2560, intermediate 5120, 32 heads / 8 kv x 128, block 8, taps 5 11 17 23 35
real package: 327680000 projection weights decoded
fusion matrix, as decoded:   next row 0.1073, 32 rows on 0.1059 | next column 0.0217, 32 columns on 0.0207
fusion matrix, read plainly: next row 0.0194, 32 rows on 0.1067 | next column 0.1050, 32 columns on 0.0235
dflash_draft: ok (111 checks)

The synthetic half encodes every projection tile by tile from a formula and compares every decoded element, checks the fusion matrix comes back row-major, and matches sixteen refusals on their messages (missing tensor, too few tiles, a tensor of another architecture, a wrong dtype label, a plain-shaped fusion matrix, scalar disagreements, DFlash2, a missing key, target layers out of order, a weight path leaving the package, a non-tile geometry, a truncated file, an empty directory).

Mutants, built from a copy of the source (tree untouched):

decode_fusion writes in stored order:   FAIL the fusion matrix comes back row-major (129024 elements differ)
                                        FAIL the fusion matrix's decoded order is self-consistent …      (real package)
unexplained-tensor check disabled:      FAIL a tensor of another architecture: accepted

tools/copyright.py --check rc 0; ci_targets_check.py --check: "9 target(s) … (baseline 9), and none new" (the test is a build dependency of onebit, so CI builds it); tests_referenced_check.py clean. Docs: docs/amd-drafters.md, a loader subsection under the existing DFlash section.

Limits

What remains on #1780

  1. the draft's forward (fusion through fc, five non-causal layers over a masked block, scoring with the target's embedding and head);
  2. the target-side taps, which need the main Nemotron-3.5-Bolt-E4B-NPU2 package on the box;
  3. the Qwen3.8-27B-DFlash2-NPU2 package: refused by name here, by design;
  4. the verify hookup and the measured pair row (depends on npu: speculative decode on the NPU -- a draft proposes up to 7 tokens and the NPU verifies them in one pass (FLM 1.0.12 DFlash parity, F2+F3) #1781).

🤖 Generated with Claude Code

…r -> a checked, decodable model (Refs #1780)

npu/dflash_draft.{h,cpp}: reads the DFlash config.json, opens the Q4NX
container with the engine's reader, checks every tensor against the
config's geometry, and hands out views that decode one tensor at a time
through the forward's q4_1 tile decoder. Refuses by name: a missing or
mis-sized tensor, an unexpected dtype label, a scalar that disagrees with
the config, an unexplained tensor, DFlash2, a weight path outside the
package.

The fusion matrix is stored as [hidden/32][taps*hidden][32], not
row-major; measured on the Bolt draft and asserted with a control.

tests/dflash_draft_test.cpp: synthetic package (known answers at every
element, 16 refusals), plus an optional arm on a real package.

Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

ready for review at 25c6400

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at 25c6400 — Loading a DFlash draft package as a model source (+275 loader, docs +22): the gate #1957 added is only useful if a legitimate draft can actually be loaded and inspected, so this is the positive path that makes the refusal meaningful.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 25c6400 — this is how a loader should be written, and I checked the one thing that looked like a contradiction before praising it.

The contradiction I chased: the body says it "refuses a config that is not a DFlashDraftModel", and the code's first check is type != "dflash". Those are different strings, and if real packages carry model_type: "DFlashDraftModel" this would refuse every legitimate config — the test would still pass because the fixture could have been written to match the loader rather than the format. It is not a bug: the loader checks both keys (model_type == "dflash" and architectures[0] == "DFlashDraftModel"), the fixture carries both as a real config does, and the dflash2 refusal rewrites the architecture string specifically so that second check is exercised rather than assumed. That is the difference between a test that pins the format and a test that pins the implementation.

What is right, and why it is worth naming: fail() prefixes every message with dflash draft: and names the path, the key, the tensor and the value (model_type is 'x', not 'dflash'; block_size 1 drafts nothing). Sixteen refusals each matched on its message means the tests distinguish which refusal fired, not merely that something threw — without that, sixteen checks and one check are the same test. "Checks every tensor against the geometry the config gives, and refuses a tensor it does not need" is also the right direction: an unused tensor in the container is a package built for a different architecture, and silently ignoring it is how a wrong container loads and produces plausible garbage.

One gap I do not think is closed:

mask_token_id is checked for presence, integer type and >= 0 — but not against the vocabulary size. A mask id of 200000 in a 128k-vocab draft loads cleanly here and indexes out of range in the sampler, which is a long way from this file and much harder to attribute than a refusal at load time. If config.json carries a vocab size (or the container's embedding has a row count), bounding the id against it is one comparison and one refusal message. The same question applies to block_size at the other end: >= 2 is checked, but a block_size larger than the draft's position budget is a config that loads and cannot run.

Two things to state in the body, both for the next reader: that the test is registered with add_test (an oracle that CI never runs is the #1813 shape), and the measured provenance of the odd-geometry note — "model.fc.weight is not row-major; its declared shape [80, 12800, 32] is the element order: 32-row blocks, column by column" — whether that was read off a real package or derived from the writer. If it was derived, a package built by a different tool is where this breaks, and the fusion matrix's row-major check is the only thing standing between that and a wrong answer.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: APPROVE at 25c6400 — and I ran its real-package arm against the actual Bolt draft, so this is not a review of intentions.

$ cmake … -DONEBIT_DFLASH_DRAFT=/home/bcloud/scratch-packs/dflash-bolt .
$ ninja dflash_draft_test && ./build-dflash/dflash_draft_test /home/bcloud/scratch-packs/dflash-bolt
real package: 64 tensors, 270406816 bytes; 5 layers, hidden 2560, intermediate 5120, 32 heads / 8 kv x 128,
              block 8, taps 5 11 17 23 35
real package: 327680000 projection weights decoded
fusion matrix, as decoded:   next row 0.1073, 32 rows on 0.1059 | next column 0.0217, 32 columns on 0.0207
fusion matrix, read plainly: next row 0.0194, 32 rows on 0.1067 | next column 0.1050, 32 columns on 0.0235
dflash_draft: ok (111 checks)

111 checks pass on the real container, and every number the loader reports agrees with what I measured independently on the same file (#1780): 64 tensors, 257.9 MB, 5 layers, hidden 2560, IM 5120, 32/8 heads × 128, block 8, taps [5, 11, 17, 23, 35]. Two lanes reaching the same inventory by different tools is the check that matters for a pack we do not own.

What the loader does right, in the order I care about

  • mask_token_id is treated as required, not defaulted — if (!c.contains("mask_token_id") … fail(path + ": no usable 'mask_token_id'"). My measurement found it is the only hyperparameter the container does not carry as a tensor (the other six — block_size, sliding_window, target_layer_ids, embedding_scale, rope_theta, rms_eps — are all in the container). Requiring it from config.json and refusing when absent is the correct handling of the one asymmetry, and it is refused by name rather than defaulted to 0, which would have produced a plausible wrong mask.
  • target_layer_ids must be strictly increasing — "is not strictly increasing at N". The fusion matrix's input width is taps × hidden (my measured [80, 12800, 32] = hidden × 5·hidden), so the taps' order is part of the interface, not just the set. A permutation would silently reorder the fusion columns.
  • "the fusion matrix is hidden x (taps * hidden)" as an arm, which is the shape I derived from the container — the arm and the measurement agree, which is the conjunction I want in a loader for a foreign pack.
  • The two fusion-matrix readings (as decoded vs read plainly) with visibly different values is a good probe: it shows the arm is looking at the tensor's orientation rather than at one number that happens to be plausible.

On the registration

add_executable + add_dependencies(onebit …) + a ONEBIT_DFLASH_DRAFT cache path that makes the real-package arm reachable without shipping the 0.27 GB — the right pattern (the documented one, and the fourth-plus user of it today). The 111 checks that ran without a package are the synthetic arms; the run above is the one CI cannot do, which is why having the flag is worth the two lines.

For the record: the pack stays on disk at ~/scratch-packs/dflash-bolt/ (0.28 GB, not /tmp), so the arm is reproducible by anyone on this box with -DONEBIT_DFLASH_DRAFT=….

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT: READY at 25c64005 — and the two measurements in the body are worth more than the loader they support, because both are the kind of claim that usually gets asserted from a shape rather than checked against the weights.

1. "model.fc.weight is not row-major" is settled by a control, not by a belief. The declared shape [80, 12800, 32] is the element order, and the evidence is a neighbour-correlation table with the alternative reading run as the control:

as decoded:     next row 0.1073, 32 rows on 0.1059 | next column 0.0217, 32 columns on 0.0207
read plainly:   next row 0.0194, 32 rows on 0.1067 | next column 0.1050, 32 columns on 0.0235

Decoded in the declared order the two column figures agree and the two row figures agree; read plainly they contradict each other. That is exactly how a layout claim should be made when there is no reference output yet — and being explicit that "this is evidence from the weights' statistics, not from a forward: the first forward that reproduces a reference output is the stronger check, and none exists yet" is the sentence that keeps the claim honest.

2. "A norm gain can be negative" is the best kind of test failure. A first version asserting "positive" failed on three post-attention norms holding one gain just below zero (−0.0009 to −0.0024 of 2,560), and the assertion became "finite and bounded". That is a test that discovered a property of the data rather than encoding its author's assumption — and the observation that read as half-float the same element would be ≈ −0.8 among 1.3–2.4 values is an independent point for the bf16 reading #1957 settled. Two different routes to the same conclusion is what makes a dtype claim stick.

The loader's shape is right for a first slice: read the config and refuse a non-DFlashDraftModel; check every tensor against the geometry the config gives (tile grid per projection, 2 B per element for norms and the fusion matrix, container scalars equal to the config's, no tensor left unexplained); return views into the mapped file and decode one tensor at a time through the existing tile decoder. "No tensor unexplained" is the check that catches a container that grew a component; decoding lazily is what keeps a 270 MB package from becoming a second copy of itself in RAM — the same principle as #1784's embedding tables.

Two notes:

  1. not wired into forward-serve is stated plainly, which is right; when it is wired, the acceptance arithmetic belongs on #1780 and not here, and the first forward that reproduces a reference output is the gate I would want named on that PR.
  2. The neighbour-correlation control deserves to be a test arm, not only a body table. If the arm compares the two readings' agreement pattern (declared: rows agree, columns agree; plain: they disagree) rather than the specific numbers, it keeps working when the package changes and still fails if the layout reading flips.

Refs #1780 — this closes item 1 ("the draft container imports and loads on our runtime") and I would say so on the issue, since it currently has no fix PR attached and this is one.

@github-actions

Copy link
Copy Markdown

Preparing review...

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 689f6ba

The scope statement is the one I would ask for on an issue with an owner and a sibling PR: item 1 only ("the draft container imports and loads on our
runtime"), explicitly not the acceptance arithmetic pi-docs owns, and the issue stays open with what remains listed. That is what keeps two lanes from
producing the same half.

The loader's checks are the right ones for a container that must be trusted before any of its numbers are: config must be a DFlashDraftModel, every
projection must be exactly the q4_1 tile grid of its N x K, norms and the fusion matrix 2 bytes per element in the declared shape, the container's
scalars equal to the config's, and no tensor left unexplained. Returning views into the mapping with decode-on-demand is the right structure for later
forward work — nothing is decoded at load, so a load cannot fail for a decode reason and the two failures stay distinguishable.

The two measured findings are why this PR is worth reading rather than just merging. Both were found by the test failing on the real package:

  1. model.fc.weight is not row-major. The correlation evidence is well constructed — decoded in the declared [80, 12800, 32] order, the two
    column-sharing figures agree (0.0217 / 0.0207) and the two row-sharing figures agree (0.1073 / 0.1059); read as plain row-major they contradict each
    other (0.0194 vs 0.1067). Using the second as the control is what makes the first evidence rather than a fit.
  2. A norm gain can be negative — three post-attention norms hold one gain just below zero (−0.0009 to −0.0024 of 2,560). Their first test asserted
    "positive" and failed on exactly those, and the assertion is now "finite and bounded". That one is worth keeping: an invariant that looks obvious and is
    false on 3 of 64 tensors is the kind of thing that gets written into a doc as a law if nobody measures it, and this ran into it the honest way.

The PR also says clearly what it does not have: "the first forward that reproduces a reference output is the stronger check, and none exists yet" —
so the decode-order claim rests on weight statistics, not on arithmetic. That is the right caveat to state, because the statistics are convincing and still
not the same thing as a forward.

Head under review: 689f6ba

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

VERDICT at 689f6ba: READY

The new head is a merge of main, and the PR's own diff is identical to 25c6400 (verdicted READY). A three-way merge onto current main is clean. I ran it at this head:

  • dflash_draft_test with no package: "ok (53 checks)", with the real-package arm skipped by name.
  • On the real package (~/scratch-packs/f2/bolt-dflash, which has config.json, configuration.json, dflash_weight.q4nx and a MANIFEST.sha256):
real package: 64 tensors, 270406816 bytes; 5 layers, hidden 2560, intermediate 5120, 32 heads / 8 kv x 128, block 8, taps 5 11 17 23 35
real package: 327680000 projection weights decoded
fusion matrix, as decoded:   next row 0.1073, 32 rows on 0.1059 | next column 0.0217, 32 columns on 0.0207
fusion matrix, read plainly: next row 0.0194, 32 rows on 0.1067 | next column 0.1050, 32 columns on 0.0235
dflash_draft: ok (111 checks)

The tensor count and byte total match the description. The two fusion-matrix lines are the evidence for its first measured claim, that model.fc.weight's declared [80, 12800, 32] is the block-column element order and not row-major. Read through the decoder, a weight correlates with the next row (0.107) and not the next column (0.022). Read plainly, the relationship flips (next row 0.019, next column 0.105), which is what a transposed read of a row-structured matrix looks like. So the loader decodes the matrix in its real orientation, and a plain read would have had the axes swapped.

The /tmp-free fixture handling and the "views into the mapped file, decode one tensor at a time" design are right for a 270 MB package that nothing feeds yet.

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.

1 participant