fix(qwen3): build causal mask batch-independently (#3582) by pjdurden · Pull Request #3586 · huggingface/candle

pjdurden · 2026-06-06T23:31:58Z

Motivation

Closes #3582.

Qwen3 produces incorrect output for batch size > 1. Model::causal_mask builds
the additive mask buffer with tgt * (tgt + offset) elements — independent of the
batch — but then shapes it (b, 1, tgt, tgt + offset). Tensor::from_slice with a
fully-specified shape doesn't validate the element count (the blanket
ShapeWithOneHole impl ignores it), so for b > 1 the tensor claims b× the
elements that actually exist in storage. Batch row 0 reads the correct mask; every
row after it reads past the buffer and is masked incorrectly. The mask only feeds
the standard matmul attention path via scores.broadcast_add(m), so this surfaces
there (e.g. Metal) and disappears at b = 1 — matching the report.

Modifications

Extracted the mask construction into a build_causal_mask free function that
shapes the mask (1, 1, tgt, tgt + offset) and lets the existing
broadcast_add apply it across the batch. This is correct for any b and also
drops the redundant per-batch allocation.
Model::forward now calls build_causal_mask directly; the b it used to pass
into the mask is no longer needed.

Tests

Added regression test causal_mask_is_batch_independent_and_broadcasts,
asserting the mask carries a leading batch dim of 1 and, broadcast onto a
2-sequence batch, produces an identical causal mask for both rows. It fails on
the old (b, 1, …) shape and passes after the fix.
cargo test -p candle-transformers --lib passes (14 tests).
cargo fmt --check and cargo clippy are clean.

pjdurden · 2026-06-06T23:32:08Z

While fixing this I noticed the same (b, 1, tgt, …) pattern — a batch-independent mask buffer shaped with the batch dimension — in several other models: qwen3_moe, quantized_qwen3, quantized_qwen3_moe, glm4_new, quantized_glm4, smol/smollm3, and z_image/text_encoder. They look like the same copy-pasted bug and would hit the same way for batch size > 1.

I kept this PR scoped to qwen3 (the reported model) for easy review, but I'm happy to fix the rest — either folded into this PR or as a follow-up. Just let me know which you'd prefer.

Copilot

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Fixes a batch-dependent corruption in Qwen3’s additive causal attention mask by constructing a batch-independent mask and relying on broadcasting, with a regression test to prevent reintroduction.

Changes:

Extracted causal mask construction into a batch-independent build_causal_mask helper (shape (1, 1, tgt, tgt + offset)).
Updated Model::forward to use the new helper and avoid building a (b, ...)-shaped mask.
Added a regression test validating correct broadcasting across batch.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Address review feedback on huggingface#3586: the sliding-window check computed `(i + offset) as i64 - j as i64 <= w as i64`, casting to signed to allow a negative result when `j > i + offset`. Rearranged to `j + w >= i + offset` — equivalent, but stays in `usize` with no signed casts and no subtraction underflow. Add a sliding-window regression test; the prior tests only covered the no-window path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

pjdurden · 2026-06-10T00:48:34Z

@ivarflakstad gentle nudge when you get a chance. Small qwen3 causal mask fix for batch greater than 1, with a regression test.

pjdurden · 2026-06-27T08:28:35Z

@ivarflakstad @LaurentMazare last ping on this one, then I'll leave it be. The fix is small and self-contained: the qwen3 causal mask was built with the batch dim baked into the buffer shape, so for batch > 1 every sequence after the first got the wrong mask. The PR shapes it batch-independently and adds a regression test that fails on main and passes here.

CI is green and the Copilot threads are resolved. If you'd rather I also fold in the same fix for the sibling models I listed above (qwen3_moe, quantized_qwen3, quantized_qwen3_moe, glm4_new, quantized_glm4, smollm3, z_image text_encoder), happy to push that in a follow-up so this PR stays scoped to qwen3. Either way, glad to rebase or tweak anything that's blocking a merge.

AlpineVibrations · 2026-06-29T14:41:06Z

any updates on this. seams like a good fix for all the related models...? nice one.

pjdurden · 2026-07-28T20:43:44Z

no maintainer feedback yet, still just the copilot pass. this one is deliberately scoped to qwen3 so it stays easy to review.

the same batch-shaped mask buffer is in qwen3_moe, quantized_qwen3, quantized_qwen3_moe, glm4_new, quantized_glm4, smollm3 and z_image text_encoder. happy to send those as a follow-up once a maintainer confirms the approach here.

