vllm-project / vllm-project/aibrix
Extend AIBrix kv cache to omni scenario
- Dominant language
- Go
- Stars
- 5.1k
- Forks
- 697
- Avg merge
- 1d 19h
- Merged PRs (30d)
- 104
Description
### π Feature Description and Motivation
https://docs.vllm.ai/projects/vllm-omni/ has its own connector base
```
OmniConnectorBase (abstract base)
βββ SharedMemoryConnector β single-node, uses POSIX shared memory
βββ MooncakeConnector β multi-node, uses Mooncake distributed store
βββ YuanrongConnector β multi-node, uses Datasystem KV client
```
Data Flow (KV Cache Transfer)
```
The full flow is managed by OmniKVTransferManager (kv_transfer_manager.py):
Stage 0 (Thinker/LLM) Stage 1 (Talker/Diffusion)
βββββββββββββββββββββ ββββββββββββββββββββββββββ
1. Prefill completes
2. _extract_kv_cache()
GPU blocks β CPU tensors
3. connector.put(kv_data)
β serialize β mooncake.put()
4. connector.get(kv_data)
β mooncake.get() β deserialize
5. apply_kv_cache_to_request()
CPU tensors β GPU
```
The AIBrix connector only needs to handle storing and retrieving byte blobs by key
### Use Case
Enable the kv cache transfer for omni distributed pipeline
### Proposed Solution
_No response_
Contributor guide
Assessment
This issue has not been assessed yet.