Skip to content

[AMD][DSV4][SWA-BITEXACT 1/5] Extension points for strict SWA state capture/restore on unified_kv + HiCache - #35303

Open
amd-danli103 wants to merge 2 commits into
sgl-project:mainfrom
amd-danli103:swa-bitexact-1n
Open

amd-danli103 wants to merge 2 commits into
sgl-project:mainfrom
amd-danli103:swa-bitexact-1n

Conversation

@amd-danli103

@amd-danli103 amd-danli103 commented Aug 18, 2026

Copy link
Copy Markdown
Contributor

Motivation

The configuration this stack addresses is DeepSeek-V4 served on the unified_kv attention backend with prefix cache enabled. There the swa and the c4 / c4-indexer overlap state live in per-request rings that the radix tree does not hand back on a prefix hit, so HiCache re-prefills a tail to rebuild them and a cache hit is not bit-identical to a cold run. Nothing outside that combination is involved. RFC #34562 proposes making that reuse bit-exact by riding the window and its coupled state on the radix nodes.

#32214 carries the whole implementation in one diff. Per review feedback it is being re-landed as a five-part stack. This is part 1/5: the extension points the later parts hang off.

This part is inert by construction. It adds an env flag with no readers, a server arg with no readers, a dataclass field with no readers, and hook methods whose base implementations are pass. No existing code path changes behavior, with or without the flag.

Modifications

  • AttentionBackend: three no-op hooks — capture_swa_windows (prefill), capture_swa_windows_decode and capture_compress_state_windows_decode (decode). Their docstrings state the ordering contract an implementation must honor. They are declared as base methods rather than probed with hasattr at the call site, so the contract lives in the base class and every backend inherits a defined answer.
  • BasePrefixCache.restore_swa_windows(): no-op, driven once from prepare_for_extend after req_pool_idx is assigned and before the first forward reads the ring. A cache that re-prefills the window instead of handing back a durable one has nothing to put there.
  • MatchPrefixParams.for_reuse (default False): distinguishes cross-request reuse matches from self-match lookups such as the one in cache_unfinished_req, which must keep trusting the device-only validators for a request's own freshly computed nodes. Three scheduler call sites set it; nothing reads it yet.
  • Call sites: two decode hooks in ModelRunner, placed outside the cuda graph after the decode forward wrote the ring and before the next step overwrites it; one prefill hook in deepseek_v4.MQALayer, on the flat post-norm+rope KV before the device ring overwrites it.
  • SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE (default off) and --hicache-swa-offload-page-stride (default 1), both documented and both unread until part 2/5.

Accuracy Tests

No accuracy impact is reachable from this part: the flag and for_reuse have no readers and every hook body is pass. test/registered/mem_cache/test_swa_bitexact_extension_points.py asserts exactly that — the hooks exist on the base classes, return None, and mutate nothing — so the inertness is enforced rather than asserted in prose.

Accuracy numbers for the feature as a whole are in #32214 and in the RFC. The first part that can move outputs is 3/5, and each such part carries its own numbers.

Speed Tests and Profiling

Two no-op virtual calls per decode step, outside the cuda graph, and one per MQA prefill layer. No measurable impact.

Stack

part contents
1/5 (this PR) extension points, flag, server arg
2/5 HiCache substrate: SWA host pool, L3 key namespace, startup guards
3/5 capture of SWA windows and c4/indexer overlap state
4/5 window riding on radix nodes and the strict reuse gate
5/5 c4/indexer state riding; completes the feature

#32214 stays open as the full-context reference until the stack is in.

A fork PR can only target main, so each later PR's "Files changed" also contains the parts before it until those land. Reviewing the last commit of each PR gives just that part's change.

Checklist

Review and Merge Process

  1. Ping Merge Oncalls to start the process. See the PR Merge Process.
  2. Get approvals from CODEOWNERS and other reviewers.
  3. Trigger CI tests with comments or contact authorized users to do so.
    • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
  4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

CI States

Latest PR Test (Base): ❌ Run #32969058333
Latest PR Test (Extra): ❌ Run #32969057910
Latest PR Test (AMD ROCm 7.2): ❌ Run #32969058022

…store

Declare the hooks a bit-exact SWA prefix reuse needs from shared code, all with
no-op defaults, plus the flag and server arg that gate it. Nothing reads them
yet: tree_core ignores MatchPrefixParams.for_reuse, both capture hooks and the
restore hook land on the base-class no-ops, and the flag has no wiring, so this
part cannot change behaviour.

Part 1 of 5, see sgl-project#34562.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants