Summary

Implement a Python-based interleaved scheduler (KVCacheV2Scheduler) for KVCacheManagerV2 that unifies capacity scheduling and micro-batch scheduling into a single pass, replacing the previous KVCacheV2DummyScheduler + BindCapacityScheduler/BindMicroBatchScheduler two-stage path for V2.

Key changes:

• New KVCacheV2Scheduler (scheduler_v2.py): merges KV cache allocation (via resize()) and token budget assignment into one unified scheduling loop. Supports:

  • MAX_UTILIZATION eviction policy (suspend-to-host) with tail-eviction strategy
  • Self-eviction for generation requests that cannot allocate
  • Block reuse integration for context requests
  • FCFS chunked prefill with configurable chunk unit size
  • PEFT (LoRA) page accounting
  • Draft model (MTP) KV cache coordination
  • LoRA-aware request sorting

• KVCacheManagerV2 scheduling API (resource_manager.py): exposes fine-grained allocation methods called by the scheduler:

  • prepare_context() — create _KVCache, handle block reuse lookup, resume from suspended state
  • resize_context() — resize KV cache to cover context tokens
  • try_allocate_generation() — allocate one additional KV slot for generation (with resume)
  • suspend_request() — suspend a request's KV cache to host tier
  • is_request_active() — check if a request has a live, non-suspended KV cache

• Draft KV cache manager refactor (resource_manager.py): extract _prepare_draft_resources() that mirrors the main manager's allocations for MTP draft layers, since the V2 scheduler only manages the primary KV cache

• PyExecutor integration (py_executor.py): when V2 scheduler is active, skip executor-level _terminate_requests and _pause_requests for paused requests (the scheduler manages suspend/resume directly)

• Bug fixes:

  • Fix BAD_PAGE_INDEX (-1) causing CUDA_ERROR_ILLEGAL_ADDRESS in Flash MLA kernel — replace with 0 for unused page slots
  • Cap get_num_available_tokens by GPU-only capacity when max_tokens is explicitly set
  • Add assertions blocking Star attention and multi-beam with V2

Test plan

• New unit tests: tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py covering scheduling logic, eviction, chunked prefill, PEFT, draft tokens, and edge cases
• New integration tests: tests/integration/defs/kv_cache/test_kv_cache_v2_scheduler.py (+542 lines) covering LLaMA accuracy with V2 KV cache + chunked context + MTP + block reuse
• Existing accuracy ITs parametrized with v2_kv_cache=True/False variants
• CI test-db entries added for B200, DGX-B200, and H100

Summary by CodeRabbit

Release Notes

• New Features

  • Enhanced KV cache scheduling strategy with improved memory management and resource efficiency
  • Added support for draft token generation during inference
  • Implemented GPU memory-aware caching for better resource utilization
  • Expanded support for adapter-based inference workflows

• Tests

  • Added comprehensive scheduling validation test suite
  • Extended test coverage with new caching configuration options