Unify full→SWA index translation in init_forward_metadata; drop pool caches by ch-wan · Pull Request #27091 · sgl-project/sglang

ch-wan · 2026-06-02T22:52:06Z

Motivation

The full→SWA index translation (translate_loc_from_full_to_swa) was triggered in two inconsistent styles:

Write path (DSV4 KV store): translated out_cache_loc lazily at the first SWA layer via DeepSeekV4TokenToKVPool.get_cached_swa_loc (env-gated by SGLANG_OPT_CACHE_SWA_TRANSLATION), with the result cached across layers.
Read path (attention): SWAKVPool translated on demand and memoized with a (data_ptr, numel) cache, kept consistent with a invalidate_loc_cache() call sprinkled through allocators, cuda-graph runners, spec draft runners, and model_runner.

This unifies the design: compute the translation once per forward in the attention backend's metadata init and store it on the backend metadata, then drop the lazy/per-layer caches and the invalidate_loc_cache machinery that existed to keep them coherent across cuda-graph capture/replay.

Modifications

DSV4 write path: computed once in init_forward_metadata_in_graph, cached on metadata. The SWA-translated KV-store write target is computed in the in-graph init op and stored on DSV4AttnMetadata.swa_out_cache_loc. Because the compute is recorded inside the cuda graph, replay re-reads the live out_cache_loc buffer (spec-v2 and DP padding rebind out_cache_loc after the out-graph prep). Multi-step draft decode mirrors the eager init's per-step slice. The field lives in assign_fields (recomputed every forward, not copied across replays) and _CP_GLOBAL_FIELDS. Applied to both deepseek_v4_backend.py and deepseek_v4_backend_hip_radix.py.
Consumers resolve through get_swa_out_cache_loc(forward_batch). All five KV-store sites (model fused stores ×2, _compute_kv_to_cache, backend store_cache ×2 twins) call one resolver: it returns the cached value only when provably current (non-idle and length matches out_cache_loc), and otherwise falls back to store-time translation — exactly the pre-cache behavior. This keeps every path correct even when the in-graph init never runs: eager IDLE (forward_idle skips attn init but the nextn model still stores dummy tokens), the draft-extend graph runners (EAGLEDraftExtendCudaGraphRunner / multi-layer / frozen-kv only run the out-graph prep), and any batch re-padded after init.
Eager init timing is unchanged from main. An earlier revision deferred the eager DP draft-extend init to after prepare_mlp_sync_batch, but that fed every backend's init a padded batch state NSA/DSA cannot consume (deep_gemm schedule-meta batch mismatch on the DeepSeek-V3.2 --dp 8 --enable-dp-attention EAGLE test), so it was reverted. The pre-pad staleness that motivated it (the deepep 4×B200 DSV4-Flash "expected 64 but got 62" crash) is instead handled by the resolver's length check, which falls back to store-time translation when DP padding rebinds out_cache_loc after init.
int32 conversion happens where metadata is built, not in translate_*. translate_loc_from_full_to_swa / translate_loc_to_hisparse_device return plain int64 lookups. Backends whose kernels require int32 convert at the metadata-build site: flash_mla (swa_out_cache_loc, swa_page_indices), FA3 (swa_page_table — eager paths previously passed int64 and crashed flash_attn_with_kvcache), trtllm-gen (_maybe_translate_swa — int64 was silently misread as int32, collapsing gpt-oss GPQA accuracy), and dsa/indexer hisparse page tables. int64-tolerant backends (triton, flashinfer, torch-native) take the mapping dtype as-is.
Removed the now-redundant caching: get_cached_swa_loc, SGLANG_OPT_CACHE_SWA_TRANSLATION, the SWAKVPool (data_ptr, numel) cache, and invalidate_loc_cache — including all call sites in cuda_graph_runner, piecewise_cuda_graph_runner, breakable_cuda_graph_runner, the eagle/mtp draft runners, model_runner, and triton's update_sliding_window_buffer. The per-layer gather is captured correctly by cuda graph, so the invalidate-before-capture dance is no longer needed.
Tests: deleted four test/manual/core tests that asserted the removed cache machinery; updated test_swa_alloc_extend_page_estimation.py and the dsv4_attention test kit for the new helper signatures; added direct unit coverage for the four resolver branches (no-metadata, current-cached, stale-shape, idle).

Accuracy Tests

Run against this branch (PYTHONPATH pointed at the worktree so the changes are actually exercised):

