huggingface / huggingface/candle
qwen3: causal_mask builds data for one batch element but shapes the tensor (b, 1, tgt, tgt+offset) — garbage/NaN output for batch > 1
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Bug
In `candle-transformers/src/models/qwen3.rs`, `Model::causal_mask` generates mask data for a single batch element but declares the tensor shape with the full batch dimension:
```rust
let mask: Vec<_> = (0..tgt)
.flat_map(|i| {
(0..(tgt + offset)).map(move |j| { /* 0. or -inf */ })
})
.collect(); // len == tgt * (tgt + offset)
Tensor::from_slice(&mask, (b, 1, tgt, tgt + offset), &self.device)? // needs b * tgt * (tgt + offset)
```
`Tensor::from_slice` does not validate the length for fully-concrete shapes (`ShapeWithOneHole` for `Into` ignores the data length), so for `b > 1` this creates a tensor whose storage is smaller than its shape. On CUDA the attention kernels then read out of bounds, and forward passes return garbage or NaN for most rows of the batch.
## Symptoms
Batched inference (e.g. embedding extraction with equal-length sequences, batch 64, no padding) produces:
- some rows correct (those that happen to read valid memory),
- some rows NaN,
- some rows with cosine ~0.1-0.4 against the single-batch reference.
Rows processed with batch size 1 are always correct, which hides the bug in typical single-sequence generation use.
Verified on candle 0.11.0, CUDA backend (RTX 5090, CUDA 12.8, no flash-attn feature, standard matmul attention path). The same construction is present on current `main`.
## Fix
The mask is identical for every batch row, so the cheapest fix is to shape it `(1, 1, tgt, tgt + offset)` and let `broadcast_add` in `forward_standard_attn` stretch it across the batch:
```rust
Tensor::from_slice(&mask, (1, 1, tgt, tgt + offset), &self.device)?.to_dtype(self.dtype)
```
(Alternatively repeat the data `b` times, but broadcasting avoids the copy.)
We are running with exactly this one-line change vendored and batched embeddings now match per-text embeddings at cosine ~1.0.
A secondary hardening suggestion: `Tensor::from_slice` silently accepting a short slice for a concrete shape is what let this produce garbage instead of an error — a length check there would have surfaced this immediately.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in candle-transformers/src/models/qwen3.rs at Model::causal_mask and inspect how the mask is used by forward_standard_attn. Ensure the mask storage matches its declared shape while preserving batch broadcasting, then verify batched Qwen3 inference against equivalent single-sequence results, including the reported batch-greater-than-one case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 2/5
- Estimated time
- 1-3 hours
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 85/100