feat: configure bidirectional pd kv transfer

Author: S1ro1

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

@@ -172,6 +172,69 @@ class MultiNodeInferenceDeploymentConfig(BaseInferenceDeploymentConfig):
     ] = "consistent_hash"
 
 
+class NixlTransportConfig(BaseModel):
+    """Configures NIXL KV transfer for disaggregated inference deployments."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    type: Literal["nixl"] = "nixl"
+
+    enable_bidirectional: Annotated[
+        bool,
+        Field(
+            description=(
+                "Whether Prefill workers can pull Decode-side KV through NIXL for later requests "
+                "in the same conversation."
+            ),
+        ),
+    ] = False
+    num_threads: Annotated[
+        int,
+        Field(ge=1, description="Number of NIXL connector threads."),
+    ] = 1
+    kv_recompute_threshold: Annotated[
+        int,
+        Field(
+            ge=0,
+            description=(
+                "Minimum number of remote Decode-side KV tokens required before a Prefill worker pulls "
+                "KV through NIXL instead of recomputing locally. Passed to NixlConnector extra config."
+            ),
+        ),
+    ] = 64
+    abort_timeout_seconds: Annotated[
+        int,
+        Field(
+            gt=0,
+            description=(
+                "Seconds vLLM NIXL waits for the peer to fetch held KV blocks before aborting and freeing them. "
+                "Exported as NIXL_ABORT_TIMEOUT and vLLM's VLLM_NIXL_ABORT_REQUEST_TIMEOUT."
+            ),
+        ),
+    ] = 480
+    router_cache_ttl_seconds: Annotated[
+        int | None,
+        Field(
+            gt=0,
+            description=(
+                "Seconds vllm-router keeps Decode-side KV metadata for bidirectional P/D reuse. "
+                "Defaults to 95% of abort_timeout_seconds."
+            ),
+        ),
+    ] = None
+
+    @model_validator(mode="after")
+    def validate_router_cache_ttl(self):
+        if self.router_cache_ttl_seconds is None:
+            self.router_cache_ttl_seconds = int(self.abort_timeout_seconds * 0.95)
+        if self.router_cache_ttl_seconds >= self.abort_timeout_seconds:
+            raise ValueError(
+                "router_cache_ttl_seconds must be less than abort_timeout_seconds "
+                f"({self.router_cache_ttl_seconds} >= {self.abort_timeout_seconds})"
+            )
+        return self
+
+
 class DisaggregatedInferenceDeploymentConfig(BaseInferenceDeploymentConfig):
     """Configures a disaggregated prefill/decode inference deployment.
 
@@ -211,6 +274,11 @@ class DisaggregatedInferenceDeploymentConfig(BaseInferenceDeploymentConfig):
         str, Field(description="Routing policy for the vllm-router (e.g. 'consistent_hash', 'round_robin').")
     ] = "consistent_hash"
 
+    kv_transport_config: Annotated[
+        NixlTransportConfig,
+        Field(description="KV transport settings for disaggregated P/D deployments."),
+    ] = NixlTransportConfig()
+
     prefill_env_overrides: Annotated[
         dict[str, str],
         Field(description="Extra environment variables exported only on prefill nodes."),

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

@@ -770,8 +770,9 @@ def auto_setup_lora(self):
 
     @model_validator(mode="after")
     def auto_setup_session_headers(self):
-        """Ensure X-Session-ID header is always set for sticky DP-aware routing at the inference router."""
+        """Ensure stable routing headers are set for inference routers."""
         self.orchestrator.client.extra_headers_from_state.setdefault("X-Session-ID", "example_id")
+        self.orchestrator.client.extra_headers_from_state.setdefault("X-Conversation-ID", "trajectory_id")
         return self
 
     @model_validator(mode="after")

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "torchaudio",
     "torchdata>=0.11.0",
     "transformers",
-    "vllm>=0.20.2",
+    "vllm==0.21.0",
     "wandb>=0.26.1",
     "ring-flash-attn>=0.1.8",
     "prime>=0.6.4",
@@ -176,6 +176,7 @@ override-dependencies = [
 [tool.uv.exclude-newer-package]
 # we want latest vllm, remove next patch
 vllm = false
+tokenspeed-mla = false
 flash_attn_3 = false
 # PrimeIntellect-published on PyPI (trusted publisher)
 prime = false
@@ -229,10 +230,10 @@ dion = { git = "https://github.com/samsja/dion.git", rev = "d891eeb" }
 transformers = { git = "https://github.com/huggingface/transformers.git", rev = "c1c3424" }
 flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", subdirectory = "flash_attn/cute", rev = "96bd151" }
 pydantic-config = { git = "https://github.com/samsja/pydantic_config.git", branch = "main" }
-vllm-router = { url = "https://github.com/PrimeIntellect-ai/router/releases/download/v0.1.22/vllm_router-0.1.22-cp38-abi3-manylinux_2_28_x86_64.whl" }
+vllm-router = { git = "https://github.com/PrimeIntellect-ai/router.git", rev = "23af7bb" }
 vllm = [
-    { url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_x86_64.whl", marker = "platform_machine == 'x86_64'" },
-    { url = "https://github.com/vllm-project/vllm/releases/download/v0.20.2/vllm-0.20.2+cu129-cp38-abi3-manylinux_2_31_aarch64.whl", marker = "platform_machine == 'aarch64'" },
+    { url = "https://files.pythonhosted.org/packages/73/6d/9b78990c9fabc70c7731de6af246a420156dc019f66b48da7c86f509c132/vllm-0.21.0-1-cp38-abi3-manylinux_2_24_x86_64.whl", marker = "platform_machine == 'x86_64'" },
+    { url = "https://files.pythonhosted.org/packages/ac/58/564b64d17dde6dc31faae836f98313538c152edf88e2a4fb43b9d551a635/vllm-0.21.0-1-cp38-abi3-manylinux_2_24_aarch64.whl", marker = "platform_machine == 'aarch64'" },
 ]
 deep-ep = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_ep-1.2.1+29d31c0-cp312-cp312-linux_x86_64.whl" }
 deep-gemm = { url = "https://github.com/PrimeIntellect-ai/prime-rl/releases/download/v0.5.0/deep_gemm-2.5.0+891d57b-cp312-cp312-linux_x86_64.whl" }

skills/config/SKILL.md

@@ -172,6 +172,29 @@ In TOML, an empty section header does the same:
 [ckpt]  # enables checkpointing with defaults
 ```
 
+### Disaggregated inference
+
+For `[deployment] type = "disaggregated"`, P/D NIXL transfer knobs live under `deployment.kv_transport_config`:
+
+```toml
+[deployment]
+type = "disaggregated"
+
+[deployment.kv_transport_config]
+type = "nixl"
+enable_bidirectional = true
+num_threads = 1
+kv_recompute_threshold = 64
+abort_timeout_seconds = 480
+router_cache_ttl_seconds = 456
+```
+
+`enable_bidirectional` defaults to `false`. When it is false, the Slurm templates pass `--pd-kv-cache-ttl-secs 0` to vllm-router so Decode-side KV metadata is not reused.
+`router_cache_ttl_seconds` can be omitted; it defaults to 95% of `abort_timeout_seconds` and must remain lower than the abort timeout.
+The Slurm templates export `abort_timeout_seconds` as both `NIXL_ABORT_TIMEOUT` and vLLM's `VLLM_NIXL_ABORT_REQUEST_TIMEOUT`.
+
+P/D NIXL deployments need UCX 1.19 or newer for H200 CUDA buffer registration. The Slurm templates add `$PROJECT_DIR/third_party/ucx` to `LD_LIBRARY_PATH`.
+
 ## Key files
 
 - `src/prime_rl/utils/config.py` — re-exports `BaseConfig` and `cli` from pydantic_config

src/prime_rl/entrypoints/inference.py

@@ -60,6 +60,7 @@ def write_slurm_script(config: InferenceConfig, config_path: Path, script_path:
             decode_port=config.deployment.decode_port,
             router_port=config.deployment.router_port,
             router_policy=config.deployment.router_policy,
+            kv_transport_config=config.deployment.kv_transport_config,
             data_parallel_rpc_port=config.data_parallel_rpc_port,
             use_deep_gemm=config.use_deep_gemm,
             prefill_env_overrides=config.deployment.prefill_env_overrides,

src/prime_rl/entrypoints/rl.py

@@ -437,8 +437,10 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) ->
             num_decode_replicas=infer_deploy.num_decode_replicas,
             gpus_per_node=config.deployment.gpus_per_node,
             router_port=infer_deploy.router_port,
+            router_policy=infer_deploy.router_policy,
             prefill_port=infer_deploy.prefill_port,
             decode_port=infer_deploy.decode_port,
+            kv_transport_config=infer_deploy.kv_transport_config,
             inference_tp=config.inference.parallel.tp,
             inference_data_parallel_rpc_port=config.inference.data_parallel_rpc_port,
             use_deep_gemm=config.inference.use_deep_gemm,

src/prime_rl/inference/patches.py

@@ -897,9 +897,9 @@ def monkey_patch_dp_engine_core_pause_resume_deadlock():
     - on resume, wake every DP rank and force an immediate global unfinished
       sync instead of waiting for the normal 32-step cadence
 
-    This keeps the upstream pause-side fix from
-    https://github.com/vllm-project/vllm/pull/37024 and extends it with the
-    resume-side wave-state fix.
+    This also bypasses vLLM's two-phase DP pause implementation
+    (https://github.com/vllm-project/vllm/pull/39366), which makes resume
+    reject states that our weight-update flow can validly hit.
     """
     from vllm.config import ParallelConfig
     from vllm.v1.core.sched.interface import PauseState
@@ -909,7 +909,8 @@ def monkey_patch_dp_engine_core_pause_resume_deadlock():
 
     _base_add_request = EngineCore.add_request
     _base_handle_client_request = EngineCoreProc._handle_client_request
-    _base_resume_scheduler = DPEngineCoreProc.resume_scheduler
+    _base_pause_complete = EngineCoreProc._pause_complete
+    _base_resume_scheduler = EngineCoreProc.resume_scheduler
 
     def _patched_add_request(self, request: Request, request_wave: int = 0):
         _base_add_request(self, request, request_wave)
@@ -930,8 +931,15 @@ def _patched_handle_client_request(self, request_type, request):
         else:
             _base_handle_client_request(self, request_type, request)
 
+    def _patched_pause_complete(self) -> bool:
+        self.pending_pause = False
+        self.ignore_start_dp_wave = False
+        return _base_pause_complete(self)
+
     def _patched_resume_scheduler(self):
         was_paused = self.scheduler.pause_state != PauseState.UNPAUSED
+        self.pending_pause = False
+        self.ignore_start_dp_wave = False
         _base_resume_scheduler(self)
         if was_paused:
             self.engines_running = True
@@ -948,6 +956,7 @@ def _patched_has_global_unfinished_reqs(self, local_unfinished: bool) -> bool:
 
     DPEngineCoreProc.add_request = _patched_add_request
     DPEngineCoreProc._handle_client_request = _patched_handle_client_request
+    DPEngineCoreProc._pause_complete = _patched_pause_complete
     DPEngineCoreProc.resume_scheduler = _patched_resume_scheduler
     DPEngineCoreProc._has_global_unfinished_reqs = _patched_has_global_unfinished_reqs
 

src/prime_rl/templates/inference.sbatch.j2

@@ -36,6 +36,8 @@ export PREFILL_PORT={{ prefill_port }}
 export DECODE_PORT={{ decode_port }}
 export ROUTER_PORT={{ router_port }}
 export RPC_PORT={{ data_parallel_rpc_port }}
+export NIXL_ABORT_TIMEOUT={{ kv_transport_config.abort_timeout_seconds }}
+export VLLM_NIXL_ABORT_REQUEST_TIMEOUT={{ kv_transport_config.abort_timeout_seconds }}
 {%- elif num_nodes > 1 %}
 export ROUTER_PORT={{ router_port }}
 export BACKEND_PORT={{ backend_port }}
@@ -171,15 +173,15 @@ srun bash -c '
     export VLLM_NIXL_SIDE_CHANNEL_PORT=5600
 
 {%- if kv_offload %}
-    PREFILL_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"'
+    PREFILL_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":{{ kv_transport_config.num_threads }},"bidirectional_kv_xfer":{{ "true" if kv_transport_config.enable_bidirectional else "false" }},"kv_recompute_threshold":{{ kv_transport_config.kv_recompute_threshold }}}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"'
 {%- else %}
-    PREFILL_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"'
+    PREFILL_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":{{ kv_transport_config.num_threads }},"bidirectional_kv_xfer":{{ "true" if kv_transport_config.enable_bidirectional else "false" }},"kv_recompute_threshold":{{ kv_transport_config.kv_recompute_threshold }}}}'"'"'
 {%- endif %}
 
 {%- if kv_offload %}
-    DECODE_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"'
+    DECODE_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":{{ kv_transport_config.num_threads }},"bidirectional_kv_xfer":{{ "true" if kv_transport_config.enable_bidirectional else "false" }},"kv_recompute_threshold":{{ kv_transport_config.kv_recompute_threshold }}}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"'
 {%- else %}
-    DECODE_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"'
+    DECODE_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":{{ kv_transport_config.num_threads }},"bidirectional_kv_xfer":{{ "true" if kv_transport_config.enable_bidirectional else "false" }},"kv_recompute_threshold":{{ kv_transport_config.kv_recompute_threshold }}}}'"'"'
 {%- endif %}
 
     DECODE_COMPILE_CFG='"'"'{"cudagraph_mode":"FULL_DECODE_ONLY"}'"'"'
@@ -250,6 +252,7 @@ srun bash -c '
             --host 0.0.0.0 \
             --port $ROUTER_PORT \
             --intra-node-data-parallel-size {{ dp_per_node }} \
+            --pd-kv-cache-ttl-secs {{ kv_transport_config.router_cache_ttl_seconds if kv_transport_config.enable_bidirectional else 0 }} \
             --worker-startup-timeout-secs 4200 \
             --log-level debug \
             >> $ROUTER_LOG 2>&1 &

src/prime_rl/templates/multi_node_rl.sbatch.j2

@@ -41,6 +41,8 @@ export NODES_PER_PREFILL_REPLICA=$((NUM_PREFILL_NODES / NUM_PREFILL_REPLICAS))
 export NODES_PER_DECODE_REPLICA=$((NUM_DECODE_NODES / NUM_DECODE_REPLICAS))
 export PREFILL_PORT={{ prefill_port }}
 export DECODE_PORT={{ decode_port }}
+export NIXL_ABORT_TIMEOUT={{ kv_transport_config.abort_timeout_seconds }}
+export VLLM_NIXL_ABORT_REQUEST_TIMEOUT={{ kv_transport_config.abort_timeout_seconds }}
 {%- else -%}
 export BACKEND_PORT={{ backend_port }}
 export INFERENCE_ENABLE_EXPERT_PARALLEL={{ "1" if inference_enable_expert_parallel else "0" }}
@@ -223,15 +225,15 @@ if [ "$SLURM_PROCID" -lt "$NUM_INFER_NODES" ]; then
     export VLLM_NIXL_SIDE_CHANNEL_PORT=5600
 
 {%- if kv_offload %}
-    PREFILL_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"'
+    PREFILL_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":{{ kv_transport_config.num_threads }},"bidirectional_kv_xfer":{{ "true" if kv_transport_config.enable_bidirectional else "false" }},"kv_recompute_threshold":{{ kv_transport_config.kv_recompute_threshold }}}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"'
 {%- else %}
-    PREFILL_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"'
+    PREFILL_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":{{ kv_transport_config.num_threads }},"bidirectional_kv_xfer":{{ "true" if kv_transport_config.enable_bidirectional else "false" }},"kv_recompute_threshold":{{ kv_transport_config.kv_recompute_threshold }}}}'"'"'
 {%- endif %}
 
 {%- if kv_offload %}
-    DECODE_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"'
+    DECODE_KV_CFG='"'"'{"kv_connector":"MultiConnector","kv_role":"kv_both","kv_connector_extra_config":{"connectors":[{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":{{ kv_transport_config.num_threads }},"bidirectional_kv_xfer":{{ "true" if kv_transport_config.enable_bidirectional else "false" }},"kv_recompute_threshold":{{ kv_transport_config.kv_recompute_threshold }}}},{"kv_connector":"OffloadingConnector","kv_role":"kv_both","kv_connector_extra_config":{"cpu_bytes_to_use":{{ kv_offload_cpu_bytes }}}}]}}'"'"'
 {%- else %}
-    DECODE_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":1}}'"'"'
+    DECODE_KV_CFG='"'"'{"kv_connector":"NixlConnector","kv_role":"kv_both","kv_connector_extra_config":{"num_threads":{{ kv_transport_config.num_threads }},"bidirectional_kv_xfer":{{ "true" if kv_transport_config.enable_bidirectional else "false" }},"kv_recompute_threshold":{{ kv_transport_config.kv_recompute_threshold }}}}'"'"'
 {%- endif %}
     DECODE_COMPILE_CFG='"'"'{"cudagraph_mode":"FULL_DECODE_ONLY"}'"'"'
 
@@ -296,12 +298,13 @@ if [ "$SLURM_PROCID" -lt "$NUM_INFER_NODES" ]; then
         REPLICA_ROUTER_ARGS=$(echo "$ALL_ROUTER_ARGS" | cut -d"|" -f$((REPLICA_IDX + 1)))
 
         vllm-router \
-            --policy consistent_hash \
+            --policy {{ router_policy }} \
             --vllm-pd-disaggregation \
             $REPLICA_ROUTER_ARGS \
             --host 0.0.0.0 \
             --port $ROUTER_PORT \
             --intra-node-data-parallel-size {{ dp_per_node }} \
+            --pd-kv-cache-ttl-secs {{ kv_transport_config.router_cache_ttl_seconds if kv_transport_config.enable_bidirectional else 0 }} \
             --worker-startup-timeout-secs 4200 \
             >> $ROUTER_LOG 2>&1 &
     fi