[Feat] KVShard: Integrate GLM DSA/MTP with PD and CPU HiCache

docs/docs/advanced_features/pd_disaggregation.mdx

@@ -74,6 +74,132 @@ destination registration. Older decode workers do not enforce this check. Keep
 `early` on both sides until every worker has been upgraded, then enable the new
 policy on both sides.
 
+## GLM DSA KV sharding with CPU HiCache
+
+On prefill workers, `--enable-kv-cache-sharding` distributes cached pages across
+the attention-CP group. Each owner stores and transfers its target/draft latent
+KV and DSA indexer pages. Optional CPU HiCache retains the same page ownership
+during backup and reload. Decode workers receive complete caches and do not
+enable KV sharding.
+
+Use this combination on Linux with NVIDIA CUDA and hardware supported by the
+TRTLLM DSA backends. The `nvidia/GLM-5.2-NVFP4` example requires Blackwell GPUs
+with enough memory for both target and draft weights on each side.
+
+| Setting | Supported configuration |
+| --- | --- |
+| Prefill parallelism | Interleave CP greater than one; pipeline parallel size one |
+| Attention and KV | TRTLLM DSA prefill/decode, `fp8_e4m3`, page size 64 |
+| DSA indexer | Uncompressed `index_kpool=1`; no cache layer splitting |
+| P/D transfer | Mooncake, with staging buffers disabled |
+| Optional MTP | GLM DSA single-layer NextN, EAGLE, top-k one |
+| Optional CPU HiCache | `cache` mode, `kernel` I/O, `layer_first` or `page_first` |
+
+KV-sharded serving selects the Python radix tree when no tree backend is
+explicitly configured. L3 storage (including runtime attachment), external cache
+linkers, HiSparse and two-batch overlap are unsupported. Prefill CUDA graphs,
+P-side speculative graphs and generic FlashInfer autotuning are disabled for
+this path. Scratch buffers for target and draft are charged against the KV
+memory budget, so the usable cache capacity also depends on context length.
+
+The following single-host GB300/NVL72 example uses GPUs 0–1 for prefill and 2–3
+for decode. Install the [Mooncake transfer engine](#requirements) first. Both
+workers use fabric NVLink with Mooncake's NVLink allocator, which makes their
+KV buffers exportable. Enabling fabric transport alone does not make ordinary
+CUDA allocations remotely accessible. For other hardware, configure the
+appropriate [RDMA or NVLink transport](#nvlink-transport-configuration).
+
+Run prefill and decode in separate terminals on the same host.
+
+### Prefill worker (GPUs 0–1)
+
+```bash
+export SGLANG_DISAGG_STAGING_BUFFER=0
+export SGLANG_MOONCAKE_CUSTOM_MEM_POOL=NVLINK
+export MC_FORCE_MNNVL=true
+
+python3 -m sglang.launch_server \
+  --model-path nvidia/GLM-5.2-NVFP4 \
+  --trust-remote-code \
+  --tp-size 2 \
+  --context-length 16384 \
+  --max-running-requests 1 \
+  --mem-fraction-static 0.85 \
+  --page-size 64 \
+  --kv-cache-dtype fp8_e4m3 \
+  --dsa-prefill-backend trtllm \
+  --dsa-decode-backend trtllm \
+  --speculative-algorithm EAGLE \
+  --speculative-num-steps 3 \
+  --speculative-eagle-topk 1 \
+  --speculative-num-draft-tokens 4 \
+  --disaggregation-transfer-backend mooncake \
+  --disaggregation-bootstrap-port 8998 \
+  --enable-metrics \
+  --enable-cache-report \
+  --disaggregation-mode prefill \
+  --port 30001 \
+  --nccl-port 29501 \
+  --enable-kv-cache-sharding \
+  --enable-prefill-cp \
+  --attn-cp-size 2 \
+  --cp-strategy interleave \
+  --chunked-prefill-size 2048 \
+  --max-total-tokens 4096 \
+  --enable-hierarchical-cache \
+  --hicache-ratio 4 \
+  --hicache-write-policy write_through \
+  --hicache-io-backend kernel \
+  --hicache-mem-layout page_first
+```
+
+### Decode worker (GPUs 2–3)
+
+```bash
+export SGLANG_DISAGG_STAGING_BUFFER=0
+export SGLANG_MOONCAKE_CUSTOM_MEM_POOL=NVLINK
+export MC_FORCE_MNNVL=true
+
+python3 -m sglang.launch_server \
+  --model-path nvidia/GLM-5.2-NVFP4 \
+  --trust-remote-code \
+  --tp-size 2 \
+  --context-length 16384 \
+  --max-running-requests 1 \
+  --mem-fraction-static 0.85 \
+  --page-size 64 \
+  --kv-cache-dtype fp8_e4m3 \
+  --dsa-prefill-backend trtllm \
+  --dsa-decode-backend trtllm \
+  --speculative-algorithm EAGLE \
+  --speculative-num-steps 3 \
+  --speculative-eagle-topk 1 \
+  --speculative-num-draft-tokens 4 \
+  --disaggregation-transfer-backend mooncake \
+  --disaggregation-bootstrap-port 8998 \
+  --enable-metrics \
+  --enable-cache-report \
+  --disaggregation-mode decode \
+  --port 30002 \
+  --nccl-port 29502 \
+  --base-gpu-id 2 \
+  --max-total-tokens 16384 \
+  --cuda-graph-max-bs-decode 1
+```
+
+Connect the two endpoints through the [P/D router](#router-integration).
+To check CPU-cache reuse, repeat a prompt after enough unrelated requests have
+evicted its GPU prefix: `cached_tokens_details.host` and the prefill
+`sglang:load_back_tokens_total` counter must increase. A device-cache hit alone
+does not exercise the host reload path.
+
+When comparing with replicated or deduplicated KV, keep the model revision,
+hardware, P/D topology, memory budget, request order, concurrency and warm-up
+policy fixed. Report device/host cache-hit tokens, TTFT percentiles, throughput,
+peak memory and answer correctness. Additional cache capacity alone does not
+establish a latency or throughput gain; use a representative replay and the
+[benchmarking guide](/docs/developer_guide/benchmark_and_profiling).
+
 ## Profiling in PD Disaggregation Mode
 
 When you need to profile prefill or decode workers in PD disaggregation mode, please refer to the [Profile In PD Disaggregation Mode](../developer_guide/benchmark_and_profiling#profile-in-pd-disaggregation-mode) section in the Benchmark and Profiling guide. Due to torch profiler limitations, prefill and decode workers must be profiled separately using dedicated command-line options.

docs/docs/advanced_features/server_arguments.mdx

@@ -2903,7 +2903,7 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec
     </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-kv-cache-sharding`</td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Shard at-rest KV cache storage at page granularity across the group that currently replicates it (attention-CP ranks for GQA models under prefill CP; attention-TP ranks for MLA models), multiplying the group's unique KV capacity. Prefill workers of PD disaggregation only. With NCCL 2.31.2 or newer and supported NVIDIA CUDA hardware, eligible single-node shard groups request zero-SM copy-engine prefix all-gather on a dedicated NCCL communicator. Older NCCL versions keep the SM-kernel all-gather. This path needs cuMem-backed NCCL memory, so on single-node deployments the server sets <code>NCCL_CUMEM_ENABLE=1</code> for every NCCL communicator of all server processes unless <code>NCCL_CUMEM_ENABLE</code> is already set; it does not under <code>--enable-memory-saver</code> or <code>SGLANG_MOONCAKE_CUSTOM_MEM_POOL</code>, which keep the SM-kernel all-gather. Multi-node deployments keep NCCL's setting (MNNVL fabrics default it to 1); export <code>NCCL_CUMEM_ENABLE=1</code> there to use this path for shard groups within one node. Set <code>SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER=0</code> to disable this path. Leave <code>NCCL_CTA_POLICY</code> unset: a global value overrides both the gather communicator's zero-SM policy and the default policy of forward collectives.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Shard at-rest KV cache storage at page granularity across the group that currently replicates it (attention-CP ranks under prefill CP; attention-TP ranks for MLA models without prefill CP), multiplying the group's unique KV capacity. Prefill workers of PD disaggregation only. See <a href="/docs/advanced_features/pd_disaggregation#glm-dsa-kv-sharding-with-cpu-hicache">GLM DSA KV sharding with CPU HiCache</a> for DSA/MTP requirements and a launch example. With NCCL 2.31.2 or newer and supported NVIDIA CUDA hardware, eligible single-node shard groups request zero-SM copy-engine prefix all-gather on a dedicated NCCL communicator. Older NCCL versions keep the SM-kernel all-gather. This path needs cuMem-backed NCCL memory, so on single-node deployments the server sets <code>NCCL_CUMEM_ENABLE=1</code> for every NCCL communicator of all server processes unless <code>NCCL_CUMEM_ENABLE</code> is already set; it does not under <code>--enable-memory-saver</code> or <code>SGLANG_MOONCAKE_CUSTOM_MEM_POOL</code>, which keep the SM-kernel all-gather. Multi-node deployments keep NCCL's setting (MNNVL fabrics default it to 1); export <code>NCCL_CUMEM_ENABLE=1</code> there to use this path for shard groups within one node. Set <code>SGLANG_ENABLE_KV_SHARD_ZERO_SM_GATHER=0</code> to disable this path. Leave <code>NCCL_CTA_POLICY</code> unset: a global value overrides both the gather communicator's zero-SM policy and the default policy of forward collectives.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
     </tr>

python/sglang/srt/arg_groups/kv_shard_hook.py

@@ -14,7 +14,11 @@
     resolving_view,
     use_mla_backend,
 )
-from sglang.srt.configs.model_config import is_deepseek_dsa
+from sglang.srt.configs.model_config import (
+    get_dsa_index_kpool,
+    get_dsa_index_kpool_compress,
+    is_deepseek_dsa,
+)
 from sglang.srt.environ import envs
 from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase
 from sglang.srt.runtime_context import attn_dp_enabled_of
@@ -26,6 +30,41 @@
 def validate_kv_shard_attention_backend(server_args: Any, cfg: Any = None) -> None:
     cfg = cfg or resolving_view(server_args)
     prefill_backend, _ = attention_backends_of(resolved_view(server_args))
+    if prefill_backend == "dsa":
+        hf_config = model_config_of(server_args).hf_config
+        if (
+            cfg.device != "cuda"
+            or not cfg.enable_prefill_cp
+            or cfg.attn_cp_size <= 1
+            or cfg.cp_strategy != "interleave"
+            or cfg.pp_size != 1
+        ):
+            raise ValueError(
+                "DSA KV sharding requires CUDA interleave prefill CP > 1 and PP1."
+            )
+        if (
+            cfg.dsa_prefill_backend != "trtllm"
+            or cfg.dsa_decode_backend != "trtllm"
+            or cfg.kv_cache_dtype != "fp8_e4m3"
+            or cfg.page_size != 64
+        ):
+            raise ValueError(
+                "DSA KV sharding requires TRTLLM prefill/decode with fp8_e4m3 "
+                "KV and page size 64."
+            )
+        if (
+            cfg.disaggregation_transfer_backend != "mooncake"
+            or cfg.enable_dsa_cache_layer_split
+            or get_dsa_index_kpool(hf_config) != 1
+            or get_dsa_index_kpool_compress(hf_config)
+        ):
+            raise ValueError(
+                "DSA KV sharding requires Mooncake, uncompressed index_kpool=1 "
+                "and no DSA cache layer splitting."
+            )
+        return
+    if is_deepseek_dsa(model_config_of(server_args).hf_config):
+        raise ValueError("DSA KV sharding requires the dsa attention backend.")
     if prefill_backend == "fa3":
         return
     if prefill_backend != "trtllm_mla":
@@ -64,9 +103,19 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None)
             "page, which is incompatible with local decode."
         )
     if cfg.speculative_algorithm is not None:
-        raise ValueError(
-            "--enable-kv-cache-sharding does not support speculative decoding."
-        )
+        model_config = model_config_of(server_args)
+        # PP1 is checked once by the DSA backend gate below.
+        if (
+            cfg.speculative_algorithm != "EAGLE"
+            or model_config.hf_config.architectures[0] != "GlmMoeDsaForCausalLM"
+            or model_config.num_nextn_predict_layers != 1
+            or cfg.speculative_eagle_topk != 1
+            or cfg.enable_multi_layer_eagle
+        ):
+            raise ValueError(
+                "KV sharding supports speculative decoding only for GLM DSA "
+                "single-layer NextN MTP with EAGLE and topk=1."
+            )
     if cfg.dllm_algorithm is not None:
         raise ValueError(
             "--enable-kv-cache-sharding does not support diffusion language "
@@ -80,6 +129,10 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None)
             "not support "
             f"--radix-cache-backend={cfg.radix_cache_backend!r}."
         )
+    # KV sharding defaults to Python even when the general tree default is Rust.
+    # Preserve an explicit backend choice for validation below.
+    if not envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.is_set():
+        envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.set("python")
     tree_core_backend = envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.get()
     # A disabled radix cache never matches or inserts, so its tree core never
     # sees a rotation base; chunked continuations keep the base on the request.
@@ -120,12 +173,24 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None)
             "(--dcp-size) are mutually exclusive users of the widened "
             "allocator index space."
         )
-    if cfg.enable_hierarchical_cache:
+    if cfg.enable_hierarchical_cache and cfg.hicache_storage_backend is not None:
         raise ValueError(
-            "--enable-kv-cache-sharding does not support "
-            "--enable-hierarchical-cache yet: HiCache backup/load assumes "
-            "each rank holds full pages."
+            "--enable-kv-cache-sharding supports HiCache host offload but not "
+            "--hicache-storage-backend yet: the storage path reads and writes "
+            "whole pages per rank, not each rank's own stripe."
         )
+    if cfg.enable_hierarchical_cache and is_deepseek_dsa(
+        model_config_of(server_args).hf_config
+    ):
+        if cfg.hicache_host_memory_mode != "cache":
+            raise ValueError("CP-sharded DSA HiCache requires host memory mode cache")
+        if cfg.hicache_io_backend != "kernel" or cfg.hicache_mem_layout not in (
+            "layer_first",
+            "page_first",
+        ):
+            raise ValueError(
+                "CP-sharded DSA HiCache requires kernel I/O with layer_first or page_first"
+            )
     if cfg.enable_lmcache:
         raise ValueError(
             "--enable-kv-cache-sharding does not support --enable-lmcache."
@@ -184,11 +249,6 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None)
         )
 
     if use_mla_backend(server_args):
-        if is_deepseek_dsa(model_config.hf_config):
-            raise ValueError(
-                "--enable-kv-cache-sharding does not support DSA models yet "
-                "(indexer buffers are not striped)."
-            )
         if cfg.attn_cp_size <= 1:
             if attn_dp_enabled_of(cfg):
                 raise ValueError(
@@ -264,15 +324,20 @@ def handle_kv_cache_sharding(server_args: Any, gpu_mem: Optional[float] = None)
             page_size,
         )
 
-    if cfg.cuda_graph_config.prefill.backend != Backend.DISABLED:
+    graph_config = cfg.cuda_graph_config
+    if graph_config.prefill.backend != Backend.DISABLED:
         logger.warning(
             "Prefill CUDA graph (%s) is incompatible with "
             "--enable-kv-cache-sharding; disabling it.",
             cfg.cuda_graph_config.prefill.backend,
         )
-        declarations["cuda_graph_config"] = with_phase(
-            cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED
-        )
+        graph_config = with_phase(graph_config, Phase.PREFILL, backend=Backend.DISABLED)
+    if cfg.speculative_algorithm is not None:
+        # P-side draft-extend uses the request's shard plan; ordinary speculative
+        # graph capture has neither that plan nor logical-page allocations.
+        graph_config = with_phase(graph_config, Phase.DECODE, backend=Backend.DISABLED)
+    if graph_config != cfg.cuda_graph_config:
+        declarations["cuda_graph_config"] = graph_config
 
     if declarations:
         declare_resolution(

python/sglang/srt/disaggregation/common/conn.py

@@ -516,13 +516,16 @@ def _should_skip_cp_replicated_state_transfer(self) -> bool:
         Prefill CP materializes global token order before writing state pools, so
         every CP rank holds the same state. When all CP ranks transfer their KV
         shards, only rank 0 needs to send that state. Cache layer split is the
-        exception because each CP rank owns different state layers.
+        exception because each CP rank owns different state layers. Page-sharded
+        DSA indexer state likewise needs every owner; Mooncake filters its page
+        lists before transfer instead of suppressing whole ranks.
         """
         return (
             self.attn_cp_size > 1
             and self.attn_cp_rank != 0
             and not get_parallel().enable_dsa_cache_layer_split
             and not get_parallel().enable_cp_tp_group_sharing
+            and getattr(self, "kv_shard_size", 1) <= 1
         )
 
     def requires_dcp_relayout(self, dst_dcp_size: int, dst_dcp_rank: int) -> bool: