Fix single-request generation being killed early on hybrid-SWA models under DFLASH/DSPARK speculative decoding

[ormandj]

Motivation

spec_prepare_for_decode routes dflash-family algorithms (DFLASH, DSPARK) to DFlashDraftInputV2.prepare_for_decode, which never calls maybe_evict_swa and never advances decode_batch_idx — which the eviction gate requires to be >= 1. eagle_prepare_for_decode does both. On hybrid-SWA models the sliding-window eviction is therefore unreachable under DFLASH/DSPARK: SWA KV is retained for every generated token, so a single request exhausts the SWA pool after roughly swa_full_tokens_ratio * max_total_num_tokens generated tokens and is retracted by the scheduler mid-generation.

Reproduced on DeepSeek-V4-Flash-0731 (2x RTX PRO 6000, SM120, TP=2, DSPARK block size 5, swa_full_tokens_ratio=0.1, max_total_num_tokens=778496). #swa token tracks generated tokens 1:1, and the request is retracted at exactly 77,824 tokens — align_page_size(int(778496 * 0.1)), the full SWA pool — while the full pool is only 10% used:

#full token: 77824, full token usage: 0.10, #swa token: 77824, swa token usage: 1.00
KV cache pool is full. Retract requests. #new_tokens_gained: 77824

Long prompts are unaffected — a 131k-token prefill holds ~8.5K SWA tokens. Only generated tokens accumulate.

The same missing call also skips the DSV4 compress-state drain (maybe_evict_dsv4_state, documented "called every decode step from ScheduleBatch"), so the NPU + DSV4 + DFLASH combination leaks state-pool slots for the same reason.

Modifications

speculative/spec_utils.py: in the dflash-family branch, mirror eagle_prepare_for_decode — call batch.maybe_evict_swa() and advance decode_batch_idx. Eviction runs before the tick so the overlap-scheduler gate (decode_batch_idx >= 1) still skips the first decode iteration, exactly as on the eagle path.

test/registered/unit/spec/test_decode_bookkeeping_ownership.py: record the new bookkeeping owner in _OWNER_SITES, as that guard requires for a genuinely new owner.

test/registered/unit/speculative/test_spec_prepare_swa_eviction.py (new): asserts the dflash branch evicts, ticks, evicts before ticking, keeps advancing across iterations (so the counter stays a clock rather than a flag), and that the eagle path is unchanged.

Accuracy Tests

Eviction frees only SWA slots behind the window — max(sliding_window, page_size) back from the sequence on the radix path, page-floored pre_len - sliding_window on the chunk-cache path. SWA layers attend within the window, and both the eagle and non-spec paths already call this same code every decode step (eagle_utils.py, mem_cache/allocation.py). No output change observed.

Speed Tests and Profiling

The ignore_eos request that previously died at 77,809 tokens completed 100,000 tokens with the fix, at a sustained ~159 tok/s and zero retractions. #swa token sampled repeatedly across the run stayed between 512 and 1,024 (~1% of the SWA pool), against 77,824 at the point of failure before the fix.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. The new test fails on the pre-fix code and passes with the fix.
• Update documentation according to Write documentations. No user-facing documentation change.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Prepared with AI assistance.

---

CI States

Latest PR Test (Base): ❌ Run #31830154702
Latest PR Test (Extra): ❌ Run #31830154341

[ormandj]

Per MAINTAINER.md, CODEOWNERS has requested the Speculative decoding owners. Flagging the KV Cache oncalls as well — @ispobock @xiezhq-hermann — since the change is in srt/speculative but its correctness depends on the SWA eviction machinery in srt/mem_cache (free_swa_out_of_window_slots, SWATokenToKVPoolAllocator.free_swa) being safe to invoke from the dflash-family path. I could not add you as reviewers directly (no permission), hence the mention.

Two points that may help review:

1. Cross-platform. The change makes maybe_evict_swa() run on the dflash-family path on every platform. maybe_evict_dsv4_state is a no-op unless the allocator has c4_state_attn_allocator (NPU only), so CUDA/ROCm are unaffected by that part. On NPU + DSV4 + DFLASH the state drain becomes reachable where it previously was not — by inspection every primitive is monotonic and idempotent (free_swa filters > 0 then zeroes full_to_swa_index_mapping; free_swa_out_of_window_slots only acts when swa_evicted_seqlen advances; _free_state_range returns early below its watermark), and DSV4NPUTokenToKVPoolAllocator.free documents that the prefix is expected to have already been returned by the eviction path. I have no NPU hardware and could not verify this at runtime — a check from someone who can would be welcome.

2. test_decode_bookkeeping_ownership.py. The new bookkeeping owner is recorded in _OWNER_SITES as that guard requires. cc @fzyzcjy as its author, in case the two entries should be expressed differently.

Non-hybrid-SWA models are unaffected: the whole body is behind tree_cache.supports_swa(), which is False by default.