NWOR: Global Deferred KV Cache with GPU integration and tests

[yuz207]

Summary

• Introduces a global DeferredWriteManager coordinating NWOR across KV-cache backends and the GPU model runner
• Adds acceptance-mask-based commit/cancel flow for KV-cache writes during speculative decoding
• Adds tests for deferred writer behavior and acceptance mask logic
• Adds environment control to disable NWOR via VLLM_DISABLE_NWOR and NWOR mode via VLLM_NWOR_MODE

Changes

• Core NWOR components
  • Add vllm/v1/kv_cache/deferred.py implementing DeferredWriteManager, record_or_write_kv_cache, get/set global manager, ShouldFallback, and related helpers
  • Add vllm/v1/kv_cache/init.py to expose DeferredWriteManager and related APIs
• Backward-compatible integration across KV-cache backends
  • Updated attention backends to route KV-cache writing through record_or_write_kv_cache instead of direct reshape_and_cache calls:
    • vllm/v1/attention/backends/flash_attn.py
    • vllm/v1/attention/backends/flashinfer.py
    • vllm/v1/attention/backends/flex_attention.py
    • vllm/v1/attention/backends/rocm_aiter_fa.py
    • vllm/v1/attention/backends/rocm_aiter_unified_attn.py
    • vllm/v1/attention/backends/tree_attn.py
    • vllm/v1/attention/backends/triton_attn.py
    • vllm/v1/attention/backends/xformers.py
  • The new path records layer_id, scales, and slot mappings, enabling NWOR staging to be applied consistently across backends
• GPU model runner NWOR integration
  • Add NWOR window lifecycle in vllm/v1/worker/gpu_model_runner.py:
    • _maybe_begin_nwor_window: starts a NWOR window based on SpecDecodeMetadata.num_draft_tokens
    • _finalize_nwor_window: builds an acceptance mask (via _build_nwor_acceptance_mask) and commits or flushes accordingly
    • _cleanup_nwor: resets NWOR state after finishing a window
    • _build_nwor_acceptance_mask: computes which drafted tokens are accepted given sampled tokens, handling per-request draft counts and device alignment
  • Integrate a global DeferredWriteManager to coordinate staging and fallback behavior during speculative decoding
• Tests
  • Added tests/v1/test_deferred_writer.py with scenarios:
    • test_deferred_manager_commit_partial_acceptance: verifies partial acceptance commits staged writes and updates metrics
    • test_deferred_manager_cancel_flush_writes_all: ensures cancel/flush of all staged writes works and writes are flushed
    • test_build_acceptance_mask_matches_expected: validates acceptance masking logic against a synthetic SpecDecodeMetadata and sampled tokens
    • test_nwor_disabled_env: verifies NWOR can be disabled via environment variable
• Config/Env
  • vllm/envs.py: add VLLM_DISABLE_NWOR environment variable to disable NWOR when set
  • vllm/envs.py: add VLLM_NWOR_MODE environment variable to select NWOR mode (e.g., "stage" or "immediate")
• Scheduler/Outputs/Metrics
  • vllm/v1/core/sched/scheduler.py: propagate nwor_stats into spec_decoding_stats flow
  • vllm/v1/metrics/stats.py: include nwor_stats field in SchedulerStats
  • vllm/v1/outputs.py: extend ModelRunnerOutput with nwor_metrics
• Migration/Compatibility
  • No API surface changes for existing KV-cache paths outside NWOR; when NWOR is not active, writes fall back to the original path via record_or_write_kv_cache, preserving existing behavior
  • The integration points are opt-in via the NWOR window lifecycle; existing workflows without an active NWOR window should experience no behavioral changes

Why this change

• NWOR enables safe speculative decoding by buffering KV-cache writes until a final acceptance decision is made, reducing the risk of KV-cache corruption when speculative steps are rejected
• Centralized/deferred KV-cache handling improves consistency across backends and simplifies future enhancements
• The global DeferredWriteManager coordinates staging and fallback behavior across async/sync code paths, including the GPU model runner lifecycle

Compatibility and impact

• No API surface changes for existing code paths outside NWOR; when NWOR is not active, writes fall back to the original path via record_or_write_kv_cache, preserving existing behavior
• The integration points are designed to be opt-in via the NWOR window lifecycle; existing workflows without an active NWOR window should experience no behavioral changes
• Environment variable VLLM_DISABLE_NWOR allows disabling NWOR at runtime for testing or fallback scenarios

How to test

• Run unit tests, focusing on NWOR-related tests:
  • tests/v1/test_deferred_writer.py
• Validate that existing attention backends still compile and run, with the new write path used only when NWOR is active
• Manual validation of NWOR flow in a scenario with speculative decoding enabled and a draft/acceptance loop

Reviewer notes

• The NWOR backend integration relies on a global DeferredWriteManager when a window is active. Review the lifecycle management to ensure proper begin/commit/cancel semantics across async/ sync code paths
• Ensure PyTorch versions with FakeTensor support interact correctly with _is_fake_tensor checks in the deferred module
• If unexpected fallsbacks occur, they should surface via the ShouldFallback pathway and trigger a safe flush of writes

Migration/Deprecation

• This change adds NWOR capabilities and does not deprecate existing KV-cache paths. NWOR can be enabled by ensuring a decode window is active and spec_decode_metadata is provided