test/registered/attention/unittests/ — 185 passed, 21 skipped, 548 subtests passed (includes cuda-graph SWA decode/verify, DSV4 capture/replay, C4/C128, dsa, dense, mla, production eagle draft / draft-extend runners, and the new resolver tests).
test/registered/unit/mem_cache/ (swa) — 614 passed.
CI: the deepep 4×B200 DeepSeek-V4-Flash EAGLE test (the only environment exercising eager-idle and eager-DP draft-extend) and the gpt-oss / trtllm accuracy suites pass on the final revision.

Speed Tests and Profiling

No dedicated benchmark. The change removes per-iteration cache bookkeeping and the invalidate_loc_cache calls; DSV4 now does the SWA write-target translation once per forward in the metadata init rather than lazily per first-SWA-layer. Paths that fall back to store-time translation (eager idle, draft-extend graph runners) match the pre-cache behavior.

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.

🤖 Generated with Claude Code

CI States

Latest PR Test (Base): 🚫 Run #26912813528
Latest PR Test (Extra): ❌ Run #26912814327

gemini-code-assist · 2026-06-02T22:52:10Z

Warning

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

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 1363d59a24

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

ch-wan · 2026-06-02T23:06:04Z

/tag-and-rerun-ci

ch-wan · 2026-06-03T18:28:48Z

/rerun-test test_deepseek_v4_flash_fp4_b200.py

github-actions · 2026-06-03T18:29:18Z

Results for /rerun-test test_deepseek_v4_flash_fp4_b200.py:

🚀 4-gpu-b200 (1 test): ❌ View workflow run

cd test/ && python3 registered/models_e2e/test_deepseek_v4_flash_fp4_b200.py

ch-wan · 2026-06-03T20:14:12Z

/rerun-test test_deepseek_v4_flash_fp4_b200.py

This reverts commit 3ea4b29861ee15c9da0b1ecc77e85a98efed7860. Deferring the eager draft-extend init to after prepare_mlp_sync_batch changed the init-time batch state for EVERY backend, and NSA/DSA cannot consume the padded batch: core fields are padded (batch_size reassigned, seq_lens padded) while spec_info / extend_seq_lens_cpu are not, so the deep_gemm fp8_paged_mqa_logits schedule metadata disagrees with its companion tensors ("_batch_size == batch_size" assert on the DeepSeek-V3.2 --dp 8 --enable-dp-attention EAGLE test). Restore main's pre-pad init timing. The DSV4 store crash this commit originally addressed is fully handled by get_swa_out_cache_loc: the pre-pad cached value fails the length check against the re-padded out_cache_loc and the store falls back to store-time translation. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

PR sgl-project#27091 (Unify full->SWA index translation; drop pool caches) drops all `token_to_kv_pool.invalidate_loc_cache()` calls because the SWA loc-cache itself was removed from the memory pools. cg-refactor inherited those call sites from the pre-split runner files; reapplied the deletion to our split files: - model_runner.py: the decode-replay call site - runner/decode_cuda_graph_runner.py: capture run_once - speculative/eagle_draft_extend_cuda_graph_runner.py: capture run_once - speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py: capture run_once The modify/delete conflicts on legacy `breakable_cuda_graph_runner.py` / `piecewise_cuda_graph_runner.py` were resolved by keeping the deletion (cg-refactor split those into runner/ + runner_backend/). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

