SIGN IN SIGN UP

fix(models): centralize the additive causal mask, build it batch-independently (#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.
P
Prajjwal Chittori committed
3b7af4a5ba7cf85b6c852680410913f6418e7487
Parent: 6f74e7c
Committed by GitHub <noreply@github.com> on 8/13/2026, 2:01:36 PM