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

Open Beginner friendly
#3,764 0 comments 0 reactions 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.