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

[pjdurden]

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]

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.

[pjdurden]

Switched to j + w >= i + offset — equivalent to the original check but stays in usize, so no signed casts and no subtraction underflow when j > i + offset. Also added a sliding-window regression test; the existing tests only covered the no-window path.

[pjdurden]

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