[DSA] k-pool: 256-token logical page so index page id = logical page id by kpham-sgl · Pull Request #42178 · sgl-project/sglang

kpham-sgl · 2026-10-02T03:07:28Z

Motivation

Fixes #41494: GLM-5.3-Flash (DSA with index_kpool=4) loses the NIAH needle at 16K when prefixes are shared. Supersedes #41645.

A pooled index key covers 4 consecutive tokens, i.e. 4 physical KV pages of 64. Before this PR the radix tree page was 64, so a prefix hit could end inside a pooled-key group. The resuming request then wrote its pooled keys into the group-head page owned by another request's cached prefix, and later readers of that prefix saw the wrong keys.

Modifications

Logical page = 64 × index_kpool (256) for the allocator, radix tree, PD and HiCache. KV kernels and the DSA backend's real_page_size keep 64-token physical pages. Logical page L owns physical pages 4L .. 4L+3.
Index page id = logical page id. Each index page holds 64 pooled keys (slots_per_page). Pooled key j of a request lives at index slot req_to_token[r, 4j] // 4, and the index page table is page_table_64[:, ::4] // 4. Every group-head lookup in the k-pool kernels and planners now divides by POOL_SIZE.
Index pool sized per pooled key: one key per index_kpool tokens, which frees about 9% more KV capacity on H200 (see below). The HiCache host declaration is sized the same way.
--dcp-size > 1 with index_kpool > 1 is rejected at startup, because a pooled key spans positions owned by different DCP ranks.
Every change is a no-op for index_kpool == 1. DeepSeek-V4's pool falls back to page_size (_index_page_slots).

Accuracy Tests

H200×8, TP8, tilelang DSA, bf16 KV, same flags for both builds.

	main	this PR
NIAH 16K, 20 samples, random needle-depth order (RULER-style)	11/20	20/20
NIAH 16K, 20 samples, increasing depth	19/19, then crash (see below)	20/20
Prefix-hit lengths	multiples of 64 (960, 2496, ...)	multiples of 256
`max_total_num_tokens`	3,531,072	3,861,760 (+9.4%)
`test_glm53_flash_h200.py` LowLatency (EAGLE) GSM8K, repeated runs	0.942 / 0.930 / 0.926 (mean 0.933)	0.934 / 0.924 / 0.924 / 0.936 (mean 0.930)
LowLatency accept length / bs1 speed	—	4.14 / pass
HighThroughput (attn-DP 8) GSM8K	—	0.944

The GSM8K difference is within run-to-run noise (one run's std ≈ 0.011); main also dips below the 0.93 threshold.

Unit tests: the new test_kpool_index_addressing.py fails 3/3 on main and passes here. The existing k-pool kernel and unit tests pass with the updated page geometry.

Known issues, also present on `main`

HiCache does not restore the indexer for hybrid mamba + DSA models: _MambaStrategy never assembles the indexer host pool and logs HybridLinearKVPool declares host pools ['indexer'] that _MambaStrategy did not assemble. NIAH after a cache flush loses needles on host-restored prefixes. This PR does not change that path.
Intermittent CUDA illegal memory access in decode at about 16K+ tokens (copy_done.synchronize). Intermittent on both main and this PR, about 1 in 3 shuffled-NIAH runs with EAGLE, and also seen once without EAGLE on main; not tied to a specific sample or length.

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.

CI States

Latest PR Test (Base): 🚫 Run #37269693427
Latest PR Test (Extra): ⏳ Run #37289097492
Latest PR Test (AMD ROCm 10): ❌ Run #37269693516

For DSA with index_kpool > 1 (GLM-5.3-Flash), the allocator / radix tree page is now 64 * index_kpool (256) while KV kernels keep 64-token physical pages. Pooled key j of a request lives at index slot loc(j * index_kpool) // index_kpool, so index page id == logical page id. Prefix-cache hits are therefore whole pooled-key groups, and a request never writes pooled keys into a page owned by another request's cached prefix (sgl-project#41494). The index pool shrinks to one pooled key per index_kpool tokens, freeing KV capacity. DCP with index_kpool > 1 is rejected at startup. All changes are no-ops for index_kpool == 1. Co-authored-by: Cursor <cursoragent@cursor.com>

Co-authored-by: Cursor <cursoragent@cursor.com>

ispobock · 2026-10-02T15:29:05Z

ref: #40134

Co-authored-by: Cursor <cursoragent@cursor.com>

ispobock · 2026-10-04T15:30:35Z

Comparison with #40134: #40134 (comment)

kpham-sgl · 2026-10-05T06:45:03Z

/rerun-test test/registered/e2e/models/test_glm53_flash_h200.py

github-actions · 2026-10-05T06:45:28Z

🚀 8-gpu-h200 (1 test): ✅ View workflow run

cd test/ && python3 registered/e2e/models/test_glm53_flash_h200.py

ispobock · 2026-10-05T09:05:02Z

-                overrides["page_size"] = 64
-                logger.warning("Setting page size to 64 for DeepSeek DSA.")
+                index_kpool = get_dsa_index_kpool(hf_config)
+                if index_kpool > 1 and getattr(cfg, "dcp_size", 1) > 1:


dcp_size is a declared server arg, use cfg.dcp_size directly instead of getattr.

ispobock · 2026-10-05T09:05:02Z

+                        f"index_kpool={index_kpool}: a pooled index key spans "
+                        "positions owned by different DCP ranks."
+                    )
+                page_size = 64 * index_kpool


