Repository navigation
[Feat] KVShard: Add attention and indexer support for page-sharded DSA caches - #44280
Merged
hzh0425 merged 7 commits intoOct 22, 2026
Merged
Conversation
Shunkangz
requested review from
1am9trash,
BBuf,
DarkSharpness,
Fridge003,
HaiShaw,
HydraQYH,
Qiaolin-Yu,
YAMY1234,
Ying1123,
alphabetc1,
celve,
hanming-lu,
hebiao064,
hnyls2002,
huangtingwei9988,
hubertlu-tw,
hzh0425,
ispobock,
kkHuang-amd,
merrymercy,
rainj-me,
xiezhq-hermann,
yizhang2077 and
yuan-luo
as code owners
October 14, 2026 12:21
- Sharded indexer-K writes take main's act_quant + set_index_k_scale_buffer fallback, i.e. the sharded pool's store_quantized; revert the Triton act_quant use_fused_store_math flag. - Translate the DSA page table after the CP request slice, so the logical indexer page table needs no second slice. - Drop the no-op needs_cpu_seq_lens override and skip the pool-capacity check for scratch-addressed sharded extends. - Keep the sharded DSA pool's assert text; its docstring now only notes that model-runner selection is still pending. - Tests: replace the mock routing tests and the Triton parity test with one real-kernel writer test, and keep a single 4-GPU real-indexer run. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Factor the sharded index-key write plan (stage every token in scratch, persist owned tokens locally) out of store_quantized into write_targets, and run fused_store_index_k_cache / fused_k_indexer_norm_rope_store over those targets. Sharded indexer keys now carry the same FP32 scales and bytes as an unsharded pool instead of the ue8m0 act_quant fallback, with no Triton change. The writer test compares against the unsharded CUDA writer byte for byte. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Conflicts with main's DSA decode context parallelism (sgl-project#43735): - dsa_indexer.py _store_index_k_cache: the fused writer loop over _index_k_write_targets now passes pool.index_page_size, as main's single call does. It equals page_size on a sharded pool (index_kpool == 1, no DCP). - dsa_backend.py __init__: supports_mha_one_shot needs both dcp_size == 1 (main) and no KV-shard pool (this branch). sgl-project#43735 also broke two hand-built fixtures in test_page_sharded_dsa_attention.py. Both are patched the way main patched its own: - The writer test builds DSATokenToKVPool without a published topology, and the constructor now reads attn_dcp_size. Fix: get_parallel().override(attn_dcp_size=1) in setUp. - The __new__-built backend now needs dcp_size and num_dcp_q_heads, which _forward_trtllm reads in the 4-GPU TRT-LLM case (SM100 only). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
Author
|
/tag-run-ci-label |
Contributor
|
Normal startup still needs two integration points for DSA page sharding:
The new tests construct the sharded pool directly, so they do not exercise this startup path. The pool docstring already marks selection as pending; if this is intentionally deferred, could the PR description identify the follow-up that will wire these entry points? |
Contributor
Author
Thank you for the comments. The final E2E integration will be done with next PR. In this PR, I will only modify the attention part. |
Contributor
Author
|
/rerun-failed-ci |
5 of 6 tasks
Shunkangz
changed the title
This comment has been minimized.
This comment has been minimized.
Shunkangz
changed the title
Shunkangz
changed the title
Preserve sharded index-key write targets with the renamed DSA cache kernel and update attention tests to the relocated act_quant module.
6 of 8 tasks
hzh0425
approved these changes
Oct 21, 2026
Collaborator
|
/rerun-failed-ci |
Collaborator
The failed ci is unreleated
|
zcnrex
added a commit
to zcnrex/sglang
that referenced
this pull request
Oct 23, 2026
Resolve conflicts with bf16 V4.1 index logits (sgl-project#44051, sgl-project#45880), grouped verify rows (sgl-project#45251), the two-level-only prefill graph (sgl-project#45783) and KV-shard index-K writes (sgl-project#44280): the adapter takes a logits_dtype and maps [batch, next_n] paged queries and schedules to grouped varlen rows on SM100. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
5 tasks done
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
This PR integrates DSA attention and indexer operations with page-sharded KV caches for GLM-5.2. It separates logical page IDs from scratch addresses and routes fused indexer writes to both scratch and owner-local persistent storage while preserving existing quantization behavior. DSA and indexer fixtures and tests are consolidated into the attention test module, keeping general page-sharding coverage separate.
Modifications
Accuracy Tests
Speed Tests and Profiling
Checklist
Review and Merge Process
/tag-and-rerun-ci,/tag-run-ci-label,/rerun-failed-ciCI States
Latest PR Test (Base): ⏳ Run #39713318009
Latest PR Test (Extra): ⏳ Run #39713318067
Latest PR Test (AMD ROCm 10): ❌ Run #39713317998