vllm-project / vllm-project/aibrix

[RFC]: vLLM v0.19 patch + dynamic packaging refactor

Open
#2,104 3 comments 0 reactions 1 assignee Claimed by @DwyaneShi View on GitHub
Dominant language
Go
Stars
5.1k
Forks
694
Avg merge
1d 20h
Merged PRs (30d)
98

Description

### Summary

We propose two coordinated changes to the AIBrix kvcache integration:

1. **Add a patch for vLLM v0.19.0**, following the existing pattern (`vllm_v0.8.5`, `v0.9.1`, `v0.10.2`, `v0.14.0`).
2. **Refactor the patch packaging strategy** so that everything that does not strictly require touching vLLM core files moves into the `aibrix_kvcache` Python package and registers itself dynamically via vLLM's public APIs (`KVConnectorFactory.register_connector`).

The combined change reduces the per-version patch from ~4,057 lines to ~125 lines (a ~97% reduction) and makes future vLLM upgrades dramatically cheaper to support.

We are willing to do the work and submit the PR. We would like maintainer feedback on direction before starting:

- Is anyone already working on a v0.19 patch?
- Is the dynamic packaging refactor aligned with the project's direction?
- Should both land in the same PR or be split?

cc/ @DwyaneShi @Jeffwan

### Motivation

### Why v0.19