CI (base-c-test-8-gpu-h200 / test_dsa_glm5_dp_mtp: GLM-5 dsa backend, --dp 8, EAGLE NextN) crashed in the draft-extend forward with deep_gemm '_batch_size == batch_size' from the DSA indexer's fp8_paged_mqa_logits schedule_meta. This is exactly the failure #27091 documented and reverted: re-planning a DP-padded batch feeds NSA/DSA a padded state its indexer schedule_meta cannot consume. The two replan_equivalent=True opt-in sites (eagle_info_v2 draft-extend, multi_layer_eagle_worker per-step) did precisely that post-pad re-plan. Drop the opt-in: draft-extend keeps its marked pre-pad metadata (no re-plan), matching the proven skip_attn_backend_init=True behavior; trtllm-MLA's defensive backstop continues to cover pre-pad-then-padded staleness. The plan-record + replan_equivalent mechanism stays but now has no production caller (dormant until a backend can tolerate a post-pad re-plan). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Same DSA fix as the eagle_info_v2 / mlv1 sites: a DP-padded re-plan of the draft-extend breaks the DSA indexer's schedule_meta (#27091). Keep the marked pre-pad metadata. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Incoming changes from origin/main include: - PR sgl-project#27192 retired DecodeInputBuffers/PrefillInputBuffers in favor of CudaGraphBufferRegistry. cg-refactor still uses both: the typed DecodeInputBuffers wrapper backs the registry's source=... adoption. Resolved the runner/decode_cuda_graph_runner.py conflict to keep cg-refactor's DecodeInputBuffers.create(...) + share_buffers() pattern while pulling in upstream's new _allocate_decode_buffers helper and its NgramEmbeddingInfo / envs references. - PR sgl-project#26676 moved SWATokenToKVPoolAllocator to mem_cache.allocator.swa; updated kv_canary/api.py and flashinfer_backend.py imports accordingly while keeping cg-refactor's cuda_graph_config imports. - PR sgl-project#27091 SWA loc-cache deletions in model_runner.py already merged in the previous merge round; the modify/delete conflicts on the legacy breakable_/piecewise_cuda_graph_runner.py files were resolved by keeping cg-refactor's deletion (split into runner / runner_backend). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Sync unified-KV attention branch with main (109 commits). Conflicts resolved in: - deepseek_v4_backend_hip_radix.py: keep unified methods (_attach_unified_kv_decode_streams/_forward_unified_kv) AND main's new get_swa_out_cache_loc (sgl-project#27091 unified full->SWA translation). - deepseek_v4.py: non-unified SWA-store now uses backend.get_swa_out_cache_loc (main sgl-project#27091) instead of the removed pool get_cached_swa_loc; unified path keeps its own ring addressing. - deepseek_v4_memory_pool.py: keep get_unified_kv; drop _should_cache_swa/ cached_loc pool caches (removed by main sgl-project#27091). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

merrymercy · 2026-06-08T21:49:28Z

+                    # The compress kernel requires an int32 write location.
+                    out_loc = compress_kv_pool.translate_loc_to_hisparse_device(
+                        out_loc
+                    ).to(torch.int32)


we should prefer int64 for any indices.

PR #27091 promised to compute the full->SWA out_cache_loc translation once per forward in the attention backend's metadata init and pass it to the KV store, but it only did so for the DeepSeek-V4 path. The general SWAKVPool path kept re-translating out_cache_loc inside SWAKVPool.set_kv_buffer once per SWA layer (the (data_ptr, numel) memo cache having been removed without a replacement). Finish the unification for every SWAKVPool-using backend: - SWAKVPool.set_kv_buffer takes a pre-translated swa_loc and uses it directly for SWA layers; it asserts swa_loc is provided (no internal translate). - Each backend translates out_cache_loc once per forward and stores it on its forward metadata as swa_out_cache_loc, mirroring how swa_page_table is cached. Eager paths compute it in init_forward_metadata; cuda-graph paths bind a pre-allocated buffer that is refilled from the live out_cache_loc before each replay (so replay re-reads live slots), never a reassigned tensor. This includes FA's topk>1 target-verify branch. - Covered backends: FlashAttention, Triton, FlashInfer, trtllm-gen, aiter, xpu (consistent by construction, no separate cuda-graph metadata path), and musa (inherits FlashAttention). Adds CPU unit coverage for SWAKVPool.set_kv_buffer (swa_loc passthrough + assert). Validated e2e on gpt-oss-20b (FA3 + SWAKVPool + cuda graph) and via the swa/{triton,flashinfer,torch_native} attention unittests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

