huggingface / huggingface/candle
broadcast_matmul to handle stride-0 broadcast dimensions
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
While implementing CPU Grouped-Query Attention, I encountered a limitation with `broadcast_matmul` and stride-0 dimensions. I'm taking a different approach (fused SIMD kernels) for my use case, but wanted to log this for others who may benefit from the fix.
Currently `broadcast_as` creates stride-0 view and `matmul` rejects stride-0 as "non-contiguous". So, I needed to use the workaround: `.broadcast_as(...).contiguous()` but this physically expands memory.
```rust
// GQA: Q has 16 heads, K/V have 8 heads (2 groups)
let q = ...; // [1, 8, 2, 2, 128]
let k = ...; // [1, 8, 1, 128, 2] ← size-1 dim should broadcast
let scores = q.broadcast_matmul(&k)?; // Error: "non-contiguous lhs"
```
I expected it to broadcast dim 2 of rhs (1 → 2) during matmul
But instead I got an error after internal broadcast creates stride-0
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by tracing the broadcast_matmul, broadcast_as, and matmul entry points to see where stride-0 dimensions become rejected as non-contiguous. Reproduce the shown grouped-query attention shape, then add coverage for broadcasting a size-1 rhs dimension and verify that broadcast_matmul works without requiring contiguous memory expansion.
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
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100