huggingface / huggingface/candle
Optimized `cumsum` operation
- Dominant language
- Rust
- Stars
- 21k
- Forks
- 1.8k
- Avg merge
- 16h 42m
- Merged PRs (30d)
- 25
Description
Currently the `cumsum` operation internally uses a matmul! This can be incredibly slow unexpectedly, and of course use a lot of memory. I propose adopting one of the MLX kernels for this, namely based on a [`scan` operation](https://github.com/ml-explore/mlx/blob/0cae0bdac83bbf5b3d1da3ca53f1f7eb95981d30/mlx/backend/metal/kernels/scan.h#L236). When using this for top-p sampling, we can really remove the necessity to do sampling on the CPU and instead focus on the GPU. I measure a **15% performance increase** by avoiding sampling and top-p/top-k on the GPU!
Implementing this would involve porting the Metal kernels, implementing it for the CPU, and CUDA kernels.
@LaurentMazare do you think this would be interesting for a PR?
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by locating Candle's cumsum operation and the CPU, Metal, and CUDA backend entry points. Compare the relevant implementation with the referenced MLX scan kernel, then verify that cumsum uses the new backend kernels and that top-p sampling can remain on the GPU without the current matmul overhead.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- machine-learning, performance
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 25/100