[feat] Add suppport for mooncake connector for P/D; custom P/D init kwargs

[SumanthRH]

What does this PR do?

Extends prefill/decode (P/D) disaggregated serving in two ways.

First, it adds role-specific engine init kwargs (prefill_init_kwargs / decode_init_kwargs) so prefill and decode servers can be launched with different vLLM args (e.g. a different all2all_backend or kv_transfer_config per role).

Second, it adds a Mooncake P2P KV-transfer path alongside the existing NIXL path, including bootstrap-server-based router discovery and a runtime patch that keeps a MultiConnector (Mooncake P2P + store) stack from tripping a Prometheus-stats assert in vLLM. It also ships a new mooncake optional-dependency extra, CPU unit tests for the new config/CLI/router logic, and docs.

Changes

Config + validation

• New InferenceEngineConfig.prefill_init_kwargs / decode_init_kwargs (both Dict[str, Any], default empty) in skyrl/train/config/config.py. Pass-through vLLM engine kwargs applied per role.
• validate_inference_engine_cfg (skyrl/train/utils/utils.py) enforces that these require enable_pd=true, are mutually exclusive with engine_init_kwargs, and — new — runs a fail-fast completeness check: since role kwargs replace engine_init_kwargs entirely, each of prefill_init_kwargs/decode_init_kwargs must be set and must itself carry a kv_transfer_config. This fails at config-validation time instead of per-role at serve-setup time in get_pd_cli_args.

Role-specific CLI-arg plumbing

• get_pd_cli_args (inference_servers/utils.py) gains role / role_init_kwargs: applies role kwargs onto a deep copy of the base args before reading kv_transfer_config.
• create_inference_servers (inference_servers/setup.py) builds separate prefill_cli_args / decode_cli_args via get_config_as_dict(...); each server group uses its own args.
• kv_role relaxed to setdefault("kv_role", "kv_both") so Mooncake's push flow can set explicit kv_producer/kv_consumer.

Mooncake connector detection

• New get_pd_p2p_connector_name(kv_config) resolves the P2P connector from a bare NixlConnector/MooncakeConnector or a MultiConnector wrapping exactly one P2P + optional store connectors; raises ValueError otherwise. SUPPORTED_PD_P2P_CONNECTORS = ("NixlConnector", "MooncakeConnector"). Replaces the old "only NixlConnector" string check.

Router bootstrap ports

• build_router_args (inference_servers/utils.py) gains prefill_bootstrap_ports and pd_kv_connector. For Mooncake, prefill_urls becomes list(zip(prefill_urls, prefill_bootstrap_ports)) and kv_connector="mooncake" is passed to RouterArgs.
• create_inference_servers computes prefill_bootstrap_ports = NIXL_SIDE_CHANNEL_BASE_PORT + i*servers_per_group + j across prefill servers when the resolved P2P connector is MooncakeConnector, and sets pd_kv_connector="mooncake".

Server-actor bootstrap port

• VLLMServerActor (inference_servers/vllm_server_actor.py) sets VLLM_MOONCAKE_BOOTSTRAP_PORT = nixl_side_channel_base + server_idx for Mooncake (matching the router's per-server value); the NIXL side-channel path (find_and_reserve_port + _setup_nixl_side_channel) is unchanged. Also calls apply_multi_connector_stats_patch() before building the engine.

MultiConnector stats patch

• New skyrl/backends/skyrl_train/patches/vllm/patch_multi_connector_stats.py: a runtime monkeypatch of vLLM's MultiKVConnectorPromMetrics.observe that skips child connectors which report KVConnectorStats but have no registered Prometheus metrics class (e.g. MooncakeConnector) instead of asserting. Without it, a MultiConnector (Mooncake P2P + store) stack crashes AsyncLLM.output_handler on the first stats flush and hangs in-flight requests.

Packaging

• New mooncake optional-dependency extra in pyproject.toml: mooncake-transfer-engine; sys_platform == 'linux' — the off-the-shelf PyPI wheel providing the Mooncake python bindings and mooncake_master binary (store connector + P2P PD transfer). Linux-only and intentionally not in the uv conflicts set (it does not clash with the fsdp/megatron/jax backend extras). uv.lock updated accordingly.

Docs

• .claude/docs/inference.md P/D Disaggregation section rewritten to document the supported P2P transfer connectors (NixlConnector pull-based vs MooncakeConnector push-based, bare or inside a MultiConnector with a store connector), the bootstrap-port/kv_connector=mooncake router mode, the stats patch, and the role-specific prefill_init_kwargs/decode_init_kwargs rules.
• docs/content/docs/configuration/config.mdx: adds prefill_init_kwargs / decode_init_kwargs to the config listing and field reference (per-role pass-through kwargs, mutual exclusion with engine_init_kwargs, both-required + per-role kv_transfer_config).

Test Plan

Tests added (CPU, no GPU):

• tests/train/test_config.py — role-kwargs config validation: test_role_init_kwargs_conflict_with_engine_init_kwargs_rejected, test_role_init_kwargs_require_enable_pd, test_partial_role_init_kwargs_rejected (partial/missing kv_transfer_config rejected), test_valid_pd_role_init_kwargs_passes.
• tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py:
  • TestGetPDP2PConnectorName — test_bare_nixl, test_bare_mooncake, test_multiconnector_resolves_single_p2p, test_multiconnector_zero_p2p_raises, test_multiconnector_two_p2p_raises, test_unsupported_bare_connector_raises.
  • TestGetPDCLIArgs — test_role_init_kwargs_applied, test_kv_role_defaults_to_kv_both, test_kv_role_preserved_when_set, test_missing_kv_transfer_config_raises.
  • test_build_vllm_cli_args_succeeds_on_gpu_less_host — marked @pytest.mark.vllm (imports vllm.platforms); it needs the vllm package (only present under --extra fsdp/--extra megatron) and was NOT run in the CPU-only env (errors with ModuleNotFoundError: No module named 'vllm' there, so it is deselected below).
• tests/backends/skyrl_train/inference_servers/test_vllm_router.py — Mooncake pd_mode router args: test_pd_mode_mooncake (prefill URLs zipped with bootstrap ports, kv_connector="mooncake"), test_pd_mode_bootstrap_ports_length_mismatch_raises, test_pd_mode_nixl_default_has_none_bootstrap_port.

uv run --isolated --extra dev --extra skyrl-train pytest \
  tests/train/test_config.py::test_role_init_kwargs_conflict_with_engine_init_kwargs_rejected \
  tests/train/test_config.py::test_role_init_kwargs_require_enable_pd \
  tests/train/test_config.py::test_partial_role_init_kwargs_rejected \
  tests/train/test_config.py::test_valid_pd_role_init_kwargs_passes \
  tests/backends/skyrl_train/inference_servers/test_build_vllm_cli_args.py \
  tests/backends/skyrl_train/inference_servers/test_vllm_router.py \
  -m "not vllm"

Manual / GPU validation (multi-node serve):

• NIXL P/D regression: uv run --isolated --extra dev --extra fsdp pytest tests/backends/skyrl_train/gpu/gpu_ci/integrations/test_pd_routing.py -v -s.
• Mooncake P/D serve on a multi-node GPU cluster: launch with enable_pd=true and role-specific prefill_init_kwargs/decode_init_kwargs carrying a MooncakeConnector (optionally a MultiConnector with MooncakeStoreConnector); verify the router discovers each prefill's bootstrap server, requests complete, and the tolerant MultiKVConnectorPromMetrics.observe patch log line appears without the output-handler crash/hang.