[wenscarl]

DeepSeek-DSA models mark most layers shared in indexer_types: those layers reuse the previous producer layer's top-k indices and never read an index-K of their own. GLM-5.2 has 78 layers, of which 21 are producers and 57 are shared.

skip_topk_layers already lets the device pool skip allocating index-K for the shared layers — index_key_cache gives those layers 0 pages, and the compress tail likewise. But _should_elide_dsa_index_k() refused whenever HiCache or PD disaggregation was enabled:

return (
    not memory_config.enable_hisparse
    and not is_draft_worker
    and not memory_config.enable_hierarchical_cache      # <-- here
    and not memory_config.enable_unified_cache_external_linker
    and get_disagg().disaggregation_mode == "null"       # <-- and here
)

That covers most serving deployments, so on a real deployment the device buffers for all 57 shared layers stayed allocated.

Both conditions existed for the same reason: HiCache and the PD transfer walk the per-layer buffer list, and eliding a layer leaves a zero-sized entry in it.

#38426 handles the HiCache half by building the indexer host pool from the
producer-layer set. This PR handles the PD transfer half and removes the two conditions.

Modifications

Dropping the conditions is one change; the other is making the DSA transfer tolerate the zero-sized entries.

That part is subtle because the two peers do not register the same layer span. A PP prefill stage registers only its own layers while decode may register the whole model, so the elision patterns are offset, not different. _send_kvcache_generic already realigns the two via get_mla_kv_ptrs_with_pp (which slices dst[prefill_start_layer : prefill_start_layer + n_src]), but that happens after the per-entry inspection. So this aligns the spans up front with the same _mla_kv_entry_span_with_pp helper, and then:

• drops entries elided on both peers, rather than issuing zero-length DMAs;
• skips the entry (no bytes move) when it is elided on exactly one peer. For DSA this is expected: e.g. a dense HiSparse decode elides index-K for shared-topk layers while prefill still allocates it, so nothing needs to be transferred for that layer. Other state types are unaffected and still go through the stride check.

Re-slicing inside _send_kvcache_generic is then a no-op: its fast path returns unchanged once the two pointer lists have equal length.

What breaks without the transfer-side half

Relaxing the two conditions alone is not enough. Once non-producer layers are elided, their device entries become zero-sized, and set_transfer_blocks() has no item_len == 0 case:

src_addr = src_ptr + int(prefill_index[0]) * item_len   # item_len == 0 -> src_ptr
dst_addr = dst_ptr + int(decode_index[0]) * item_len
length   = item_len * len(prefill_index)                # 0
transfer_blocks.append((src_addr, dst_addr, length))    # queued anyway

So every elided layer contributes a zero-length block, pointing at the empty allocation, for every transfer. That is silent: no error is raised, the blocks are simply handed to the transfer engine. This PR filters those entries out before the DMA instead.

The ordering matters too, and is easy to get wrong. A PP prefill stage registers only its own layers while decode may register the whole model, so the two elision patterns are offset, not different. Any per-entry comparison has to happen after the spans are aligned — otherwise a PP4 stage's 20 entries are compared against the decode's 78 and nothing lines up. _send_kvcache_generic already performs that alignment via get_mla_kv_ptrs_with_pp, but it does so after the per-entry inspection point, which is why this PR aligns up front with the same _mla_kv_entry_span_with_pp helper.

We hit exactly this ordering problem on an internal tree carrying an equivalent elision patch plus a stricter layout check: with PP prefill and elision forced on, 0 of 4,958 warmup requests completed in 4 hours, every one parked in flight, because the check compared 20 prefill entries against 78 decode entries and rejected them. The same topology without elision ran to completion, since with no zero-length entries the check is skipped and the existing prefill_start_layer slicing handles PP correctly.

Accuracy Tests

No accuracy evaluation was run.

The change should not alter what is computed: shared layers never read their own index-K, so not allocating it is expected to be a no-op for model output. That argument is structural — it rests on dsa_layer_skips_topk() and the existing skip_topk_layers path, not on an experiment run here. An accuracy comparison with elision on vs off, on a DSA model with HiCache and PD disaggregation enabled, is the right check and I have not run one.

Benchmarking and Profiling

The win is KV cache capacity, and only that. The same index-K is computed either way; eliding the non-producer layers' device buffers hands that memory back to the KV pool, so the same GPU footprint holds more tokens. There is no throughput claim here, and the measurements below say so explicitly.

KV capacity

Measured at startup from max_total_num_tokens, which is deterministic and does not depend on load.

GLM-5.2-NVFP4 on 12× VR200 (sm_107), disaggregated: prefill is one PP=4 worker (layer partition 20,20,20,18 over 78 layers), decode is one TP=DP=EP=8 worker with dp-attention and DeepEP low-latency. kv_cache_dtype=fp8_e4m3, mem_fraction_static=0.8/0.85, HiCache 128 GiB write-back on prefill, EAGLE speculative decoding on both roles, CUDA graphs on for decode.

role	#38426 only	+this PR
prefill max_total_num_tokens	6,360,640	7,096,512	+11.6%
decode max_total_num_tokens	2,525,184	2,872,320	+13.7%

Both figures reproduced exactly across three i

The decode figure also reproduces on differentmodel, kv_cache_dtype=fp8_e4m3, at an identical 96.33 GB KV footprint per decode rank, 1,661,504 → 1,889,920 tokens, again +13.7%.

Throughput: unchanged

Three runs per arm, agentic-coding trace at concurrency 112, 900 s measurement window, recomputed from the raw per-request
export rather than a summary:

arm	req/s/GPU (3 runs)	mean	range
#38426 only	0.1138 / 0.1145 / 0.1172	0.1152	[0.1138, 0.1172]
+this PR	0.1089 / 0.1125 / 0.1183	0.1132	[0.1089, 0.1183]
The ranges overlap and the 1.7% mean differencspread, so no throughput effect is detectableeither way on this workload.
That is the expected result for this configuration, and worth stating plainly: in these runs the prefill KV pool sa4–10% utilisation while the host-side HiCache PU KV capacity was not the binding constraintand extra capacity had nothing to buy. The capacity win should only turn into throughput on workloads where GPU KV capacity is the limit.

Reproducing

Requires the zero-length-entry fix in this PR. Without it the elided layers' zero-length placeholders are handed to
Mooncake, batch_register fails with

transfer_engine_impl.cpp: Transfer Engine doesry region

its return value is not checked, and every KV emory that was never registered.

Capacity is the cheap check and needs no benchy and compare one log line:

python3 -m sglang.launch_server \
  --model-path <GLM-5.2-NVFP4> --trust-remote-code \
  --kv-cache-dtype fp8_e4m3 --quantization mod
  --tp 8 --dp 8 --enable-dp-attention \                                                                            --enable-hierarchical-cache --hicache-size 1
  --hicache-write-policy write_back --hicache-io-backend direct \                                                  --hicache-mem-layout page_first_direct \
  --disaggregation-mode decode --disaggregation-transfer-backend mooncake                                        
grep -o 'max_total_num_tokens=[0-9]*' <server log>                                          

HiCache and --disaggregation-mode are what matter: those are the two conditions this PR removes from _should_elide_dsa_index_k, so with either one ly and the delta is zero. The model must be aDeepSeek-DSA one with shared entries in indexer_types — GLM-5.2 has 21 producer layers of 78, which is where the double-digit percentage comes from.

For the throughput comparison, run the same setained load for a fixed window and comparecompleted requests per GPU, with at least three runs per arm — a single run is not enough to separate a real effect from run-to-run variance at this scale.
NOTE: this PR was established firstly on #38426** (producer-layer HiCache host pool) which later closed.

Stacked on #38426. Semantically required, not textually: the diff applies cleanly without it, but relaxing enable_hierarchical_cache while the indexer host pool is still sized for every layer would leave it backing device buffers that no longer exist.

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add or update documentation as needed. (the _should_elide_dsa_index_k docstring now records why the two conditions existed and what removed them)
• Add unit tests as needed. (no existing test covers DSA elision with HiCache or PD; pointers on where such a test belongs are welcome)
• Update CODEOWNERS if relevant. (not relevant)

---

CI States

Latest PR Test (Base): ⏳ Run #38373698945
Latest PR Test (Extra): ❌ Run #38373698523
Latest PR Test (AMD ROCm 10): ⏳ Run #38373698892

[nvpohanh]

[by Claude Code] Severity: nit | Confidence: High

Issue

The PR body still describes the old design. It says this PR is stacked on #38426 and semantically requires it, but #38426 is closed. It also says the transfer aligns spans up front with _mla_kv_entry_span_with_pp and raises on one-sided elision. ed92a34 removed that block and now skips one-sided pairs.

Fix

Update the body to the current design: dense kv_args, 0-length skips at Mooncake registration and in layers_params, one-sided pairs skipped in the length checks, and the Ascend PD guard. Remove the #38426 dependency, since main's owned_device_layers covers the HiCache side.

[nvpohanh]

[by Claude Code] Severity: functional | Confidence: High

Issue

This branch conflicts with main in mooncake/conn.py, inside this check. #42051 deleted both inline item-length loops that this commit edits. Main now checks lengths in _state_stride_mismatch(), and _validate_peer_state_layout() also runs it at registration. If the conflict is resolved by taking main's side, the zero-length skip is lost.

Example

On CPU, I ran main's _state_stride_mismatch() with an elided prefill and a dense HiSparse decode. It returns StateType.DSA item length mismatch for positional entry 1: prefill=0 decode=64. Main then rejects that decode at registration, and every request to it fails.

Fix

Rebase onto main and move both skips into the two loops of _state_stride_mismatch(). Add a CPU test that calls it with one-sided zeros.

[wenscarl]

Based to main and _state_stride_mismatch() in mooncake/conn.py now skips an entry when either side's item length is 0. Also added CPU test that calls _state_stride_mismatch() with one-sided zeros in both directions on both pairing paths, plus a negative case that must still fail. It fails without the skips with exactly your message (prefill=0 decode=64). I removed the older test that only re-tested the list comprehension.

[nvpohanh]

@wenscarl Could you address the final round of reviews? After this, I will start the CI and ask SGLang for review. Thanks!

[Lzy17]

Thanks for the PR. Remaining items:

1. mori: skip 0-length buffers in _register_local_buffers as well, keeping state_mem_descs index-aligned. RegisterRdmaMemoryRegionAuto aborts when ibv_reg_mr fails and IOEngine::RegisterMemory has no size==0 guard, so a failing ibv_reg_mr(len=0) would abort mori PD at startup. I'll confirm on ionic and report back.
2. Update the description: it says one-sided elision raises (the code skips it), and that this is stacked on #38426 (closed; the compaction is already on main).
3. Limit the one-sided skip in _state_stride_mismatch to StateType.DSA; it currently applies to every flat state type.
4. NPU: _should_elide_dsa_index_k only guards PD. With enable_hierarchical_cache removed, NPU arch35 + HiCache now gets the compact layout. Please confirm its host pool supports that, or guard it.

[nvpohanh]

@wenscarl please address comments from @Lzy17 . After that, I will trigger and babysit CI

[wenscarl]

Thanks for the PR. Remaining items:

1. mori: skip 0-length buffers in _register_local_buffers as well, keeping state_mem_descs index-aligned. RegisterRdmaMemoryRegionAuto aborts when ibv_reg_mr fails and IOEngine::RegisterMemory has no size==0 guard, so a failing ibv_reg_mr(len=0) would abort mori PD at startup. I'll confirm on ionic and report back.
2. Update the description: it says one-sided elision raises (the code skips it), and that this is stacked on [HiCache] Compact DSA indexer host storage to producer layers #38426 (closed; the compaction is already on main).
3. Limit the one-sided skip in _state_stride_mismatch to StateType.DSA; it currently applies to every flat state type.
4. NPU: _should_elide_dsa_index_k only guards PD. With enable_hierarchical_cache removed, NPU arch35 + HiCache now gets the compact layout. Please confirm its host pool supports that, or guard it.
   Thanks @Lzy17

• Mori: _register_local_buffers now skips 0-length state buffers and keeps a None placeholder, so state_mem_descs stays index-aligned with state_item_lens. The pack/unpack helpers and teardown handle None. The transfer path already skips a zero-length source before it touches the descriptor. Please confirm on ionic when you get a chance. I couldn't test mori PD here.
• PR description: Updated. It now says one-sided elision is skipped rather than raising, and I removed the "stacked on [HiCache] Compact DSA indexer host storage to producer layers #38426" line since that PR is closed and the compaction is already on main.
• Mooncake _state_stride_mismatch: The one-sided zero-length destination skip is now limited to StateType.DSA, through a new allow_dst_elision argument on is_elided_entry. A zero-length source is still always skipped. The default is unchanged, so nixl and mori behave as before.
• NPU: _should_elide_dsa_index_k now returns False on NPU when HiCache is enabled, so NPU HiCache keeps the non-compact layout. I haven't verified whether the NPU host pool supports the compact layout. If it does, we can drop the guard in a follow-up.

[nvpohanh]

/tag-and-rerun-ci

[nvpohanh]

[by Claude Code] I've started babysitting the NVIDIA CI of this PR. I'll check it every 30 minutes until all NV pipelines pass. If a relevant job fails, I may rerun it, merge the latest main into this branch, push a small test-only fix, or post a note here.

Anyone can stop me at any time: just post a comment asking me to stop (for example, stop babysitting). I'll stop at my next check, within 30 minutes, and take no further action on this PR.

[nvpohanh]

/rerun-failed-ci

[nvpohanh]

[by Claude Code] All NV pipelines have passed.

[nvpohanh]

[by Claude Code] I've started babysitting the NVIDIA CI of this PR. I'll check it every 30 minutes until all NV pipelines pass. If a relevant job fails, I may rerun it, merge the latest main into this branch, push a small test-only fix, or post a note here.

Anyone can stop me at any time: just post a comment asking me to stop (for example, stop babysitting). I'll stop at my next check, within 30 minutes, and take no further action on this PR.