[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
Open
amd-danli103 wants to merge 2 commits into
amd-danli103 wants to merge 2 commits into
Conversation
amd-danli103
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
alphabetc1,
hanming-lu,
hebiao064,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
August 18, 2026 09:18
5 tasks
…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.
amd-danli103
force-pushed
the
swa-bitexact-1n
branch
from
August 25, 2026 06:14
756627c to
3d8a74d
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
The configuration this stack addresses is DeepSeek-V4 served on the
unified_kvattention 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_decodeandcapture_compress_state_windows_decode(decode). Their docstrings state the ordering contract an implementation must honor. They are declared as base methods rather than probed withhasattrat 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 fromprepare_for_extendafterreq_pool_idxis 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(defaultFalse): distinguishes cross-request reuse matches from self-match lookups such as the one incache_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.ModelRunner, placed outside the cuda graph after the decode forward wrote the ring and before the next step overwrites it; one prefill hook indeepseek_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_reusehave no readers and every hook body ispass.test/registered/mem_cache/test_swa_bitexact_extension_points.pyasserts exactly that — the hooks exist on the base classes, returnNone, 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
#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
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #32969058333
Latest PR Test (Extra): ❌ Run #32969057910
Latest PR Test (AMD ROCm 7.2): ❌ Run #32969058022