Scheduler-Cache Interface Design

Introduction

In the previous lesson, you successfully built the core inference loop, integrating your StallFreeScheduler with a model to create a continuous batching engine. This was a major step, but our scheduler has a critical blind spot: it plans batches based on token budgets and batch sizes, completely unaware of the most precious resource in LLM inference—GPU memory for the KV cache. This can lead to the scheduler creating a plan that the system can't execute, resulting in out-of-memory (OOM) errors.

This lesson addresses that problem directly. Your goal is to design the data structures and the interface for communication between the scheduler and a paged KV cache manager. We will define the formal "contract" that allows the scheduler to become memory-aware. Instead of blindly scheduling requests, it will learn to query for available memory, request block allocations, and release memory when it's no longer needed. This interface is the control plane for memory management in a high-throughput inference engine.

1. The Role of the KV Cache Manager

Before we can design the interface, let's solidify our understanding of what the KV cache manager is responsible for. In a system with PagedAttention, the KV cache for a single request isn't a monolithic, contiguous block of memory. It's a collection of fixed-size physical blocks, and a block table maps the logical sequence of tokens to these potentially non-contiguous physical blocks.

The KVCacheManager is the component that orchestrates this entire system. To see how it works conceptually, let's revisit a key segment from the vLLM presentation.

The KVCacheManager's primary responsibilities are:

• Managing the global pool of all physical KV cache blocks.
• Maintaining a list of free blocks.
• Handling requests from the scheduler to allocate blocks for new tokens.
• Handling commands to free blocks from completed requests.
• Managing sharing of blocks through reference counting (e.g., for prefix caching or parallel sampling).

Our task is to define the API that allows the scheduler to access these functions safely and efficiently.

2. Designing the Core Data Structures and Interface

A well-designed interface hides implementation details. The scheduler shouldn't need to know about the internal workings of the block pool or reference counting. It should operate on abstract handles and clear commands. Production systems like vLLM achieve this with carefully designed data structures and class APIs.

Let's examine the vLLM source documentation, which serves as an excellent blueprint for our design.

From this documentation, we can derive the key components of our design:

Data Structures

• Request: The object representing a user's request, containing token IDs, sampling parameters, and state (e.g., number of tokens processed). This is what the scheduler manages.
• KVCacheBlocks: An abstract container. The scheduler receives this from the KVCacheManager after a successful allocation. Its primary purpose is to be passed to the model executor, which uses it to find the physical memory locations for the PagedAttention kernel.
• BlockTable: A mapping (e.g., a list or dictionary) for each request, holding the IDs of the physical blocks that store its KV cache. This is managed inside the KVCacheManager but conceptually associated with each request.

The following diagrams provide a visual mental model for these structures.

Interface (API)

The KVCacheManager class should expose the following core methods to the Scheduler:

• get_num_free_blocks() -> int: Allows the scheduler to query the amount of available memory before making a decision.
• allocate_slots(request: Request, num_tokens_to_alloc: int) -> Optional[KVCacheBlocks]: The main allocation function. The scheduler asks for space for a certain number of tokens for a specific request. It returns a KVCacheBlocks handle on success or None on failure (insufficient memory).
• free(request: Request) -> None: When a request is finished, the scheduler calls this to return all associated blocks to the free pool.
• (Optional) can_allocate(request: Request, num_tokens_to_alloc: int) -> bool: A helper that can be more efficient than get_num_free_blocks by calculating the exact block requirement for a request.

To see how these pieces fit together in a dynamic system, let's turn to a high-level walkthrough.

3. Implementation Exercise: Creating a Memory-Aware Scheduler

Now it's time to put this design into practice. You will modify the StallFreeScheduler from the last lesson to make it memory-aware. I will provide a mock KVCacheManager that implements the interface we've just designed. Your task is to update the schedule_next_batch method to use this manager, ensuring no allocation happens without first securing the required memory.

Your Task:
Modify the StallFreeScheduler class in the code below. The key areas to change are marked with # TODO.

1. In the decode-scheduling loop, you must check if you can_allocate 1 token and then allocate_slots for it.
2. In the prefill-scheduling loops (for both continuing and new requests), you must do the same: check if you can_allocate the chunk and then allocate_slots for it.
3. When a request is completed and removed, you must call kv_cache_manager.free() to release its resources.

import torch
from collections import deque
from dataclasses import dataclass, field
from typing import List, Deque, Dict, Optional




# --- Data Structures (from previous lesson) ---

@dataclass
class Request:
    id: int
    prompt_tokens: List[int]
    max_new_tokens: int
    tokens_processed: int = 0
    output_tokens: List[int] = field(default_factory=list)
    is_finished: bool = False

    def is_prefill_done(self) -> bool:
        return self.tokens_processed >= len(self.prompt_tokens)

    def is_complete(self) -> bool:
        return self.is_finished or (self.is_prefill_done() and len(self.output_tokens) >= self.max_new_tokens)

