diff --git a/docs/inference.md b/docs/inference.md index ab8a05ebb90..66b59bda02d 100644 --- a/docs/inference.md +++ b/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. diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index 6855ea356fd..48c3af6b858 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/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 diff --git a/src/prime_rl/templates/inference.sbatch.j2 b/src/prime_rl/templates/inference.sbatch.j2 index 4e513d8e8d7..a762ae7ab65 100755 --- a/src/prime_rl/templates/inference.sbatch.j2 +++ b/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" %} diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 71627e9aa97..d325cb2eeb8 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/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" %} diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index f718ad156d7..c4e9d60214f 100644 --- a/tests/unit/test_configs.py +++ b/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