vllm-project / vllm-project/aibrix

Extend AIBrix kv cache to omni scenario

Open
#1,955 0 comments 0 reactions 1 assignee Claimed by @DwyaneShi View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.