@dataclass
class ScheduledSequence:
    request_id: int
    tokens: List[int]
    is_decode: bool
    position_ids_start: int




# --- Mock KVCacheManager (implements our designed interface) ---

class MockKVCacheManager:
    """A mock KV cache manager to simulate memory allocation."""
    def __init__(self, num_blocks: int, block_size: int):
        self.num_total_blocks = num_blocks
        self.block_size = block_size
        self.free_blocks = num_blocks
        self.request_blocks: Dict[int, int] = {}  # req_id -> num_blocks

    def _get_num_required_blocks(self, num_tokens: int) -> int:
        return (num_tokens + self.block_size - 1) // self.block_size

    def can_allocate(self, request: Request, num_tokens_to_alloc: int) -> bool:



        # Check if this request is new or has existing blocks
        current_blocks = self.request_blocks.get(request.id, 0)
        current_tokens = current_blocks * self.block_size
        



        # Calculate tokens that are already covered by existing blocks
        # This is a simplification; a real manager would be more precise
        tokens_in_last_block = (current_tokens - request.tokens_processed) if request.tokens_processed < current_tokens else 0
        
        effective_tokens_to_alloc = max(0, num_tokens_to_alloc - tokens_in_last_block)
        
        required_new_blocks = self._get_num_required_blocks(effective_tokens_to_alloc)
        return self.free_blocks >= required_new_blocks

    def allocate_slots(self, request: Request, num_tokens_to_alloc: int) -> bool:
        if not self.can_allocate(request, num_tokens_to_alloc):
            return False
            
        current_blocks = self.request_blocks.get(request.id, 0)
        current_tokens = current_blocks * self.block_size
        tokens_in_last_block = (current_tokens - request.tokens_processed) if request.tokens_processed < current_tokens else 0
        effective_tokens_to_alloc = max(0, num_tokens_to_alloc - tokens_in_last_block)
        
        required_new_blocks = self._get_num_required_blocks(effective_tokens_to_alloc)
        
        self.free_blocks -= required_new_blocks
        self.request_blocks[request.id] = self.request_blocks.get(request.id, 0) + required_new_blocks
        print(f"    [Cache] Allocated {required_new_blocks} blocks for Req {request.id}. Free blocks: {self.free_blocks}")
        return True # In a real system, would return KVCacheBlocks handle

    def free(self, request: Request):
        if request.id in self.request_blocks:
            num_freed_blocks = self.request_blocks.pop(request.id)
            self.free_blocks += num_freed_blocks
            print(f"    [Cache] Freed {num_freed_blocks} blocks from Req {request.id}. Free blocks: {self.free_blocks}")
            
    def get_num_free_blocks(self) -> int:
        return self.free_blocks




# --- Your Task: Modify the Scheduler ---

class StallFreeScheduler:
    def __init__(self, token_budget: int, max_batch_size: int, kv_cache_manager: MockKVCacheManager):
        self.token_budget = token_budget
        self.max_batch_size = max_batch_size
        self.waiting_queue: Deque[Request] = deque()
        self.active_pool: Dict[int, Request] = {}



        # The scheduler now holds a reference to the cache manager
        self.kv_cache_manager = kv_cache_manager

    def add_request(self, request: Request):
        self.waiting_queue.append(request)
        
    def schedule_next_batch(self) -> List[ScheduledSequence]:
        next_batch: List[ScheduledSequence] = []
        current_tokens = 0




        # 1. Evict completed requests and FREE their memory
        completed_ids = [req_id for req_id, req in self.active_pool.items() if req.is_complete()]
        for req_id in completed_ids:



            # TODO: Free the KV cache for the completed request.
            self.kv_cache_manager.free(self.active_pool[req_id])
            del self.active_pool[req_id]
        



        # 2. Prioritize ongoing decodes
        for req in list(self.active_pool.values()): # Use list to allow modification
            if req.is_prefill_done() and not req.is_complete():



                # TODO: Check if we can allocate memory for 1 decode token.
                can_alloc = self.kv_cache_manager.can_allocate(req, 1)
                
                if current_tokens + 1 <= self.token_budget and len(next_batch) < self.max_batch_size and can_alloc:



                    # TODO: If we can, formally allocate the slot.
                    if self.kv_cache_manager.allocate_slots(req, 1):
                        pos_start = len(req.prompt_tokens) + len(req.output_tokens)



                        # The last token to use for decode is the one just generated.
                        # If no output tokens yet, it means prefill just finished. We need a dummy token.
                        last_token = req.output_tokens[-1] if req.output_tokens else -1 # Placeholder
                        next_batch.append(ScheduledSequence(request_id=req.id, tokens=[last_token], is_decode=True, position_ids_start=pos_start))
                        current_tokens += 1
                



        # 3. Handle partially completed prefills
        for req in list(self.active_pool.values()):
            if not req.is_prefill_done():
                remaining_prompt = len(req.prompt_tokens) - req.tokens_processed
                chunk_size = min(remaining_prompt, self.token_budget - current_tokens)
                



                # TODO: Check if we can allocate memory for the prefill chunk.
                can_alloc = self.kv_cache_manager.can_allocate(req, chunk_size)

                if chunk_size > 0 and len(next_batch) < self.max_batch_size and can_alloc:



                    # TODO: If we can, formally allocate the slots.
                    if self.kv_cache_manager.allocate_slots(req, chunk_size):
                        start = req.tokens_processed
                        end = req.tokens_processed + chunk_size
                        next_batch.append(ScheduledSequence(request_id=req.id, tokens=req.prompt_tokens[start:end], is_decode=False, position_ids_start=start))
                        current_tokens += chunk_size
                        req.tokens_processed += chunk_size
                break # Only one prefill at a time




        # 4. Admit new requests
        # Use a copy of the queue to iterate since we might modify the original
        for req in list(self.waiting_queue):
            if len(self.active_pool) >= self.max_batch_size or current_tokens >= self.token_budget:
                break
            
            chunk_size = min(len(req.prompt_tokens), self.token_budget - current_tokens)
            



            # TODO: Check if we can allocate memory for the new request's prefill chunk.
            can_alloc = self.kv_cache_manager.can_allocate(req, chunk_size)

            if chunk_size > 0 and can_alloc:



                # TODO: If we can, pop the request, add to active pool, and allocate slots.
                self.waiting_queue.popleft()
                self.active_pool[req.id] = req
                if self.kv_cache_manager.allocate_slots(req, chunk_size):
                    next_batch.append(ScheduledSequence(request_id=req.id, tokens=req.prompt_tokens[:chunk_size], is_decode=False, position_ids_start=0))
                    current_tokens += chunk_size
                    req.tokens_processed += chunk_size
        
        return next_batch


