Repository navigation
fix(models): correct batched causal mask under-sizing tensor storage - #3780
Open
aashish-thapa wants to merge 2 commits into
Open
aashish-thapa wants to merge 2 commits into
aashish-thapa wants to merge 2 commits into
Conversation
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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_maskfills the mask for one batch element (tgt * (tgt + offset)values) but shapes the tensor(b, 1, tgt, tgt + offset).from_slicedoesn't check length against a fully-concrete shape, so forb > 1the storage isbx too small and the nextbroadcast_addreads 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:
Dropping the batch dim makes
bdead, so it's gone from eachcausal_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/smollm3has the identical code and builds it anywhere:The mask holds 16 elements but its shape claims 32. On GPU that same read is silent and returns NaN.
Test.
forwardonly builds the mask off-CPU, so I extracted qwen3'scausal_maskto 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.