Implement safe chunked KV writes for CUDA graph capture

[yuz207]

Summary

• Implement safe, chunked KV writes to support CUDA graph capture
• Introduce per-segment acceptance tracking and slicing helpers
• Extend GPUModelRunner to return acceptance counts alongside mask

Changes

KV Cache

• Add _slice_scale_segment to support per-segment scale slicing based on entry length
• Extend DeferredWriteManager with:
  • _req_start_offsets to map draft tokens to request offsets
  • commit now accepting a sequence of accepted counts instead of a boolean mask
  • Per-entry segmentation logic to compute accepted segments across all draft requests
  • Optimized path for full-entry, single-segment acceptance
  • Per-segment slicing for keys, values, slots, and scales, with proper device handling
  • Robust exception handling that triggers fallback with metrics updates
  • Clearing of request offsets on flush/clear

GPU Model Runner

• _build_nwor_acceptance_mask now returns a tuple (mask, counts)
• Vectorized path updated to propagate and return acceptance counts
• _scv_vectorized_mask now returns (mask_work, accepted_counts) when available
• Commit path updated to pass accepted_counts to DeferredWriteManager.commit

Tests

• Update tests to reflect new API:
  • Use manager.commit([1]) instead of a boolean mask
  • Validate that _build_nwor_acceptance_mask returns (mask, counts) and that counts match expectations
  • Extend tests to verify per-segment acceptance counts in vectorized paths

Why

• Enables safe CUDA graph capture by supporting partial and segmented writes without risking incorrect slices or race conditions
• Provides explicit acceptance tracking across multiple draft requests, increasing robustness in graph-captured execution

API Changes

• DeferredWriteManager.commit signature changed from commit(mask: Tensor) to commit(accepted_counts: Sequence[int])
• _build_nwor_acceptance_mask now returns (mask: Tensor, counts: list[int]) instead of a sole mask

Backward Compatibility

• Internal refactors align with CUDA graph capture goals; external call sites within this repo have been updated accordingly. If any external callers relied on the old signature, they will need to adapt to the new accepted_counts interface.

Testing Plan

• Run unit tests for KV cache deferred writer: pytest tests/v1/test_deferred_writer.py
• Validate that commit calls with [1] produce expected writes and that counts are correctly propagated
• Validate _build_nwor_acceptance_mask returns both mask and counts and that counts align with draft token segmentation

Notes for reviewers

• The new per-segment logic ensures safe writes even when CUDA graph capture requires partial acceptance across multiple draft requests
• Some parts of the implementation introduce additional branch paths; expected behavior remains correct including fallback behavior on failure

🌿 Generated by Terry

---

ℹ️ Tag @terragon-labs to ask questions and address PR feedback

📎 Task: https://www.terragonlabs.com/task/7ed05cf2-ef51-40c2-b908-a395d47b0386