Converge the two SWA predicates, and stop conditioning the capture sink on the pool

[ch-wan]

Stack — third of four. Depends on #37511 → #37512 (its base branch). One commit of its own; #37560 sits on top.

Motivation

Two findings from auditing the KVIndexTranslator series (#35245 → #34613 → #37307). Both predate that series in effect; both were given their current form by it.

The same question had two answers. A backend decides "does this pool have a full→swa index mapping" through _resolve_swa_kv_pool, which keys on BaseSWAKVPool. The translator decided it again, for the write loc, keying on SWAKVPool:

self._swa_write_loc_from_full = (
    token_to_kv_pool.translate_loc_from_full_to_swa
    if isinstance(token_to_kv_pool, SWAKVPool)   # too narrow
    else None
)

translate_loc_from_full_to_swa is an @abstractmethod on BaseSWAKVPool, so every BaseSWAKVPool has it. The narrow spelling drops DeepSeekV4TokenToKVPool, which is a BaseSWAKVPool and not a SWAKVPool — a backend that resolves such a pool as SWA passes its outer guard and then hands None to copy_.

The capture sink was conditioned on the pool, and three backends had no sink at all. Three sites (flashinfer, trtllm_mha, flashattention) read:

if in_capture and self.kv_index_translator.is_translating:
    swa_out_cache_loc[:n].zero_()
else:
    swa_out_cache_loc[:n].copy_(sliding_window_write_loc_for(out_cache_loc))

Four combinations for two wanted behaviours, and only the unified one got the sink. The capture forward does write KV, so under a static SWA pool the else branch derived a real slot from a capture batch's dummy out_cache_loc and wrote KV there.

Modifications

• The translator keys on BaseSWAKVPool — the class that declares the capability, and what update_sliding_window_buffer in the triton backend already used.
• All six sites that write cuda_graph_swa_out_cache_loc now condition on in_capture alone. Nothing depends on the value either way: graph capture bakes pointers, not buffer contents, and the same code refills on every replay-prep. Zeros route the capture write to slot 0, the reserved sink in every id space, whatever the pool is.
• Enumerating by the condition found three sites; enumerating by who writes the buffer found six. triton_backend._fill_cuda_graph_swa_out_cache_loc (called from inside the if in_capture: branch), aiter_backend.init_forward_metadata_out_graph (function-body top level, so both phases) and trtllm_mha._write_ragged_verify_graph_metadata had no sink at all -- they derived a live slot from a capture batch's dummy out_cache_loc and wrote KV there unconditionally. in_capture is threaded into the first and third; aiter already had it in scope.
• The comments were part of the problem — they said "zeros are the page-0 sink" while sitting above a branch that computes a live one. They now state why the phase is the whole condition.

Accuracy Tests

Neither fix has an end-to-end run, because the pools whose behaviour they change are not reachable here: DSv4 (trtllm_mha rejects MLA models) and hybrid-SWA under the static pool.

What is checkable is collateral damage on the configurations that do run:

config	result	band
Kimi-Linear TP2 + unified + flashinfer_mla	0.9200	0.9125–0.9175
Falcon-H1-1.5B + unified + fa3	0.7925	0.7625–0.7900
gpt-oss-20b + unified + triton (SWA, 2 runs)	0.950 / 0.945	0.9387 mean over 4 base runs

The gpt-oss row covers the new triton sink specifically: hybrid-SWA on triton, so _fill_cuda_graph_swa_out_cache_loc runs on both phases. aiter is ROCm and unreachable here.

Both one question outside the band's edge, which is inside the ±2 that three runs of one unchanged server reproduce here. test/registered/unit/mem_cache/: 2127 passed.

Speed Tests and Profiling

No performance effect. The sink change removes a per-capture gather on the static-SWA path (once per graph, not per step); the predicate change is an isinstance at construction.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests. — both fixes are reachable only on pools this box cannot run; a test would have to fake the pool, which the repo's style rules out.
• Update documentation according to Write documentations. — no user-facing surface changes.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

A third finding, deliberately not fixed

ROCm gfx950 DeepSeek MHA FP8 (forward_mha.py, forward_mha_rocm.py) runs filter_dcp_local_kv_indices + translate_dcp_read_ids inside forward_normal_prepare — a per-layer method — over page_table_1_flattened, which is a per-step field. Both are pure functions of it, so the same tensor is recomputed once per layer.

Every clean hoist needs either a mutable field on the frozen DSAMetadata or module-level memo state, and the path wants ROCm gfx950 + FP8 KV + unified + DCP + prefill, none of which is reachable on this box. I would rather leave it named than ship an untested change to it.

---

CI States

Latest PR Test (Base): 🚫 Run #33697277863
Latest PR Test (Extra): 🚫 Run #33697277668
Latest PR Test (AMD ROCm 7.2): 🚫 Run #33697277873

[chatgpt-codex-connector]

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review	Status	Commit	Review trigger
📝 Code Review	✅ Completed 2026-09-02T06:30:20.077965Z	62ba19a	PR opened

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.