fix(dcp): GLM-5.3-Flash norope c-KV sharding under DCP (Hopper/tilelang only) - #36989
Closed
junliu-mde wants to merge 4 commits into
Closed
junliu-mde wants to merge 4 commits into
junliu-mde wants to merge 4 commits into
Conversation
github-actions
Bot
added
quant
junliu-mde
changed the title
…using virtual locs raw On DCP the allocator issues virtual locs in [0, size * dcp_world_size). GLM-5.3-Flash is the first qk_rope_head_dim == 0 model to run DCP, and its c-KV write goes through set_mla_kv_buffer_kernel_norope, which lacked the owner-filter + divide the rope variant has. The read side (decode page table, fused topk output, extend) also used virtual locs raw. Below max_total_num_tokens this is accidentally self-consistent (replicated writes, LSE merge of identical partials); past it, reads and writes go out of bounds silently (expandable_segments absorbs the MMU faults) and quality degrades state-dependently. Four changes: 1. mla_buffer.py: norope kernel gains the same DCP localization as the rope variant (loc % dcp == rank filter, loc // dcp addressing); the wrapper takes dcp_localize so the replicated draft pool (which indexes virtual locs by design) can opt out. 2. memory_pool.py: MLATokenToKVPool.dcp_localized_writes drives the kernel flag; the torch set_kv_buffer DCP branch now also divides (it filtered without dividing); HybridLinearKVPool delegates the flag. 3. dsa_backend.py decode/verify: owner-rule page-table localization (loc % dcp == rank -> loc // dcp, else -1); the existing cross-rank LSE merge reconstitutes full topk attention. 4. dsa_backend.py extend: route sparse extend to the RAGGED topk transform and gather the rank-sharded prefix c-KV into the per-request-contiguous layout RAGGED indices already address (all_gather_kv_cache_for_mha_extend). Also scale the flattened page-table capacity assert by dcp_world_size: its consumer (the index-K buffer) is virtual-sized, so virtual locs are legal there. kv_cache_configurator.py sizes the DSA index-K buffer by max_total * dcp / loc_space_scale (it stores every token on every rank for topk scoring but was sized per-rank) and opts the replicated draft pool out of localized writes. Fixes sgl-project#36886
junliu-mde
force-pushed
the
fix/glm53-dcp-sharded-ckv
branch
from
August 29, 2026 07:46
334ba44 to
b9e5540
Compare
Author
|
Lol, looks like I targeted the wrong branch on my first PR, so now it's flooded with labels. |
junliu-mde
changed the title
junliu-mde
marked this pull request as ready for review
August 29, 2026 09:07
junliu-mde
requested review from
Fridge003,
Qiaolin-Yu,
Ying1123,
alphabetc1,
hanming-lu,
hebiao064,
hnyls2002,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
August 29, 2026 09:07
junliu-mde
requested review from
BBuf,
DarkSharpness,
HaiShaw,
HydraQYH,
celve and
yuan-luo
as code owners
August 29, 2026 09:07
Fridge003
force-pushed
the
xinyuan/glm-5.3-flash-support
branch
from
August 31, 2026 08:35
cfb1c83 to
493322b
Compare
5 tasks
Leoyzen
added a commit
to Leoyzen/sglang
that referenced
this pull request
Sep 7, 2026
…5.3-Flash (SM90) Port of sgl-project#36904 (beastllama, closed unmerged) re-based onto main, plus the tilelang return_lse removed by sgl-project#37487's GLM-5.3-Flash decode-CP removal, plus DCP integration for glm-dcp-0907. - tilelang_kernel: dispatch sparse_mla_fwd_decode_partial_fp8 on CUDA (block_I=32/threads=128 tiles sized for the fp8 smem class, SM count via _cuda_sm_count); restore return_lse on sparse_attention_fwd_kernel_v1/v2 (caller-allocated base-2 LSE kernel arg, written only when traced) and extend sparse_mla_fwd_decode_combine with a global base-2 LSE output so the partial+combine fp8 path can also return LSE. - overrides (_check_tilelang_dsa_fp8_kv): allow CUDA + tilelang + fp8_e4m3 (SM89+ required, both DSA backends must be tilelang); DCP allowed since the raw-fp8 writer now applies the DCP owner filter. - kv_cache_configurator (calculate_mla_kv_cache_dim): route the raw 512B/token MLA pool layout on CUDA when both DSA backends are tilelang, keyed on the backend pair instead of the platform. - memory_pool (_write_mla_kv_buffer): key the raw fused-quant write branch on layout (not _is_hip); integrate with our DCP-aware write dispatch — DCP-widened locs go through set_mla_kv_buffer_dcp_sharded_triton_fp8_quant. - mla_buffer: add DCP_RANK/DCP_WORLD_SIZE owner filter to set_mla_kv_buffer_fp8_quant_kernel (mirrors the bf16 kernel) and a set_mla_kv_buffer_dcp_sharded_triton_fp8_quant wrapper. This closes the gap that made sgl-project#36904 refuse tilelang-fp8 + DCP. - dsa_backend: tilelang/tilelang added to the DCP backend-pair whitelist (SM80+, bf16 or raw-fp8 KV; SM89+ belt-and-braces check for fp8); _forward_tilelang gains return_lse with the DCP contract (base-2 LSE, [T,H,D]/[T,H] fp32, fixup_zero_kv_rows for owner-less rows — same contract as flashmla_kv/trtllm); supports_mha_one_shot disabled for the CUDA raw-fp8 tilelang layout (sgl-project#36904). - test: port of sgl-project#36904's CUDA raw-fp8 tilelang probe (one-hot exactness + spread-case fp8 budget with scrambled-index negative control). - test(hicache): register the pre-existing kv_cache_dtype storage-key unit test with CI and fix its bare pytest.main (pre-commit requirement). Provenance: internal deployment branch glm-dcp-0907; refs sgl-project#36904, sgl-project#37487 (removal reverted for tilelang), sgl-project#36989 (c-KV localization pattern).
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
3 participants
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.
Note
Targets the #36507 branch (
xinyuan/glm-5.3-flash-support, based at 033446b) so the fix rides along with GLM-5.3-Flash support.TL;DR
GLM-5.3-Flash is the first
qk_rope_head_dim == 0model to run DCP. Its c-KV write path (set_mla_kv_buffer_kernel_norope) and the DSA sparse read paths all use the allocator's virtual locs ([0, size * dcp_world_size)) raw against per-rank-sized buffers. Belowmax_total_num_tokensthis is accidentally self-consistent and outputs are correct; past it, reads/writes silently corrupt KV (expandable_segmentsabsorbs the MMU faults) and output quality degrades state-dependently. Full report with repro protocol: #36886.Changes
mla_buffer.py): norope kernel gains the same localization the rope variant already has (loc % dcp == rankfilter,loc // dcpaddressing). Wrapper takesdcp_localize; the replicated draft pool (virtual-sized, indexes virtual locs by design) opts out.memory_pool.py):dcp_localized_writesflag onMLATokenToKVPool(delegated byHybridLinearKVPool); the torchset_kv_bufferDCP branch now divides after filtering (it filtered only).dsa_backend.py): owner-rule page-table localization (owned →loc // dcp, else-1); the existing cross-rank LSE merge reconstitutes full topk attention from the per-rank partials.dsa_backend.py): no LSE merge exists on this path, so gather the rank-sharded prefix c-KV into the per-request-contiguous layout (all_gather_kv_cache_for_mha_extend) and route topk through the RAGGED transform — its flattenedcu_seqlens_k-offset indices address exactly that layout. Also scale the flattened-page-table capacity assert bydcp_world_size(its consumer, the index-K buffer, is virtual-sized, so virtual locs are legal there; the per-rank bound fail-stops the scheduler on legal traffic).kv_cache_configurator.py): the DSA index-K buffer stores every token on every rank for topk scoring but was sized per-rank; size itmax_total * dcp / loc_space_scale. This is the crash half of [Bug] DCP on GLM-5.3-Flash (norope MLA): virtual locs used raw against per-rank KV buffers — index-K OOB crash + state-dependent quality corruption past max_total_num_tokens #36886.Scope: validated on Hopper only
All validation ran on H100 with
dsa_prefill_backend=tilelang,dsa_decode_backend=tilelang,dsa_topk_backend=sgl-kernel, bf16 KV. The write-kernel and pool changes are backend-agnostic, but the read-side integration (page-table localization, RAGGED gathered extend) was only exercised on the tilelang path — flashmla_sparse / trtllm variants and Blackwell are untested here and should be treated as out of scope.Validation (H100 tp8 + dcp8, bf16, tilelang, MTP on)
max_total_num_tokens; patched does not.Why quality was correct below the watermark
With writes unlocalized every rank writes every token at the virtual loc (replication), and reads use the same virtual locs, so reads hit exactly what was written. The decode LSE merge combines 8 identical partials, which is mathematically lossless. Corruption only starts when a virtual loc exceeds the per-rank buffer end — hence fresh restarts look healthy and long soaks degrade.
CI States
Latest PR Test (Base): ❌ Run #33365284758
Latest PR Test (Extra): ❌ Run #33365284414
Latest PR Test (AMD ROCm 7.2): ❌ Run #33365284663