feat(inference): experimental HiSparse on P/D decode

docs/inference.md

@@ -15,6 +15,7 @@ This page covers the inference configuration and the supported features/deployme
 - [Adaptive Concurrency](#adaptive-concurrency)
 - [Advanced Configuration](#advanced-configuration)
     - [KV Cache Offload](#kv-cache-offload)
+    - [HiSparse](#hisparse)
     - [Optimized P/D disaggregation deployment](#optimized-pd-disaggregation-deployment)
     - [Other vLLM features](#other-vllm-features)
     - [Router Replay](#router-replay)
@@ -263,6 +264,19 @@ num_bytes = 0
 
 Mooncake keys a block by model, parallel rank and content hash only. Roles that share the store must therefore use the same KV layout (tensor parallel size, `kv_cache_dtype`, block size): do not change these in only one role's `*_vllm_overrides`.
 
+### HiSparse
+
+> **Experimental.** Off by default and not used by any example config. Validated so far: P/D with HiSparse on decode, router replay and sampling replay on, on a small random DSA model (statistical checks only). Not validated yet: a real GLM-5.x long-prompt (40k+) quality A/B, aggregated HiSparse against the trainer forward pass, and preemption or failed KV loads with HiSparse.
+
+HiSparse (vLLM, DSA sparse-MLA models such as GLM-5.x) keeps the sparse-MLA KV in a pinned host pool. Decode attention reads per-request GPU buffers that hold the indexer top-k rows. This frees GPU memory for many more concurrent decode sequences.
+
+```toml
+[inference.hisparse]
+host_pool_gib = 160   # pinned host RAM per GPU worker
+```
+
+prime-rl adds vLLM's `HiSparseConnector` to the KV transfer config. Under disaggregated P/D it applies to decode instances only: the launcher writes one engine config per role (`inference-prefill.json`, `inference-decode.json`) and drops HiSparse from the prefill one. Budget decode-node RAM as `GPUs per node × host_pool_gib`. Prefill KV offload works with it, but decode cannot also offload KV: set `inference.deployment.decode_kv_cache_offload = "None"`. Native offload rejects HiSparse's host-resident KV layout, and Mooncake would store that layout under the same keys as prefill's regular layout, so prefill fails to load it. Without P/D, native offload is rejected and Mooncake offload works. vLLM checks HiSparse's requirements at startup: the V2 model runner (vLLM selects it), no pipeline or decode context parallelism, and a model with `index_topk`. vLLM 0.32 also rejects `cudagraph_mode = "FULL"`; the P/D decode launch uses `FULL_DECODE_ONLY`, which it accepts. Size decode `max_num_seqs` explicitly (e.g. 96 per rank): the GPU buffers scale with it.
+
 ### Optimized P/D disaggregation deployment
 
 For optimal P/D disaggregation deployment, we automatically set the decode `all2all_backend` to `deepep_low_latency` and the prefill `all2all_backend` to `deepep_high_throughput`. We currently don't support customizing all2all backends for P/D disaggragation out of the box. You can do this by overriding the slurm template only.

packages/prime-rl-configs/src/prime_rl/configs/inference.py

@@ -263,6 +263,13 @@ def to_connector_dict(self) -> dict[str, Any]:
 ]
 
 
+class HiSparseConfig(BaseConfig):
+    """vLLM HiSparse for DSA sparse-MLA models (e.g. GLM-5.x): the sparse-MLA KV lives in a pinned host pool and decode attention reads per-request GPU buffers of the indexer top-k rows."""
+
+    host_pool_gib: float = Field(..., gt=0)
+    """Pinned host pool per GPU worker, in GiB (``HiSparseConnector`` ``host_pool_gib``). Budget ``GPUs per node × host_pool_gib`` against node RAM."""
+
+
 # Known llm-d EPP scorer plugins (used to guard the ``scorers`` map against typos).
 KNOWN_SCORERS = frozenset(
     {
@@ -447,6 +454,9 @@ class InferenceConfig(BaseConfig):
     kv_cache_offload: KVCacheOffloadConfig | None = None
     """KV cache offload for inference workers, as composable CPU/disk tiers. Discriminated on ``type``: ``native`` (vLLM ``OffloadingConnector``/``TieringOffloadingSpec``, self-contained) or ``mooncake`` (per-node Mooncake distributed store). Disaggregated P/D combines the chosen connector with NIXL through ``MultiConnector``."""
 
+    hisparse: HiSparseConfig | None = None
+    """Enable vLLM HiSparse (experimental). Under disaggregated P/D it applies to decode instances only, and decode must not offload KV (``deployment.decode_kv_cache_offload = None``). vLLM checks its requirements at startup (a DSA model, the V2 model runner, no pipeline or decode context parallelism)."""
+
     use_pd_kv_transfer: bool = False
     """Auto-set for disaggregated P/D: emit the NIXL transfer connector. Persisted into the per-node config (which drops ``deployment``) so the connector is still built per worker. Not meant to be set by hand."""
 
@@ -489,6 +499,20 @@ def validate_llmd_no_routed_experts(self):
             )
         return self
 
+    @model_validator(mode="after")
+    def validate_hisparse_decode_offload(self):
+        """A HiSparse decode instance keeps its sparse-MLA KV in a host-resident layout: native offload rejects it, and Mooncake stores it under the same keys as prefill's regular layout, so prefill fails to load it."""
+        if self.hisparse is None:
+            return self
+        if self.kv_cache_offload_by_role.get("decode") is not None:
+            raise ValueError(
+                "HiSparse runs on decode, which cannot also offload KV: set "
+                "inference.deployment.decode_kv_cache_offload = None (prefill offload is supported)."
+            )
+        if isinstance(self.kv_cache_offload_by_role.get("all"), NativeKVCacheOffloadConfig):
+            raise ValueError("HiSparse does not work with native KV cache offload: use Mooncake or no offload.")
+        return self
+
     @model_validator(mode="after")
     def validate_router_deployment(self):
         """The llm-d router (EPP + Envoy) is launched by the SLURM templates only; multi-node deployments need a router to front the per-rank engines."""
@@ -559,18 +583,20 @@ def kv_cache_offload_by_role(self) -> dict[str, KVCacheOffloadConfig | None]:
         return {role: self.kv_cache_offload if offload == "inherit" else offload for role, offload in overrides.items()}
 
     def for_pd_role(self, role: Literal["prefill", "decode"]) -> "InferenceConfig":
-        """Engine config for one role of a disaggregated deployment, with that role's KV cache offload."""
+        """Engine config for one role of a disaggregated deployment: the role's KV cache offload, and HiSparse on decode only."""
         assert self.deployment.type == "disaggregated"
         config = self.model_copy(deep=True)
         config.kv_cache_offload = self.kv_cache_offload_by_role[role]
+        if role == "prefill":
+            config.hisparse = None
         return config
 
     def build_kv_transfer_config(self) -> dict[str, Any] | None:
-        """Build the single vLLM ``kv_transfer_config`` from the transfer + offload connectors.
+        """Build the single vLLM ``kv_transfer_config`` from the transfer, offload and HiSparse connectors.
 
-        Disaggregated P/D always uses NIXL for prefill→decode transfer. KV cache offload (if
-        configured) contributes its own connector. When both are present they are composed via
-        ``MultiConnector``. Returns None when neither applies.
+        Disaggregated P/D always uses NIXL for prefill→decode transfer. KV cache offload and
+        HiSparse (if configured) contribute their own connectors. Several connectors are composed
+        via ``MultiConnector``. Returns None when none applies.
         """
         connectors: list[dict[str, Any]] = []
         if self.use_pd_kv_transfer:
@@ -587,6 +613,14 @@ def build_kv_transfer_config(self) -> dict[str, Any] | None:
                 # Decode instances also store the KV they generate, so the next turn's prefill can load it.
                 connector["kv_connector_extra_config"]["save_decode_cache"] = True
             connectors.append(connector)
+        if self.hisparse is not None:
+            connectors.append(
+                {
+                    "kv_connector": "HiSparseConnector",
+                    "kv_role": "kv_both",
+                    "kv_connector_extra_config": {"host_pool_gib": self.hisparse.host_pool_gib},
+                }
+            )
 
         if not connectors:
             return None

src/prime_rl/templates/inference.sbatch.j2

@@ -192,7 +192,7 @@ srun --kill-on-bad-exit=1 bash -s <<'LAUNCH_SH'
 {% include '_launch_rank.sh.j2' %}
 {% include '_launch_router.sh.j2' %}
 
-    # NIXL/RDMA registers pinned host memory (KV offload pools).
+    # NIXL/RDMA registers pinned host memory (KV offload pools, the HiSparse host pool).
     ulimit -l unlimited
 {%- if mooncake_offloads %}
 {% include "_mooncake_store.sh.j2" %}

src/prime_rl/templates/multi_node_rl.sbatch.j2

@@ -338,7 +338,7 @@ if [ "$SLURM_PROCID" -lt "$NUM_INFER_NODES" ]; then
 {%- endfor %}
     LOCAL_IP=$(hostname -I | awk "{print \$1}")
 
-    # NIXL/RDMA registers pinned host memory (KV offload pools).
+    # NIXL/RDMA registers pinned host memory (KV offload pools, the HiSparse host pool).
     ulimit -l unlimited
 {%- if mooncake_offloads %}
 {% include "_mooncake_store.sh.j2" %}

tests/unit/test_configs.py

@@ -1174,6 +1174,21 @@ def test_pd_role_connectors():
     assert [c["kv_connector"] for c in decode] == ["NixlConnector", "MooncakeStoreConnector"]
     assert decode[1]["kv_connector_extra_config"]["save_decode_cache"]
 
+    with pytest.raises(ValidationError, match="decode_kv_cache_offload"):
+        InferenceConfig(deployment=deployment, slurm={}, kv_cache_offload=mooncake, hisparse={"host_pool_gib": 160})
+    deployment["decode_kv_cache_offload"] = None
+    config = InferenceConfig(
+        deployment=deployment, slurm={}, kv_cache_offload=mooncake, hisparse={"host_pool_gib": 160}
+    )
+    decode = config.for_pd_role("decode").build_kv_transfer_config()["kv_connector_extra_config"]["connectors"]
+    assert [c["kv_connector"] for c in decode] == ["NixlConnector", "HiSparseConnector"]
+    assert config.for_pd_role("prefill").hisparse is None
+
+    with pytest.raises(ValidationError, match="native KV cache offload"):
+        InferenceConfig(
+            kv_cache_offload={"type": "native", "cpu": {"num_bytes": 1000}}, hisparse={"host_pool_gib": 160}
+        )
+
 
 CUSTOM_RENDERER_SOURCE = """
 from typing import Literal