huggingface / huggingface/candle
Qwen2: prepare_attention_mask is bidirectional (no causal triangle) when attn_mask provided
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Severity: CRITICAL — silently wrong outputs under batched eval / training
`Qwen2Model::prepare_attention_mask` at `candle-transformers/src/models/qwen2.rs:326-339` constructs the attention mask using **only the padding signal**, with no causal triangle. When an `attn_mask` is provided to `Model::forward()`, the model becomes a bidirectional encoder — every non-padding position can attend to every other non-padding position, including **future tokens**.
**Why this is silent.** Inference of a single un-padded prompt hits the `None` branch in `forward()` and is correct. The bug only triggers under batched eval / training / any caller that passes `attn_mask` for variable-length padding. Symptom: eval `token_acc = 100%` because the model trivially "predicts" tokens it can already see.
**Verified present at candle main HEAD `1d7e9274591ef7eb101d965b71243a41994214d2`** (committed 2026-05-26).
## Reproducer
https://github.com/nxrobins/candle-bug-6-qwen2-causal-mask
Standalone crate pinned to commit `7c7a8c570e6a16ea24cc30a7501c8ffbbdb51680` (the "Bump candle version to 0.10.2" commit, since 0.10.2 isn't tagged in the repo). One-line repro:
\`\`\`bash
git clone https://github.com/nxrobins/candle-bug-6-qwen2-causal-mask && cd candle-bug-6-qwen2-causal-mask && cargo run --release
# exit 1: BUG PRESENT (6 zeros above the diagonal of a 4×4 mask)
# exit 0: bug appears fixed
\`\`\`
Since `prepare_attention_mask` is `pub(crate)`, the reproducer replicates the upstream function byte-for-byte with attribution and ships a positive-control corrected implementation. Both \`cargo run\` and \`cargo test\` (2 tests) pass; no CUDA/Metal required.
## Suggested fix
Combine the padding mask with a lower-triangular causal mask via elementwise AND before converting to the \`0 / -inf\` attention bias:
\`\`\`rust
fn prepare_attention_mask(&self, attn_mask: &Tensor) -> Result {
let (b_sz, sql_len) = attn_mask.dims2()?;
let mut mask: Vec = vec![];
for b in 0..b_sz {
mask.push(attn_mask.i((b, ..))?.expand((1, 1, sql_len, sql_len))?);
}
let padding_mask = Tensor::cat(&mask, 0)?;
// Lower-triangular causal mask.
let causal: Vec = (0..sql_len)
.flat_map(|i| (0..sql_len).map(move |j| if j <= i { 1.0 } else { 0.0 }))
.collect();
let causal_2d = Tensor::from_slice(&causal, (sql_len, sql_len), &self.device)?;
let causal_bool = causal_2d
.reshape((1, 1, sql_len, sql_len))?
.broadcast_as(padding_mask.shape())?
.to_dtype(padding_mask.dtype())?;
// A cell is allowed iff BOTH padding and causal allow it.
let combined = (padding_mask * causal_bool)?;
let on_true = combined.zeros_like()?.to_dtype(self.dtype)?;
let on_false = Tensor::new(f32::NEG_INFINITY, &self.device)?
.broadcast_as(combined.shape())?
.to_dtype(self.dtype)?;
combined.where_cond(&on_true, &on_false)
}
\`\`\`
This was applied locally to a vendored copy of \`qwen2.rs\` and verified end-to-end against a Python reference: \`token_acc\` returned to expected ~35% on Qwen2.5-0.5B (was 100% with the bug).
## Likely affected sibling models
Any model in \`candle-transformers/src/models/\` with a similar \`prepare_*_mask\` helper that doesn't include causality. \`git grep -lE \"fn prepare_(attention|padding)_mask\" candle-transformers/src/models/\` returns the audit list. Happy to audit + file per-model PRs if useful.
## Context
Surfaced during Phase 0.5b validation of an external project against a Python reference implementation. Full writeup with the four bugs that surfaced together: https://github.com/nxrobins/wave-agent/blob/main/docs/upstream-candle-bugs.md
Contributor guide
No contributing guide indexed for this repository
Research direction
Start with candle-transformers/src/models/qwen2.rs:326-339 and compare prepare_attention_mask with the standalone reproducer and its Python reference. Run cargo test or cargo run --release in candle-bug-6-qwen-causal-mask; the work is done when the reproducer exits 0 and batched attention no longer permits future-token access.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 70/100