[AMD] Fix stale SWA ring buffer on radix prefix reuse for DeepSeek-V4 with unified_kv backend

[amd-danli103]

Motivation

The DeepSeek-V4 unified_kv backend under ROCm stores SWA KV in a per-request ring buffer, addressed by req_pool_idx * window + pos % window. Unlike the compressed/global KV, this ring is not content-addressed and is never written into the radix tree.

Radix prefix caching reuses matched prefix pages across requests. For the index-addressed compressed KV this is safe (content-stable). But the SWA ring slots covering the reused-prefix region still hold whatever the previous occupant of that req_pool slot wrote — i.e. stale SWA from an unrelated request.

When a request reuses a cached prefix and its decode sliding window (the trailing window = 128 tokens) reaches back into that reused-prefix region, the SWA path reads stale ring contents, leading to wrong attention output.

This affects only the unified_kv layout. The default index-addressed SWA pool (triton/tilelang backend) is content-stable and unaffected.

Modifications

• BasePrefixCache.swa_reprefill_tail_tokens() -> int: new base method returning 0 (no-op for all layouts).
• SWARadixCache.swa_reprefill_tail_tokens(): override returning sliding_window_size only when the unified_kv_triton backend is active on HIP (is_unified_kv_triton()), else 0.
• Scheduler prefix-match paths (schedule_batch.py, schedule_policy.py): cap the radix match length by input_len - reprefill_tail, so the trailing sliding window is held back from prefix reuse and re-prefilled into this request's own ring. The decode window therefore reads freshly-written data.

The scheduler cap is generic — it just calls tree_cache.swa_reprefill_tail_tokens(), so any cache implementing the method benefits, and it is a strict no-op for every other layout/backend (base returns 0).

Scope / activation: the fix engages only when all of the following hold:

• SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton on HIP and the model has a sliding window.
• radix prefix caching is enabled (when disabled the cache is ChunkCache/SWAChunkCache, which inherits the base 0 — and with no prefix reuse there is no stale-ring read to fix);
• --enable-hierarchical-cache off and SGLANG_ENABLE_UNIFIED_RADIX_TREE off (these route to UnifiedRadixCache, owned by [AMD] Enable unified-KV HiCache on DeepSeek-V4 #29417) and the experimental C++ radix tree off;

In every other configuration swa_reprefill_tail_tokens() returns 0 and this is a complete no-op.

Accuracy Tests

GSM8K does not surface this bug (that's also why we didn't notice this bug before):

The root cause only manifests when several conditions hold at the same time. It's this conjunction that makes the bug so easy to miss:

1. Prefix reuse — a request's prefix is served from the radix cache (a real cache hit).
2. Short uncached tail — the tokens after the reused prefix number fewer than the sliding window W (tail < W), so during decode the sliding-window attention still needs keys/values that live inside the reused prefix.
3. Stale ring — the per-request SWA slots that the reused prefix maps to were, in the meantime, overwritten with different content by another request that recycled the same req_pool slot. The SWA ring is addressed by req_pool_idx * W + pos % W and is never stored in the radix tree, so it is not content-stable.

So GSM8K structurally cannot surface this bug: it breaks (2) and rarely meets (3). Empirically it's unchanged before/after (0.945–0.950, within run-to-run noise) — a no-regression guard, not a detector.

Determinism harness (the actual detector)

At temperature=0, prefix-cache reuse must not increase output divergence beyond the model's inherent non-determinism floor. We make this rigorous and self-calibrating: for each uncached tail length we send K=64 identical prompts and report how many fall outside the majority cluster.

• tail >= window (128): the decode window cannot reach the reused region → these rows are the non-determinism floor (in-run control).
• tail < window: the decode window reaches into the reused prefix → the stale-ring read is exposed.

Result (dsv4 Pro, unified_kv backend, TP8/DP8, --page-size 256, radix on, MTP):

tail	BEFORE (3rounds)	AFTER (3rounds)
16	0 / 0 / 0	0
32	4 / 6 / 17	0
64	0 / 0 / 0	1/0/0
96	0 / 0 / 0	0
127	2 / 1 / 0	0
160 / 220	0	0–1

Conclusion: Within tail<W, which specific tail fires is scheduling-dependent; tail=32 reproduces most reliably. The headline is not any single cell but: every tail<W row can exceed the ≥W floor BEFORE, and all return to the floor AFTER.

How to reproduce?

1. Launch server with dsv4 unified_kv backend + radix cache enabled

export TRITON_CACHE_AUTOTUNING=1
export TRITON_CACHE_DIR=/sgl-workspace/.triton_cache/
export SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton

export SGLANG_OPT_FP4_MOE_K_ALIGN=128
export SGLANG_DSV4_INDEXER_K_CACHE_PRESHUFFLE=1

export SGLANG_DEFAULT_THINKING=1
export SGLANG_DSV4_REASONING_EFFORT=max
export SGLANG_OPT_DEEPGEMM_HC_PRENORM=false
export SGLANG_USE_AITER=1
export SGLANG_USE_ROCM700A=1
export SGLANG_OPT_USE_FUSED_COMPRESS=true
export SGLANG_OPT_FP8_WO_A_GEMM=false
export SGLANG_OPT_USE_JIT_INDEXER_METADATA=false
export SGLANG_OPT_USE_TOPK_V2=false
export SGLANG_OPT_USE_AITER_INDEXER=true
export SGLANG_OPT_USE_TILELANG_INDEXER=false
export SGLANG_OPT_USE_TILELANG_MHC_PRE=false
export SGLANG_OPT_USE_TILELANG_MHC_POST=false
export SGLANG_FP8_PAGED_MQA_LOGITS_TORCH=1
export SGLANG_OPT_USE_FUSED_COMPRESS_TRITON=true

export SGLANG_OPT_USE_MULTI_STREAM_OVERLAP=false
export SGLANG_ROCM_USE_MULTI_STREAM=false

export AITER_BF16_FP8_MOE_BOUND=0
MODEL=/data/models/DeepSeek-V4-Pro
python3 -m sglang.launch_server \
    --model-path /data/models/DeepSeek-V4-Pro \
    --trust-remote-code \
    --swa-full-tokens-ratio 0.1 \
    --tp 8 \
    --dp 8 \
    --speculative-algorithm EAGLE \
    --speculative-num-steps 3 \
    --speculative-num-draft-tokens 4 \
    --speculative-eagle-topk 1 \
    --enable-dp-attention \
    --enable-cache-report \
    --attention-backend dsv4 \
    --max-running-request 1024 \
    --page-size 256 \
    --chunked-prefill-size 32768 \
    --mem-fraction-static 0.9 \
    --port 30001 \
    --disable-shared-experts-fusion \
    --tool-call-parser deepseekv4 \
    --reasoning-parser deepseek-v4 \
    --enable-prefill-delayer

2. The trigger is scheduling-dependent, so the magnitude varies run-to-run; run the script several times.
   python3 swa_ring_determinism.py 30001 /path/to/DeepSeek-V4-Pro

#!/usr/bin/env python3
# Deterministic reproduction for the DeepSeek-V4 unified_kv SWA-ring reuse bug.
#
# This harness is reference-free and self-calibrating:
#   * For each `tail`, send K identical prompts and report how many fall OUTSIDE
#     the majority output cluster (`off_majority`) and the number of distinct
#     outputs (`distinct`).
#   * `tail <  W` : the decode window reaches back into the REUSED prefix -> the
#     stale-ring read is exposed  (BUG-EXPOSED rows).
#   * `tail >= W` : the decode window cannot reach the reused region -> these rows
#     measure the inherent non-determinism FLOOR (in-run control).
#
# The churn filler MUST be >= W tokens: the SWA ring is addressed by
# `req_pool_idx * W + pos % W`, so a filler shorter than W only overwrites the
# first few slots and fails to pollute the window the reused prefix maps to.
#
import sys, requests, concurrent.futures as cf
from collections import Counter
from transformers import AutoTokenizer

PORT = int(sys.argv[1]) if len(sys.argv) > 1 else 30001
MODEL = sys.argv[2] if len(sys.argv) > 2 else "/data/models/DeepSeek-V4-Pro"
BASE = f"http://127.0.0.1:{PORT}"
W = 128          # sliding window, hard-coded to 128
K = 64           # identical requests per tail

tok = AutoTokenizer.from_pretrained(MODEL, trust_remote_code=True)
FILLER = ("The history of mathematics spans thousands of years and many cultures. "
          "Ancient civilizations developed counting, geometry, and astronomy. ")

def gen(prompt, mx=60):
    j = requests.post(f"{BASE}/generate", json={
        "text": prompt,
        "sampling_params": {"temperature": 0.0, "max_new_tokens": mx},
    }, timeout=600).json()
    return j["text"], j.get("meta_info", {}).get("cached_tokens", -1)

def build(tag, n):
    ids = tok(tag + " " + FILLER, add_special_tokens=False)["input_ids"]
    out = []
    while len(out) < n:
        out.extend(ids)
    return tok.decode(out[:n])

def churn(n=500):                       # distinct >=W-token requests recycle req_pool
    with cf.ThreadPoolExecutor(max_workers=32) as ex:   # slots and overwrite the ring
        list(ex.map(lambda i: gen(build(f"C{i}", 256), 1), range(n)))

def run(tail):
    P = build("FIXEDTAG", 512 + tail)   # distinctive prefix; uncached tail = `tail`
    gen(P); churn()                     # warm prefix into radix, then pollute ring slots
    with cf.ThreadPoolExecutor(max_workers=K) as ex:
        outs = [t for t, _ in ex.map(lambda i: gen(P), range(K))]
    c = Counter(outs)
    off_majority = K - c.most_common(1)[0][1]
    zone = "<W  (bug-exposed)" if tail < W else ">=W (floor/control)"
    print(f"tail={tail:3d} {zone:20s} distinct={len(c):2d}  off_majority={off_majority}/{K}")

if __name__ == "__main__":
    print(f"K={K} identical reqs per tail; temp=0; churn filler=256 (>=W)")
    for tail in [16, 32, 64, 96, 127, 160, 220]:
        run(tail)

Speed Tests and Profiling

The cap re-prefills the trailing sliding window (128) of an otherwise-cached prefix, rounded up to page granularity — i.e. bounded by window + page_size, a constant independent of context/prefix length.

Measured with generated-shared-prefix (shared system prompt = 2048, 64 groups × 16 prompts = 1024 req, output = 256, TP8/DP8, unified_kv, radix on, page-size 256, concurrency 64; steady-state of 3 runs; --cache-report on):
Client: python3 -m sglang.bench_serving --port 30001 --dataset-name generated-shared-prefix --gsp-num-groups 64 --gsp-prompts-per-group 16 --gsp-system-prompt-len 2048 --gsp-question-len 64 --gsp-output-len 256 --max-concurrency 64 --cache-report

tail ≥ W (question_len=128) — cap is a no-op:

metric	BEFORE	AFTER
Total cached tokens	2.09M	2.09M (identical)
Output tok/s	3200	3206
Mean TTFT (ms)	510	520
Mean TPOT (ms)	17.2	17.2

tail < W (question_len=64) — cap active (this is the fix's worst case):

metric	BEFORE	AFTER	perf diff
Total cached tokens	~2.09M	~1.85M	−11.7% (≈1 page/hit re-prefilled)
Output tok/s	3247	3188	−1.8%
Mean TTFT (ms)	494	525	+31 ms
Mean TPOT (ms)	17.12	17.35	+1.3%

The drop in cached tokens (2.09M → 1.85M) confirms the fix path is actually exercised; even so the throughput cost is ≤2% and TTFT +~30 ms. In real traffic only a fraction of requests have tail < W, so the aggregate overhead is smaller.

Relationship with #29417 (HiCache for unified_kv)

This PR and #29417 are complementary and intentionally non-overlapping in ownership:

	This PR	#29417
Cache class	SWARadixCache (radix-only)	UnifiedRadixCache (HiCache host-offload)
Scope	unified_kv + radix, HiCache off	unified_kv + HiCache
SWA tail fix	SWARadixCache override	UnifiedRadixCache override (its own path)

• BasePrefixCache returns 0, so UnifiedRadixCache inherits a no-op here and gets its own override in [AMD] Enable unified-KV HiCache on DeepSeek-V4 #29417 — no functional conflict.
• Both PRs touch the same scheduler cap sites. Since the cap is generic (tree_cache.swa_reprefill_tail_tokens()), whichever merges first provides it; the other should drop its duplicate scheduler edit and keep only its cache-class override. @1am9trash

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 #28841909733
Latest PR Test (Extra): 🚫 Run #28847534328

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[HaiShaw]

/tag-and-rerun-ci

[amd-danli103]

@amd-bot ci-status

[amd-bot]

@amd-danli103

CI Status for PR #30339

Merge verdict: ⛔ Not ready — do not merge on green. Two independent problems: (1) PR CI is incomplete — the AMD pipeline is still running (3 stage-c-*-8-gpu-amd mi325 shards queued) and the Base pipeline fast-fail-skipped all downstream base-b-* jobs after base-b-test-1-gpu-large (4) failed. (2) The changed code is not exercised by PR CI at all — the behavioral change is a no-op unless SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton on HIP, which only the nightly-only AMD DeepSeek-V4 suites set. Every executed failure I inspected is unrelated to this PR (infra/pre-existing), so there is no red attributable to the PR — but "green" would prove nothing here.

Caution

This PR's changed code is not exercised by any PR-CI test. The new logic in swa_radix_cache.swa_reprefill_tail_tokens() returns non-zero only when is_unified_kv_triton() is true (is_hip() and SGLANG_HACK_FLASHMLA_BACKEND == "unified_kv_triton"). The only tests that set that env var are test/registered/amd/test_deepseek_v4_*.py, all registered nightly=True (suites nightly-amd-8-gpu-mi35x-deepseek-v4-*), which do not run on PR CI. For every PR-CI test, swa_reprefill_tail_tokens() returns 0 and the diff is a pure no-op. Green does NOT verify this fix. Before merge, run one of the nightly DeepSeek-V4 unified_kv suites (e.g. test/registered/amd/test_deepseek_v4_pro_fp4.py on an 8-GPU MI35x host with SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton), ideally with a radix-prefix-reuse scenario that reproduced the stale SWA ring.

Caution

PR CI is incomplete. AMD (PR Test (AMD)) is still running — 3 stage-c-test-large-8-gpu-amd (linux-mi325-8gpu-sglang) shards are queued; not tested yet. In Base (PR Test Base), wait-for-base-b + base-b-test-1-gpu-large (4) fast-failed, so base-b-test-1-gpu-small (6/7/8), base-b-test-1-gpu-large (6/7), and base-b-test-2-gpu-large (2/3) were skipped, not tested. Wait for AMD to finish and address/rerun the base-b root cause before relying on this pipeline.

Changed files: schedule_batch.py (+8), schedule_policy.py (+7/-1), mem_cache/base_prefix_cache.py (+3), mem_cache/swa_radix_cache.py (+21)

Executed CI failure attribution: AMD: 3 executed failures (0 related) + 3 shards still pending · Others: 6 root failures (0 related) + fast-fail cascades collapsed. Every executed failure is on a code path the PR does not touch.

AMD Executed Failures

Job	Test File	Test Function	Error	Related?	Why
stage-c-test-large-8-gpu-amd-mi35x (mi35x-gpu-8, 1)	N/A	N/A	ERROR: VRAM usage exceeds threshold (5%)	🟢	Runner pre-check — leftover VRAM on the host, not a test; PR doesn't touch GPU allocation
stage-b-test-1-gpu-small-amd-mi35x	N/A	N/A	ERROR: VRAM usage exceeds threshold (5%)	🟢	Same runner-infra pre-check
stage-b-test-2-gpu-large-amd	test/registered/model_loading/test_load_weights_from_remote_instance.py	test_load_weights_from_remote_instance	AssertionError: sgl_dp_1_dst_params rank 1 — Parameters not close	🟢	Remote-weight-loading DP correctness; no SWA/radix/schedule path involved

Still pending (not tested): stage-c-test-large-8-gpu-amd (linux-mi325-8gpu-sglang, 0/1/3).

Other Executed Failures

Job	Test File	Test Function	Error	Related?	Why
base-b-test-1-gpu-large (4)	test/registered/kernels/test_dsa_indexer.py	test_topk_fused_backends_equivalence	AssertionError: topk_v2_plan must be preprocessed per forward	🟢	DSA indexer topk kernel; unrelated to SWA radix cache. This is the base-b fast-fail root cause
build-test (Arm64)	test/registered/cpu/test_norm.py	TestFusedQKGemmaRMSNorm::test_fused_qk_gemma_rmsnorm[_with_gate]	AttributeError: 'sgl_kernel' object has no attribute 'fused_qk_gemma_rmsnorm_with_gate_cpu'	🟢	Missing CPU kernel op in the build; PR touches no sgl_kernel/norm code
build-test (xeon-gnr, base-b-test-cpu)	test/registered/cpu/test_norm.py	same as above	same AttributeError	🟢	Same missing-CPU-kernel issue (same cluster)
stage-b-test-1-gpu-xpu	test/registered/xpu/test_intel_xpu_backend.py	(server startup)	timeout after 1200s	🟢	XPU backend timeout; unrelated to AMD SWA path
stage-b-test-4-npu-a3	test_ascend...	TestAscendW4A4::test_gsm8k	test failure + RPC failed / GnuTLS recv error (modelscope download)	🟢	NPU-specific + network infra
multimodal-gen-test-2-npu-a3	sglang/multimodal_gen/test/server/ascend/test_server_2_npu.py	TestDiffusionServerTwoNpu::test_diffusion_generation[flux_2_image_t2i_2npu]	test failure + download errors	🟢	NPU diffusion server; unrelated
stage-b-test-16-npu-a3	N/A	N/A	RPC failed; HTTP 504 … cache-service pypi	🟢	Pure infra: pip cache download failure
single-node-poc (ascend qwen3_6_27b…)	N/A	N/A	HttpClient timeout	🟢	Infra timeout

Collapsed fast-fail / aggregator jobs (not independent failures): wait-for-base-b, pr-test-finish, finish, pr-test-npu-finish, pr-test-amd-extra-finish, and the skipped base-b-* shards (6/7/8, large 6/7, 2gpu 2/3) — all cascades of the failures above.

Details / what to do before merge

• Verify the fix for real (highest priority). PR CI never runs the changed path. Run a nightly AMD DeepSeek-V4 unified_kv suite on MI35x with SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton, exercising a radix-prefix-reuse case, and confirm the stale-SWA-ring symptom is gone. Consider whether a lighter PR-CI-runnable regression test (even a unit test asserting swa_reprefill_tail_tokens() behavior / the key_limit capping) could be added so this doesn't silently regress.
• Wait for AMD to finish. The 3 queued stage-c-*-8-gpu-amd (mi325) shards must complete; do not read the AMD pipeline as green until then. The two VRAM pre-check failures are runner hygiene — a rerun on a clean host should clear them.
• base-b / build / XPU / NPU failures are pre-existing or infra, none attributable to this PR (all searched: no open PR references topk_v2_plan; the fused_qk_gemma_rmsnorm_with_gate_cpu op was added for HIP in closed #27656 and is simply absent on the CPU/Arm build). If full base-b signal is needed, rerun after the DSA-indexer and CPU-norm issues are fixed on main, or apply bypass-fastfail (sparingly) to force downstream base-b jobs to run.

---

Generated by amd-bot using Claude Code CLI

[amd-danli103]

/rerun-failed-ci

[amd-danli103]

hi @HaiShaw , per discussion, I did some manual tests locally on an 8*MI355X server for this PR.

Nightly test_deepseek_v4_pro_fp4.py with radix cache ENABLED

Config = registered nightly test/registered/amd/test_deepseek_v4_pro_fp4.py env (SGLANG_DEFAULT_THINKING=1, SGLANG_DSV4_REASONING_EFFORT=max, SGLANG_USE_ROCM700A=0, SGLANG_DP_USE_GATHERV=1, SGLANG_DSV4_FP4_EXPERTS=true, AITER_BF16_FP8_MOE_BOUND=0; --tp 8 --page-size 256 --swa-full-tokens-ratio 0.1 --chunked-prefill-size 8192 --mem-fraction-static 0.90), but with radix cache ENABLED——To actually exercise the changed path, I enabled radix caching.

• GSM8K accuracy (8-shot, 1319 q) — no regression: both above the 0.92 nightly threshold; delta within run-to-run noise.

Accuracy	Invalid
BEFORE (no fix)	0.946
AFTER (fix on)	0.950

• SWA-ring determinism harness (K=64 identical prompts, temp=0, 3 rounds) —— fixed
  tail<W(128) = decode window reaches the reused prefix (bug-exposed);
  tail>=W = control/floor.

tail	off_majority BEFORE (R1/R2/R3)	off_majority AFTER (R1/R2/R3)
16	0 / 0 / 0	0 / 0 / 0
32	13 / 12 / 9	0 / 0 / 0
64	4 / 1 / 3	0 / 0 / 0
96	0 / 0 / 0	0 / 0 / 0
127	0 / 0 / 0	0 / 0 / 0
160	0 / 0 / 0 (floor)	0 / 0 / 0
220	0 / 0 / 0 (floor)	0 / 0 / 0
off_majority = # of the 64 outputs outside the majority cluster.

Conclusion: BEFORE, tail<W rows (esp. tail=32) diverge well above the >=W floor; AFTER, every row returns to the floor (0/64).

Nightly test_deepseek_v4_pro_fp4.py with radix cache DISABLED

Pls Note: the registered nightly suites launch with --disable-radix-cache, under which this fix is a strict no-op, behavior is identical to baseline.

• GSM8K accuracy (8-shot, 1319 q) — no regression

Accuracy	Invalid
AFTER (fix on)	0.954

[HaiShaw]

gated by is_unified_kv_triton