- Several upstream KV-cache compression and quantization PRs (e.g. vllm-project/vllm#38479) target vLLM `main` (0.19.x line). Without an AIBrix patch for that line, AIBrix users cannot prototype these features.
- Adoption of newer vLLM releases is accelerating; staying behind by 5 minor versions creates drift that gets harder to close over time.
- We need v0.19 internally to evaluate next-generation models (Qwen3.5 family, hybrid Mamba/attention) on top of AIBrix offloading. We will do the port either way and would prefer to contribute it upstream.

### Why also propose a refactor

We analyzed the v0.14 patch (`vllm_v0.14.0-aibrix-kvcache.patch`, 4,057 lines, 9 modified files). Of those 9 files, only **3 actually require touching vLLM core code** (the hooks in the hot path inside `gpu_model_runner.py`, `kv_connector_model_runner_mixin.py`, and `base.py`). The other 6 can be moved into the `aibrix_kvcache` Python package using vLLM's existing public APIs.

This is the same model that LMCache, Mooncake and NixlConnector already follow — they live as importable modules inside the vLLM tree (or external packages) and register themselves at import time.

| Component in current v0.14 patch | Lines today | Lines after refactor | How |
|---|---|---|---|
| Type1 connector (new file) | ~1,400 | 0 in patch | Move to `aibrix_kvcache.connectors.type1` |
| Type2 connector (new file) | ~765 | 0 in patch | Move to `aibrix_kvcache.connectors.type2` |
| PDReuse connector (new file) | ~1,657 | 0 in patch | Move to `aibrix_kvcache.connectors.pd_reuse` |
| `factory.py` (registration) | ~22 | 0 | `KVConnectorFactory.register_connector()` at package import |
| `envs.py` (AIBRIX_* vars) | ~24 | 0 | Read with `os.environ` from the package |
| `kv_transfer_metrics.py` | ~50 | 0 | Implement inside the AIBrix package |
| `base.py` (handle_preemptions sig) | ~25 | ~25 | Stays as patch (base API change) |
| `gpu_model_runner.py` (new hook in hot path) | ~90 | ~90 | Stays as patch |
| `kv_connector_model_runner_mixin.py` | ~10 | ~10 | Stays as patch |
| **TOTAL** | **~4,057** | **~125** | **97% reduction** |

### Concrete maintenance benefit

Today, each new vLLM release requires re-resolving conflicts in 4,000+ lines of patch. With the refactor, only ~125 lines of hot-path patches need re-locating. New vLLM releases become trivial to support — and the `aibrix_kvcache` package becomes mostly version-independent.

This also unblocks evaluation of newer models (Qwen3.5) and KV compression PRs for users on the v0.19 line, but the benefit is general across the AIBrix user base.

### Proposed Change

### Proposed package structure

```text
aibrix_kvcache/
├── __init__.py # auto-registers the 3 connectors via factory
├── envs.py # AIBRIX_* env vars (os.environ)
├── metrics.py # was kv_transfer_metrics.py inside vLLM
├── connectors/
│ ├── __init__.py
│ ├── type1.py # was aibrix_offloading_connector_type1.py
│ ├── type2.py # was aibrix_offloading_connector_type2.py
│ └── pd_reuse.py # was aibrix_pd_reuse_connector.py
└── integration/
└── vllm/
└── patches/
└── vllm_v0.19.0-aibrix-kvcache.patch # ~125 lines, hot-path hooks only
```

### Auto-registration at import

```python
# aibrix_kvcache/__init__.py
from vllm.distributed.kv_transfer.kv_connector.factory import (
KVConnectorFactory,
)

for name, module, cls in [
("AIBrixOffloadingConnectorV1Type1",
"aibrix_kvcache.connectors.type1", "Type1Connector"),
("AIBrixOffloadingConnectorV1Type2",
"aibrix_kvcache.connectors.type2", "Type2Connector"),
("AIBrixPDReuseConnector",
"aibrix_kvcache.connectors.pd_reuse", "PDReuseConnector"),
]:
KVConnectorFactory.register_connector(name, module, cls)
```

User experience becomes simply: `pip install aibrix-kvcache` plus the (much smaller) standard AIBrix patch.

### Concrete findings from our delta exploration (vLLM 0.14 → 0.19)

We've cloned both repos and validated the deltas. Sharing the data so the discussion is grounded:

- `factory.py` in v0.19: 58-line diff vs v0.14, mostly registers new connectors (FlexKV, SimpleCPUOffload, Mooncake reorganized). Trivial to port.
- `kv_connector/v1/base.py`: 178-line diff. Real changes: new `KVConnectorWorkerMetadata` abstract class, `handle_preemptions` signature change (`set[str]` → `KVConnectorMetadata`), new `build_connector_worker_meta()`. Moderate adaptation needed in our 3 connectors.
- `kv_transfer_metrics.py`: **removed** in v0.19, merged into `kv_connector/v1/metrics.py`. Patch chunk needs reubication or removal.
- `gpu_model_runner.py`: 4,209-line diff with 154 hunks between v0.14 and v0.19. Our patch only modifies 5 small hunks (around lines 834, 887, 950, 982, 3134 in v0.14). Main effort is locating the equivalent positions in the refactored v0.19 file.
- `kv_connector_model_runner_mixin.py`: 65-line diff. Several methods removed (`maybe_setup_kv_connector`, `maybe_wait_for_kv_save`, `get_finished_kv_transfers`), new `defer_finalize` parameter. Manageable.
- Attention backends (`flash_attn.py`, `flashinfer.py`) **not touched** by AIBrix patch — and they moved from `vllm/attention/backends/` to `vllm/v1/attention/backends/`. Doesn't affect us.

### Backward compatibility

Existing patches (v0.8.5, v0.9.1, v0.10.2, v0.14.0) **are not modified**. They keep their inline connectors. Only the new v0.19.0 patch uses the externalized package layout.

This avoids regression risk on supported versions and lets us validate the new packaging in production before discussing migration of older patches.

### Implementation plan

If maintainers approve, we propose to land the work in this order:

1. **Refactor first.** Move the 3 connectors + envs + metrics into the
`aibrix_kvcache` package. Existing patches (v0.8.5–v0.14) keep working
unchanged. This lands as a standalone PR so it can be reviewed in isolation.

2. **Add the v0.19 patch.** Once the refactor is merged, add
`vllm_v0.19.0-aibrix-kvcache.patch` (~125 lines, hot-path hooks only),
adapted to the new `handle_preemptions` signature.

3. **CI + validation.** Extend the CI matrix to build images with
`VLLM_VERSION=v0.19.0`. We will validate against Qwen3.5 models on L4
GPUs (AWQ + hybrid Mamba/attention coverage) and benchmark against the
v0.14 baseline before requesting final review.

We are happy to land all three steps as separate PRs or squashed into one,
whichever maintainers prefer.

### Note on positioning vs `SimpleCPUOffloadConnector`

We noticed vLLM v0.19 ships a native `SimpleCPUOffloadConnector` and an `offloading/` framework (~1,300 lines). This partially overlaps with AIBrix L1 use cases. AIBrix still differentiates with L2 (remote KVCache cluster), PDReuse, and the broader control-plane integration — but we'd appreciate the team's view on positioning before we publish the PR description.

### Questions for maintainers

1. Is anyone currently working on a v0.19 patch? If so, we'll happily contribute as reviewers/co-authors instead.
2. Does the dynamic packaging refactor align with the project direction?
3. Should the refactor land in the same PR as v0.19, or split into two PRs?
4. Any concerns about the package auto-registering connectors at import time? We can make registration explicit if preferred.
5. Should we coordinate with the maintainers behind `SimpleCPUOffloadConnector` to clarify positioning?

### Alternatives Considered

### Alternative 1: Static patch for v0.19, no refactor

The simplest path: copy `vllm_v0.14.0-aibrix-kvcache.patch` → `vllm_v0.19.0-aibrix-kvcache.patch`, resolve conflicts, ship it.

**Why not chosen:** keeps the per-version maintenance cost identical to today (~4,000 lines of patch to re-resolve every release). Misses the opportunity to fix the underlying packaging issue while we already have the relevant context. We would still propose this as a fallback if maintainers reject the refactor.

### Alternative 2: Skip v0.19, wait for the next stable

We could wait for vLLM v0.20 or v0.21 hoping the API stabilizes.

**Why not chosen:** there is no signal that the relevant APIs (`base.py`, `gpu_model_runner.py`) are about to stabilize. Each vLLM release brings churn in the same files. Waiting just delays the inevitable.

### Alternative 3: Monkey-patch at runtime instead of `.patch` files

Replace `GPUModelRunner.execute_model` at import time from the AIBrix package, eliminating the need for any `.patch` file.

**Why not chosen:** this is technically possible but extremely fragile. Each vLLM release modifies `execute_model` substantially (4,209-line diff between v0.14 and v0.19). A monkey-patch would silently break or, worse, run with stale logic. Static patches at least fail loudly when the target file changes.

### Alternative 4: Push all hot-path hooks upstream into vLLM first

The cleanest long-term solution is to convince vLLM to expose a public hook (e.g. `KVConnectorBase_V1.before_update_states`) that AIBrix and similar projects can override. If accepted, the v0.19 patch would shrink from ~125 lines to ~0.

**Why not chosen as the primary path:** vLLM upstream changes take weeks-to-months of design discussion. We don't want to block AIBrix's v0.19 support on that timeline. We will pursue this in parallel as a separate RFC to vLLM (companion to this proposal). If/when it lands, the AIBrix patch can be slimmed further in a follow-up.

### Alternative 5: Use Python entry_points for connector registration

Instead of explicit `register_connector()` calls inside `aibrix_kvcache/__init__.py`, we could declare entry points in `setup.py` so vLLM auto-discovers the connectors.

**Why not chosen for the first iteration:** entry-point auto-discovery has subtle import-order and packaging gotchas, and vLLM does not currently use this pattern for its own connectors. Starting with explicit registration is simpler, more debuggable, and consistent with how LMCache and Mooncake are wired today. We could migrate to entry_points in a follow-up if maintainers prefer.

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.