Skip to content

fix(loader): refuse wrong-cardinality tensors before allocating or reading - #1824

Open
Xore wants to merge 1 commit into
JustVugg:devfrom
Xore:fix/tensor-cardinality-read-side
Open

Xore wants to merge 1 commit into
JustVugg:devfrom
Xore:fix/tensor-cardinality-read-side

Conversation

@Xore

@Xore Xore commented Oct 1, 2026

Copy link
Copy Markdown

Splits out the loader-hardening part of #1564, as requested. Rebuilt from scratch on current dev (4521832a) so it carries no unrelated changes.

Problem

load_t() sized tensors by reading whatever the file contained. A truncated or mismatched container could therefore be read as a config-sized tensor, and the resulting out-of-bounds access happened during inference rather than at load time, where it can be refused cleanly.

Change

Pass the expected element count to the loader at every call site in colibri.c, inkling.c and olmoe.c, and refuse a mismatch before allocation or read. Upstream's embed_norm_name() handling in inkling.c is preserved.

if (n != want) { /* refuse: shape does not match the config */ }

The guard runs before allocation, so no correctly shaped container is rejected and no short one is ever read at config width.

Verification

Rebuilt onto current dev; no upstream behaviour is removed (the only deletions are loader call sites rewritten to pass the expected cardinality).

  • make colibri inkling olmoe and the build/segment/{inkling,olmoe}.o objects — clean
  • make test-c — full suite passes
  • New: test_colibri_read_trust, test_inkling_read_trust, test_olmoe_read_trust — all pass
  • Real model, real CUDA, pinned oracle, on the real 35B MoE container:
int4                      Matching tokens: 16/16
int8  CUDA_EXPERT_GB=18   Matching tokens: 16/16
int8  CUDA_EXPERT_GB=14   Matching tokens: 16/16
int8  CUDA_EXPERT_GB=10   Matching tokens: 16/16

Identical to pristine dev on the same containers.

Notes

This replaces the loader-hardening portion of #1564. The rest of that PR is still being split out separately; this one stands alone and can be reviewed and merged on its own.

@JustVugg

JustVugg commented Oct 1, 2026 •

Copy link
Copy Markdown
Owner

Thanks for splitting this out. The hardening is right: ld()/load_t() sized the buffer from the file while the forward pass indexes it with config dims, so a short tensor was read past its end at inference.

We checked every expected count against the forward pass on dev, including the GLM MTP norms (enorm, hnorm, shared_head.norm), the indexer k_norm, inkling's router (n_experts + n_shared), its sconv kernels (conv_k) and rel_logits_proj (d_rel x ext). They all match. Locally on top of dev:

  • glm_tiny oracle: teacher forcing 32/32, greedy 20/20;
  • olmoe tiny oracle: 8/8;
  • the three new tests pass, and they exercise the real loaders since they include the engine sources.

One thing blocks it. Since #1758, dev generates header dependencies with -MMD, and tests.test_makefile_deps fails on the three new rules because they list headers by hand. They should name only the source and the engine, like the neighbouring rules:

tests/test_colibri_read_trust$(EXE): tests/test_colibri_read_trust.c colibri.c $(VK_OBJ)
tests/test_inkling_read_trust$(EXE): tests/test_inkling_read_trust.c inkling.c $(INK_CUDA_OBJ) $(METAL_OBJ)
tests/test_olmoe_read_trust$(EXE): tests/test_olmoe_read_trust.c olmoe.c

(keep whatever recipe line each rule already has). A small note on the description: the 35B MoE runs exercise qwen36, which this PR does not touch, so the tiny oracles above are the relevant evidence for colibri, inkling and olmoe.

Your CI has not run yet because this is a first contribution; we will approve it after the update and merge when it is green.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants