michaelfeil / michaelfeil/infinity

Move `.detach().cpu()` into `encode_core`, and option to use cuda streams

Open
#155 5 comments 0 reactions 0 assignees View on GitHub
help wanted
Dominant language
Python
Stars
2.9k
Forks
206
PR merge metrics
No merged PRs in 30d

Description

In `encode_post` of `SentenceTransformerPatched` we have
```
embeddings = out_features.detach().cpu().to(torch.float32)
```
On GPU if I'm understanding correctly, `.cpu()` triggers a device to host synchronization which will have to wait for all the computation to finish, and may be blocking for quite some time.

In batch handler, however, `encode_post` runs in async function `_postprocess_batch`, which isn't a good place for blocking code.

Is there any reason for this implementation?

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by inspecting encode_post in SentenceTransformerPatched, encode_core, and the async _postprocess_batch path mentioned in the issue. Check where detach().cpu().to(torch.float32) runs and how CUDA stream use would affect batching. Done requires an agreed design and verification that embedding results remain correct without blocking the async postprocessing path.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
backend, performance
Issue type
Refactor
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.