Skip to content

refactor(model): unify Qwen + DeepSeek weight staging behind one declarative pipeline - #174

Merged
superxf merged 15 commits into
hw-native-sys:mainfrom
lterrac:refactor/weight-pipeline
Aug 26, 2026
Merged

refactor(model): unify Qwen + DeepSeek weight staging behind one declarative pipeline#174
superxf merged 15 commits into
hw-native-sys:mainfrom
lterrac:refactor/weight-pipeline

Conversation

@lterrac

@lterrac lterrac commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Steps 1-4 of #163, plus the sidecar parity: the family-neutral pieces, DeepSeekV4's layer contract expressed as data, and the whole-model slab geometry. Additive except for one delegation, and every step gated by a differential test rather than by inspection.

@ndleslx — you confirmed on the issue that you have no branch for this, so this is the whole refactor: the family-neutral pipeline, both families expressed as rule tables, Qwen's lazy staging, and the legacy path retired.

It was two PRs (#174 + #177) until #177's deletions made that a bad split — the first added a legacy rename and two differential suites that the second deleted, so a reviewer read work that cancelled. Combined here; #177 is closed as superseded, and its CI history is there if the intermediate state is ever of interest.

What is here

common/weights/store.py LazySafetensorsStore — the checkpoint index and the grouped reads, lifted out of DeepSeekV4WeightStore, which now subclasses it
common/weights/spec.py the rule schema: LayerContext, and five rule kinds
common/weights/shard.py Replicate and ExpertParallel
common/weights/packer.py pack_layer, the generic per-layer evaluator
common/weights/stacker.py StackGroup plus allocate_slabs / destinations_for / copy_packed_layer / stack_layers
deepseek/weight_spec.py all 49 DeepSeekV4 layer weights as data
tests the parity harness, plus the differential

pack_deepseek_v4_layer_weights now evaluates that table, and load_stacked_layer_weights drives the generic stacker instead of three DeepSeek-shaped helpers. Signatures, output names and order, dtypes and error wording are unchanged.

Slab geometry becomes data: a StackGroup is a set of weights plus the ordered layer ids contributing to it, and a layer's position in that list is the slice it occupies — which is how a group covers a subset of layers (CSA holds its layers contiguously in first-appearance order, which is what the fused kernels index and what the sidecars on disk were written from). Two old properties kept deliberately: nothing calls torch.cat, since concatenating would hold the model twice at the peak (~346 GB it need not pay), and a group with no layers allocates nothing.

The parity harness comes first, on purpose

pypto-lib fuses every layer into one kernel, so the whole-model stacked layout is an output contract and "byte-identical" has to be checkable, not intended. tests/unit/model/conftest.py adds a (shape, dtype, sha256) fingerprint; tests/unit/model/deepseek/conftest.py writes a synthetic checkpoint whose tensor set comes from deepseek_v4_layer_weight_names rather than a copied list, so it tracks the contract instead of drifting from it.

The tests are written as sensitivity rather than as green checks — each one fails on a specific regression the generic path could introduce:

  • reproducibility (without it the rest is noise);
  • layer order — a mis-ordered slab keeps every shape and dtype, so only content shows it;
  • group provenance — perturbing a compress_ratio==4 layer must move the CSA slabs and must not touch the HCA ones, and vice versa;
  • the rank axis and contiguity that alloc_stacked_tensor needs for the resident upload;
  • the full 49-name differential, over all three attention kinds, on both the direct and the destination path;
  • slab placement, over five layouts — mixed, no groups at all, CSA-only, HCA-only, interleaved — because placement is the half of the contract a shape check cannot see: every slab keeps its shape and dtype no matter which layer landed where, so wrong geometry yields correctly-shaped wrong weights;
  • the prepacked sidecar, in both directions: the payload this writer produces against the one the helpers produced, and a legacy-written file read back through the (untouched) reader and compared against a fresh pack. Plus the fingerprint recomputed against its own definition — it decides whether published files stay valid, and a changed payload would make every sidecar in the fleet look stale at once — and a check that a valid sidecar is consumed rather than silently repacked.

Three things worth flagging for review

The legacy packer and the three old stacker helpers are still in the file, the packer renamed _pack_deepseek_v4_layer_weights_legacy. Each differential compares two independent implementations; pointing both sides at the new code would make every assertion tautologically true. They go away with the suites, at the cleanup step.

safetensors does not write byte-identical files, and the obvious sidecar assertion is therefore wrong. It serializes its metadata map in nondeterministic order — eight writes of one dict in a single process gave two different key orders — so a byte-for-byte file comparison failed 2 runs in 3, in the header, for a reason unrelated to this refactor. Found by re-running rather than by reading: it passed the first time. The test compares the header's tensor entries and offsets, the metadata as a dict, and the payload bytes; the reasoning is in the docstring so nobody "fixes" it by touching the writer.

One behavioural trap, found while wiring the stacker rather than by reading it: moving the progress log into the pack callback silently dropped layer 0's line, because layer 0 takes the template-copy path and never reaches the packer. stack_layers grew an on_layer_done hook that fires once per layer whichever path it took.

The direct and destination paths cast differently — explicit .to(dtype) on one, implicitly inside copy_() on the other. #163 names this as a byte-identity trap. They agree for the conversions in use, and there is now a test asserting it rather than assuming it.

The compressor/indexer tensors use production dimensions in the fixture. The packer validates the active branch against fixed model dimensions and zero-fills the inactive branch at the same sizes, and slabs are allocated from layer 0's template — so a toy-sized CSA tensor cannot match the placeholder. The fixture imports those constants instead of copying them.

What it buys

Peak host memory during Qwen3 startup drops from ~2× the model to ~1×. The eager path kept three representations alive at the peak — the whole state dict from _load_safetensors_dir, the cast LayerWeights built from it, and the stacked destinations. Now each layer's raw tensors are read, written into their slab slice, and dropped before the next layer is read; only the destinations persist.

What could not move, and why

Executor.lookup_embeddings reads model.embed_tokens at request time, not at compile time, so the embedding cannot be released after staging — the issue flags this and it holds. The globals therefore stay eager. _load_safetensors_subset reads just those few names through safe_open rather than load_file, so reaching one tensor does not materialize the whole shard around it.

RuntimeModel.layers is now left empty for this loader, and the executor reads what it needs through extra (model_dir, weight_map). Nothing else in the tree read layers — that was checked, not assumed.

Removed as dead

_load_safetensors_dir, the eager LayerWeights loop, and _release_layer_weights — the last existed to blank out per-layer tensors after copying, which lazy reading makes unnecessary.

Step 8: the legacy path is gone

Net −1057 lines. Deleted now that the differentials guarding them have run green here and on #174: _pack_deepseek_v4_layer_weights_legacy plus the five helpers only it used, the three hand-written stacker helpers, and the two differential suites themselves. weight_loader.py goes from 1589 to 1097 lines.

Keeping the differentials longer was tempting and wrong: they compare two implementations, so the moment one is deleted the suite compares nothing. The proof is in the record and git log has both sides. What remains is the property that outlives the refactor — the sidecar this writer produces round-trips through the reader and matches a fresh pack of the same checkpoint.

LayerWeights and RuntimeModel.layers are gone for real. I had deferred that on the grounds that six open PRs construct a RuntimeModel; checking rather than assuming, none of them pass layers= (#114's two matches are num_layers=model.config.num_hidden_layers). Five call sites dropped the argument.

New docs/dev/model/weight-staging.md covers the pipeline and, more usefully, the invariants that are not visible in the code: rule order being sidecar contract, why nothing in the stacker calls torch.cat, why the stack axis differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head cannot pad with zeros. Both model docs point at it with their own specifics.

Step 7, and one thing the issue asks for that I did not do

RuntimeModel.layers is now defaulted and empty for every loader, and LayerWeights is documented as deprecated with a pointer to what replaced it. Both are kept rather than deleted: #168, #114, #152, #145, #144 and #132 are open and some construct a RuntimeModel, so removing the field would break them for a gain that can wait.

--num-layers-override loses its runtime_model.layers[:n] slice, which had quietly become a no-op — staging reads config.num_hidden_layers and pulls that many layers from the checkpoint, so replacing the config is the override now.

The stage_weights hook is deliberately absent. #163 proposes it between _create_runner and init_kv_cache, because Qwen's DistributedWorker forks inside init_kv_cache and staging must precede that. Checking the real order in PyptoExecutor.register_model: _compile_model runs first and Qwen stages inside it, well ahead of _create_runner and the fork — so the constraint the hook exists to satisfy already holds.

That leaves it as an architectural tidy-up, separating "compile kernels" from "stage weights" into named phases, and it would touch common/runner/model_runner.py plus both npu_runner.py — exactly where those six PRs are. Adding surface to the most contested files in the repo to formalise an ordering that already holds seemed the wrong trade; happy to be told otherwise, and it is a small change to make once they land.

Verification

223 tests pass across tests/unit; test_model_components.py passes unmodified, which is the gate #163 sets for this step. ruff, check-headers and check-english-only clean.

No perf change: legacy vs rules on the same input is −1.8% to −8.0% (i.e. marginally faster, at the noise level) across the three attention kinds. That is only the dispatch overhead — the byte movement is identical by construction, since both go through the same policy.

The hardware gates (DeepSeek 8-card, Qwen3 accuracy) are not run locally, but they are meaningful again now that #173 pins the pypto revision CI builds against — that pin sits before the upstream stall (hw-native-sys/pypto#2354), so a red gate on this PR would be about this PR. This was a draft until #173 landed precisely because that was not true.

Net effect

26 files, +3337/-764. weight_loader.py goes from 1589 to 1097 lines; LayerWeights,
RuntimeModel.layers, the hand-written packer and the three stacker helpers are gone.

224 unit tests pass (four fewer than before the rebase: main's #176 deleted
test_offline.py and test_device_sampling.py along with LLMEngine).

CI: the DeepSeek prefix-cache failure is not from this branch

If unit-tests comes back red on test_deepseek_v4_http_completion_matches_expected_text[k1-prefix-cache], that is #183, not this PR. That case issues two identical requests and asserts the second returns the same text as the first; the direction of the mismatch flips between runs, main failed the same assertion on 65b69fef and went green on 9dc3bf8e without anyone touching that path, and the base rate across recent runs is 3 in 9.

It cannot be this branch: the assertion compares two responses from the same weights, so a staging bug would make both wrong and the expected_text checks — including the non-prefix-cache K=1 case in the same run — would fail too. They pass.

@coderabbitai

coderabbitai Bot commented Aug 20, 2026

Copy link
Copy Markdown

Important

  • 🔍 Trigger review

This repository does not receive automatic reviews because it has fewer than 10 stars.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: d1375031-8bb6-4c05-be30-da69739676ee


Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@lterrac
lterrac marked this pull request as ready for review August 20, 2026 11:31
@lterrac
lterrac requested review from ndleslx and removed request for ndleslx August 20, 2026 11:51
lterrac added a commit to lterrac/pypto-serving that referenced this pull request Aug 21, 2026
…he new one

hw-native-sys#163 step 8, which closes the issue. Net -1057 lines.

**Deleted, now that the differentials that guarded them have run green in CI:**
`_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used
(`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`,
`_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three
hand-written stacker helpers, and the two differential suites themselves.
`weight_loader.py` goes from 1589 to 1097 lines.

Keeping them longer was tempting and wrong: a differential compares two
implementations, so the moment one is deleted the suite is comparing nothing. The
proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on
this branch — and `git log` has both sides if a byte-level question ever comes back.
What remains is the property that outlives the refactor: the sidecar written by this
writer round-trips through the reader and matches a fresh pack of the same checkpoint.

**`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated.
The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` —
did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are
`num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in
tests and the loader dropped the argument.

**Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages,
the six rule kinds, and the invariants that are not obvious from the code — rule order
being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis
differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head
cannot pad with zeros. Both model docs gain a section pointing at it with their own
specifics.

The two traps found by running the tests rather than reading them are recorded there
too, so they are not rediscovered: safetensors serializing its metadata in
nondeterministic order (which makes a whole-file assertion flaky), and a synthetic
checkpoint needing production dimensions for the compressor/indexer weights.

228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which
is the gate hw-native-sys#163 sets.
lterrac added a commit to lterrac/pypto-serving that referenced this pull request Aug 21, 2026
…he new one

hw-native-sys#163 step 8, which closes the issue. Net -1057 lines.

**Deleted, now that the differentials that guarded them have run green in CI:**
`_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used
(`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`,
`_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three
hand-written stacker helpers, and the two differential suites themselves.
`weight_loader.py` goes from 1589 to 1097 lines.

Keeping them longer was tempting and wrong: a differential compares two
implementations, so the moment one is deleted the suite is comparing nothing. The
proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on
this branch — and `git log` has both sides if a byte-level question ever comes back.
What remains is the property that outlives the refactor: the sidecar written by this
writer round-trips through the reader and matches a fresh pack of the same checkpoint.

**`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated.
The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` —
did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are
`num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in
tests and the loader dropped the argument.

**Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages,
the six rule kinds, and the invariants that are not obvious from the code — rule order
being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis
differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head
cannot pad with zeros. Both model docs gain a section pointing at it with their own
specifics.

The two traps found by running the tests rather than reading them are recorded there
too, so they are not rediscovered: safetensors serializing its metadata in
nondeterministic order (which makes a whole-file assertion flaky), and a synthetic
checkpoint needing production dimensions for the compressor/indexer weights.

228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which
is the gate hw-native-sys#163 sets.
@lterrac
lterrac force-pushed the refactor/weight-pipeline branch from 58259ab to a15e69d Compare August 21, 2026 09:41
lterrac added a commit to lterrac/pypto-serving that referenced this pull request Aug 21, 2026
…he new one

hw-native-sys#163 step 8, which closes the issue. Net -1057 lines.

**Deleted, now that the differentials that guarded them have run green in CI:**
`_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used
(`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`,
`_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three
hand-written stacker helpers, and the two differential suites themselves.
`weight_loader.py` goes from 1589 to 1097 lines.

Keeping them longer was tempting and wrong: a differential compares two
implementations, so the moment one is deleted the suite is comparing nothing. The
proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on
this branch — and `git log` has both sides if a byte-level question ever comes back.
What remains is the property that outlives the refactor: the sidecar written by this
writer round-trips through the reader and matches a fresh pack of the same checkpoint.

**`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated.
The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` —
did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are
`num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in
tests and the loader dropped the argument.

**Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages,
the six rule kinds, and the invariants that are not obvious from the code — rule order
being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis
differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head
cannot pad with zeros. Both model docs gain a section pointing at it with their own
specifics.

The two traps found by running the tests rather than reading them are recorded there
too, so they are not rediscovered: safetensors serializing its metadata in
nondeterministic order (which makes a whole-file assertion flaky), and a synthetic
checkpoint needing production dimensions for the compressor/indexer weights.

228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which
is the gate hw-native-sys#163 sets.
@lterrac lterrac changed the title refactor(model): stage DeepSeekV4 layer weights from a declarative rule table refactor(model): unify Qwen + DeepSeek weight staging behind one declarative pipeline Aug 21, 2026
lterrac added a commit to lterrac/pypto-serving that referenced this pull request Aug 21, 2026
…he new one

hw-native-sys#163 step 8, which closes the issue. Net -1057 lines.

**Deleted, now that the differentials that guarded them have run green in CI:**
`_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used
(`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`,
`_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three
hand-written stacker helpers, and the two differential suites themselves.
`weight_loader.py` goes from 1589 to 1097 lines.

Keeping them longer was tempting and wrong: a differential compares two
implementations, so the moment one is deleted the suite is comparing nothing. The
proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on
this branch — and `git log` has both sides if a byte-level question ever comes back.
What remains is the property that outlives the refactor: the sidecar written by this
writer round-trips through the reader and matches a fresh pack of the same checkpoint.

**`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated.
The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` —
did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are
`num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in
tests and the loader dropped the argument.

**Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages,
the six rule kinds, and the invariants that are not obvious from the code — rule order
being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis
differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head
cannot pad with zeros. Both model docs gain a section pointing at it with their own
specifics.

The two traps found by running the tests rather than reading them are recorded there
too, so they are not rediscovered: safetensors serializing its metadata in
nondeterministic order (which makes a whole-file assertion flaky), and a synthetic
checkpoint needing production dimensions for the compressor/indexer weights.

228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which
is the gate hw-native-sys#163 sets.
@lterrac
lterrac force-pushed the refactor/weight-pipeline branch 2 times, most recently from 9d9075d to f166bb5 Compare August 24, 2026 10:12
lterrac added a commit to lterrac/pypto-serving that referenced this pull request Aug 24, 2026
…he new one

hw-native-sys#163 step 8, which closes the issue. Net -1057 lines.

**Deleted, now that the differentials that guarded them have run green in CI:**
`_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used
(`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`,
`_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three
hand-written stacker helpers, and the two differential suites themselves.
`weight_loader.py` goes from 1589 to 1097 lines.

Keeping them longer was tempting and wrong: a differential compares two
implementations, so the moment one is deleted the suite is comparing nothing. The
proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on
this branch — and `git log` has both sides if a byte-level question ever comes back.
What remains is the property that outlives the refactor: the sidecar written by this
writer round-trips through the reader and matches a fresh pack of the same checkpoint.

**`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated.
The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` —
did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are
`num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in
tests and the loader dropped the argument.

**Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages,
the six rule kinds, and the invariants that are not obvious from the code — rule order
being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis
differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head
cannot pad with zeros. Both model docs gain a section pointing at it with their own
specifics.

The two traps found by running the tests rather than reading them are recorded there
too, so they are not rediscovered: safetensors serializing its metadata in
nondeterministic order (which makes a whole-file assertion flaky), and a synthetic
checkpoint needing production dimensions for the compressor/indexer weights.

228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which
is the gate hw-native-sys#163 sets.

def _load_safetensors_dir(model_dir: Path) -> dict[str, torch.Tensor]:
"""Load all safetensors shards from a local Hugging Face directory."""
def _load_safetensors_subset(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

We can make the ModelLoader a common base class for both qwen/deepseek to subclass it

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.

Good call — that is the same move this PR made for the stores (LazySafetensorsStore) and did not make for the loaders. Done in 33b3411.

SafetensorsDirectoryLoader now holds what is genuinely shared: supports_format was byte-identical in both, and both need a config.json before any of their own checks mean anything, so can_load becomes that precondition plus a _recognises hook per family.

I left load out rather than making it a template method, and it is worth saying why in case you disagree: the two produce different things from a directory — one stages Hugging Face layers lazily, the other validates a quantized checkpoint contract and returns metadata — so a shared skeleton would need a hook per step, which is more coupling than the duplication it removes. Happy to go further if you had a specific shape in mind.

One asymmetry kept deliberately and now commented: the DeepSeekV4 detection treats an unreadable config.json as "not mine" rather than an error, so the registry moves on to the next loader, while the Hugging Face one accepts any directory carrying safetensors.

@@ -1197,393 +1053,45 @@ def pack_deepseek_v4_layer_weights(
destinations: Mapping[str, torch.Tensor] | None = None,
prefix: str | None = None,
) -> DeepSeekV4PackedLayerWeights:

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

why we still kept this no-op func?

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.

Not a no-op, but the docstring made it look like one — it led with the history of the hand-written packer it replaced instead of what it does. Rewritten in 33b3411.

It is DeepSeekV4's binding to the generic evaluator: it turns the family's arguments into a LayerContext, selects the rank policy and the expert placement, supplies the synthetic-weight factories (the Hadamard index, the MTP smooth rows), and carries the two error templates that the existing diagnostics — and the tests matching on them — depend on. Both production callers (load_packed_layer_weights and load_mtp_weights) would otherwise repeat that wiring, and the MTP path would additionally have to know to thread prefix="mtp.0" through it.

If the concern is that the name now oversells it, I am happy to rename — but the argument surface is what the rest of the DeepSeekV4 code already calls, including tools/prepack_deepseek_v4.py, so I would rather not change the signature in this PR.

lterrac added a commit to lterrac/pypto-serving that referenced this pull request Aug 24, 2026
…n common

Review on hw-native-sys#174 asked for a common base class for the Qwen and DeepSeekV4 loaders, which
is the same move this PR made for the *stores* and had not made for the loaders.

`SafetensorsDirectoryLoader` now holds the parts that are genuinely shared:
`supports_format` was byte-identical in both, and both must find a `config.json` before
any of their own checks mean anything — so `can_load` becomes that precondition plus a
`_recognises` hook each family implements.

`load` deliberately stays out of it. The two produce different things from a directory —
one stages Hugging Face layers lazily, the other validates a quantized checkpoint
contract and hands back metadata — and a shared skeleton would need a hook per step,
which is more coupling than the duplication it removes.

The DeepSeekV4 detection keeps its own shape rather than being folded into the base: it
requires a shard index and a config that names DeepSeekV4, and treats an unreadable config
as "not mine" rather than an error, so the registry moves on to the next loader. That
asymmetry with the Hugging Face loader (which accepts any directory carrying safetensors)
is intentional and now has a comment saying so.

Also rewrote `pack_deepseek_v4_layer_weights`'s docstring, which the same review read as a
no-op wrapper — fairly, because it led with the history of the hand-written packer it
replaced instead of what it does. It is the family's binding to the generic evaluator: it
builds the `LayerContext`, selects the rank and expert-placement policies, supplies the
synthetic-weight factories, and carries the error templates its callers' diagnostics are
matched against. Both production callers would otherwise repeat that wiring, and the MTP
path would additionally have to thread `prefix="mtp.0"` through it.

257 tests pass; `ruff` and the repo hooks clean.
lterrac added a commit to lterrac/pypto-serving that referenced this pull request Aug 24, 2026
…n common

Review on hw-native-sys#174 asked for a common base class for the Qwen and DeepSeekV4 loaders, which
is the same move this PR made for the *stores* and had not made for the loaders.

`SafetensorsDirectoryLoader` now holds the parts that are genuinely shared:
`supports_format` was byte-identical in both, and both must find a `config.json` before
any of their own checks mean anything — so `can_load` becomes that precondition plus a
`_recognises` hook each family implements.

`load` deliberately stays out of it. The two produce different things from a directory —
one stages Hugging Face layers lazily, the other validates a quantized checkpoint
contract and hands back metadata — and a shared skeleton would need a hook per step,
which is more coupling than the duplication it removes.

The DeepSeekV4 detection keeps its own shape rather than being folded into the base: it
requires a shard index and a config that names DeepSeekV4, and treats an unreadable config
as "not mine" rather than an error, so the registry moves on to the next loader. That
asymmetry with the Hugging Face loader (which accepts any directory carrying safetensors)
is intentional and now has a comment saying so.

Also rewrote `pack_deepseek_v4_layer_weights`'s docstring, which the same review read as a
no-op wrapper — fairly, because it led with the history of the hand-written packer it
replaced instead of what it does. It is the family's binding to the generic evaluator: it
builds the `LayerContext`, selects the rank and expert-placement policies, supplies the
synthetic-weight factories, and carries the error templates its callers' diagnostics are
matched against. Both production callers would otherwise repeat that wiring, and the MTP
path would additionally have to thread `prefix="mtp.0"` through it.

257 tests pass; `ruff` and the repo hooks clean.
@lterrac
lterrac force-pushed the refactor/weight-pipeline branch from 33b3411 to 6432384 Compare August 24, 2026 13:10
lterrac added a commit to lterrac/pypto-serving that referenced this pull request Aug 24, 2026
…n common

Review on hw-native-sys#174 asked for a common base class for the Qwen and DeepSeekV4 loaders, which
is the same move this PR made for the *stores* and had not made for the loaders.

`SafetensorsDirectoryLoader` now holds the parts that are genuinely shared:
`supports_format` was byte-identical in both, and both must find a `config.json` before
any of their own checks mean anything — so `can_load` becomes that precondition plus a
`_recognises` hook each family implements.

`load` deliberately stays out of it. The two produce different things from a directory —
one stages Hugging Face layers lazily, the other validates a quantized checkpoint
contract and hands back metadata — and a shared skeleton would need a hook per step,
which is more coupling than the duplication it removes.

The DeepSeekV4 detection keeps its own shape rather than being folded into the base: it
requires a shard index and a config that names DeepSeekV4, and treats an unreadable config
as "not mine" rather than an error, so the registry moves on to the next loader. That
asymmetry with the Hugging Face loader (which accepts any directory carrying safetensors)
is intentional and now has a comment saying so.

Also rewrote `pack_deepseek_v4_layer_weights`'s docstring, which the same review read as a
no-op wrapper — fairly, because it led with the history of the hand-written packer it
replaced instead of what it does. It is the family's binding to the generic evaluator: it
builds the `LayerContext`, selects the rank and expert-placement policies, supplies the
synthetic-weight factories, and carries the error templates its callers' diagnostics are
matched against. Both production callers would otherwise repeat that wiring, and the MTP
path would additionally have to thread `prefix="mtp.0"` through it.

257 tests pass; `ruff` and the repo hooks clean.
@lterrac
lterrac force-pushed the refactor/weight-pipeline branch from 6432384 to c5ddc89 Compare August 24, 2026 13:11
@lterrac

lterrac commented Aug 24, 2026

Copy link
Copy Markdown
Collaborator Author

unit-tests is red on c5ddc89 for an upstream reason, and retrying will not clear it.

The DeepSeek accuracy cases (k1-fused, k1-prefix-cache) fail with DeepSeek server exited before becoming healthy (code=3). The server dies in pypto's parser before it loads anything:

File ".../pypto/language/parser/type_resolver.py", line 53, in _implicit_tile_view_defaults
  default_fractal = ir.TileView().fractal
TypeError: __init__(): incompatible function arguments.
ParserSyntaxError: Failed to parse function '_hc_pre_separate'

main @ 013cabe8 fails identically (run of 03:01 today), so this is not the PR — it reaches us because #178 dropped PYPTO_PIN and CI now builds pypto HEAD. Filed as hw-native-sys/pypto#2503.

Note this is a different failure from the intermittent prefix-cache assertion in #183; that one is not what is firing here.

platform-build and pre-commit are green. The unit tests for the staging code added here pass locally; happy to paste that run if it helps a review proceed while pypto#2503 is open.

@lterrac

lterrac commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator Author

Correcting my comment above: the attribution was wrong, and the issue I linked is closed as a duplicate.

The ir.TileView() TypeError is not a native-side change. nanobind 3.0.0 was released on 2026-08-24, and nanobind 3.0 stopped accepting the null holder that TileView::start_offset uses as its default — so an all-defaulted overload no longer binds with zero arguments. That is why unrelated branches went red within minutes of each other. pypto diagnosed and capped it in hw-native-sys/pypto#2492.

What is still true: retrying this PR will not clear it, and it is not #183's flake.

What I got wrong: it is no longer purely upstream. pypto's cap lives in pyproject.toml under [build-system] requires, and our setup-ci-job action installs pypto with --no-build-isolation, so that block is never consulted — the extension compiles against the nanobind we install, which is uncapped. This run pulled nanobind-3.0.0. So the remaining fix is on our side, not theirs.

I will put up the cap for our action separately; this PR needs it to go green but does not otherwise depend on it.

…/weights

First step of hw-native-sys#163, additive: the checkpoint reader is family-neutral, so it moves
out of DeepSeek's weight loader and into a shared base the Qwen path can adopt in
turn. `LazySafetensorsStore` owns the index (`filename_for`, `path_for`, `require`,
`__contains__`) and the reads (`load_tensor`, `load_many`, grouped one open per
shard); `DeepSeekV4WeightStore` subclasses it and keeps what is actually DeepSeek's
— which names a checkpoint must expose, and how global and per-layer tensors are
packed for the fused kernels. No behaviour change: 75 lines leave the loader, 21
arrive, and the loader's 1589-line pack path is untouched.

Two seams were deliberate rather than incidental:

- The three error messages are class attributes (`missing_name_error`,
  `missing_names_error`, `missing_shard_error`) instead of hard-coded strings, so
  DeepSeek keeps the diagnostics its users already recognise, word for word. The
  wording is part of a family's contract too, and a shared base that silently
  rewords it would be a regression nobody notices until they grep a log.
- The default opener is resolved through `_default_open_fn()`, a method, not a
  module global read from the base. `DeepSeekV4WeightStore` overrides it to return
  this module's `_default_safe_open`, which carries the DeepSeekV4-specific import
  diagnostic and is what `test_model_components.py` monkeypatches by name. Reading
  the base module's global instead would have left that patch pointing at nothing
  and the test passing for the wrong reason.

`tests/unit/model/common/test_weight_store.py` covers the generic store directly:
shard grouping (interleaved names must not reopen a shard), caller order with
duplicates dropped, the three message templates, a missing shard file, a real
safetensors round trip, and a guard on the per-construction opener resolution above.

Gate for this step, from hw-native-sys#163: `tests/unit/model/deepseek/test_model_components.py`
passes unmodified — 81 tests, and 187 across `tests/unit`.
Step 1 of hw-native-sys#163, and deliberately first: the refactor has to keep the stacked slabs
byte-identical, because pypto-lib fuses every layer into one kernel and treats the
whole-model layout as an output contract. That is only enforceable if a test can
fail on it, so the harness lands before any pipeline change.

Three pieces:

- `tests/unit/model/conftest.py` — `fingerprint_tensors`, a byte-level fingerprint
  (`name -> (shape, dtype, sha256)`). Hashing the bytes rather than comparing with
  `torch.equal` is what catches a dtype change that happens to round-trip through
  the same values.
- `tests/unit/model/deepseek/conftest.py` — a synthetic on-disk checkpoint. Which
  tensors a layer needs comes from `deepseek_v4_layer_weight_names`, not from a
  copied list, so the fixture tracks the contract and fails loudly when it grows.
- `tests/unit/model/deepseek/test_weight_pack_parity.py` — four properties over the
  real load -> pack -> stack path with the sidecar bypassed.

The tests are chosen so each one fails on a specific regression the generic stacker
could introduce: reproducibility (without it the rest is noise), sensitivity to
layer order (a mis-ordered slab keeps every shape and dtype, so only content shows
it), group provenance (perturbing a compress_ratio==4 layer must move the CSA slabs
and must not touch the HCA ones, and vice versa), and the rank axis plus contiguity
that `alloc_stacked_tensor` requires for the resident upload.

Two things the harness had to encode rather than smooth over, both found by running
it:

- The compressor/indexer tensors cannot be toy-sized. The packer validates the
  active branch against fixed model dimensions and zero-fills the inactive branch
  at those same dimensions, and every slab is allocated from layer 0's template —
  so a toy-sized CSA tensor cannot match the template placeholder and the stack
  fails on a shape mismatch. The fixture imports those constants instead of copying
  them, and stores the two tensors the packer transposes in their pre-transpose
  orientation.
- Each build needs its own directory. With a shared one the second checkpoint
  overwrote the first's shards, so the order-sensitivity test compared a checkpoint
  with itself and passed while proving nothing. That is exactly the failure mode a
  parity harness must not have, and it is why the sensitivity tests exist at all.

Not asserted here: "the stacker never calls `torch.cat`". Banning it across this
path would also ban `deepseek_v4_hadamard_idx`, which legitimately builds its matrix
with `torch.cat` for every CSA layer; the ban only makes sense where the packer is
mocked out, which is where `test_model_components.py` already keeps it.

191 tests pass across `tests/unit`.
…table

Step 3 of hw-native-sys#163, first slice, and additive: nothing calls the new path in production
yet. What lands is the schema, the evaluator, the rank policy, DeepSeekV4's core rule
table, and the differential test that proves the two agree — so the rewire that comes
next is a delegation rather than a leap.

- `common/weights/spec.py` — `LayerContext` and `LayerWeightRule`. Only `transpose`
  and `reshape_groups` are expressible as data, because both are re-orientations of
  the same bytes; anything computing new values stays in a family's code rather than
  hiding arithmetic behind a field that looks declarative.
- `common/weights/shard.py` — `Replicate`, holding both replication paths: into a
  slab slice (which is what keeps the stacker off `torch.cat`) or into a fresh
  expanded buffer.
- `common/weights/packer.py` — `pack_layer`, which walks the rules in order and skips
  a rule whose destination is absent, so a caller can stage one group of a multi-group
  layout without packing the rest.
- `deepseek/weight_spec.py` — the 25 core weights as data, present for every layer
  whatever its attention kind. The optional compressor/indexer branches, the router
  and the routed experts stay in code: each needs a rule kind this schema does not
  express yet, and inventing those kinds before their parity test exists is how a
  byte-identity regression gets in.

Rule order is contract, not style: the slab allocator lays out whole-model tensors in
the order the packed mapping is built, and the prepacked sidecar records the resulting
name-to-offset map, so reordering entries would invalidate every sidecar on disk. The
first test therefore asserts the rule order equals the hand-written table's order.

The other four are the differential hw-native-sys#163 asks for — old and new importable side by
side, compared directly rather than against a golden file:

- byte-for-byte equality on the direct path;
- byte-for-byte equality on the destination path, which the legacy packer needs full
  destinations for, so the reference pack supplies them and the comparison is taken on
  the subset the rules cover;
- direct vs destination agreement, which is the trap hw-native-sys#163 names explicitly: the direct
  path casts with `.to(dtype)` while the destination path casts implicitly inside
  `copy_()`. They agree for the conversions in use, and this asserts it instead of
  assuming it;
- the subset-of-destinations behaviour above.

196 tests pass across `tests/unit`.
…eights

Finishes the declarative half of hw-native-sys#163 step 3. Still additive — production calls the
hand-written packer — but the table is now whole, so the rewire that follows can be a
single delegation instead of a half-declarative intermediate state nobody can review.

Three rule kinds the core slice did not need, each added because a weight class exists
that cannot be honestly expressed without it:

- `OptionalWeightRule` — compressor/indexer weights, present for one attention kind and
  **zero-filled at a fixed shape** for the others. The inactive branch is written, not
  skipped: every layer must present the same kernel signature. The shape comes from the
  model, not the checkpoint, which is also why a synthetic checkpoint has to use
  production dimensions for these.
- `DefaultedWeightRule` — router weights the checkpoint may omit. `required_when` names
  the context flag that decides whether an absent source is an error or simply "this
  layer does not use that router mode", so the requirement cannot quietly decay into a
  zero fill.
- `ExpertWeightRule` + `ExpertParallel` — routed experts, sharded rather than replicated.
  Rank ownership is injected from the loader's own `deepseek_v4_local_expert_ids` rather
  than reimplemented as "contiguous block": the placement is a model decision, and a
  policy that assumed it would silently reshuffle a family that numbers experts
  differently.

`SyntheticWeightRule` covers the Hadamard index, whose factory is looked up by key so
the rule stays data. It sits between `csa_weights_proj` and `csa_inner_wkv` because that
is where the hand-written table puts it, and the order is the sidecar's offset map.

Dimensions are literals or the name of a `LayerContext` field (`resolve_shape`), which is
what keeps a config-dependent shape — `gate_bias` at `n_routed_experts` — expressible as
data rather than as a lambda smuggled into a table.

The differential now covers the whole contract: all 49 names, parametrised over the three
attention kinds, on both the direct and the destination path, plus the router case where
the layer carries `gate_bias` and the `tid2eid` placeholder is the interesting half, plus
a test that a required source still raises when absent. Name order is asserted too, not
just contents.

204 tests pass across `tests/unit`.
Completes hw-native-sys#163 step 3: `pack_deepseek_v4_layer_weights` now builds its 49 weights by
evaluating the rule table instead of by hand. Signature, output names and order,
dtypes and diagnostics are unchanged — the last of those deliberately, since the
wording is what users grep for, so the family's two "missing raw DeepSeekV4 ..."
messages became templates the evaluator formats rather than strings buried in it.

The hand-written implementation stays, renamed `_pack_deepseek_v4_layer_weights_legacy`
and documented as the parity reference until hw-native-sys#163's cleanup step. That is not
hesitation: the differential suite compares two independent implementations, and
pointing both of its sides at the public function would have made every assertion in
it tautologically true. Deleting it belongs with the step that also deletes the suite.

Two tests were added because the rewire creates a gap the existing ones cannot see.
The wrapper picks the context, the rank policy, the expert policy and the synthetic
factories itself, so a wrong expert policy or a forgotten factory would yield
plausible output that a test calling `pack_layer` directly never exercises. The new
cases compare the public entry point against the legacy packer across all three
attention kinds, and assert the family's error wording survives.

208 tests pass across `tests/unit`.
hw-native-sys#163 step 4. `load_stacked_layer_weights` now drives `common/weights/stacker.py`
instead of three DeepSeek-shaped helpers, and the slab geometry becomes data: a
`StackGroup` is a set of weights plus the *ordered* layer ids that contribute to it,
and a layer's position in that list is the slice it occupies. That is what lets a
group cover a subset of layers — the CSA group holds its layers contiguously in
first-appearance order, which is what the fused kernels index and what the prepacked
sidecars on disk were written from.

Two properties of the old code kept deliberately, because both are load bearing:

- **Nothing calls `torch.cat`.** Slabs are allocated once from the template layer's
  shapes and every later layer is packed straight into a view of its own slice.
  Concatenating instead would hold the whole model twice at the peak — ~346 GB for
  DeepSeek V4 that it does not have to pay.
- **A group with no layers allocates nothing**, so a model using none of an attention
  kind reserves no slabs for it.

The FWD group declares `members=None`, meaning "everything the other groups do not
claim, in the packer's own order". A new weight therefore joins it automatically, the
way `fwd_names` used to be derived, rather than needing to be listed in two places.

One behavioural trap found while wiring it: moving the progress log into the pack
callback silently dropped layer 0's line, because layer 0 takes the template-copy path
and never reaches the packer. `stack_layers` grew an `on_layer_done` hook that fires
once per layer whichever path it took.

`test_weight_stacker_parity.py` compares the generic stacker against the hand-written
helpers over five layouts — mixed, no groups at all, CSA-only, HCA-only, and
interleaved — because placement is the half of this contract a shape check cannot see:
every slab keeps its shape and dtype no matter which layer landed where, so wrong
geometry produces correctly-shaped wrong weights. Each layer is given distinct values
so a misplacement changes the bytes. Plus the view-aliasing property, the empty-group
case, and the two rejections that keep the catch-all unambiguous.

The three old helpers are unreferenced now but stay until the cleanup step, for the
same reason the legacy packer does: they are the other side of this differential.

218 tests pass across `tests/unit`.
…actor

The sidecars already published are the artifacts most exposed by hw-native-sys#163: each one is a
whole-model stacked payload that took ~40 minutes to build, and the refactor is only
safe if it neither changes what gets written nor rejects what was written before.
Both directions are now covered:

- the payload the current stacker writes, against the payload the hand-written helpers
  wrote — tensor entries, offsets, metadata values and every payload byte;
- a legacy-written sidecar read back through the current reader and compared against a
  fresh pack of the same checkpoint, which is the direction that protects files on
  disk;
- a currently-written sidecar actually being consumed rather than silently repacked,
  asserted by making the shard path fail the test if it is reached;
- the fingerprint recomputed against its own definition, since it is what decides
  whether published files stay valid — a changed payload would invalidate every
  sidecar in the fleet at once, silently, by making them all look stale;
- staleness still detected, via a source shard whose mtime moved.

**A byte-for-byte file comparison would have been the obvious assertion and it is
wrong.** safetensors serializes its metadata map in nondeterministic order: eight
writes of one dict in a single process produced two different key orders, so the first
version of this test failed 2 runs in 3, in the header, for a reason that has nothing
to do with this refactor. Anyone "fixing" that by touching the writer would be chasing
a phantom. The test now compares the header's tensor entries, the metadata as a dict,
and the payload bytes — which is what a reader depends on — and the reasoning is
recorded in the docstring so the trap is not rediscovered.

Found by re-running the file rather than by reading it: it passed the first time.

223 tests pass across `tests/unit`, and this file five times in a row.
The two pieces hw-native-sys#163 lists that the Qwen migration cannot start without. Both are
additive: DeepSeek now declares its staging policy instead of implying it, and nothing
else changes behaviour.

**`common/weights/pipeline.py`** owns the staging order, because peak host memory is
decided there and the two families need opposite answers. DeepSeek stages serially —
packing one layer allocates ~8 GB of intermediates, 256 routed experts each stacked and
rank-replicated, so overlapping multiplies the peak and contends on bandwidth instead of
hiding latency. Qwen wants a pool, its layers being small enough that read latency
dominates. Neither is a sensible default, so `StagingPolicy` is a parameter.

Two details in it are load bearing rather than cosmetic:

- `workers=1` takes the no-pool path, not a pool of one, so a family that must not
  overlap cannot be made to by a scheduling accident.
- each pooled worker pins torch to one thread for its duration and restores it after.
  Without that, N staging threads each fan out into torch's own pool and oversubscribe
  the machine — the copies then run *slower* than serially, which reads as "threading
  did not help" rather than as a misconfiguration.

`stack_layers` now takes the policy and routes its loop through it, so the module is
used rather than decorative. Overlapping is safe there by construction — each layer
writes a disjoint slice of each slab — but safety is not speed, which is the caller's
call to make.

**`GlobalWeightRule`** covers embedding, LM head and final norm, whose three conditions
are all silent when wrong:

- `fallback_source` — a tied checkpoint ships no `lm_head` and the embedding stands in.
  Falling back is correct; inventing zeros is not.
- `pad_to_multiple` — the fused LM head hard-codes a padded vocabulary, so the weight
  has to grow to meet it. Those rows are never selected, but they are matmul operands.
- `pad_fill` — and it differs between the two: the embedding pads with zeros, the LM
  head pads by **replicating row 0**. Zero rows in an LM head give every padded token
  the same finite logit rather than an impossible one, so the mistake survives review
  and surfaces as sampling noise. Both behaviours are copied from the Qwen executor,
  which is where they exist today.

Twenty tests, including the ones that would otherwise be assumed: that `workers=1`
really runs on the calling thread, that a pool really overlaps (via a barrier, which a
serial runner would deadlock on), that a failing layer propagates rather than leaving a
half-written slab, that torch's thread count is restored, and that a layer's tensors are
unreachable once its write returns — checked with weak references, since "we do not keep
it" is the kind of claim that quietly stops being true.

One thing removed rather than added: a `finally: raw = None` that looked like it was
releasing each layer. It was theatre — the mapping is a local of one call and dies with
it — so the docstring now says that instead.

243 tests pass across `tests/unit`.
hw-native-sys#163 step 5. The draft layer needed no new machinery for its 49 layer weights —
`LayerContext` already carries the prefix, so `mtp.0` goes through the same table as
`layers.N` — and its twelve extras, previously a hand-written dict literal, are now a
rule table of their own.

Ten of them are replicated reads. The two `_smooth` tensors are synthesized all-ones
rows, and they carry the one trap here: the factory returns a **single row** and the
rank policy replicates it to `[ranks, hidden]`. A factory returning the rank-shaped
tensor directly — which is what the literal built — would come back as
`[ranks, ranks, hidden]`, still all ones, so every value check would pass while the
shape was wrong for the kernel. There is a test for exactly that shape.

The fixture grew the draft layer, which it did not have before, so none of this was
covered until now. Its names are derived the way `load_mtp_weights` derives them —
`deepseek_v4_layer_weight_names` with `layers.0` rewritten — so the fixture cannot
drift from the loader it is meant to exercise.

Three differential tests: the twelve extras against a reproduction of the old dict
literal, the synthesized pair's shape and values, and the full `load_mtp_weights` path
end to end through the store, asserting the extras land last and in rule order.

246 tests pass across `tests/unit`.
hw-native-sys#163 step 6, first half: Qwen3's checkpoint contract, stacked layout and globals as
data, with a differential against the executor's own staging loop. Additive — the
executor still runs its loop; the metadata-only loader flip comes next.

Two family differences had to become parameters of the shared code rather than
assumptions inside it, and neither is cosmetic:

- **Stack axis.** A DeepSeekV4 weight leads with the rank axis its upload shards on and
  stacks its layers on axis 1; a rank-less Qwen weight stacks on axis 0. Get it wrong
  and the slab is the right size holding a transposed model, which no shape check
  catches. `allocate_slabs` / `destinations_for` / `stack_layers` now take `stack_axis`,
  and use `narrow` rather than a hard-coded slice.
- **No rank axis at all.** Added `NoShard` beside `Replicate`, as a distinct policy
  rather than `Replicate(ranks=1)`: the former yields `[*shape]` and the latter
  `[1, *shape]`, and a slab built from the wrong one has an extra axis nothing indexes.

`allocate_slabs` also takes the allocator, because Qwen's slabs must be shared memory —
its upload reads them from a forked child — while DeepSeekV4's need not be.

Two schema additions, both for behaviour that is silent when wrong:

- `flatten_to_row` on layer rules: a 1-D norm gamma has nothing to stack on, and
  reshaping it to `[1, dim]` is what makes it stackable at all.
- `default_fill="ones"` on `DefaultedWeightRule`, for a Qwen3 checkpoint that ships no
  QK norms. That is a model variant rather than a fault, and the neutral gamma is ones —
  zeros would annihilate the activations the gamma scales instead of leaving them
  unscaled, changing the model's output without changing a single shape.
- `LayerContext.dims`, a map rather than more fields, so `head_dim` is available to a
  rule without one family's config vocabulary accumulating on a type every family
  shares.

**A positional-argument bug the differential caught, worth recording.** Inserting
`flatten_to_row` after `default_shape` silently shifted DeepSeekV4's positional
`required_when`: the string `"include_gate_bias"` landed in `flatten_to_row`, truthy, so
`gate_bias` came out `[2, 1, 4]` instead of `[2, 4]` — same values, same dtype, extra
axis. The field moved to the end of the dataclass and those call sites now pass by
keyword, with a comment saying why.

Nine Qwen tests: the eleven slabs byte-for-byte against a reproduction of
`_stage_stacked_decode_weights`, absent QK norms defaulting to ones, the transposed
projections and `[1, dim]` gammas, no slab carrying a rank axis, a reordered-layer
sensitivity check, and the four global-weight behaviours including the tied head and the
two padding fills.

255 tests pass across `tests/unit`.
…int eagerly

hw-native-sys#163 step 6b: `HuggingFaceDirectoryLoader` becomes metadata-only for the per-layer
weights, and the Qwen executor stages them one layer at a time through the rule table.

Peak host memory drops from ~2x the model to ~1x. The eager path kept three
representations alive at the peak: the whole state dict from `_load_safetensors_dir`,
the cast `LayerWeights` built from it, and the stacked destinations. Now a layer's raw
tensors are read, written into their slab slice, and dropped before the next layer is
read; only the destinations persist.

The globals stay eager, and that is a constraint rather than a choice:
`Executor.lookup_embeddings` reads `model.embed_tokens` at **request** time, so the
embedding cannot be released after staging. `_load_safetensors_subset` reads just those
few names through `safe_open` rather than `load_file`, so reaching one tensor does not
materialize the shard holding it.

Removed as dead once nothing called them: `_load_safetensors_dir`, the eager
`LayerWeights` loop, and `_release_layer_weights` — the last of which existed to blank
out per-layer tensors after copying, which lazy reading makes unnecessary.

**The staging function is now exercised end to end, which it was not before.** The
earlier Qwen tests compared the rule table against a reproduction of the executor's old
loop, which only proves the table matches my reading of that loop. Four new tests drive
the real `_stage_stacked_decode_weights` against a Hugging Face checkpoint written to
disk and compare it to that same reproduction: the eleven slabs byte-for-byte, the
absent-QK-norm variant, that every slab is shared memory (the upload reads them from a
forked child, so private memory would be a silent correctness bug), and that a model
loaded without the metadata fails with an actionable message rather than a `KeyError`.

259 tests pass across `tests/unit`.
files the rest of it would touch.

`RuntimeModel.layers` is now defaulted and empty for every loader in the tree, and
`LayerWeights` is documented as deprecated with a pointer to what replaced it. Both are
kept rather than deleted on purpose: hw-native-sys#168, hw-native-sys#114, hw-native-sys#152, hw-native-sys#145, hw-native-sys#144 and hw-native-sys#132 are open, and
some construct a `RuntimeModel`. Removing the field would break them for no gain that
cannot wait — it goes when they have landed.

`--num-layers-override` in the Qwen example loses its `runtime_model.layers[:n]` slice,
which had quietly become a no-op: staging reads `config.num_hidden_layers` and pulls
exactly that many layers from the checkpoint, so replacing the config *is* the override
now. The knob still works; there is simply nothing eager left to truncate.

**The `stage_weights` hook is deliberately not added, and the reason is worth recording
because the issue asks for it.** hw-native-sys#163 proposes it "between `_create_runner` and
`init_kv_cache`", motivated by Qwen's `DistributedWorker` forking inside
`init_kv_cache` — so staging must happen before that. Checking the actual order in
`PyptoExecutor.register_model`, it already does: `_compile_model` runs first, and Qwen
stages inside it, well ahead of `_create_runner` and the fork. The constraint the hook
exists to satisfy is met without it.

That leaves the hook as an architectural tidy-up — separating "compile kernels" from
"stage weights" into named phases — and it would touch `common/runner/model_runner.py`
plus both `npu_runner.py`, which is exactly where the six open PRs are. Adding surface
to the most contested files in the repo to formalise a phase ordering that already holds
is a poor trade this week. Worth doing after they land, with the timing rationale
restated then rather than assumed.
…he new one

hw-native-sys#163 step 8, which closes the issue. Net -1057 lines.

**Deleted, now that the differentials that guarded them have run green in CI:**
`_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used
(`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`,
`_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three
hand-written stacker helpers, and the two differential suites themselves.
`weight_loader.py` goes from 1589 to 1097 lines.

Keeping them longer was tempting and wrong: a differential compares two
implementations, so the moment one is deleted the suite is comparing nothing. The
proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on
this branch — and `git log` has both sides if a byte-level question ever comes back.
What remains is the property that outlives the refactor: the sidecar written by this
writer round-trips through the reader and matches a fresh pack of the same checkpoint.

**`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated.
The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` —
did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are
`num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in
tests and the loader dropped the argument.

**Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages,
the six rule kinds, and the invariants that are not obvious from the code — rule order
being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis
differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head
cannot pad with zeros. Both model docs gain a section pointing at it with their own
specifics.

The two traps found by running the tests rather than reading them are recorded there
too, so they are not rediscovered: safetensors serializing its metadata in
nondeterministic order (which makes a whole-file assertion flaky), and a synthetic
checkpoint needing production dimensions for the compressor/indexer weights.

228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which
is the gate hw-native-sys#163 sets.
…n common

Review on hw-native-sys#174 asked for a common base class for the Qwen and DeepSeekV4 loaders, which
is the same move this PR made for the *stores* and had not made for the loaders.

`SafetensorsDirectoryLoader` now holds the parts that are genuinely shared:
`supports_format` was byte-identical in both, and both must find a `config.json` before
any of their own checks mean anything — so `can_load` becomes that precondition plus a
`_recognises` hook each family implements.

`load` deliberately stays out of it. The two produce different things from a directory —
one stages Hugging Face layers lazily, the other validates a quantized checkpoint
contract and hands back metadata — and a shared skeleton would need a hook per step,
which is more coupling than the duplication it removes.

The DeepSeekV4 detection keeps its own shape rather than being folded into the base: it
requires a shard index and a config that names DeepSeekV4, and treats an unreadable config
as "not mine" rather than an error, so the registry moves on to the next loader. That
asymmetry with the Hugging Face loader (which accepts any directory carrying safetensors)
is intentional and now has a comment saying so.

Also rewrote `pack_deepseek_v4_layer_weights`'s docstring, which the same review read as a
no-op wrapper — fairly, because it led with the history of the hand-written packer it
replaced instead of what it does. It is the family's binding to the generic evaluator: it
builds the `LayerContext`, selects the rank and expert-placement policies, supplies the
synthetic-weight factories, and carries the error templates its callers' diagnostics are
matched against. Both production callers would otherwise repeat that wiring, and the MTP
path would additionally have to thread `prefix="mtp.0"` through it.

257 tests pass; `ruff` and the repo hooks clean.
@lterrac
lterrac force-pushed the refactor/weight-pipeline branch from c5ddc89 to 119f2ef Compare August 25, 2026 08:17
@lterrac

lterrac commented Aug 25, 2026

Copy link
Copy Markdown
Collaborator Author

Correction to my last paragraph: no PR from me is needed. The cap already landed on main as #191 (975cbd6), which I had missed — my upstream/main was stale at 013cabe8.

Rebased onto d76111f and force-pushed (c5ddc89119f2ef); the 14 commits are unchanged apart from the new base. CI is re-running.

@superxf superxf left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Requesting changes for one process-wide Torch thread-count race; see the inline comment. The local unit and lint suites otherwise pass.


def _run(layer_id: int) -> int:
if policy.pin_torch_threads:
previous = torch.get_num_threads()

@superxf superxf Aug 25, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Restore the process-wide Torch thread count outside the worker pool

torch.set_num_threads() changes process-wide state, so saving and restoring it independently in each worker races. For example, worker A saves 8 and sets 1; worker B then saves 1; A restores 8; and B finally restores 1. The serving process is then left at one Torch thread after staging. I reproduced this deterministically with two event-ordered workers (torch_threads_after=1). The previous Qwen implementation avoided this by setting the count once before entering the pool and restoring it once after all workers had joined. Please keep the save/set/restore around the whole pool and add a test that forces overlapping workers to exit in this order.

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.

Confirmed and fixed in 1b7e292 — the race is real and the reasoning is exactly right. The pin now wraps the whole pool, under try/finally so a failing layer cannot leak it either, which is what the previous Qwen implementation did.

One correction to your test request, because I tried it your way first and it does not work. The interleaving cannot be forced from stage: each worker reads the count before its callback runs, so by the time an event inside stage can fire, both workers have already saved. I wrote that test with two event-ordered workers, ran it against the pre-fix code, and it passed — worthless as a regression test.

So it asserts the protocol instead: exactly one pin and one restore for the whole pool, whatever the worker count.

real_get = torch.get_num_threads
sets: list[int] = []
monkeypatch.setattr(torch, "set_num_threads", sets.append)
stage_layers([0, 1, 2, 3], stage=lambda _: None, policy=StagingPolicy(workers=2))
assert sets == [1, real_get()]

Verified in both directions: against the pre-fix code it fails with 8 set_num_threads calls for 4 layers instead of 2; after the fix it passes. A second test covers the raising path.

If you would still rather have the ordered-worker version alongside it, I can add it as documentation of the scenario — but it should not be the regression test, since it does not fail on the bug.

…orker

`torch.set_num_threads` is process-wide, so saving and restoring it inside each
worker races with itself: worker A saves 4 and sets 1, worker B then saves *1*,
A restores 4, and B -- finishing last -- restores 1. Staging then leaves the whole
serving process single-threaded for everything that follows. Reported in review
with a deterministic reproduction; the previous Qwen implementation set the count
once before entering the pool, which this restores.

The pin now wraps the pool, under try/finally so a failing layer cannot leak it
either.

On the test: the interleaving cannot be forced from `stage`, because each worker
reads the count before its callback runs. My first attempt did exactly that and
passed against the broken code -- worthless as a regression test. This asserts the
protocol instead: exactly one pin and one restore for the whole pool, whatever the
worker count. Verified to fail on the pre-fix code (8 set_num_threads calls for
4 layers instead of 2) and pass after. A second test covers the raising path.
@lterrac
lterrac force-pushed the refactor/weight-pipeline branch from 1b7e292 to c1b64c4 Compare August 25, 2026 13:11
@superxf
superxf merged commit 9db7e20 into hw-native-sys:main Aug 26, 2026
4 checks passed
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.

3 participants