michaelfeil / michaelfeil/infinity
Move `.detach().cpu()` into `encode_core`, and option to use cuda streams
- 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