huggingface / huggingface/candle

Tensor::to_scalar very high latency

Open
#2,239 6 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 description:
When converting a cuda tensor, which is indexed from an other large tensor to CPU memory scalar, the time consumption of this process is unusually long.

Minimal reproducible example:
```rust
use candle_core::{Device, IndexOp, Tensor};
use std::time::Instant;

fn time_it(func: impl FnOnce() -> T) -> T {
let start = Instant::now();
let result = func();
let duration = start.elapsed();
println!("{:?}", duration);
result
}

fn main() {
let cuda_dev = Device::new_cuda(0).unwrap();
let cpu_dev = Device::Cpu;

let shape = [10000, 10000];
let tensor = Tensor::rand(0f32, 1f32, &shape, &cpu_dev).unwrap();
println!("Tensor on CPU: {:?}", tensor);

let tensor = time_it(|| tensor.to_device(&cuda_dev).unwrap());
println!("Tensor on GPU: {:?}", tensor.layout());

let cell = time_it(|| tensor.i((42, 42)).unwrap());

let scalar = time_it(|| cell.to_scalar::().unwrap());
dbg!(scalar);
}
```

In the most recent crate version v0.5.1, the line `let scalar = time_it(|| cell.to_scalar::().unwrap());` will have the identical time consumption of line `let tensor = time_it(|| tensor.to_device(&cuda_dev).unwrap());`.

In my investigation, the problem is because candle's tensor uses a storage + layout approach to represent slice, this approach in most of the time is effective, but it's inefficient when copying a small slice from cuda mem to CPU mem because it will still copy the full storage and use layout to offset to the slice you want. this creates massive latency(Before my attempt to fix ~300ms, after my attempt to fix ~17us).

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by running the Rust minimal reproducible example in the issue and inspect Candle's tensor storage and layout handling for CUDA-to-CPU scalar conversion. The issue's stated done condition is that converting an indexed scalar should avoid copying the full storage and take substantially less time than the reported roughly 300 ms.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
backend, performance
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.