huggingface / huggingface/candle

Optimized `cumsum` operation

Open
#2,948 1 comment 1 reaction 0 assignees View on GitHub
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.