```grasp
{
  "type": "exercise",
  "id": "5ff6958f-0085-402d-811b-db2a314a7674"
}

--- Example Usage (Simplified from last lesson) ---

if name == 'main':

# Setup a cache manager with 100 blocks of size 16
kv_manager = MockKVCacheManager(num_blocks=100, block_size=16)
scheduler = StallFreeScheduler(token_budget=1024, max_batch_size=8, kv_cache_manager=kv_manager)




# Add some requests
prompts = [
    "The capital of France is a city known for its beautiful art and history.", # 14 tokens
    "The best programming language for machine learning is arguably Python because of its vast ecosystem.", # 17 tokens
    "Large language models are transforming the field of artificial intelligence by enabling new applications.", # 16 tokens
    "To be or not to be, that is the question. Whether 'tis nobler in the mind to suffer the slings and arrows of outrageous fortune." # 28 tokens
]
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2")

for i, prompt in enumerate(prompts):
    prompt_tokens = tokenizer.encode(prompt)
    scheduler.add_request(Request(id=i, prompt_tokens=prompt_tokens, max_new_tokens=5))




# Simulate the inference loop
iteration = 0



# In this mock run, we just schedule and see the state changes. We don't execute a model.
while len(scheduler.active_pool) > 0 or len(scheduler.waiting_queue) > 0:
    iteration += 1
    print(f"\n--- Iteration {iteration} ---")
    
    batch = scheduler.schedule_next_batch()
    
    if not batch:
        print("Scheduler returned no batch. Ending run.")
        break
        
    print(f"Scheduled {len(batch)} sequences:")
    for seq in batch:
        req = scheduler.active_pool[seq.request_id]
        print(f"  - Req {seq.request_id}: {'Decode' if seq.is_decode else 'Prefill'} {len(seq.tokens)} tokens. Req state: {req.tokens_processed}/{len(req.prompt_tokens)} prompt done.")
        



        # Simulate completion of decode steps
        if seq.is_decode:
            req.output_tokens.append(123) # Dummy output token

### Conclusion

Excellent work. By modifying the scheduler to communicate with the `KVCacheManager`, you have fundamentally improved the robustness and intelligence of our inference engine. The scheduler is no longer "blind"; it now formally negotiates for memory resources before committing to a plan.

**Key Takeaways:**

*   **Separation of Concerns:** The scheduler is responsible for *policy* (which requests to run), while the `KVCacheManager` is responsible for the *mechanism* of memory allocation. A clean interface between them is essential.
*   **The API Contract:** The core of the interface consists of three main operations: querying for available resources (`can_allocate`), requesting an allocation (`allocate_slots`), and releasing resources (`free`).
*   **Failure Handling:** A critical part of the design is handling allocation failure. When the manager returns `None` or `False`, the scheduler must gracefully skip the request for the current cycle, preventing an OOM crash.

**Preview of the Next Lesson:**

We have successfully defined the interface *to* the `KVCacheManager` and used it to make our scheduler memory-aware. Now, it's time to look inside the black box. In our next lesson, we will begin implementing the `KVCacheManager` itself. You will **implement a paged KV cache allocator** that manages the pool of non-contiguous physical memory blocks, turning the abstract API calls we designed today into concrete memory management operations.