Implementing Custom KV Cache Eviction Policies in vLLM for Multi-Tenant Chatbots — code snippets

snippet-1.py

from fastapi import Request, HTTPException
from typing import Dict, Any
import jwt

# Secret key used for decoding JWT headers passed by the API Gateway
JWT_SECRET = "production_super_secret_key_rotation_99"

class TenantContextMiddleware:
    async def __call__(self, request: Request, call_next):
        # Extract JWT from Authorization header
        auth_header = request.headers.get("Authorization")
        if not auth_header or not auth_header.startswith("Bearer "):
            # Default to free tier for unauthenticated internal diagnostic requests
            request.state.tenant_id = "anonymous_free"
            request.state.priority = 2  # Free Tier
            request.state.lease_seconds = 0
            return await call_next(request)

        try:
            token = auth_header.split(" ")[1]
            payload = jwt.decode(token, JWT_SECRET, algorithms=["HS256"])
            
            # Map claims to scheduler metadata
            request.state.tenant_id = payload.get("tenant_id", "anonymous_free")
            tier = payload.get("tier", "free")
            
            if tier == "premium":
                request.state.priority = 0
                request.state.lease_seconds = 300
            elif tier == "standard":
                request.state.priority = 1
                request.state.lease_seconds = 60
            else:
                request.state.priority = 2
                request.state.lease_seconds = 0
                
        except jwt.PyJWTError:
            raise HTTPException(status_code=401, detail="Invalid tenant credentials")

        return await call_next(request)

snippet-2.py

import time
from typing import List, Optional
from dataclasses import dataclass

@dataclass
class TenantMetadata:
    tenant_id: str
    priority: int
    lease_seconds: int
    last_active_time: float

class CustomSequenceGroupMetadata:
    def __init__(self, request_id: str, tenant_metadata: TenantMetadata):
        self.request_id = request_id
        self.tenant = tenant_metadata
        self.allocated_blocks: List[int] = []

    def update_activity(self):
        self.tenant.last_active_time = time.time()

    def is_lease_active(self) -> bool:
        elapsed = time.time() - self.tenant.last_active_time
        return elapsed < self.tenant.lease_seconds

snippet-3.py

import collections
from typing import Dict

class TenantCacheTracker:
    def __init__(self):
        # Map tenant_id -> set of block IDs allocated
        self.tenant_to_blocks: Dict[str, set] = collections.defaultdict(set)
        # Map request_id -> CustomSequenceGroupMetadata
        self.request_metadata: Dict[str, CustomSequenceGroupMetadata] = {}

    def register_request(self, request_id: str, metadata: CustomSequenceGroupMetadata):
        self.request_metadata[request_id] = metadata

    def deregister_request(self, request_id: str):
        if request_id in self.request_metadata:
            del self.request_metadata[request_id]

    def record_block_allocation(self, tenant_id: str, block_id: int):
        self.tenant_to_blocks[tenant_id].add(block_id)

    def record_block_deallocation(self, tenant_id: str, block_id: int):
        self.tenant_to_blocks[tenant_id].discard(block_id)

    def get_tenant_block_count(self, tenant_id: str) -> int:
        return len(self.tenant_to_blocks[tenant_id])

snippet-4.py

from typing import List, Tuple
from vllm.core.scheduler import Scheduler
from vllm.sequence import SequenceGroup, SequenceStatus

def select_preemption_candidates(
    running_seq_groups: List[SequenceGroup],
    tracker: TenantCacheTracker
) -> List[SequenceGroup]:
    """
    Evaluates running sequences and returns them sorted by suitability for preemption.
    The sequence at the beginning of the returned list will be preempted first.
    """
    candidates: List[Tuple[int, bool, float, SequenceGroup]] = []
    
    for seq_group in running_seq_groups:
        req_id = seq_group.request_id
        meta = tracker.request_metadata.get(req_id)
        
        if not meta:
            # Fallback for unclassified system traffic: prioritize preemption
            candidates.append((3, False, 0.0, seq_group))
            continue
            
        priority = meta.tenant.priority
        lease_active = meta.is_lease_active()
        last_active = meta.tenant.last_active_time
        
        # Sort key construction:
        # 1. Non-leased sequences are preempted before leased sequences.
        # 2. Higher priority value (Free = 2) is preempted before lower priority value (Premium = 0).
        # 3. Older active sessions are preempted first (LRU behavior within the tier).
        # We use a tuple for lexical sorting.
        
        lease_sort_key = 0 if not lease_active else 1
        priority_sort_key = -priority  # We want higher numbers (lower priority) first
        
        candidates.append((lease_sort_key, priority_sort_key, last_active, seq_group))

    # Sort based on criteria: non-leased first, lowest priority tier first, oldest first
    candidates.sort(key=lambda x: (x[0], x[1], x[2]))
    
    return [item[3] for item in candidates]

snippet-5.py

class TenantAwareScheduler(Scheduler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.tenant_tracker = TenantCacheTracker()

    def _schedule(self) -> Tuple[List[SequenceGroup], List[SequenceGroup], List[SequenceGroup]]:
        """
        Custom scheduling cycle. Overrides default preemption sorting before delegating
        to the allocation state machine.
        """
        # Obtain default active queues
        running = list(self.running)
        
        # Check if we need to preempt due to physical block starvation
        # If block space allocator cannot accommodate run queue, sort running queue
        # by preemption preference
        if self.block_manager.get_num_free_gpu_blocks() < self.cache_config.block_size:
            sorted_candidates = select_preemption_candidates(running, self.tenant_tracker)
            
            # Reconstruct the scheduler queue with preferred preemption targets placed at the head
            # so that parent class helper `_preempt` picks them first
            self.running.clear()
            for seq_group in sorted_candidates:
                self.running.append(seq_group)
                
        # Delegate core allocation execution to the optimized backend
        return super()._schedule()

snippet-6.py

from prometheus_client import Counter, Gauge

# Telemetry metrics for multi-tenant cache behavior
KV_CACHE_EVICTION_COUNTER = Counter(
    "vllm_tenant_kv_cache_evictions_total",
    "Number of KV cache blocks evicted by custom policy",
    labelnames=["tenant_id", "priority", "preemption_reason"]
)

KV_CACHE_SWAP_GAUGE = Gauge(
    "vllm_tenant_active_swapped_blocks",
    "Current number of memory blocks swapped to CPU RAM per tenant",
    labelnames=["tenant_id"]
)

def record_eviction_telemetry(seq_group: SequenceGroup, reason: str, block_count: int):
    # Retrieve request metadata
    req_id = seq_group.request_id
    # Record to Prometheus counters
    KV_CACHE_EVICTION_COUNTER.labels(
        tenant_id=getattr(seq_group, "tenant_id", "unknown"),
        priority=getattr(seq_group, "priority", 2),
        preemption_reason=reason
    ).inc(block_count)

snippet-8.py

def determine_preemption_mode(allocated_blocks: int) -> str:
    # 1 block = 16 tokens. 128 blocks = 2048 tokens.
    # Below 2048 tokens, it is computationally cheaper to recalculate than to swap.
    if allocated_blocks < 128:
        return "recompute"
    return "swap"