`Model::causal_mask` built the additive mask buffer with `tgt*(tgt+offset)` elements (independent of the batch) but shaped it `(b, 1, tgt, tgt+offset)`. For `b > 1` the tensor claims b× the elements actually present, so every batch row but the first reads past the buffer and is masked incorrectly, producing wrong output for batched forwards (the bug disappears at b=1). This is hit on the standard matmul attention path (e.g. Metal), as reported in huggingface#3582. Extract the mask construction into a `build_causal_mask` free function that shapes the mask `(1, 1, tgt, tgt+offset)` and relies on the existing `broadcast_add` to apply it across the batch. Add a CPU regression test asserting the mask is batch-independent and broadcasts to a causal, per-row-identical mask. The same `(b, 1, tgt, ...)` pattern exists in several sibling models (qwen3_moe, quantized_qwen3{,_moe}, glm4_new, quantized_glm4, smollm3, z_image/text_encoder); happy to fix those in this PR or a follow-up.

Address review feedback on huggingface#3586: the sliding-window check computed `(i + offset) as i64 - j as i64 <= w as i64`, casting to signed to allow a negative result when `j > i + offset`. Rearranged to `j + w >= i + offset` — equivalent, but stays in `usize` with no signed casts and no subtraction underflow. Add a sliding-window regression test; the prior tests only covered the no-window path.

ivarflakstad

Thanks for the contribution! :)

ivarflakstad · 2026-07-29T09:06:58Z

+/// Builds the additive causal attention mask of shape `(1, 1, tgt, tgt + offset)`.
+///
+/// The mask values depend only on the query/key positions, never on the batch,
+/// so it is built with a leading batch dim of `1` and broadcast across the batch
+/// by `broadcast_add` in attention. Shaping this `b`-independent buffer as
+/// `(b, 1, tgt, tgt + offset)` claims `b×` the elements actually present, so every
+/// batch row but the first reads past the buffer and is masked incorrectly
+/// (see https://github.com/huggingface/candle/issues/3582).


The first line is fine.

The section below is not necessary. We don't need an explanation about why the old approach was incorrect - it'll just be confusing when the old approach is gone, because then it's just a rant about something unrelated which is no longer there.

Anyway now I'm ranting. Just remove the section and it's fine :)

ivarflakstad · 2026-07-29T09:08:45Z