add a comment on why * index_kpool: a logical page has to own a whole pooled index group, so a prefix hit never ends inside one.

ispobock · 2026-10-05T09:05:02Z

+                        "positions owned by different DCP ranks."
+                    )
+                page_size = 64 * index_kpool
+                overrides["page_size"] = page_size


validate_mamba_extra_buffer checks mamba_track_interval % page_size before this override sets page_size, so --mamba-track-interval 64 or 128 now passes startup. maybe check it here next to the dcp check.

ispobock · 2026-10-05T09:05:02Z

+        assert self.page_size % index_kpool == 0, (
+            f"page_size {self.page_size} must be a multiple of index_kpool {index_kpool}"
+        )
+        self.slots_per_page = self.page_size // index_kpool


nit: slots_per_page does not say which buffer or unit. maybe index_page_size (pooled keys per index page). fine as a follow-up rename.

ispobock · 2026-10-05T09:05:02Z

        0, page_table_64.shape[-1], pool_size, device=page_table_64.device
    )
-    return page_table_64[..., idx]
+    return page_table_64[..., idx] // pool_size


nit: pool_size reads like a memory pool size, it is the k-pool ratio. maybe kpool, same follow-up.

ispobock · 2026-10-05T09:05:02Z

        )
        self.dsa_index_topk = get_dsa_index_topk(hf_config)
        self.dsa_index_kpool = get_dsa_index_kpool(hf_config)
+        self.real_page_size = model_runner.page_size // self.dsa_index_kpool


nit: maybe physical_page_size instead of real_page_size, matching the local in DSATokenToKVPool.__init__ and the physical/virtual page naming DCP already uses. same follow-up.

ispobock · 2026-10-05T09:07:42Z

Verified on GLM-5.3-Flash. A prompt is cached up to the middle of a 256-token group, then two requests extend it with different suffixes. The scored continuation copies a span from the first request's suffix, so it depends on the index keys the second request overwrites on main.

copy mean logprob    before 2nd request    after 2nd request
main                 -0.33 to -0.18        -8.16 to -7.94
this PR              -0.32 to -0.17        -0.30 to -0.20

HiCache host restore of the indexer on hybrid models is still missing, as noted in the known issues. With #40915 on top of this PR:

copy mean logprob    device            after host restore
this PR              -0.28 to -0.23    -11.82 to -10.94
with #40915          -0.30 to -0.22    -0.28 to -0.22

kpham-sgl · 2026-10-05T19:19:07Z

/rerun-failed-ci

Since sgl-project#42178, slots_per_page = page_size // index_kpool, but the sharded pool translates indexer locations in token units. With index_kpool > 1 that would address the wrong rows, so fail at construction instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

github-actions Bot added deepseek hicache Hierarchical Caching for SGLang jit-kernel memory-pool labels Oct 2, 2026

Drop the k-pool getattr fallbacks and the arithmetic-only test

8d98818

Co-authored-by: Cursor <cursoragent@cursor.com>

ispobock added run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) labels Oct 2, 2026

kpham-sgl and others added 2 commits October 2, 2026 18:16

Merge branch 'main' into kpham/dsa-kpool-index-by-loc

c3e3455

Give hand-built DSA pools the index page the host decl reads.

f1731ec

