[AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated by yuttian1 · Pull Request #38192 · sgl-project/sglang

yuttian1 · 2026-09-06T07:38:18Z

Motivation

Under unified-KV, the DSV4 kernels address SWA and C4 state as per-request
rings, but the pool sizer still charged both as token-scaled pools. It reserved
memory the runtime never allocates, so full-token capacity came out well below
what the card can hold.

#30315 fixed the sizing but was reverted in #38163 because it was not fully
gated. This reland keeps the fix and puts every branch behind the unified-KV
gate; the non-unified FP8 path is untouched.

Modifications

pool_configurator: under unified-KV, SWA and C4 state are sized per request
slot (_fixed_swa_bytes, _fixed_c4_state_bytes); the remainder goes to
full KV.
deepseek_v4_memory_pool: C4 state pool becomes num_req_slots * c4_ring_size, resolved from the same gate as the sizer so the two cannot
drift.
swa allocator: swa_available_size() reports ring capacity in tokens
(free req slots x per-slot cost) instead of page counts.
schedule_policy: a request is charged one ring slot at first admission
rather than a per-chunk token budget; continuations are not charged again.
c_plan.cuh: the C4 plan uses request-ring addressing
(req_pool_idx * ring_size + pos % ring_size).

Results

Full-token capacity per DP worker:
6,848,000 -> 12,569,600 (+83.55%).
The increase comes from accounting C4 state and the SWA ring as
request-scoped allocations instead of token-scaled pools.
GSM8K accuracy:
95.2% -> 95.1% (1 question difference; no observed accuracy regression).
Serving benchmark on 8x MI355X:
no measurable throughput, TTFT, or TPOT regression within run-to-run variance
at concurrency 512 (replicated); concurrency 256 is a single run per arm.

Validation

GSM8K: 1319 questions
Serving benchmark: TP8 + DP8, 8x MI355X
Added regression tests for request-ring addressing, request-slot reuse,
unified/non-unified gating, and token-cap monotonicity.
CI: pr-gate fails on the current head, which skips the downstream suites.
Not yet diagnosed.

Caveat

The unified-KV BF16 layout is pre-existing and is outside the scope of this PR.

CI States

Latest PR Test (Base): ❌ Run #34116192994
Latest PR Test (Extra): ✅ Run #34116192774
Latest PR Test (AMD ROCm 7.2): ❌ Run #34116193084

…g" (sgl-project#30315) This reverts commit f5819b0 (sgl-project#38163), restoring sgl-project#30315 verbatim. The follow-up commit gates every remaining shared-path change behind the unified-KV switch.

The reverted sgl-project#30315 claimed "no behavior change on the non-unified (fp8) path", but three changes reached shared code with no gate at all: 1. deepseek_v4_memory_pool.py: `c4_state_pool_size = max(caller, num_req_slots * ring)` ran *before* `_unified_kv` was resolved, so the fp8 path also paid the request-addressed floor -- raising its footprint and lowering max_total_num_tokens. The gate is now resolved before any sizing, and unified sets the size exactly instead of via max(); fp8 keeps the caller-supplied SWA-addressed size byte-for-byte. 2. schedule_policy.py: the `swa_needed >= rem_swa_tokens` -> `>` relaxation at both admission sites applied to every hybrid-SWA model. It is correct only where rem_swa_tokens is an exact ring-slot capacity, so it is now conditional; the legacy SWA-token path keeps its conservative `>=`. 3. invariant_checker.py: the unconditional `logical_attn_allocator` unwrap changed which object the invariant was asserted against on every DSV4-HiSparse deployment, not just unified ones. Also tightens the two duck-typed `clear_c4_req_states` call sites (allocation.py, disaggregation/decode.py) to check `_unified_kv` explicitly rather than relying on the hook's internal early return, and skips the `newly_allocated` list comprehension entirely off the unified path. The gate is identity-compared (`is True`) so a duck-typed test stub that auto-creates attributes cannot accidentally select the unified path. Drops the unused `host_hit_len` / `storage_hit_len` parameters sgl-project#30315 added to `_update_prefill_budget`; no call site ever passed them. test_dsv4_c4_state_lifecycle.py: the unified lifecycle test built a bare MagicMock and relied on duck-typing alone, so it now declares `_unified_kv` like the pool it stands in for. Adds the matching negative test that the reset never fires off the unified path.

The reverted PR carried edits that change no behavior and only widen the review surface. Removing them shrinks the diff from 113 hunks / +860/-166 to 92 hunks / +815/-97 with no functional delta. - Restore 14 import blocks that were collapsed to one line. The repo's own ruff-format (v0.15.1, magic trailing comma) keeps the multi-line form, so the collapse did not come from repo tooling. - Collapse two byte-identical Triton branches. `elif use_req_ring:` had the same body as `if compress_ratio == 128:` in attn.py and compressor.py; both operands are tl.constexpr, so `or` folds at compile time. This is already the form compress_hip.py uses in the same series. - Drop a local `attn_head_dim` in _fixed_c4_state_bytes that duplicates the `self.attn_head_dim` the series itself introduces. - Restore seven comments in c_plan.cuh that were re-wrapped to ~80 columns. The governing .clang-format sets ColumnLimit: 120 and clang-format 20.1.7 reports both the old and the new text as clean, so the re-wrap was gratuitous; two of them broke a column-aligned address table.

alloc_extend, alloc_decode and new_pages_available all take a unified early return; alloc_extend_swa_tail did not, while new_pages_available stopped consulting num_swa_pages on that path. The vestigial paged SWA allocator is therefore called with no capacity gate, and a bare assert turns its exhaustion into a scheduler crash instead of a None backoff. Reachable via DecodePreallocQueue, which routes DSV4 page_size > 1 preallocation through this method.

Every unified-KV gate now spells the check `... is True` instead of relying on truthiness. `_unified_kv` is a plain bool on the real pool, so this is a no-op in production, but a duck-typed stub (a bare MagicMock auto-creates a truthy attribute) would otherwise select the unified path in tests. Six sites still used the truthy form: the allocator's master switch in SWATokenToKVPoolAllocator.__init__, the four CompressorHip ring-address branches, the pool-stats observer, the SWA leak-invariant skip, the swa_max_total_num_tokens reconciliation, and the decode-bottleneck diagnostic. SWATokenToKVPoolAllocator.debug_print() reported swa_attn_allocator's available size, which on the unified path is the vestigial paged allocator and contradicts swa_available_size(). It now calls swa_available_size(), so the two agree; off the unified path the two are the same expression. Also condenses four comment blocks added earlier in this series that ran past the two-line limit in .claude/rules/comment-style.md.

The two compressor entry points still gated use_req_ring on bare truthiness. Both take a DeepSeekV4TokenToKVPool, so read _unified_kv directly; the same holds for the four compress_hip sites, which sit after an isinstance assert. The XPU compress plan builder has no use_req_ring parameter, so a unified-KV plan would silently fall back to the SWA-paged layout. Assert instead. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…bute _compute_dsv4_sizes now reads self._unified, but the white-box fixture builds DSV4PoolConfigurator via object.__new__ and only set five attributes, so test_dsv4_accepts_pool_above_floor raised AttributeError. Set it, and cover the unified branch where the c4 state pool is sized per request instead. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…ies on calculate_pool_sizes subtracts the three request-scoped fixed pools from the byte budget; calculate_pool_sizes_from_max_tokens cannot, because it takes a token count and subtracting bytes there would double-count. That is only safe while every constraint in config_from_budget is a min(), which nothing checked. Assert it, state the caller contract, and add a regression test that caps a budget-derived token count and requires the total footprint to shrink. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Apply .claude/rules/comment-style.md to the branch's added comments: drop function-top preambles whose facts are already stated at the lines they constrain, drop docstrings on private helpers, and condense multi-line rationale to the single non-recoverable fact. Kept: cross-file mirrors (swa_ring_size, CompressStatePool row math), layout and unit facts, magic-number derivations, the kernel write-pad aliasing bound, and the token-count-vs-byte-budget caller contract. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

yuttian1 · 2026-09-06T13:05:54Z

Additional benchmark detail: the c=512 comparison used 7 baseline, 6 PR, and 3
baseline-control runs. The baseline-control variance is larger than the observed
PR delta, so the result is reported as "no measurable regression within
run-to-run variance".

yuttian1 · 2026-09-06T16:03:58Z

Re: memory accounting

The 11.63 GiB difference between the two arms is accounted for. The reverted
baseline charges SWA and C4 state as token-scaled pools, although the unified
runtime allocates both as request-scoped rings. This PR changes the accounting
to match the actual request-ring layout and uses the reclaimed capacity for
additional full KV tokens.

The allocation ledger matches the observed PR footprint within allocator
fragmentation. The remaining BF16 unified-KV layout cost is pre-existing and is
not changed by this PR.

A startup accounting assertion or follow-up regression test would still be
useful, but there is no unexplained PR-specific memory residual.

HaiShaw · 2026-09-06T20:38:51Z

cc @1am9trash @Duyi-Wang @amd-danli103 @hnyls2002 @xiezhq-hermann @ispobock @hebiao064 @ShangmingCai

ShangmingCai

Disaggregation part LGTM

…oc hook

…of fallback

…evictable

HaiShaw · 2026-09-07T20:12:24Z

NV CI clean, (XPU) / stage-b-test-1-gpu-xpu stuck not relevant.

…lly gated (sgl-project#38192) Co-authored-by: hnyls2002 <lsyincs@gmail.com> Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>

yuttian1 added 4 commits September 6, 2026 01:00

Reland "[AMD][DSV4] Fix unified-KV pool sizing and SWA ring accountin…

c653f2e

…g" (sgl-project#30315) This reverts commit f5819b0 (sgl-project#38163), restoring sgl-project#30315 verbatim. The follow-up commit gates every remaining shared-path change behind the unified-KV switch.

github-actions Bot added deepseek jit-kernel memory-pool labels Sep 6, 2026

yuttian1 and others added 5 commits September 6, 2026 08:15

yuttian1 marked this pull request as ready for review September 6, 2026 16:36

yuttian1 requested review from ByronHsu, Fridge003, HaiShaw, Qiaolin-Yu, Ying1123, alphabetc1, hanming-lu, hebiao064, hnyls2002, huangtingwei9988, hzh0425, ispobock, merrymercy, xiezhq-hermann and yizhang2077 as code owners September 6, 2026 16:36

github-actions Bot added the run-ci label Sep 6, 2026

HaiShaw assigned hnyls2002 Sep 6, 2026

amd-danli103 mentioned this pull request Sep 7, 2026

[AMD][DSV4] feat: enable fp8 two-pool unified_kv on gfx950 #37413

Merged

5 tasks

ShangmingCai reviewed Sep 7, 2026

View reviewed changes

Comment thread python/sglang/srt/disaggregation/decode.py

ShangmingCai approved these changes Sep 7, 2026

View reviewed changes

typed swa_req_ring flag instead of getattr; c4 reset via req pool all…

c062146

…oc hook

alphabetc1 mentioned this pull request Sep 7, 2026

[CI] Fix the DSpark dp-tier unit test fixture after #34919 #38314

Merged

hnyls2002 added 5 commits September 7, 2026 03:21

shared get_swa_ring_size helper

fbbb2bd

revert comment-only kernel hunks; ring mode asserts req pool instead …

0fa9334

…of fallback

hoist req_ring_state in compress_hip; drop is True on bool

0e595f4

trim comments

c2e9b8d

ring mode: size_swa is the ring total; PD swa budget ignores phantom …

8e005cf

…evictable

hnyls2002 approved these changes Sep 7, 2026

View reviewed changes

hnyls2002 added bypass-fastfail run-ci-extra labels Sep 7, 2026

Merge branch 'main' into reland/dsv4-unified-c4-request-ring

d09055a

HaiShaw approved these changes Sep 7, 2026

View reviewed changes

HaiShaw merged commit 570087c into sgl-project:main Sep 7, 2026
185 of 206 checks passed

This was referenced Sep 8, 2026

[AMD] Avoid the FP8 wo_a path when the weight is BF16 #38453

Draft

[CI Monitor] Daily Report - 2026-09-08 bingxche/sglang-ci-bot#189

Open

ankith117 mentioned this pull request Sep 8, 2026

[AMD][DSV4] Skip the paged SWA page return under the per-request ring #38571

Merged

5 tasks

YAMY1234 mentioned this pull request Sep 8, 2026

[DSv4] Integrate TRT-LLM DSv4 Attention for SM100/103 #30805

Merged

5 tasks

bingxche mentioned this pull request Sep 9, 2026

[CI Monitor] Daily Report - 2026-09-09 bingxche/sglang-ci-bot#190

Open

amd-danli103 mentioned this pull request Sep 9, 2026

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

Open

5 tasks

curnane-lab mentioned this pull request Sep 9, 2026

[Refactor] Price hybrid SWA pools from a declarative pricing table #38644

Open

5 tasks

bingxche mentioned this pull request Sep 10, 2026

[CI Monitor] Daily Report - 2026-09-10 bingxche/sglang-ci-bot#191

Open

xiezhq-hermann mentioned this pull request Sep 13, 2026

[HiCache] Rework the buffer-mode storage prefetch pipeline and retry bookkeeping #39283

Open

6 tasks

Navigation Menu

Uh oh!

[AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated - #38192

[AMD][DSV4] Reland unified-KV pool sizing and SWA ring accounting, fully gated#38192
HaiShaw merged 16 commits into
sgl-project:mainfrom
yuttian1:reland/dsv4-unified-c4-request-ring

yuttian1 commented Sep 6, 2026 •

edited by github-actions Bot

Loading

Uh oh!

yuttian1 commented Sep 6, 2026 •

edited

Loading

Uh oh!

yuttian1 commented Sep 6, 2026 •

edited

Loading

Uh oh!

HaiShaw commented Sep 6, 2026

Uh oh!

Uh oh!

ShangmingCai left a comment

Uh oh!

HaiShaw commented Sep 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Uh oh!

Conversation

yuttian1 commented Sep 6, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Results

Validation

Caveat

CI States

Uh oh!

yuttian1 commented Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

yuttian1 commented Sep 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Re: memory accounting

Uh oh!

HaiShaw commented Sep 6, 2026

Uh oh!

Uh oh!

ShangmingCai left a comment

Choose a reason for hiding this comment

Uh oh!

HaiShaw commented Sep 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

yuttian1 commented Sep 6, 2026 •

edited by github-actions Bot

Loading

yuttian1 commented Sep 6, 2026 •

edited

Loading

yuttian1 commented Sep 6, 2026 •

edited

Loading