…caches (sgl-project#27091) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

ch-wan requested review from Fridge003, HaiShaw, Qiaolin-Yu, Ying1123, alphabetc1, hanming-lu, hebiao064, hnyls2002, hzh0425, ispobock, kpham-sgl, merrymercy, pyc96, xiezhq-hermann and yizhang2077 as code owners June 2, 2026 22:52

github-actions Bot added deepseek hicache Hierarchical Caching for SGLang labels Jun 2, 2026

chatgpt-codex-connector Bot reviewed Jun 2, 2026

View reviewed changes

Comment thread python/sglang/srt/layers/attention/deepseek_v4_backend.py Outdated

Comment thread python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py Outdated

github-actions Bot added the run-ci CI: run the baseline test suite on this PR label Jun 2, 2026

ch-wan force-pushed the cheng/swa-translate-int64 branch from 2d85528 to 866b5a0 Compare June 3, 2026 00:45

sglang-bot mentioned this pull request Jun 3, 2026

CUDA Coredump Tracker #26340

Open

github-actions Bot added the blackwell SM100/SM120 label Jun 3, 2026

ch-wan force-pushed the cheng/swa-translate-int64 branch 2 times, most recently from 4ef5862 to 5deade1 Compare June 3, 2026 08:25

ch-wan removed the run-ci CI: run the baseline test suite on this PR label Jun 3, 2026

ch-wan mentioned this pull request Jun 3, 2026

[refactor] Retire DecodeInputBuffers / PrefillInputBuffers in favor of CudaGraphBufferRegistry #27192

Merged

5 tasks

ch-wan merged commit c9ca56d into main Jun 3, 2026
182 of 202 checks passed

ch-wan deleted the cheng/swa-translate-int64 branch June 3, 2026 23:12

kkHuang-amd mentioned this pull request Jun 4, 2026

[AMD][WA] force to use gate_mode interleaved to fix tp2/tp4/tp8 acc issue #27201

Merged

5 tasks

ch-wan mentioned this pull request Jun 4, 2026

Replace skip_attn_backend_init with a batch-carried attention plan marker (+ staleness re-plan) #27193

Merged

5 tasks

amd-bot mentioned this pull request Jun 4, 2026

[CI Monitor] Daily Report - 2026-06-04 bingxche/sglang-ci-bot#93

Open

amd-bot mentioned this pull request Jun 5, 2026

[CI Monitor] Daily Report - 2026-06-05 bingxche/sglang-ci-bot#94

Open

amd-bot mentioned this pull request Jun 7, 2026

[CI Monitor] Daily Report - 2026-06-07 bingxche/sglang-ci-bot#96

Open

merrymercy reviewed Jun 8, 2026

View reviewed changes

ch-wan mentioned this pull request Jun 9, 2026

[SWA] Cache full→SWA out_cache_loc per forward across attention backends #27617

Merged

5 tasks

ch-wan mentioned this pull request Jun 11, 2026

[DSV4] Use int64 for compressor out_loc tensors #27973

Merged

4 tasks

nvpohanh mentioned this pull request Jun 15, 2026

Performance regression on 8x B300 Llama-4-Scout after PR #27091 (AIPerf output throughput -1.85%) #28300

Closed

3 tasks

This was referenced Jun 26, 2026

[Bug] Gemma-4-26B-A4B NVFP4 quality drops #26518

Closed

Avoid dynamic Q quantization in trtllm_mha #29423

Merged

luoroger37 mentioned this pull request Jul 23, 2026

[HiCache][DSV4] Complete SWA pressure and host-pool fixes bytedance-iaas/sglang#659

Merged

Chronostasys pushed a commit to MindLab-Research/sglang that referenced this pull request Aug 24, 2026

Unify full→SWA index translation in init_forward_metadata; drop pool …

d66703e

…caches (sgl-project#27091) Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Navigation Menu

Unify full→SWA index translation in init_forward_metadata; drop pool caches - #27091

Unify full→SWA index translation in init_forward_metadata; drop pool caches#27091
ch-wan merged 9 commits into
mainfrom
cheng/swa-translate-int64

ch-wan commented Jun 2, 2026 •

edited by github-actions Bot

Loading

Uh oh!

gemini-code-assist Bot commented Jun 2, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Uh oh!

Uh oh!

Uh oh!

ch-wan commented Jun 2, 2026

Uh oh!

ch-wan commented Jun 3, 2026

Uh oh!

github-actions Bot commented Jun 3, 2026 •

edited

Loading

Uh oh!

ch-wan commented Jun 3, 2026

Uh oh!

Uh oh!

merrymercy Jun 8, 2026

Uh oh!

ch-wan Jun 11, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

ch-wan commented Jun 2, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Speed Tests and Profiling

Checklist

CI States

Uh oh!

gemini-code-assist Bot commented Jun 2, 2026

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

ch-wan commented Jun 2, 2026

Uh oh!

ch-wan commented Jun 3, 2026

Uh oh!

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

Uh oh!

Uh oh!

ch-wan commented Jun 3, 2026

Uh oh!

Uh oh!

merrymercy Jun 8, 2026

Choose a reason for hiding this comment

Uh oh!

ch-wan Jun 11, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ch-wan commented Jun 2, 2026 •

edited by github-actions Bot

Loading

github-actions Bot commented Jun 3, 2026 •

edited

Loading