…vent cross-stream data race

In overlap scheduling (MTPv2), `process_batch_result(N-1)` runs on the
default stream concurrently with `forward(N)` on the forward stream.
When a request finishes, `release_kv_cache` immediately returns its
`req_pool_idx` to the free list.  A new request can then recycle that
pool index and `prepare_for_decode` overwrites the `req_to_token` row on
the default stream while `forward(N)` still reads it — causing an
"index out of bounds" assertion in IndexKernel.cu.

Fix: defer the pool-index free by one overlap iteration.

- `ReqToTokenPool.deferred_free(req)`: withholds the pool index from
  the free list (the slot cannot be reallocated).
- `ReqToTokenPool.flush_deferred_frees()`: moves deferred slots back to
  the free list once the forward that read them has completed.
- `release_kv_cache(..., defer_pool_free=True)`: used in the decode
  result-processing path when overlap is enabled.
- `process_batch_result_decode`: flushes deferred frees right after
  `copy_done.synchronize()`, which guarantees the previous forward
  has finished reading `req_to_token`.

This is the overlap-scheduling counterpart of PR sgl-project#17850, which fixed the
same class of race for chunked prefill.

…vent cross-stream data race

In overlap scheduling (MTPv2), `process_batch_result(N-1)` runs on the
default stream concurrently with `forward(N)` on the forward stream.
When a request finishes, `release_kv_cache` immediately returns its
`req_pool_idx` to the free list.  A new request can then recycle that
pool index and `prepare_for_decode` overwrites the `req_to_token` row on
the default stream while `forward(N)` still reads it — causing an
"index out of bounds" assertion in IndexKernel.cu.

Fix: defer the pool-index free by one overlap iteration.

- `ReqToTokenPool.deferred_free(req)`: withholds the pool index from
  the free list (the slot cannot be reallocated).
- `ReqToTokenPool.flush_deferred_frees()`: moves deferred slots back to
  the free list once the forward that read them has completed.
- `release_kv_cache(..., defer_pool_free=True)`: used in the decode
  result-processing path when overlap is enabled.
- `process_batch_result_decode`: flushes deferred frees right after
  `copy_done.synchronize()`, which guarantees the previous forward
  has finished reading `req_to_token`.

This is the overlap-scheduling counterpart of PR sgl-project#17850, which fixed the
same class of race for chunked prefill.

Previously the scheduler explicitly freed a chunked request's req_pool
slot in get_next_batch_to_run before allocating the next batch, then
relied on alloc(num_reqs) to hand out a fresh slot to that same request
on the next round. Under chunked prefill with the pool near capacity,
the freed slot could be re-allocated to a different request whose
matched-prefix indices were then written to the same row of
req_to_token while the previous batch's forward was still reading that
row's host buffer. JAX provides no happen-before between host numpy
mutations and an in-flight dispatch, and on the PJRT host-CPU zero-copy
path the host write is the device buffer; the result is silent KV
gather corruption.

Fix by stabilising slot ownership across batches:

- ReqToTokenPool.alloc(reqs) now treats reqs whose req_pool_idx is
  already non-None as in-place reuse and only consumes free_slots for
  the rest. Atomic on capacity miss: returns None and leaves every req
  field untouched.
- ReqToTokenPool.free(req) is req-aware; it releases the slot and
  clears req.req_pool_idx so a stale value cannot leak to a later
  caller after the slot has been re-handed out.
- The explicit free of the chunked request's slot in
  Scheduler.get_next_batch_to_run is removed. The req keeps its
  req_pool_idx; the next prepare_for_extend's alloc(reqs) reuses it.
- All call sites (ScheduleBatch.alloc_req_slots, release_req,
  ChunkCache / RadixCache / SWARadixCache.cache_finished_req) are
  switched to the req-aware API.

Mirrors the upstream sglang fix in
sgl-project/sglang#17850 and addresses the
chunked-prefill ownership item in
#992 .

New unit tests in test/mem_cache/test_req_to_token_pool.py lock in:
fresh allocation, in-place reuse, atomic capacity miss (with and
without partial reuse), free-clears-req-field, alloc/free roundtrip,
and end-to-end chunked-req survival across batches.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Previously the scheduler explicitly freed a chunked request's req_pool
slot in get_next_batch_to_run before allocating the next batch, then
relied on alloc(num_reqs) to hand out a fresh slot to that same request
on the next round. Under chunked prefill with the pool near capacity,
the freed slot could be re-allocated to a different request whose
matched-prefix indices were then written to the same row of
req_to_token while the previous batch's forward was still reading that
row's host buffer. JAX provides no happen-before between host numpy
mutations and an in-flight dispatch, and on the PJRT host-CPU zero-copy
path the host write is the device buffer; the result is silent KV
gather corruption.

Fix by stabilising slot ownership across batches:

