fix(qwen3): build causal mask batch-independently (#3582) (#3586)
* fix(qwen3): build causal mask batch-independently (#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 #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 #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 P
Prajjwal Chittori committed
6e823a4306ac1c65b88b924b762ad656a8d4c216
Parent: 21cca0b
Committed by GitHub <noreply@github.com>
on 7/30/2026, 6:00:08 AM