[AMD] Enable unified-KV HiSparse on ROCm for DeepSeek-V4 - #29168
amd-danli103 wants to merge 21 commits into
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements unified-KV HiSparse support on ROCm for DeepSeek-V4, allowing the compressed C4 KV cache to reside directly within the unified pool's rows. Key updates include removing the temporary ROCm guard, introducing HiSparseUnifiedC4DevicePool to alias the unified compressed region, updating the swap-in logic to use the linear MLA path, and preventing runtime reallocations of DSA decode metadata buffers. The review feedback highlights a critical indexing bug in compressor_v2.py where unmapped slots (indicated by -1) could corrupt the SWA ring, recommends raising a RuntimeError to guard against post-capture buffer reallocations in dsa_backend.py, and suggests initializing inherited KVCache attributes in HiSparseUnifiedC4DevicePool to prevent potential AttributeErrors.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| out_loc = ( | ||
| token_to_kv_pool.c4_kv_pool._translate_loc_to_hisparse_device( | ||
| self.forward_metadata.core_metadata.c4_out_loc | ||
| ) | ||
| ) | ||
| out_loc = out_loc + token_to_kv_pool.unified_swa_pages |
There was a problem hiding this comment.
If c4_out_loc contains -1 (representing an invalid or unmapped slot), _translate_loc_to_hisparse_device will return -1 (as the last element of the mapping is -1). Doing out_loc + token_to_kv_pool.unified_swa_pages will then result in -1 + unified_swa_pages. Since unified_swa_pages is positive, this becomes a valid positive index (e.g., unified_swa_pages - 1), which is the last slot of the SWA ring. This will cause the store kernel to write the compressed KV of the unmapped/invalid token into the SWA ring, corrupting the SWA ring's KV cache.
To prevent this, use torch.where to only add unified_swa_pages to valid (non-negative) indices, keeping -1 as -1.
out_loc = (
token_to_kv_pool.c4_kv_pool._translate_loc_to_hisparse_device(
self.forward_metadata.core_metadata.c4_out_loc
)
)
out_loc = torch.where(
out_loc >= 0,
out_loc + token_to_kv_pool.unified_swa_pages,
out_loc,
)There was a problem hiding this comment.
c4_out_loc = tl.where(c4_should_compress, raw_out_loc // 4, 0) so c4_out_loc won't contain -1. Assumption does not make sense.
| self.data_ptrs = torch.tensor( | ||
| [x.data_ptr() for x in self.kv_buffer], | ||
| dtype=torch.uint64, | ||
| device=self.device, | ||
| ) |
There was a problem hiding this comment.
HiSparseUnifiedC4DevicePool inherits from HiSparseC4DevicePool (which inherits from DeepSeekV4SingleKVPool -> KVCache), but its __init__ intentionally skips calling super().__init__ to avoid allocating a second device buffer. As a result, several attributes defined in KVCache.__init__ (such as self.custom_mem_pool, self.layer_transfer_counter, self.enable_memory_saver, and self.memory_saver_adapter) are left uninitialized. If other parts of the codebase interact with HiSparseUnifiedC4DevicePool expecting it to be a fully initialized KVCache instance, this could lead to AttributeErrors. We should explicitly initialize these missing attributes to safe defaults to ensure robust compatibility.
| self.data_ptrs = torch.tensor( | |
| [x.data_ptr() for x in self.kv_buffer], | |
| dtype=torch.uint64, | |
| device=self.device, | |
| ) | |
| self.data_ptrs = torch.tensor( | |
| [x.data_ptr() for x in self.kv_buffer], | |
| dtype=torch.uint64, | |
| device=self.device, | |
| ) | |
| self.custom_mem_pool = None | |
| self.layer_transfer_counter = None | |
| self.enable_memory_saver = False | |
| from sglang.srt.mem_cache.memory_pool import get_memory_saver_adapter | |
| self.memory_saver_adapter = get_memory_saver_adapter(False) |
9337e0e to
7c21ef9
Compare
Removed "preventing runtime reallocations of DSA decode metadata buffers" - this potential bug fix from this PR. Rationale: that fix is orthogonal to this feature. It will be submitted as a separate PR. |
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
|
hi @1am9trash or @HaiShaw , could u pls help to trigger the CI test? Thank you! |
clintg6
left a comment
clintg6
left a comment
There was a problem hiding this comment.
Verified all 19 tests pass locally: JIT 10/10, unit 5/5, TP8 e2e 2/2 (GSM8K 0.952), TP8+DP8 e2e 2/2 (GSM8K 0.950). Long-context swap path confirmed working.
All runtime changes are safely gated behind _is_hip / unified_hisparse so no impact on the NVIDIA path.
Gemini findings reviewed: the -1 index concern in compressor_v2.py is a false positive (c4_out_loc defaults to 0 via tl.where, not -1). The skipped super().init in HiSparseUnifiedC4DevicePool is intentional and safe as the uninitialized KVCache attributes are never accessed on this pool.
Hi @clintg6 , thank you so much for the thorough verification and the approval! The results aligned with my local results. |
|
@amd-bot ci-status |
CI Status for PR #29168Merge verdict: 🚫 Do not merge on CI signal — PR CI never ran. All 20 "failures" are the Caution This PR's changed code is not exercised by any PR-CI test — every test job was skipped because the gate failed. The Changed files: Executed CI failure attribution: AMD: 0 executed failures · Others: 0 executed failures — all failures are the gate-block cascade, not executed tests. No test job ran in any workflow. Root cause (single, shared across all workflows)
Every downstream stage ( Coverage gap — what is currently unverifiedThis PR adds a new memory pool (
None of these ran. Until the gate passes, there is zero signal on the changed code. What to do before merge
Generated by amd-bot using Claude Code CLI |
7c21ef9 to
194e347
Compare
|
/tag-and-rerun-ci |
|
@amd-bot ci-status |
CI Status for PR #29168Merge verdict: No executed CI failure is attributable to this PR — every red job is an unrelated pre-existing/flaky/infra failure (FP-comparator flake, perf thresholds, a pytest-missing build env, a stale test signature). The new code's unit + JIT tests passed on AMD (the target backend). BUT PR CI did not verify the actual feature: the two new end-to-end accuracy evals are Caution This PR's headline value — DeepSeek-V4 Pro FP4 HiSparse unified-KV accuracy on ROCm — is exercised only by the two new tests registered Warning NVIDIA PR CI is incomplete: in run 28346854737, Changed files: Executed CI failure attribution: AMD: 4 executed failures (0 related) · Others: 4 executed failures (0 related) · plus 2 fast-fail/gate cascades. NVIDIA AMD Executed Failures
Other Executed Failures
Details / what to do before merge
Generated by amd-bot using Claude Code CLI |
|
/rerun-failed-ci |
…nts) New MI355X 2-node 1P1D leg exercising both memory features in one PD run: HiCache on the prefill role, HiSparse on the decode role, KV transfer over MoRI, unified-KV layout. The two features live on opposite PD roles and do not conflict (HiCache is prefill-only, HiSparse is decode-only), so a single recipe per model covers both. Added for all four DSV4 variants: mi355x-fp8/dsv4pro, mi355x-fp8/dsv4flash, mi355x-fp4/dsv4pro, mi355x-fp4/dsv4flash Two launcher changes the recipes depend on: - Env ordering: a recipe's prefill/decode_extra_env is now applied after the hardcoded DSV4 env so it can override it (needed to pin SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton and the MoRI host-registration knobs). - disable_radix_cache knob: --disable-radix-cache was hardcoded into all four COMMON_FLAGS strings; HiCache L1 *is* the radix cache and the combination is a hard ValueError. The recipe field disable_radix_cache: false now drives a RADIX_FLAG (defined before the role if/else so both branches see it under set -u), so HiCache can run on the prefill role. The recipes run against a checkout tree containing three not-yet-merged PRs (SGLANG_USE_CHECKOUT_RUNTIME=1): sgl-project#29168 (unified-KV HiSparse on ROCm), sgl-project#32368 (HiSparse PD over MoRI, based on sgl-project#29168), and sgl-project#36966 (device-alias pointer fix so both features stop GPU-faulting at PD warmup on non-SVM host kernels). Not wired into nightly-configs.yaml until they land. Validation: the DSV4-Pro-FP4 config (deepseek-ai/DeepSeek-V4-Pro) was validated end-to-end on the hand-driven path earlier (HiSparse decode GSM8K 0.960, HiCache prefill GSM8K 0.950, both 0 GPU faults). Through the CI launcher, boot + PD KV transfer (generate 200) + hisparse/mori-live were confirmed; the full GSM8K gate run under checkout-runtime is pending an image-native run once the dependency PRs merge (checkout-runtime JIT-compiles aiter on first start, which desyncs decode/prefill startup and races the bench probe -- an artifact that disappears with SGLANG_USE_CHECKOUT_RUNTIME=0). Accuracy gate is a placeholder pending that number.
Follow-up to sgl-project#30315 (pool sizing) and complementary to sgl-project#29168 (unified-KV HiSparse device pool). This adds the runtime accounting so scheduling and allocation treat the unified SWA pool as a fixed per-request ring instead of a linearly-consumed token pool. On the unified path SWA is addressed by state_slot (== req_pool_idx) inside the DSV4 kernels; the paged SWA indices / full_to_swa_index_mapping are never consumed, so accounting SWA as a linear token pool over-throttles admission and decode retract. The real bound is concurrency (num_req_slots), already enforced by req_to_token_pool. All changes are gated on get_kvcache()._unified_kv; the non-unified (fp8) path is unchanged. - swa.py: account SWA as a fixed per-request ring slot (swa_ring_cost_tokens) bounded by num_req_slots (via req_to_token_pool). available_size / swa_available_size / new_pages_available / alloc_extend / alloc_decode take the unified branch and skip the vestigial paged SWA allocator + mapping. - schedule_policy.py: rem_swa_tokens / _swa_budget_for_req use the ring-based swa_available_size and charge one fixed ring slot per request; skip the linear SWA clamp and the double-charge on chunked continuations. - schedule_batch.py: unified-KV check_decode_mem diagnostic ([SWA-BOTTLENECK]). - hisparse.py: forward swa_ring_cost_tokens through the DSV4-HiSparse wrapper (the unified HiSparse C4 device pool itself is sgl-project#29168). - model_runner: wire req_to_token_pool into the SWA allocator.
Follow-up to sgl-project#30315 (pool sizing) and complementary to sgl-project#29168 (unified-KV HiSparse device pool). This adds the runtime accounting so scheduling and allocation treat the unified SWA pool as a fixed per-request ring instead of a linearly-consumed token pool. On the unified path SWA is addressed by state_slot (== req_pool_idx) inside the DSV4 kernels; the paged SWA indices / full_to_swa_index_mapping are never consumed, so accounting SWA as a linear token pool over-throttles admission and decode retract. The real bound is concurrency (num_req_slots), already enforced by req_to_token_pool. All changes are gated on get_kvcache()._unified_kv; the non-unified (fp8) path is unchanged. - swa.py: account SWA as a fixed per-request ring slot (swa_ring_cost_tokens) bounded by num_req_slots (via req_to_token_pool). available_size / swa_available_size / new_pages_available / alloc_extend / alloc_decode take the unified branch and skip the vestigial paged SWA allocator + mapping. - schedule_policy.py: rem_swa_tokens / _swa_budget_for_req use the ring-based swa_available_size and charge one fixed ring slot per request; skip the linear SWA clamp and the double-charge on chunked continuations. - schedule_batch.py: unified-KV check_decode_mem diagnostic ([SWA-BOTTLENECK]). - hisparse.py: forward swa_ring_cost_tokens through the DSV4-HiSparse wrapper (the unified HiSparse C4 device pool itself is sgl-project#29168). - model_runner: wire req_to_token_pool into the SWA allocator.
Co-authored-by: Cursor <cursoragent@cursor.com>
Main now rejects new files under test/registered/amd/ and GPU suites under unit/.
HiSparse aliases bf16 rows; fp8 splits nope/rope and needs its own adapter. Pin the live C4 out_loc remap that sits next to main's fp8_2buff store.
_init_compressed_pools was overwriting HiSparseUnifiedC4DevicePool with kv_pools[4] (None on the unified path), so the allocator assert fired at server start.
Keep unified-KV HiSparse C4 across DSV4.1 indexer paging and encoder replay. Port the nightly eval job into nightly-test-amd.yml.
|
/rerun-failed-ci |
Keep unified-KV linear swap beside MiniMax M3 HiSparse; paged FP8 dsv4 kernels stay gated on is_dsv4_paged_layout.
main's NetworkAddress.to_url() already prepends http://.
This branch has not been deployed
4 participants
Motivation
On ROCm, HiSparse currently runs only on the separate-KV DeepSeek-V4 path (
SGLANG_HACK_FLASHMLA_BACKEND=triton, packed FP8 KV layout), enabled by #26639. The unified-KV backend (SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton) — where compressed C4 KV lives in the unified pool'srows[swa_pages:]and the HiSparse hot buffer is a bf16 view into that region backed by a linear host cold pool — was blocked by a temporary startup guard.This PR enables unified-KV + HiSparse on ROCm end-to-end, and adds JIT/unit/eval coverage.
Modifications
hisparse_coordinator.py,deepseek_v4_memory_pool.py,compressor_v2.py,arg_groups/hisparse_hook.py,srt/mem_cache/allocator/hisparse.py): remove the temporary guard that blocked unified_kv_triton + HiSparse, wire the unified C4 device pool / linear host cold pool, and fix the compressed-C4 remap to account for the pre-applied swa_pages offset.test/registered/kernels/ops/kvcache/test_hisparse.py).deepseek_v4_memory_pool.py,mem_cache/allocator/hisparse.py): upstreampool_configuratoralready foldshost_to_device_ratio = Sintobytes_per_full_tokenas a1/(4S)C4 term and raisesmax_total_num_tokensaccordingly, but the unified pool ignored it. This sizes the on-device C4 region to the shrunk budget viac4_compress_pages, and stopsavailable_size()/full_available_size()from clamping capacity to the device C4 pool — on the unified path the cold remainder is host-resident, so the bound is the logical full-token pool. Net effect: freed VRAM actually reachesfull_token. Separate-KV / CUDA behaviour is unchanged.test/registered/unit/mem_cache/test_hisparse_unified_pool.py): host-mirror geometry, view aliasing of the unified compressed region, c4-only layer views, alloc/free mapping, oversubscribe and that c4_compress_pages actually shrinks the device C4 region while the c128 layers are untouched..../mi35x/..._unified_eval_mi35x.pyand..._unified_dp_eval_mi35x.py): GSM8K accuracy + a ~19k-token long-context swap-path retrieval test, on TP8 and on TP8+DP-attention.Highlight — the NVIDIA CUDA forward path is unchanged. All modifications are confined to ROCm-only branches:
_is_hipguards and HiSparse/unified-KV-gated code. (The JIT test change only removesskipif(is_hip())markers, so those tests now also run on ROCm; CUDA behavior is unaffected.)To enable this feature, you can pass e.g.
in the server launch configs.
Accuracy Tests
GSM8K few-shot accuracy alone does not validate the host↔device swap path. With
device_buffer_size=2048(compressed-token budget) and GSM8K prompts of ~1k–1.7k tokens (~250–425 compressed after the 4x C4 ratio), every request stays on the fast path (host_len <= device_buffer_size→ preloaded to the device hot buffer), so decode is all device-resident hits and the cold→hot swap-in (miss-copy) kernel is never exercised. GSM8K therefore validates: sparse top-k selection numerics, unified-KV pool wiring/aliasing, the compressor write path, and HIP-graph decode stability — but not swap correctness.Note: a nonzero #cpu token in the decode log during GSM8K is the authoritative host backup of compressed C4 (every request is backed up to host), not decode-time swap-in. The swap-in (miss-copy) path is gated per-request by
host_len > device_buffer_size; GSM8K's per-request footprint (~250–425 compressed) stays belowdevice_buffer_size=2048, so it runs entirely on the preloaded hot path.Swap correctness is established by, in increasing scope:
test/registered/kernels/ops/kvcache/test_hisparse.py):test_dsv4_swap_in_reads_paged_host_layout,*_miss_copy_layout, and the LRU/miss cases check the swap-in kernel byte-for-byte against a reference.test_b_long_context_swap): a ~19k token prompt (~4750 compressed > 2048) forces a host cold pool and a swap-in on every decode step; a passcode buried in the long context can only be retrieved if cold→hot copies land correctly. This is the authoritative end-to-end swap oracle.Accuracy of record: The numbers reported here (TP8 0.953 / TP8 DP8 0.949) were reproduced locally on
deepseek-ai/DeepSeek-V4-Pro.Tests reproduction guidance
JIT kernel (11 passed)
These now also run on ROCm (the
skipif(is_hip())markers were removed). The swap-in kernel cases (test_dsv4_swap_in_reads_paged_host_layout,*_miss_copy_layout, LRU/miss) compare the host→device cold-copy byte-for-byte against a reference, so a pass proves the swap-in kernel is numerically correct.unit test (6 passed)
Validates the unified-KV pool: host-mirror geometry, the bf16 view aliasing into
rows[swa_pages:], C4-only layer views, alloc/free index mapping, and oversubscribe handling.e2e test (2 passed each)
Each file runs two methods. What proves the feature is live:
(a) Server startup — the unified-KV HiSparse host cold pool is allocated, and the config is applied:
The
dsv4_hisparse_c4host pool only exists on the HiSparse path; its presence withunified_kv_tritonconfirms the unified hot/cold wiring is taken.(b)
test_a_gsm8k— accuracy / sparse-selection guard. Expect the method to PASS with GSM8K accuracy ≈ 0.95 (assertion gate> 0.91).(c)
test_b_long_context_swap— the swap-path guard, and the authoritative evidence that host↔device swap works. A ~19k-token prompt with a passcode buried at the top is sent; expect the method to PASS with the passcode echoed back:The corresponding server decode line is the smoking gun — a single long request whose compressed C4 footprint far exceeds
device_buffer_size=2048:#gpu token: 2049≈device_buffer_size(2048) → the device hot buffer is capped at the configured budget.#cpu token: 9566(> 0) → the remaining cold C4 tokens live in the host pool and are streamed in on every decode step.--dp 8 --enable-dp-attention).Speed Tests and Profiling
This is a capacity feature, not a latency one, and it is worth being explicit about the trade rather than reporting a throughput number that misses the point.
host_to_device_ratio = Sshrinks the GPU-resident C4 region to1/Sof the budget and mirrors the remainder on the host. Upstream_get_bytes_per_full_tokenalready carries the resulting1/(4S)C4 term, so the freed VRAM shows updirectly as a larger
max_total_num_tokens. #32368 measures 12.90M → 17.56M (+36%) on the same 1P1D topology.The cost is decode-step latency: the swap-in runs on the main stream between the indexer and attention of every C4 layer. It is latency/occupancy-bound, not bandwidth-bound —
LaunchKernelusesgrid = batch size, so a per-rank decode batch of ~20 leaves ~9% occupancy on 256 CUs, which is why #32368'sdevice_buffer_sizesweep (2048 → 16384, ~8x fewer bytes moved) made decode monotonically worse. The levers are multi-CTA sharding (#31341) andside-stream overlap (#28523), neither of which is in this PR's scope.
When to turn this on: HiSparse trades decode latency for token capacity, so it is a loss whenever the KV working set of your concurrent requests already fits in VRAM — you pay the swap-in cost and get nothing back. It pays off only past
the point where the non-HiSparse configuration can no longer admit the load, at which point the comparison is not "slower" but "runs at all".
CI coverage
PR-CI:
test/registered/kernels/ops/kvcache/test_hisparse.py(HiSparse JIT kernels; HIP skip removed so it now runs on ROCm)test/registered/unit/mem_cache/test_hisparse_unified_pool.py(unified-KV C4 device pool)These give per-PR coverage of the unified-KV + HiSparse pool/kernel paths.
Nightly e2e (8-GPU TP8 / TP8 DP8,
nightly=True):nightly-amd-8-gpu-mi35x-deepseek-v4-pro-hisparse-unifiednightly-amd-8-gpu-mi35x-deepseek-v4-pro-hisparse-unified-dpThese two suites are registered but not yet wired into any workflow, so they do not run on PR or on the nightly schedule yet. Following the precedent of the unified-KV PR #27380, the PR-CI/nightly wiring is tracked in the AMD CI coverage issue #27521 (cc @bingxche @yctseng0211). I will add them to the issue list then.
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ❌ Run #35813712809
Latest PR Test (Extra): ❌ Run #35813712476
Latest PR Test (AMD ROCm 10): ❌ Run #35813712792