- ReqToTokenPool.alloc(reqs) now treats reqs whose req_pool_idx is
  already non-None as in-place reuse and only consumes free_slots for
  the rest. Atomic on capacity miss: returns None and leaves every req
  field untouched.
- ReqToTokenPool.free(req) is req-aware; it releases the slot and
  clears req.req_pool_idx so a stale value cannot leak to a later
  caller after the slot has been re-handed out.
- The explicit free of the chunked request's slot in
  Scheduler.get_next_batch_to_run is removed. The req keeps its
  req_pool_idx; the next prepare_for_extend's alloc(reqs) reuses it.
- All call sites (ScheduleBatch.alloc_req_slots, release_req,
  ChunkCache / RadixCache / SWARadixCache.cache_finished_req) are
  switched to the req-aware API.

Mirrors the upstream sglang fix in
sgl-project/sglang#17850 and addresses the
chunked-prefill ownership item in
#992 .

New unit tests in test/mem_cache/test_req_to_token_pool.py lock in:
fresh allocation, in-place reuse, atomic capacity miss (with and
without partial reuse), free-clears-req-field, alloc/free roundtrip,
and end-to-end chunked-req survival across batches.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Previously the scheduler explicitly freed a chunked request's req_pool
slot in get_next_batch_to_run before allocating the next batch, then
relied on alloc(num_reqs) to hand out a fresh slot to that same request
on the next round. Under chunked prefill with the pool near capacity,
the freed slot could be re-allocated to a different request whose
matched-prefix indices were then written to the same row of
req_to_token while the previous batch's forward was still reading that
row's host buffer. JAX provides no happen-before between host numpy
mutations and an in-flight dispatch, and on the PJRT host-CPU zero-copy
path the host write is the device buffer; the result is silent KV
gather corruption.

Fix by stabilising slot ownership across batches:

- ReqToTokenPool.alloc(reqs) now treats reqs whose req_pool_idx is
  already non-None as in-place reuse and only consumes free_slots for
  the rest. Atomic on capacity miss: returns None and leaves every req
  field untouched.
- ReqToTokenPool.free(req) is req-aware; it releases the slot and
  clears req.req_pool_idx so a stale value cannot leak to a later
  caller after the slot has been re-handed out.
- The explicit free of the chunked request's slot in
  Scheduler.get_next_batch_to_run is removed. The req keeps its
  req_pool_idx; the next prepare_for_extend's alloc(reqs) reuses it.
- All call sites (ScheduleBatch.alloc_req_slots, release_req,
  ChunkCache / RadixCache / SWARadixCache.cache_finished_req) are
  switched to the req-aware API.

Mirrors the upstream sglang fix in
sgl-project/sglang#17850 and addresses the
chunked-prefill ownership item in
#992 .

New unit tests in test/mem_cache/test_req_to_token_pool.py lock in:
fresh allocation, in-place reuse, atomic capacity miss (with and
without partial reuse), free-clears-req-field, alloc/free roundtrip,
and end-to-end chunked-req survival across batches.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

#1008)

Previously the scheduler explicitly freed a chunked request's req_pool
slot in get_next_batch_to_run before allocating the next batch, then
relied on alloc(num_reqs) to hand out a fresh slot to that same request
on the next round. Under chunked prefill with the pool near capacity,
the freed slot could be re-allocated to a different request whose
matched-prefix indices were then written to the same row of
req_to_token while the previous batch's forward was still reading that
row's host buffer. JAX provides no happen-before between host numpy
mutations and an in-flight dispatch, and on the PJRT host-CPU zero-copy
path the host write is the device buffer; the result is silent KV
gather corruption.

Fix by stabilising slot ownership across batches:

- ReqToTokenPool.alloc(reqs) now treats reqs whose req_pool_idx is
  already non-None as in-place reuse and only consumes free_slots for
  the rest. Atomic on capacity miss: returns None and leaves every req
  field untouched.
- ReqToTokenPool.free(req) is req-aware; it releases the slot and
  clears req.req_pool_idx so a stale value cannot leak to a later
  caller after the slot has been re-handed out.
- The explicit free of the chunked request's slot in
  Scheduler.get_next_batch_to_run is removed. The req keeps its
  req_pool_idx; the next prepare_for_extend's alloc(reqs) reuses it.
- All call sites (ScheduleBatch.alloc_req_slots, release_req,
  ChunkCache / RadixCache / SWARadixCache.cache_finished_req) are
  switched to the req-aware API.

Mirrors the upstream sglang fix in
sgl-project/sglang#17850 and addresses the
chunked-prefill ownership item in
#992 .

New unit tests in test/mem_cache/test_req_to_token_pool.py lock in:
fresh allocation, in-place reuse, atomic capacity miss (with and
without partial reuse), free-clears-req-field, alloc/free roundtrip,
and end-to-end chunked-req survival across batches.

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>