+/// `(b, 1, tgt, tgt + offset)` claims `b×` the elements actually present, so every
+/// batch row but the first reads past the buffer and is masked incorrectly
+/// (see https://github.com/huggingface/candle/issues/3582).
+fn build_causal_mask(


Any specific reason to moving this to a fn instead of in the impl?

Not a problem, just wondering if there was a specific reason.

ivarflakstad · 2026-07-29T09:12:34Z

+    // Regression test for https://github.com/huggingface/candle/issues/3582:
+    // the additive causal mask is independent of the batch dimension, so it must
+    // be built with a leading batch dim of 1 and broadcast across the batch.
+    // Shaping it `(b, 1, tgt, tgt + offset)` from a `b`-independent buffer claims
+    // `b×` the elements actually present, corrupting every batch row but the first.


Suggested change

// Regression test for https://github.com/huggingface/candle/issues/3582:

// the additive causal mask is independent of the batch dimension, so it must

// be built with a leading batch dim of 1 and broadcast across the batch.

// Shaping it `(b, 1, tgt, tgt + offset)` from a `b`-independent buffer claims

// `b×` the elements actually present, corrupting every batch row but the first.

// Regression test for https://github.com/huggingface/candle/issues/3582

pjdurden · 2026-07-30T01:20:37Z

Trimmed the doc comment to the first line and applied the test comment suggestion. build_causal_mask is a free fn only so the tests can call it without building a Model. It only needed device and dtype off self, so passing those in seemed simpler than keeping it a method. Can move it back into impl Model if you prefer.

ivarflakstad

👍

…pendently Follow-up to huggingface#3586 / huggingface#3582, extending that fix to every remaining model with the same pattern and moving the shared logic into `crate::utils` alongside `build_causal_mask`, as huggingface#3437 did for the rectangular mask. Eight models filled the additive mask buffer with `tgt * (tgt + offset)` elements (independent of the batch) but shaped the tensor `(b, 1, tgt, tgt + offset)`. `Tensor::from_slice` does not validate the element count against the shape, since the blanket `impl<S: Into<Shape>> ShapeWithOneHole` discards `el_count`, so the oversized tensor is built silently and only misbehaves at use: - CPU: panics with `range end index N out of range for slice of length M` once the mask is read. - Metal/CUDA: reads past the buffer, so every batch row after the first is masked with garbage and the model returns wrong output. That is what huggingface#3582 reported. Affected: qwen3_moe, quantized_qwen3, quantized_qwen3_moe, glm4_new, quantized_glm4, smol/smollm3, smol/quantized_smollm3 and z_image/text_encoder. Unlike qwen3, where the mask path was gated to CPU-only under the `flash-attn` feature, these gate only on `l == 1`, so the broken mask was built on every multi-token forward on every backend. `utils::build_additive_causal_mask` returns `(1, 1, tgt, tgt + offset)` and is broadcast over the batch by the existing `broadcast_add`. All nine models now call it, including qwen3, which drops the per-model copy added in huggingface#3586. The sliding-window check uses `usize` arithmetic (`j + w >= i + offset`) as in huggingface#3586, avoiding the signed casts. Tests move with the helper: mask shape and batch broadcast, the KV-cache offset case, and the sliding window.

…pendently (#3879) Follow-up to #3586 / #3582, extending that fix to every remaining model with the same pattern and moving the shared logic into `crate::utils` alongside `build_causal_mask`, as #3437 did for the rectangular mask. Eight models filled the additive mask buffer with `tgt * (tgt + offset)` elements (independent of the batch) but shaped the tensor `(b, 1, tgt, tgt + offset)`. `Tensor::from_slice` does not validate the element count against the shape, since the blanket `impl<S: Into<Shape>> ShapeWithOneHole` discards `el_count`, so the oversized tensor is built silently and only misbehaves at use: - CPU: panics with `range end index N out of range for slice of length M` once the mask is read. - Metal/CUDA: reads past the buffer, so every batch row after the first is masked with garbage and the model returns wrong output. That is what #3582 reported. Affected: qwen3_moe, quantized_qwen3, quantized_qwen3_moe, glm4_new, quantized_glm4, smol/smollm3, smol/quantized_smollm3 and z_image/text_encoder. Unlike qwen3, where the mask path was gated to CPU-only under the `flash-attn` feature, these gate only on `l == 1`, so the broken mask was built on every multi-token forward on every backend. `utils::build_additive_causal_mask` returns `(1, 1, tgt, tgt + offset)` and is broadcast over the batch by the existing `broadcast_add`. All nine models now call it, including qwen3, which drops the per-model copy added in #3586. The sliding-window check uses `usize` arithmetic (`j + w >= i + offset`) as in #3586, avoiding the signed casts. Tests move with the helper: mask shape and batch broadcast, the KV-cache offset case, and the sliding window.

…uggingface#3586) * fix(qwen3): build causal mask batch-independently (huggingface#3582) `Model::causal_mask` built the additive mask buffer with `tgt*(tgt+offset)` elements (independent of the batch) but shaped it `(b, 1, tgt, tgt+offset)`. For `b > 1` the tensor claims b× the elements actually present, so every batch row but the first reads past the buffer and is masked incorrectly, producing wrong output for batched forwards (the bug disappears at b=1). This is hit on the standard matmul attention path (e.g. Metal), as reported in huggingface#3582. Extract the mask construction into a `build_causal_mask` free function that shapes the mask `(1, 1, tgt, tgt+offset)` and relies on the existing `broadcast_add` to apply it across the batch. Add a CPU regression test asserting the mask is batch-independent and broadcasts to a causal, per-row-identical mask. The same `(b, 1, tgt, ...)` pattern exists in several sibling models (qwen3_moe, quantized_qwen3{,_moe}, glm4_new, quantized_glm4, smollm3, z_image/text_encoder); happy to fix those in this PR or a follow-up. * fix(qwen3): use usize arithmetic for sliding-window mask check Address review feedback on huggingface#3586: the sliding-window check computed `(i + offset) as i64 - j as i64 <= w as i64`, casting to signed to allow a negative result when `j > i + offset`. Rearranged to `j + w >= i + offset` — equivalent, but stays in `usize` with no signed casts and no subtraction underflow. Add a sliding-window regression test; the prior tests only covered the no-window path. * docs(qwen3): trim causal mask comments per review

…pendently (huggingface#3879) Follow-up to huggingface#3586 / huggingface#3582, extending that fix to every remaining model with the same pattern and moving the shared logic into `crate::utils` alongside `build_causal_mask`, as huggingface#3437 did for the rectangular mask. Eight models filled the additive mask buffer with `tgt * (tgt + offset)` elements (independent of the batch) but shaped the tensor `(b, 1, tgt, tgt + offset)`. `Tensor::from_slice` does not validate the element count against the shape, since the blanket `impl<S: Into<Shape>> ShapeWithOneHole` discards `el_count`, so the oversized tensor is built silently and only misbehaves at use: - CPU: panics with `range end index N out of range for slice of length M` once the mask is read. - Metal/CUDA: reads past the buffer, so every batch row after the first is masked with garbage and the model returns wrong output. That is what huggingface#3582 reported. Affected: qwen3_moe, quantized_qwen3, quantized_qwen3_moe, glm4_new, quantized_glm4, smol/smollm3, smol/quantized_smollm3 and z_image/text_encoder. Unlike qwen3, where the mask path was gated to CPU-only under the `flash-attn` feature, these gate only on `l == 1`, so the broken mask was built on every multi-token forward on every backend. `utils::build_additive_causal_mask` returns `(1, 1, tgt, tgt + offset)` and is broadcast over the batch by the existing `broadcast_add`. All nine models now call it, including qwen3, which drops the per-model copy added in huggingface#3586. The sliding-window check uses `usize` arithmetic (`j + w >= i + offset`) as in huggingface#3586, avoiding the signed casts. Tests move with the helper: mask shape and batch broadcast, the KV-cache offset case, and the sliding window.

Copilot AI review requested due to automatic review settings June 6, 2026 23:31

Copilot AI reviewed Jun 6, 2026

View reviewed changes

Comment thread candle-transformers/src/models/qwen3.rs

pjdurden added 2 commits July 28, 2026 16:44

pjdurden force-pushed the fix/qwen3-causal-mask-batch-gt-1 branch from 7eb8896 to ae860e9 Compare July 28, 2026 21:45

ivarflakstad reviewed Jul 29, 2026

View reviewed changes

docs(qwen3): trim causal mask comments per review

c164cbf

pjdurden force-pushed the fix/qwen3-causal-mask-batch-gt-1 branch from bce00ee to c164cbf Compare July 30, 2026 00:51

ivarflakstad approved these changes Jul 30, 2026

View reviewed changes

ivarflakstad merged commit 6e823a4 into huggingface:main Jul 30, 2026
11 checks passed

pjdurden mentioned this pull request Aug 13, 2026

fix(models): centralize the additive causal mask, build it batch-independently #3879

Merged

Navigation Menu

Repository navigation

fix(qwen3): build causal mask batch-independently (#3582) - #3586

fix(qwen3): build causal mask batch-independently (#3582)#3586
ivarflakstad merged 3 commits into
huggingface:mainfrom
pjdurden:fix/qwen3-causal-mask-batch-gt-1

pjdurden commented Jun 6, 2026

Uh oh!

pjdurden commented Jun 6, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

pjdurden commented Jun 10, 2026

Uh oh!

pjdurden commented Jun 27, 2026

Uh oh!

AlpineVibrations commented Jun 29, 2026

Uh oh!

pjdurden commented Jul 28, 2026

Uh oh!

ivarflakstad left a comment

Uh oh!

ivarflakstad Jul 29, 2026

Uh oh!

ivarflakstad Jul 29, 2026

Uh oh!

ivarflakstad Jul 29, 2026

Uh oh!

pjdurden commented Jul 30, 2026

Uh oh!

ivarflakstad left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

pjdurden commented Jun 6, 2026

Motivation

Modifications

Tests

Uh oh!

pjdurden commented Jun 6, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Uh oh!

pjdurden commented Jun 10, 2026

Uh oh!

pjdurden commented Jun 27, 2026

Uh oh!

AlpineVibrations commented Jun 29, 2026

Uh oh!

pjdurden commented Jul 28, 2026

Uh oh!

ivarflakstad left a comment

Choose a reason for hiding this comment

Uh oh!

ivarflakstad Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

ivarflakstad Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

ivarflakstad Jul 29, 2026

Choose a reason for hiding this comment

Uh oh!

pjdurden commented Jul 30, 2026

Uh oh!

ivarflakstad left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants