Repository navigation
[Feat][GLM5.2] Add DSA Cache Layer Split under Prefill CP - #29421
Conversation
Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
…split Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
![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 introduces DSA (DeepSeek Sparse Attention) cache layer splitting across context-parallel (CP) ranks to reduce per-rank KV memory. It adds the enable_dsa_cache_layer_split configuration option and updates memory pools, disaggregation connections, and model runners to support sharding of KV and indexer cache layers, alongside prefetching and broadcasting mechanisms. Feedback on the changes identifies a critical bug in memory_pool.py where torch.distributed.broadcast incorrectly uses a group-relative rank instead of a global rank for the source process, which can cause hangs or crashes in multi-GPU setups. Additionally, a missing use_layer_broadcast_comm=True argument was found when broadcasting the index buffer, which prevents it from utilizing the dedicated layer broadcast communicator.
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.
Dovis01
changed the title
… decode indexer Under --enable-dsa-cache-layer-split, DSATokenToKVPool allocates a 0-block index_k_with_scale_buffer for layers not owned by this rank. The DSA decode indexer (_get_topk_paged) read this buffer directly via get_index_k_with_scale_buffer, so deep_gemm.fp8_paged_mqa_logits got num_kv_blocks=0, making the TMA descriptor (cuTensorMapEncodeTiled) fail with CUDA_ERROR_INVALID_VALUE. This crashed the flashinfer-autotune warmup, whose dummy DECODE forward runs the indexer on every layer (including non-owned ones) -- so removing --disable-flashinfer-autotune made the prefill server fail to start. Route get_index_k_with_scale_buffer through _get_broadcastable_index_buffer when layer sharding is enabled, mirroring get_index_k_continuous (extend path) and get_key_buffer (attention), which already broadcast non-owned layers' KV from the owner rank. Non-layer-sharded pools are unchanged. Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
…split Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
pool_configurator -> get_glm_dsa_layer_split_effective_num_layers crashes every scheduler start since sgl-project#29421 merged: it imports get_attention_cp_size / get_attention_cp_rank / get_attention_cp_group from sglang.srt.layers.dp_attention, but the get_parallel() refactor (sgl-project#30492/sgl-project#30493) retired those accessors before sgl-project#29421 landed, so the names no longer exist anywhere. Switch the three call sites to get_parallel().attn_cp_size / .attn_cp_rank / .attn_cp_group. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
This PR can trigger an ImportError on current main. However, the new code added by this PR still imports those removed APIs: # python/sglang/srt/layers/cp/utils.py
from sglang.srt.layers.dp_attention import get_attention_cp_size
from sglang.srt.layers.dp_attention import get_attention_cp_rank
# python/sglang/srt/mem_cache/dsa_cache_layer_split.py
from sglang.srt.layers.dp_attention import get_attention_cp_group |
|
Have we tested the performance? By how much will the performance decrease? |
|
only need --enable-prefill-cp is ok , --cp-strategy interleave is default open @Dovis01 what's the use for ablout this pr ? add --enable-dsa-cache-layer-split can have more kvcache than not add ? |
- README: sgl-project#31075/sgl-project#30951 already in old base (not v0.5.16 new); sgl-project#29421 is fork-only custom patch; add full list of v0.5.16 new fixes (GLM-5.2/ROCm/ spec-decoding/PD) verified via git merge-base --is-ancestor - values-1p1d-batch1.yaml: remove incompatible --enable-mixed-chunk and --enable-hierarchical-cache --hicache-ratio 2.0; update spec-num-steps 2->3; add num-continuous-decode-steps 2 and optimistic-prefill-attempts 1 to match actual pod YAMLs
…t#29421) Signed-off-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
…led layersplit memory sharding The v0.5.16 merge (74a0065, KVCacheConfigurator) left our branch's layersplit pool (CpLayerSplitKVPool via build_kv_pool_maybe_layersplit) unreachable: _init_pools has no callers post-merge, and the configurator gates its replacement (LayerSplitDSATokenToKVPool, upstream sgl-project#29421) on the NEW flag --enable-dsa-cache-layer-split, while our scripts pass the LEGACY flag --enable-dsa-prefill-cp-layersplit. Result: every rank allocated the FULL 78-layer DSA KV pool (8x intended) and prefill OOM'd at startup in every fraction (0.84/0.85/0.87/0.88/0.90 — deficit ~0.6-2.6GB, layer ~33). The upstream pool itself is NOT usable on this topology either: with it built (any flag combo), prefill froze in DeepGEMM warmup ~77% with all GPUs spinning at 100% — an NCCL collective divergence (dual-flag AND new-flag- only both deadlock; the only configuration ever validated here is the legacy pool, pre-merge). Fix — restore the validated semantics: - kv_cache_configurator._build_dsa_kv_pool: legacy flag takes precedence and routes through build_kv_pool_maybe_layersplit (CpLayerSplitKVPool); the upstream pool stays reachable for new-flag-only configs. - cp_layersplit_pool.broadcast_owner_layer_prefix: migrate off the retired dp_attention accessors (sgl-project#30653) to get_parallel() — the path had rotted invisible while dead code (ImportError surfaced on first re-wire); also keep the get_req_to_token_pool import. - cp/utils.is_glm_dsa_cache_layer_split_enabled: honor EITHER flag so KV cell sizing reflects the sharded per-rank layout (without this the profiler sizes the cell at 78 layers). Also in this commit (same hunt): negative weight_indices (-1 padding-slot sentinel) guards in the six LoRA kernels (csgmv shrink/expand + absorbed step_a/step_b q/v) — padding segments skip before any lora_ranks load. VERIFICATION (2P3D production topology, 1101 prefill + 1100/1103 decode): prefill healthy (12M tokens / 79.13GB, layersplit on, warmup clean); eviction sequence (L0->L1->L2->L3 evicts base->L0 ret->base-after) 0 bad; garbling battery (12 filler + 20 topic) x3 rounds: 0/96 junk; gateway + proxy up; public 18777 chat returns clean reasoning output.
we had tested this pr，after open layer-split，throughtout, ttft had both decrease 3x。when reached 1m，it became even worse。 I want to know the reason, you find it? @CSEEduanyu |
do you test this under the pd-diagg mode? |
Summary
Support new features for GLM5.2 Model. Add DSA cache layer-split for DeepSeek Sparse Attention under prefill context-parallel: the KV/indexer cache layers are sharded across CP ranks (each rank owns a disjoint layer range) instead of every rank holding all layers, cutting per-rank KV memory by ~74% (0.77 → 0.20 GB/rank at 8192 tokens on GLM-5.2-FP8, 78 layers / cp_size=4). Enabled via
--enable-dsa-cache-layer-split(requires--enable-prefill-cp --cp-strategy interleave).Key changes
memory_pool.py): per-rank owned layer ranges, non-owned layers allocated withnum_pages=0, cross-rank KV/index read via_broadcast_tensor_from_ownerover a dedicated CP NCCL communicator.prefill.py,mooncake/conn.py,common/conn.py,base/conn.py): shard-aware layer-range registration and transfer of owned layers; decode learns prefill's layer-split state through the bootstrap handshake (PrefillServerInfo) and pulls from all CP ranks.server_args.py): new--enable-dsa-cache-layer-splitflag, auto-disabled on non-DSA models, decode/non-PD workers, and when prefill-CP interleave isn't set._broadcast_tensor_from_ownerran on a default -disabledPyNcclCommunicator(silent no-op → garbage output); now wrapped inchange_state(enable=True).enable_all_cp_ranks_for_transferhad no layer-split term, so decode only pulled KV from prefill rank 0 (correct first token, then repetition); now all CP ranks participate on both sender and receiver sides.For better optimization support, it is necessary to rely on the following fix PR.
Submitted the Fix PR:
CI States
Latest PR Test (Base): 🚫 Run #29010063850
Latest PR Test (Extra): ⏳ Run #29010063585