Skip to content

fix(FIX-UNALIGNED-LOADERS-772): refuse to borrow safetensors bytes that sit at an odd element offset (#2558) - #2561

Open
lu-zero wants to merge 1 commit into
mudler:mainfrom
lu-zero:fix/exl3-borrow-alignment
Open

fix(FIX-UNALIGNED-LOADERS-772): refuse to borrow safetensors bytes that sit at an odd element offset (#2558)#2561
lu-zero wants to merge 1 commit into
mudler:mainfrom
lu-zero:fix/exl3-borrow-alignment

Conversation

@lu-zero

@lu-zero lu-zero commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator

fix(FIX-UNALIGNED-LOADERS-772): refuse to borrow safetensors bytes that sit at an odd element offset (#2558)

Fifth recurrence of the unaligned-safetensors-loaders class, this time on the
borrow path rather than in a loader. EXL3's suh/svh/trellis reach the CPU
kernels through BorrowStTensorBytes, which checked only size identity and
borrowed the safetensors mapping verbatim. A safetensors payload starts at
8 + , and the header length has arbitrary parity, so an
F16/BF16 tensor can begin on an odd byte. That is an ordinary file, not a
corrupt one. The result under the sanitizer lane is a fatal misaligned
uint16_t load in HadRowBlock (cpu_exl3_kernels.cpp:130), which aborts
test_qwen35_exl3 on main (run 33593242487, job 100131383741).

The gate: BorrowStTensorBytes now refuses to borrow when the incoming data
pointer is not aligned to the element size of the requested dtype, and to
2 bytes for kI8 trellis that EXL3's TileWord32 reads as uint16_t even though
SizeOf(kI8) is 1. Refusing costs only the direct-upload lever for that one
tensor, which is the function's stated fail-closed contract; the call sites
already fall back to MakeOwned + memcpy, so the loaded values are unchanged
bit for bit. The belt: vt::Exl3Gemm's check block now VT_CHECKs that
suh/svh/trellis pointers are 2-byte aligned, so a future producer of an
unaligned operand refuses by name instead of aborting under UBSan.

The same gate reaches the sibling finding #2540, whose odd pointer is the
RMSNorm gamma: the qwen3 loader loads input_layernorm.weight through
LoadBf16Direct (dense_weight_loaders.h:355), which rides the same
BorrowStTensorBytes borrow and the same MakeOwned + memcpy fallback. With the
gate in place the odd gamma is refused and copied instead, and the RmsNorm
widen at cpu_ops.cpp:557 never sees an odd pointer again. Measured on this
branch, sanitizer lane, VT_POOL_BYPASS=1: test_qwen35_exl3 1131/1131,
test_llama_embedding_fold, test_dots3_note_attn, test_muse_glimmer_text,
test_capi and test_openai_api_server all rc 0 — six of the eight
sanitize-cpu failures on main tip (run 33601071886) are this one defect
class, and this change claims them. The remaining failures are the
qwen4_exp rank-5 refusal (#2430's open caller-reconciliation question) and
whatever the lane surfaces next once these stop aborting first.

The test pins the fixture's payload parity deterministically, the way
test_loader_unaligned_offsets.cpp does: the checkpoint is written twice
(unpadded and padded by one header space), the two mapped payload offsets
REQUIRE different parity, and the odd one is kept. Red-first evidence on the
pinned fixture, before the fix, is the misaligned-load diagnostic at
cpu_exl3_kernels.cpp:130 with the odd address 0x...ef; after the fix the same
binary passes 1131/1131. Mutation evidence: neutralizing the gate turns the
suite red again through the belt's named refusal (the layers compose), and
forcing the gate to refuse every borrow keeps it green at 1131/1131, which
proves the copy fallback is value-correct and the fix does not depend on
borrow luck.

Closes #2558

The same gate is expected to close #2540 as well: WidenRowToF32's bf16 RMSNorm
gamma reaches the CPU kernels through this identical BorrowStTensorBytes
producer, one consumer away from suh/svh, and it borrows as kBF16 (elem == 2), so
an odd base now returns false and LoadBf16Direct falls through to its existing
MakeOwned + memcpy. That would take six further sanitize-cpu tests green
(test_llama_embedding_fold, test_muse_glimmer_text,
test_muse_glimmer_text_fallback, test_dots3_note_attn, test_openai_api_server,
test_capi). #2540 is deliberately NOT closed by this body: the prediction is
static tracing and the sanitize lane on this head is what has to confirm it.

NOT READY TO MERGE as of 2026-09-02. build-test-cpu is red on this head, and
two of the three failures are this gate rather than a stale base:
test_load_direct_upload and test_qwen3_5_dense_load_residency both assert
that the loader BORROWS the mapping, and the gate refuses those borrows. The
cost is not per-tensor: the direct-upload fixture's payload base is 8 + 171 =
179, so every tensor in it is refused, and a safetensors header length is
arbitrary. See the review comment on this pull request for the measurement and
for the vt::LoadUnaligned alternative, which removes the undefined behaviour
without switching off the zero-copy path.

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:zai-glm-5.3-flash [maki]

…at sit at an odd element offset (mudler#2558)

The borrow lever in BorrowStTensorBytes now refuses to borrow safetensors bytes when the data pointer is not aligned to the element size of the requested dtype. This prevents EXL3's HadRowBlock and TileWord32 from faulting on misaligned uint16_t loads under UBSan. For kI8 data that will be accessed as uint16_t (EXL3 trellis), the gate enforces 2-byte alignment even though SizeOf(kI8) == 1. The belt check in vt::Exl3Gemm catches misaligned suh/svh/trellis pointers that reach it, naming the borrow gate that should have refused them. The call sites already fall back to MakeOwned + memcpy when the borrow is refused, so this fix only costs the lever, never correctness.

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:zai-glm-5.3-flash [maki]
@localai-org-maint-bot

Copy link
Copy Markdown
Collaborator

Reviewed as part of an external-contributor sweep. The defect this fixes is real and the diagnosis is right, but please hold this one: build-test-cpu went red on it, and two of the three failures are caused by the gate itself rather than by a stale base.

14  - test_load_direct_upload            (Failed)
201 - test_qwen4_exp_layer_loop          (Failed)   <- stale base, not this PR
672 - test_qwen3_5_dense_load_residency  (Failed)

The two that are this change

tests/vllm/test_load_direct_upload.cpp:142:
  ERROR: CHECK( w.bytes.borrowed() ) is NOT correct!
    values: CHECK( false )

in TEST CASE: direct upload: a verbatim BF16 weight VIEWS the mapping, no host copy, and the sibling assertion in test_qwen3_5_dense_load_residency.cpp:296, dense load: the MODELOPT NVFP4 arm BORROWS the mapping, it does not copy.

Both assert the zero-copy mechanism, and the new gate refuses the borrow, so both now take the MakeOwned + memcpy path.

This is not a fixture accident

I computed the fixture's payload alignment rather than assuming. Header() in test_load_direct_upload.cpp is 171 bytes, so the payload base is 8 + 171 = 179, and every tensor in it is refused:

tensor dtype file offset address required result
w BF16 0 179 2 odd, refused
v BF16 16 195 2 odd, refused
f F32 32 211 4 211 % 4 = 3, refused

Nothing about 171 is unusual. A safetensors header length is arbitrary, so the parity of 8 + header_len is arbitrary too. That means the gate as written disables the direct-upload lever for every tensor in roughly half of all bf16 checkpoints, and for three quarters of f32 ones, rather than for the occasional odd tensor.

The body says refusing "costs only the direct-upload lever for that one tensor". The measurement says the blast radius is the whole checkpoint, decided by one arbitrary header length. That is a shipped capability being switched off by luck, and it is why I would rather not land this shape.

The direction I would suggest

This tree already has the seam for this exact class: vt::LoadUnaligned (include/vt/unaligned.h), which is how the first four recurrences in .agents/specs/unaligned-safetensors-loaders.md were repaired, and it is also what cpu_layernorm.cpp and vt::cpu::LoadF32 use. Fixing the consumer keeps the borrow and removes the UB, where refusing the borrow removes the UB by removing the feature.

Concretely: keep the VT_CHECK belts in Exl3Gemm, since a named refusal beats a UBSan abort, and move HadRowBlock (cpu_exl3_kernels.cpp:130) and TileWord32 (cpu_exl3_dequant.cpp:64) onto vt::LoadUnaligned instead of gating BorrowStTensorBytes. That closes #2558 without costing the zero-copy path.

If you would rather keep the gate, then the two tests above encode the current contract and would need to change deliberately, with the cost argued in the body, which is a bigger decision than this PR is currently scoped as.

One that is not yours

test_qwen4_exp_layer_loop is the rank-5 paged-KV break (#2559 / #2560), already fixed on main by 33b08e463. Your merge base is bb78d1ee8, which predates it. A rebase clears that one.

Worth keeping

Your producer-side diagnosis looks correct and is more valuable than this PR's own scope: the same BorrowStTensorBytes hands WidenRowToF32 its bf16 RMSNorm gamma, which is #2540, so one producer explains both findings. I have noted that on the PR body. windows-msvc-cpu and windows-msvc-vulkan are red tree-wide and are not yours.

@localai-org-maint-bot

Copy link
Copy Markdown
Collaborator

Follow-up, and it is good news that I want on the record before my earlier comment reads as more negative than it should.

The sanitize-cpu (address,undefined) lane has now reported on this head, and your diagnosis is confirmed. It went from eight failures to three, and every test in the misaligned-load class passed:

test_llama_embedding_fold        Passed    0.62 sec
test_muse_glimmer_text           Passed    0.70 sec
test_muse_glimmer_text_fallback  Passed    0.71 sec
test_dots3_note_attn             Passed    2.21 sec
test_openai_api_server           Passed  147.68 sec
test_capi                        Passed   38.91 sec
test_qwen35_exl3                 Passed    1.32 sec

with zero misaligned address reports in the whole job, and zero at cpu_matmul_elem.cpp:577.

That is worth stating plainly: this change closes #2558 and #2540 together, seven tests across both, and it proves they were always one defect with one producer. That was the genuinely valuable finding here, and it is not obvious from either issue.

The three that remain are:

14  - test_load_direct_upload            <- the gate's cost
672 - test_qwen3_5_dense_load_residency  <- the gate's cost
201 - test_qwen4_exp_layer_loop          <- stale base, clears on rebase

So the position is narrower than my first comment may have implied. Nothing is wrong with the analysis, and nothing is wrong with where you put the fix in the dependency chain. The only open question is the mechanism, and it is now a straight trade-off with both sides measured:

  • As written: 7 tests fixed, 2 zero-copy assertions broken, and direct upload disabled for every tensor in any checkpoint whose 8 + header_len has the wrong parity.
  • Consumer-side vt::LoadUnaligned: the same 7 fixed, zero-copy kept, more work.

Given the confirmation above, the consumer-side version is now well-targeted rather than speculative: we know the producer is the single root cause, so putting WidenRowToF32 (cpu_matmul_elem.cpp:577), HadRowBlock (cpu_exl3_kernels.cpp:130) and TileWord32 (cpu_exl3_dequant.cpp:64) on the seam should reproduce exactly this pass list without touching the borrow.

I am happy to take that on if you would rather not, since the branch has already done the hard part. Your call.

@lu-zero

lu-zero commented Sep 2, 2026

Copy link
Copy Markdown
Collaborator Author

Respinning consumer-side. Your alignment computation is what settled it: a checkpoint-wide lever switched off by one arbitrary header length is a product regression no token gate can see, and the two red tests are the tree stating that contract. The borrow stays.

The shape of the respin, on this branch:

  • BorrowStTensorBytes loses the alignment gate; VT_LOAD_DIRECT_UPLOAD keeps its meaning, and both zero-copy assertions stand.
  • The consumers move onto vt::LoadUnaligned, the seam the first four recurrences used: WidenRowToF32 (cpu_matmul_elem.h:110, inlined into the cpu_ops.cpp:557 kernel — which also retires the unowned cpu_matmul_elem.cpp:577 finding), HadRowBlock (cpu_exl3_kernels.cpp:130), and TileWord32 (cpu_exl3_dequant.cpp:64).
  • The VT_CHECK belts in vt::Exl3Gemm stay: a named refusal at the entry beats a sanitizer abort inside the kernel.
  • A rebase picks up the rank-5 paged-KV repair (33b08e463), which clears the third red you listed.
  • The fixture-parity pin in test_qwen35_exl3.cpp is reworked to pin the unaligned-consumer path rather than the refusal.

We are taking the respin ourselves rather than handing it over — the branch is ours to rebase and the worktree is already set up for the sanitizer lane — but thank you for the offer, and for the sweep. The producer-side finding stands either way: one BorrowStTensorBytes explains both #2558 and #2540, your pass list is the evidence the spec now cites, and the three-site consumer repair is exactly the shape that finding predicted.

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

Labels

None yet

Projects

None yet

2 participants