Skip to content

fix(models): correct batched causal mask under-sizing tensor storage - #3780

Open
aashish-thapa wants to merge 2 commits into
huggingface:mainfrom
aashish-thapa:fix/causal-mask-batch-broadcast
Open

aashish-thapa wants to merge 2 commits into
huggingface:mainfrom
aashish-thapa:fix/causal-mask-batch-broadcast

Conversation

@aashish-thapa

@aashish-thapa aashish-thapa commented Jul 26, 2026 •

Copy link
Copy Markdown

Fixes #3764.

Hit this batching a few equal-length sequences through a Qwen3 model for embeddings. One at a time was fine, batched gave row 0 correct and the rest NaN. Traced it to the causal mask, found #3764 already open with nobody on it, so I fixed it properly.

causal_mask fills the mask for one batch element (tgt * (tgt + offset) values) but shapes the tensor (b, 1, tgt, tgt + offset). from_slice doesn't check length against a fully-concrete shape, so for b > 1 the storage is bx too small and the next broadcast_add reads off the end. That gives NaN on CUDA and a hard panic on CPU.

The mask is the same for every row, so build it once and let the existing broadcast stretch it:

Tensor::from_slice(&mask, (1, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype)

Dropping the batch dim makes b dead, so it's gone from each causal_mask.

The same (b, 1, ...)-from-single-batch-data pattern is in eight models, all broadcasting the mask, so the one change is correct for each: qwen3, qwen3_moe, quantized_qwen3, quantized_qwen3_moe, glm4_new, quantized_glm4, smol/smollm3, z_image/text_encoder. Split into two commits: qwen3 plus the test, then the siblings.

Repro (CPU, no GPU or weights). qwen3 only builds the mask off-CPU, but smol/smollm3 has the identical code and builds it anywhere:

batch = 1 ... all finite
batch = 2 ... panicked: range end index 32 out of range for slice of length 16

The mask holds 16 elements but its shape claims 32. On GPU that same read is silent and returns NaN.

Test. forward only builds the mask off-CPU, so I extracted qwen3's causal_mask to a free fn (already the crate idiom) to make it testable on CPU, and added a test pinning the shape and checking it broadcasts to identical rows. Fails on the old shape, passes on the new. The other seven are character-identical, so one canonical test rather than eight copies. Happy to add per-model if you prefer.

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.

qwen3: causal_mask builds data for one batch element but shapes the tensor (b, 1, tgt, tgt+offset) — garbage/NaN output for batch > 1

1 participant