huggingface / huggingface/candle
conv1d with large groups (depthwise) is extremely slow on CUDA — O(groups) kernel launches
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
## Summary
`Tensor::conv1d()` with `groups > 1` decomposes the operation into `groups` individual `conv1d_single_group` calls ([conv.rs L196-203](https://github.com/huggingface/candle/blob/main/candle-core/src/conv.rs#L196-L203)):
```rust
let blocks = self.chunk(groups, 1)?; // groups chunks
let kernel = kernel.chunk(groups, 0)?; // groups chunks
let blocks = blocks.iter().zip(&kernel)
.map(|(block, kernel)| block.conv1d_single_group(kernel, ¶ms))
.collect::>>()?;
Tensor::cat(&blocks, 1)
```
For **depthwise convolution** (`groups = channels`), this creates **3 × groups CUDA kernel launches** (chunk + conv + cat). With `groups=2048` and `kernel_size=3`, that's ~6,000 kernel launches per call — each one trivially small but dominated by launch overhead.
## Impact — real-world model
Running [MioTTS LFM2](https://huggingface.co/Aratako/MioTTS-2.6B) (30-layer model, 22 conv layers, `hidden_size=2048`, `kernel_size=3`) on an RTX 5090:
| | candle `conv1d` | manual workaround |
|---|---|---|
| Per conv layer (prefill, 19 tokens) | **33 ms** | **< 0.1 ms** |
| 22 conv layers total | **726 ms** | **~2 ms** |
| Prefill (total) | **740 ms** | **19 ms** |
The conv layers were **54% of total inference time**, entirely due to kernel launch overhead.
## Workaround
For depthwise conv1d with small kernels, expanding into `narrow` + `broadcast_mul` + `add` is trivially equivalent and ~360× faster:
```rust
// input: (B, D, T+K-1) padded, weight: (D, K) squeezed
// Equivalent to conv1d(input, weight.unsqueeze(1), groups=D)
let mut out = input.narrow(2, 0, t)?
.broadcast_mul(&weight.narrow(1, 0, 1)?.unsqueeze(0)?)?;
for k in 1..kernel_size {
out = (out + input.narrow(2, k, t)?
.broadcast_mul(&weight.narrow(1, k, 1)?.unsqueeze(0)?)?)?;
}
// Total: K narrow + K mul + (K-1) add = 8 kernels for K=3
```
## Suggested fix
For the general case, rather than looping over groups, the grouped convolution could be implemented as a single batched operation in the CUDA/cuDNN backend. Specifically for depthwise conv (`groups == c_in && c_in_k == 1`), cuDNN's grouped convolution handles this natively in a single kernel call.
Alternatively, a special-case path for depthwise conv1d (similar to the workaround above) would eliminate the overhead without needing cuDNN changes.
## Environment
- candle-core 0.9.2
- CUDA 13.1 / RTX 5090
- The same issue applies to `conv2d` with large groups ([conv.rs L329-340](https://github.com/huggingface/candle/blob/main/candle-core/src/conv.rs#L329-L340))
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.