vllm - 💡(How to fix) Fix [Bug]: SimpleCPUOffloadScheduler crashes with AssertionError: Expected N hit tokens, got 0 (TOCTOU race in update_state_after_alloc) [1 participants]

AssertionError: Expected 19264 hit tokens, got 0
  File "vllm/v1/simple_kv_offload/manager.py", line 267, in update_state_after_alloc
    assert hit_length == num_external_tokens, (
        f"Expected {num_external_tokens} hit tokens, got {hit_length}"
    )

---

1. scheduler.py:619
   get_num_new_matched_tokens()
   └─ find_longest_cache_hit() → finds N tokens in CPU LRU cache
      returns hit_length=N, BUT discards cpu_hit_blocks (_ = ...)
      *** no touch() called — blocks are NOT pinned ***

2. scheduler.py:746-754
   kv_cache_manager.allocate_slots()
   └─ may call cpu_block_pool.get_new_blocks()
      └─ triggers CPU LRU eviction
         *** the N blocks found in step 1 are evicted here ***

3. scheduler.py:771
   update_state_after_alloc(request, blocks, num_external_tokens=N)
   └─ find_longest_cache_hit() again → returns 0 (blocks evicted)
      assert 0 == N  →  CRASH

---

_, hit_length = self.cpu_coordinator.find_longest_cache_hit(   # blocks discarded!
    remaining_hashes, max_hit_len
)

---

# In get_num_new_matched_tokens():
cpu_hit_blocks, hit_length = self.cpu_coordinator.find_longest_cache_hit(...)
if hit_length > 0:
    # Pin immediately to prevent LRU eviction before update_state_after_alloc()
    all_hit_blocks = [blk for grp in cpu_hit_blocks for blk in grp if not blk.is_null]
    self.cpu_block_pool.touch(all_hit_blocks)          # temporary pin
    self._pending_cpu_hits[request.request_id] = (cpu_hit_blocks, hit_length)
    return hit_length, True

# In update_state_after_alloc():
# Pop cached result — no second find_longest_cache_hit() needed
pending = self._pending_cpu_hits.pop(req_id, None)
if pending is None:
    return  # graceful fallback instead of assert
cpu_hit_blocks, actual_hit_length = pending
# ... use cpu_hit_blocks directly ...
self.cpu_block_pool.touch(cpu_blocks_to_touch)     # persistent pin for async load
self.cpu_block_pool.free_blocks(all_hit_blocks)    # release temporary pin

# In request_finished():
# Release temporary pin if request is preempted before update_state_after_alloc()
pending = self._pending_cpu_hits.pop(req_id, None)
if pending is not None:
    self.cpu_block_pool.free_blocks([...])

---

VLLM_USE_SIMPLE_KV_OFFLOAD=1 vllm serve <model> \
  --kv-offloading-size 32 \
  --no-disable-hybrid-kv-cache-manager \
  --tensor-parallel-size 2 \
  --kv-cache-dtype fp8_e4m3 \
  --max-model-len 262144