[RFC][kv_offload]: Tiering Admission Policy Design
RFC
### Motivation.
Unconditional cascade/promotion submission causes:
1. **Tier overload** (PR #50045): cascades pile up on slow tiers.
2. **Promotion contention** (PR #50014): promotions evict each other before consumption.
Both are admission decisions — gate *before* submission.
### Proposed Change.
## Interface
```python
class TieringAdmissionPolicy(ABC):
@abstractmethod
def should_admit(self, job: JobMetadata) -> bool:
"""Pure predicate — no side effects."""
@abstractmethod
def on_admitted(self, job: JobMetadata) -> None:
"""Called after job is committed to submission."""
@abstractmethod
def on_completed(self, job: JobMetadata, result: JobResult) -> None:
"""Called when a job finishes."""
@abstractmethod
def reset(self) -> None:
"""Clear internal state (reset_cache)."""
@classmethod
def build_metric_definitions(
cls, extra_config: dict[str, Any]
) -> dict[str, OffloadingMetricMetadata]:
return {}
def get_stats(self) -> "OffloadingConnectorStats | None":
return None
```
Types from PR #48798:
- `JobMetadata` = `NamedTuple(transfer_job: TransferJob, tier_idx: int)`
- `TransferJob`: `job_id`, `keys`, `block_ids: np.ndarray | None`, `is_promotion`, `req_context`
- `JobResult`: `job_id`, `success`, `successful_keys`, `transfer_time`
`should_admit` is side-effect-free for `CompositePolicy` composability. State
mutation happens only in `on_admitted`.
## Integration
### Cascades — inside `create_store_job()`
All cascade paths (manager direct + `_SecondaryTierFacingParent`) funnel here:
```python
def create_store_job(self, keys, req_context, tier_idx) -> TransferJob | None:
probe = JobMetadata(
TransferJob(job_id=self._next_job_id(), keys=keys,
block_ids=None, is_promotion=False,
req_context=req_context),
tier_idx,
)
if not self._policy.should_admit(probe):
return None # rejected — no blocks pinned
spec = self.primary_tier.prepare_read(keys, req_context)
probe.transfer_job.block_ids = spec.block_ids
self._register_job(probe)
self._policy.on_admitted(probe)
return probe.transfer_job
```
Breaking change: `ParentManager.create_store_job()` returns `TransferJob | None`.
Downstream tiers (`p2p`, `fs`, `obj`) must handle `None`.
### Promotions — inside `_initiate_promotion`
Admission per-block at lookup time; I/O submission remains batched in
`_flush_pending_promotions()`:
```python
probe = JobMetadata(
TransferJob(job_id=self._next_job_id(), keys=[key],
block_ids=None, is_promotion=True,
req_context=req_context),
tier_idx,
)
if not self._policy.should_admit(probe):
return False # LookupResult.MISS
self._policy.on_admitted(probe)
alloc = self.primary_tier.prepare_store([key], req_context)
if alloc is None:
return False
# accumulate into _pending_load_submissions...
```
### Completion / Reset
```python
self._policy.on_completed(job_metadata, completed_job) # _process_finished_jobs
self._policy.reset() # reset_cache
```
## Factory
Same lazy-import registry pattern as `SecondaryTierFactory`:
```python
class AdmissionPolicyFactory:
_registry: dict[str, Callable[[], type[TieringAdmissionPolicy]]] = {}
@classmethod
def register_policy(cls, policy_type, module_path, class_name): ...
@classmethod
def create_policy(cls, policy_config: dict) -> TieringAdmissionPolicy:
config = policy_config.copy()
policy_cls = cls.get_policy_class(config.pop("type"))
return policy_cls(**config)
```
Config: `extra_config["admission_policy"]`, default `{"type": "always_admit"}`.
Composite config:
```json
{"type": "composite", "policies": [
{"type": "cascade_backpressure", "high_water_s": 1.0, "low_water_s": 0.5},
{"type": "promotion_concurrency", "max_inflight_promotions": 256}
]}
```
## File Layout
```
vllm/v1/kv_offload/tiering/admission/
├── factory.py # AdmissionPolicyFactory + registrations
├── base.py # TieringAdmissionPolicy ABC
├── always.py # AlwaysAdmitPolicy
├── cascade_backpressure.py # CascadeBackpressurePolicy
├── promotion_concurrency.py # PromotionConcurrencyPolicy
└── composite.py # CompositePolicy
```
## Implementations
### `CascadeBackpressurePolicy` (PR #50045)
Gates cascades to overloaded tiers. Signal: EMA of submission-to-completion
wall-clock time (captures queue wait + I/O, not just `transfer_time`).
```python
class CascadeBackpressurePolicy(TieringAdmissionPolicy):
def __init__(self, high_water_s=1.0, low_water_s=0.5, alpha=0.3,
warmup_completions=3):
self._high = high_water_s
self._low = low_water_s
self._alpha = alpha
self._warmup = warmup_completions
self._ema: dict[int, float] = {}
self._under_pressure: dict[int, bool] = {}
self._completions: dict[int, int] = {}
self._submit_time: dict[int, float] = {}
def should_admit(self, job):
if job.transfer_job.is_promotion:
return True
return not self._under_pressure.get(job.tier_idx, False)
def on_admitted(self, job):
if not job.transfer_job.is_promotion:
self._submit_time[job.transfer_job.job_id] = time.monotonic()
def on_completed(self, job, result):
if job.transfer_job.is_promotion:
return
submit_t = self._submit_time.pop(job.transfer_job.job_id, None)
if submit_t is None:
return
total_latency = time.monotonic() - submit_t
tier = job.tier_idx
ema = self._ema.get(tier, 0.0)
self._ema[tier] = self._alpha * total_latency + (1 - self._alpha) * ema
count = self._completions.get(tier, 0) + 1
self._completions[tier] = count
if count < self._warmup:
return
if self._ema[tier] > self._high:
self._under_pressure[tier] = True
elif self._ema[tier] < self._low:
self._under_pressure[tier] = False
def reset(self):
self._ema.clear()
self._under_pressure.clear()
self._completions.clear()
self._submit_time.clear()
```
### `PromotionConcurrencyPolicy` (PR #50014)
Gates promotions when in-flight count exceeds cap.
```python
class PromotionConcurrencyPolicy(TieringAdmissionPolicy):
def __init__(self, max_inflight_promotions: int = 512):
self._max = max_inflight_promotions
self._inflight: int = 0
def should_admit(self, job):
if not job.transfer_job.is_promotion:
return True
return self._inflight < self._max
def on_admitted(self, job):
if job.transfer_job.is_promotion:
self._inflight += 1
def on_completed(self, job, result):
if job.transfer_job.is_promotion:
self._inflight -= 1
def reset(self):
self._inflight = 0
```
### `CompositePolicy`
```python
class CompositePolicy(TieringAdmissionPolicy):
def __init__(self, policies: list[TieringAdmissionPolicy]):
self._policies = policies
def should_admit(self, job):
return all(p.should_admit(job) for p in self._policies)
def on_admitted(self, job):
for p in self._policies:
p.on_admitted(job)
def on_completed(self, job, result):
for p in self._policies:
p.on_completed(job, result)
def reset(self):
for p in self._policies:
p.reset()
```
## Pin Lifecycle (Manager, not Policy)
Separate from admission — introduced by PR #50014. After promotion lands
(`complete_store` sets ref_cnt 0), manager pins via `prepare_load()` (ref_cnt
+1). Released on consumption or request finish. Policy controls *arrival rate*;
manager controls *protection duration*.
## Metrics
**Manager-level** (policy-agnostic, in `TieringMetricsTracker`):
- `vllm:kv_offload_tiering_cascades_dropped{tier}`
- `vllm:kv_offload_tiering_promotions_dropped{tier}`
**Policy-specific** — same hooks as `SecondaryTierManager`:
- `build_metric_definitions(cls, extra_config) -> dict[str, OffloadingMetricMetadata]`
- `get_stats(self) -> OffloadingConnectorStats | None`
`TieringOffloadingSpec.build_metric_definitions` fans out to the policy class
(same pattern as secondary tiers).
Example (`CascadeBackpressurePolicy`):
```python
@classmethod
def build_metric_definitions(cls, extra_config):
return {
"vllm:kv_offload_tiering_backpressure_active": OffloadingGaugeMetadata(
documentation="1 when tier is under backpressure"),
"vllm:kv_offload_tiering_store_latency_ema": OffloadingGaugeMetadata(
documentation="EMA of store latency (seconds)"),
}
```
## Summary
| | `CascadeBackpressurePolicy` | `PromotionConcurrencyPolicy` |
|--|--|--|
| Gates | Cascades | Promotions |
| Signal | Wall-clock EMA with hysteresis | In-flight count vs cap |
| `on_admitted` | Record submit timestamp | Increment count |
| `on_completed` | Update EMA | Decrement count |
| Post-landing | N/A | Manager pins via ref_cnt |
### Feedback Period.
_No response_
### CC List.
@bnellnm @varun-sundar-rabindranath @ronensc @AlejandroParedesLT
### Any Other Things.
_No response_
### Before submitting a new issue...
- [x] Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the [documentation page](https://docs.vllm.ai/en/latest/), which can answer lots of frequently asked questions.
4 条评论