[AMD] Strict bit-exact SWA HiCache for DeepSeek-V4 with unified_kv: SWA-window + c4/indexer state riding across L1/L2/L3

[amd-danli103]

Motivation

DeepSeek-V4's unified_kv backend on ROCm packs three KV families in one pool: the content-addressed compressed KV (C4/C128), the per-request SWA ring (addressed by req_pool_idx * window + pos % window), and the compressor's c4 / indexer overlap state. Only the compressed KV is content-stable; the SWA ring and the c4/indexer overlap state are not, and are never written into the radix tree.

Two prior PRs made reuse of such prefixes safe but only approximately correct:

• [AMD] Enable unified-KV HiCache on DeepSeek-V4 #29417 enabled unified_kv HiCache, but it offloads only the compressed KV (C4/C128). The SWA ring is not preserved, so on a hierarchical hit the trailing sliding window is re-prefilled (swa_reprefill_tail_tokens()) rather than restored.
• [AMD] Fix stale SWA ring buffer on radix prefix reuse for DeepSeek-V4 with unified_kv backend #30339 applied the same re-prefill idea to the radix-only path: cap the match by the trailing window and recompute it into the request's own ring, so decode never reads a stale ring.

Re-prefill recomputes the trailing window (and, implicitly, the c4/indexer overlap state for that region) instead of reproducing the exact values the original prefill wrote. That is enough to prevent gross stale-ring errors and to keep GSM8K-level accuracy within noise — but it is not bit-exact: a partial re-prefill cannot reconstruct the original SWA window and c4/indexer overlap state byte-for-byte, so the reused output can differ from a full fresh compute at the bit level.

This PR closes that gap. Instead of re-prefilling an approximation, it captures and offloads the true values of the SWA window and the c4/indexer overlap state across L1→L2→L3, and restores them exactly on a hierarchical hit — making unified_kv SWA prefix reuse strict and bit-exact. A hard reuse gate re-prefills only when the truth is unavailable (never a stale/approximate read).

This is an opt-in gated by SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE (default off), and it engages only when all of the following hold at once:

• running DeepSeek-V4 on ROCm/HIP;
• the unified_kv_triton attention backend is chosen (SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton);
• HiCache is enabled (--enable-hierarchical-cache).

In every other case, the feature is a strict no-op, no pool is allocated, no capture, no restore; behavior is byte-identical to main.

Modifications

Organized as three layers:

1. Substrate (Commit 1–5):

• SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE flag + --hicache-swa-offload-page-stride arg (environ.py, arg_groups/fields/memory.py).
• Dedicated pinned-host SWA-window offload pool with capture-done handshake + upstream deferred release (memory_pool_host.py).
• Two independent, tile-start-addressed L3 pools for c4 state and indexer state, packing (deepseek_v4_memory_pool.py, hybrid_cache/hybrid_pool_assembler.py).
• Version-namespaced L3 keys (hicache_storage.py, storage/nixl/hicache_nixl.py).
• Startup fail-fast guard (hybrid_pool_assembler.py): with the strict flag on, raise a clear ValueError unless --hicache-write-policy write_through, before any host pool is pinned.

2. Capture (Commit 6–8):

• Capture the SWA ring window at each completed prefill/decode page with EAGLE-safe geometry (deepseek_v4_backend_hip_radix.py, model_runner.py, deepseek_v4.py).
• Snapshot c4/indexer overlap state for the same window into the independent state pools (compress_hip.py, compressor_v2.py).

3. Reuse/restore + strict gate (Commit 9–10):

• Ride the captured window + state back to device at the correct buf_lo offset on a hierarchical hit; strict reuse gate caps the match and re-prefills any tail that isn't provably bit-exact, never a stale read; commit-coupling guard keeps window+state+full-KV committed atomically (swa_component.py, unified_radix_cache.py, base_prefix_cache.py, schedule_batch.py, schedule_policy.py).

4. Follow-ups on the strict path:

• HiCache capture/restore now addresses c4/indexer state by request slot (translate_from_req_position_to_state_loc) instead of swa_loc, matching the per-request compressor read path from [AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated #38192 — the old addressing landed a restored window on the wrong C4 rows (compress_hip.py).
• State staging is sized from one prefix's needs — ceil(ISL/(page*stride)) + 1 windows at per-rank batch 1, ISL = min(200k, context_length) — instead of scaling with cache capacity; bounded at 1.5x durable, overridable by SGLANG_SWA_HICACHE_STATE_STAGING_PAGES (hybrid_pool_assembler.py).

Note to reviewers

Scope. Off by default: no pool is allocated, no capture or restore runs.

deepseek_v4_backend_hip_radix.py and compress_hip.py are ROCm-only by construction. In the common files, the L3 key namespace and the unified_swa_hicache pool branch are gated on the flag at startup; the rest is either a new symbol that re-routes nothing (BindCapturedSWAHost, MatchPrefixParams.for_reuse, the new arg/env fields, the new host-pool methods) or a call whose callee returns early unless strict mode and the SWA host pool are both live.

The one match-path change that could otherwise have altered another configuration — clamping the FULL device anchor to the host-gated boundary on a cross-request reuse match — is opt-in per component (TreeComponent.device_anchor_needs_reuse_clamp, False everywhere except strict SWA), so a non-strict tree matches exactly what it matched before. test_reuse_does_not_clamp_when_no_component_asks_for_it pins that.

This PR can be split into three cumulative sub-PRs on request — substrate / capture / reuse (C1–5 / C6–8 / C9–10) — each gated by the same flag and a no-op by default.

Accuracy Tests

Two detectors (like #30339): a byte-exact harness is the real detector; GSM8K is only a no-regression guard.

1. Unit Tests (6 test files, 194 UTs total)

How: cd test && python3 -m pytest -q srt/mem_cache/test_swa_*.py
Expected: 194 passed — byte-exact swa+c4_state/c4_indexer_state capture/restore across L1→L2→L3, strict-gate reject, commit-coupling guard, dirty-read protection, L3-only restore.

No regression on the shared unified-cache suite. python3 -m pytest -q registered/unit/mem_cache gives the same result on this branch as on the merge-base (2115 passed / 1351 skipped / 401 subtests; the 6 failures are pre-existing in our container — 5 need a built rust ext, 1 is a HIP JIT compile failure).

2. E2E (DeepSeek-V4 unified_kv on ROCm, HiCache + strict flag)

How: launch with the gates on and SGLANG_SWA_DBG_CHECKSUM=1 (emits a per-window checksum vs a fresh device re-read), trigger prefix reuse, grep the log.

export SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton
export SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE=1     # dev; set 0 for the "off" baseline
export SGLANG_SWA_DBG_CHECKSUM=1                  # + standard dsv4 env as in #30339
python3 -m sglang.launch_server --model-path /data/models/DeepSeek-V4-Pro --attention-backend dsv4 \
  --page-size 256 --swa-full-tokens-ratio 0.1 --enable-cache-report --port 30001 \
  --enable-hierarchical-cache --hicache-write-policy write_through \
  --hicache-io-backend direct --hicache-mem-layout layer_first \
  --tp 4 --dp 4 --enable-dp-attention  2>&1 | tee server.log
# decode + EAGLE spec path (tp8/dp8): add --speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-num-draft-tokens 4 --speculative-eagle-topk 1
# L3-only variant: add --hicache-storage-backend file --hicache-storage-prefetch-policy wait_complete
#   + SGLANG_HICACHE_FILE_BACKEND_STORAGE_DIR=/ssd/dir, populate, kill, restart same dir, resend prefix

Trigger reuse — send the same long prompt N rounds; cached_tokens climbs 0 → large:

python3 - <<'PY'
import json, urllib.request
P = "This is a long shared prefix sentence. " * 4000    # >> page_size (256) tokens
for r in range(6):                                      # 6 for prefill; 20 for decode dp8 (every rank must hit)
    b = json.dumps({"model": "x", "prompt": P, "max_tokens": 8, "temperature": 0}).encode()
    u = json.loads(urllib.request.urlopen(urllib.request.Request(
        "http://127.0.0.1:30001/v1/completions", data=b,
        headers={"Content-Type": "application/json"})).read())["usage"]
    print("round", r, "cached_tokens", (u.get("prompt_tokens_details") or {}).get("cached_tokens", 0))
PY

Expected (grep server.log):

check	grep	expect
byte-exact SWA + c4/indexer	[C4-STATE-DBG] state ride bit-exact / [LB-DEV] device landing byte-exact / [SWA-DBG] restore verified	L1/L2 same-process reuse: all three > 0. L3-only cross-restart: [LB-DEV] + [C4-STATE-DBG] > 0 ([SWA-DBG] is 0 on this path — the prefill checksum is not persisted across a process restart, so the window's byte-exact verification is carried by [LB-DEV] device-landing).
no failure / no dirty read	MISMATCH|reuse_reject|BIND-MISS|AssertionError|Traceback	0
L3-only source	prefetch .*matched=[0-9]+ loaded=[0-9]+	matched=0 loaded>0

L3-only restore validated (cold restart, empty L1/L2): prefetch matched=0 loaded=2304, the first request immediately reports cached_tokens=2304, [LB-DEV] + [C4-STATE-DBG] byte-exact all pass, zero MISMATCH.

gsm8k eval:
Two arms, each run once per server profile — off (SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE=0) and dev (=1) — with the cache flushed (/flush_cache) before each pass and temperature=0, so the only variable is the strict flag. GSM8K is a no-regression guard (the byte-exact harness above is the real correctness detector); the acceptance criterion is that dev is not below off and invalid ≈ 0.

• non-thinking

python3 benchmark/gsm8k/bench_sglang.py \
  --host 127.0.0.1 --port 30001 --num-shots 8 --num-questions 1319 --parallel 256

For off (SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE=0) and dev (=1) , all above tests we got acc ~0.94(invalid ≈ 0).

• thinking

python3 -m sglang.test.run_eval \
  --eval-name gsm8k --host 127.0.0.1 --port 30001 \
  --num-threads 256 --temperature 0.0 --repeat 1 --num-shots 8 --max-tokens 8192

For off, we got 0.933, and for dev, we got 0.942(invalid 0).

Speed Tests and Profiling

Both arms HiCache on, DBG=0, --hicache-ratio 1.0 (single variable = strict flag); customer clients. Just take some of the configs as example.

scenario	metric	off	dev	note
Prefill (dp4, 50k ISL, prefix-cache0.6)	tokens/s	≈ 32.7k	≈ 30.8k	dev ≈ −5.8%
Prefill (dp4, 50k ISL, prefix-cache0.8)	tokens/s	≈ 31.4k	≈ 29.1k	dev ≈ −7.3%
Decode TPOT (dp8, EAGLE, batch 32/48/64)	ms/token	≈ 70.6 / 82.4 / 96.8	≈ 70.8 / 83.8 / 98.4	dev slightly slower, no regression

The prefill cost is the deterministic price of strict correctness (positional H2D of the SWA window + c4/indexer state on reuse, not a content-addressed page copy), applied only with the flag on.

Optional in-repo reproduction.
The table above comes from the customer perf clients. For a self-contained check that needs no external scripts, the stock bench_serving client reproduces the strict overhead directly in-repo (it is a different, smaller workload, so absolute numbers differ). Launch the same server twice with DBG=0 — once SGLANG_UNIFIED_KV_BIT_EXACT_HICACHE=1 (dev), once =0 (off) — then run against each:

python3 -m sglang.bench_serving --backend sglang --dataset-name random \
  --random-input-len 50000 --random-output-len 200 --random-range-ratio 1 \
  --num-prompts 64 --max-concurrency 64 --seed 42 --port 30001

Compare peak output token throughput: dev trails off by a small, consistent margin (measured off ≈ 1536 tok/s vs dev ≈ 1472 tok/s, ≈ −4%), the same direction as the customer-client table. Note that at this concurrency/input length the Median/P99 TTFT is dominated by request queueing (~50–100 s) and is not a reliable reuse detector — prefix-reuse correctness is verified separately by the byte-exact harness and the climbing cached_tokens above, not by this client's TTFT.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

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 #34570730373
Latest PR Test (Extra): ❌ Run #34570730288
Latest PR Test (AMD ROCm 10): ❌ Run #34570730197

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.