Co-authored-by: Cursor <cursoragent@cursor.com>

ispobock mentioned this pull request Oct 4, 2026

Store index cache per page independently and enable hicache offload for GLM-5.3-Flash #40134

Closed

Merge branch 'main' into kpham/dsa-kpool-index-by-loc

2b90fe9

kpham-sgl added the highest-priority CI: all three control labels, plus never batch-cancelled or stale-closed label Oct 5, 2026

ispobock reviewed Oct 5, 2026

View reviewed changes

kpham-sgl added run-ci-extra CI: also run the extra suite (requires run-ci) and removed run-ci-extra CI: also run the extra suite (requires run-ci) labels Oct 5, 2026

ispobock approved these changes Oct 5, 2026

View reviewed changes

kpham-sgl merged commit 61ba984 into sgl-project:main Oct 5, 2026
307 of 353 checks passed

kpham-sgl mentioned this pull request Oct 5, 2026

[DSA] Clarify pooled-page naming and ownership #42671

Merged

leonardHONG mentioned this pull request Oct 6, 2026

[Fix] Correct pooled page IDs in target verify replay #42722

Closed

mmangkad mentioned this pull request Oct 6, 2026

[DSA] Fix pooled page table refresh for target verify CUDA graph replay #42750

Merged

nvpohanh mentioned this pull request Oct 6, 2026

[CI] test_glm53_flash_b200.py (TestGLM53FlashB200DFlash2.test_gsm8k) fails on main: GSM8K ~0.87 < 0.93 since 2026-10-06 #42749

Closed

vincentzed mentioned this pull request Oct 6, 2026

[DSA] Fix decode context parallelism for RoPE sparse attention on Blackwell #42536

Closed

4 tasks done

kpham-sgl mentioned this pull request Oct 6, 2026

[DSv4] Fix shared index cache accessor page size #42832

Merged

ormandj mentioned this pull request Oct 7, 2026

[HiCache] Budget declared DSA index host pools and keep mismatched MTP draft rows separate #42918

Open

2 tasks done

kevin-mii mentioned this pull request Oct 7, 2026

[Docs] GLM-5.3-Flash: measured MI355X MXFP4 agentic recipe #43015

Merged

2 tasks done

kpham-sgl mentioned this pull request Oct 7, 2026

feat(kda): integrate FlashInfer prefill checkpoints #41400

Merged

3 of 4 tasks

ormandj mentioned this pull request Oct 13, 2026

GLM-5.3-Flash bug tracking #37524

Open

3 of 12 tasks

Navigation Menu

Repository navigation

[DSA] k-pool: 256-token logical page so index page id = logical page id - #42178

[DSA] k-pool: 256-token logical page so index page id = logical page id#42178
kpham-sgl merged 6 commits into
sgl-project:mainfrom
kpham-sgl:kpham/dsa-kpool-index-by-loc

kpham-sgl commented Oct 2, 2026 •

edited by github-actions Bot

Loading

Uh oh!

ispobock commented Oct 2, 2026

Uh oh!

ispobock commented Oct 4, 2026 •

edited

Loading

Uh oh!

kpham-sgl commented Oct 5, 2026

Uh oh!

github-actions Bot commented Oct 5, 2026 •

edited

Loading

Uh oh!

ispobock Oct 5, 2026

Uh oh!

ispobock Oct 5, 2026

Uh oh!

ispobock Oct 5, 2026

Uh oh!

ispobock Oct 5, 2026

Uh oh!

ispobock Oct 5, 2026

Uh oh!

ispobock Oct 5, 2026

Uh oh!

ispobock commented Oct 5, 2026

Uh oh!

kpham-sgl commented Oct 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

kpham-sgl commented Oct 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

Known issues, also present on main

Checklist

CI States

Uh oh!

ispobock commented Oct 2, 2026

Uh oh!

ispobock commented Oct 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kpham-sgl commented Oct 5, 2026

Uh oh!

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

Uh oh!

Uh oh!

ispobock Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ispobock Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ispobock Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ispobock Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ispobock Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ispobock Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ispobock commented Oct 5, 2026

Uh oh!

kpham-sgl commented Oct 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kpham-sgl commented Oct 2, 2026 •

edited by github-actions Bot

Loading

Known issues, also present on `main`

ispobock commented Oct 4, 2026 •

edited

Loading

github-actions Bot commented Oct 5, 2026 •

edited

Loading