feat: add P/D disaggregation output to llm-d generator

Add prefill/decode disaggregation as an output option to the llm-d
deployment generator. When pd_enabled=True, the generator produces
separate Kustomize patches for prefill and decode deployments instead
of the single patch-vllm.yaml.

All three patch variants (vllm, prefill, decode) use a single unified
template (patch-modelserver.yaml.j2) rendered with different context
(deployment_name, replica_count, extra_args) to avoid maintaining
near-identical templates that would drift. If the roles diverge
significantly, splitting is a one-step refactor.

Also adds Field(ge=1) validation on prefill_replicas/decode_replicas
in the API schema.

Signed-off-by: Amit Oren <amoren@redhat.com>

Author: amito

src/planner/api/routes/configuration.py

@@ -6,7 +6,7 @@
 
 import yaml
 from fastapi import APIRouter, Depends, HTTPException, Request, status
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from starlette.concurrency import run_in_threadpool
 
 from planner.api.dependencies import (
@@ -32,6 +32,9 @@ class GenerateDeploymentRequest(BaseModel):
     configuration: DeploymentConfiguration
     namespace: str = "default"
     stack: StackType = "vllm"
+    pd_enabled: bool = False
+    prefill_replicas: int = Field(1, ge=1, le=32)
+    decode_replicas: int = Field(1, ge=1, le=32)
 
 
 class DeploymentModeRequest(BaseModel):
@@ -53,27 +56,21 @@ def _generate_yaml_from_config(
     deployment_generator: DeploymentGenerator,
     llmd_generator: LlmdDeploymentGenerator,
     yaml_validator: YAMLValidator,
+    pd_enabled: bool = False,
+    prefill_replicas: int = 1,
+    decode_replicas: int = 1,
 ) -> dict[str, Any]:
-    """Generate YAML files from a deployment configuration.
-
-    Args:
-        config: Deployment configuration
-        namespace: Kubernetes namespace
-        stack: Deployment stack (vllm or llm-d)
-        deployment_generator: vLLM deployment generator
-        llmd_generator: llm-d deployment generator
-        yaml_validator: YAML validator
-
-    Returns:
-        Dict with deployment_id, namespace, files (file paths), and contents (YAML strings)
-
-    Raises:
-        HTTPException: If stack is unknown or YAML validation fails
-    """
+    """Generate YAML files from a deployment configuration."""
     logger.info(f"Generating deployment for model: {config.model_name} (stack={stack})")
 
     if stack == "llm-d":
-        result = llmd_generator.generate_all(config=config, namespace=namespace)
+        result = llmd_generator.generate_all(
+            config=config,
+            namespace=namespace,
+            pd_enabled=pd_enabled,
+            prefill_replicas=prefill_replicas,
+            decode_replicas=decode_replicas,
+        )
     elif stack == "vllm":
         result = deployment_generator.generate_all(config=config, namespace=namespace)
     else:
@@ -132,6 +129,9 @@ async def generate_deployment(
             deployment_generator,
             llmd_generator,
             yaml_validator,
+            pd_enabled=request.pd_enabled,
+            prefill_replicas=request.prefill_replicas,
+            decode_replicas=request.decode_replicas,
         )
 
         return DeploymentBundle(

src/planner/configuration/llmd_generator.py

@@ -70,30 +70,93 @@ def generate_all(
         self,
         config: DeploymentConfiguration,
         namespace: str = "default",
+        pd_enabled: bool = False,
+        prefill_replicas: int = 1,
+        decode_replicas: int = 1,
     ) -> dict[str, Any]:
         """Generate all llm-d deployment files.
 
         Returns a dict with: deployment_id, namespace, files, contents.
+
+        When *pd_enabled* is True, separate prefill and decode patches are
+        generated instead of the single ``patch-vllm.yaml``.
         """
+        if not 1 <= prefill_replicas <= 32:
+            msg = f"prefill_replicas must be between 1 and 32, got {prefill_replicas}"
+            raise ValueError(msg)
+        if not 1 <= decode_replicas <= 32:
+            msg = f"decode_replicas must be between 1 and 32, got {decode_replicas}"
+            raise ValueError(msg)
+
         deployment_id = generate_deployment_id(config)
         context = self._prepare_context(config, deployment_id, namespace)
 
-        configs: list[tuple[str, str, str]] = [
-            ("kustomization.yaml.j2", "modelserver/kustomization.yaml", "kustomization"),
-            ("patch-vllm.yaml.j2", "modelserver/patch-vllm.yaml", "patch_vllm"),
-            ("values.yaml.j2", "scheduler/values.yaml", "helm_values"),
+        context["pd_enabled"] = pd_enabled
+
+        # Build the list of (template, output path, key, extra_context) tuples.
+        # All model-server patches use the same unified template with different
+        # deployment_name, replica_count, and extra_args.
+        patch_template = "patch-modelserver.yaml.j2"
+
+        configs: list[tuple[str, str, str, dict[str, Any]]] = [
+            ("kustomization.yaml.j2", "modelserver/kustomization.yaml", "kustomization", {}),
         ]
 
+        if pd_enabled:
+            configs.append(
+                (
+                    patch_template,
+                    "modelserver/patch-prefill.yaml",
+                    "patch_prefill",
+                    {
+                        "deployment_name": "prefill",
+                        "replica_count": prefill_replicas,
+                        "extra_args": [
+                            "--kv-connector=nixlv2",
+                            "--kv-role=kv_producer",
+                            "--enable-chunked-prefill",
+                        ],
+                    },
+                )
+            )
+            configs.append(
+                (
+                    patch_template,
+                    "modelserver/patch-decode.yaml",
+                    "patch_decode",
+                    {
+                        "deployment_name": "decode",
+                        "replica_count": decode_replicas,
+                        "extra_args": ["--kv-connector=nixlv2", "--kv-role=kv_consumer"],
+                    },
+                )
+            )
+        else:
+            configs.append(
+                (
+                    patch_template,
+                    "modelserver/patch-vllm.yaml",
+                    "patch_vllm",
+                    {
+                        "deployment_name": "decode",
+                        "replica_count": context["replicas"],
+                        "extra_args": [],
+                    },
+                )
+            )
+
+        configs.append(("values.yaml.j2", "scheduler/values.yaml", "helm_values", {}))
+
         deployment_dir = self.output_dir / deployment_id
         (deployment_dir / "modelserver").mkdir(parents=True, exist_ok=True)
         (deployment_dir / "scheduler").mkdir(parents=True, exist_ok=True)
 
         generated_files: dict[str, str] = {}
         generated_contents: dict[str, str] = {}
 
-        for template_name, output_rel_path, config_type in configs:
+        for template_name, output_rel_path, config_type, extra_ctx in configs:
             template = self.env.get_template(template_name)
-            rendered = template.render(**context)
+            rendered = template.render(**context, **extra_ctx)
 
             output_path = deployment_dir / output_rel_path
             output_path.parent.mkdir(parents=True, exist_ok=True)

src/planner/configuration/templates/llmd/kustomization.yaml.j2

@@ -39,4 +39,9 @@ labels:
         create: true
 
 patches:
+{% if pd_enabled %}
+  - path: patch-prefill.yaml
+  - path: patch-decode.yaml
+{% else %}
   - path: patch-vllm.yaml
+{% endif %}

…ration/templates/llmd/patch-vllm.yaml.j2 …templates/llmd/patch-modelserver.yaml.j2src/planner/configuration/templates/llmd/patch-vllm.yaml.j2 renamed to src/planner/configuration/templates/llmd/patch-modelserver.yaml.j2 src/planner/configuration/templates/llmd/patch-vllm.yaml.j2 renamed to src/planner/configuration/templates/llmd/patch-modelserver.yaml.j2

@@ -1,9 +1,9 @@
 apiVersion: apps/v1
 kind: Deployment
 metadata:
-  name: decode
+  name: {{ deployment_name }}
 spec:
-  replicas: {{ replicas }}
+  replicas: {{ replica_count }}
   template:
     spec:
       containers:
@@ -13,6 +13,9 @@ spec:
             - "{{ model_id }}"
             - "--tensor-parallel-size={{ tensor_parallel }}"
             - "--port=8000"
+{% for arg in extra_args %}
+            - "{{ arg }}"
+{% endfor %}
           resources:
             requests:
               nvidia.com/gpu: "{{ gpus_per_replica }}"

tests/unit/test_llmd_generator.py

@@ -348,3 +348,195 @@ def test_generate_deployment_with_stack_vllm_is_default(
         assert response.status_code == 200
         data = response.json()
         assert "inferenceservice" in data["files"]
+
+
+@pytest.mark.unit
+class TestPDDisaggregation:
+    """Tests for P/D (prefill/decode) disaggregation output."""
+
+    def test_pd_disabled_produces_single_patch(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """Default (pd_enabled=False) has patch_vllm, no patch_prefill/patch_decode."""
+        result = llmd_generator.generate_all(sample_config)
+
+        assert "patch_vllm" in result["contents"]
+        assert "patch_prefill" not in result["contents"]
+        assert "patch_decode" not in result["contents"]
+
+    def test_pd_disabled_patch_targets_decode_deployment(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """Non-PD patch must use name 'decode' to match the llm-d base Deployment."""
+        result = llmd_generator.generate_all(sample_config)
+        parsed = yaml.safe_load(result["contents"]["patch_vllm"])
+
+        assert parsed["metadata"]["name"] == "decode"
+
+    def test_pd_enabled_produces_prefill_and_decode_patches(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """pd_enabled=True produces patch_prefill and patch_decode, no patch_vllm."""
+        result = llmd_generator.generate_all(sample_config, pd_enabled=True)
+
+        assert "patch_prefill" in result["contents"]
+        assert "patch_decode" in result["contents"]
+        assert "patch_vllm" not in result["contents"]
+
+    def test_pd_prefill_patch_has_correct_replicas(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """Prefill patch uses prefill_replicas value."""
+        result = llmd_generator.generate_all(
+            sample_config, pd_enabled=True, prefill_replicas=2
+        )
+        parsed = yaml.safe_load(result["contents"]["patch_prefill"])
+
+        assert parsed["spec"]["replicas"] == 2
+        assert parsed["metadata"]["name"] == "prefill"
+
+    def test_pd_decode_patch_has_correct_replicas(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """Decode patch uses decode_replicas value."""
+        result = llmd_generator.generate_all(
+            sample_config, pd_enabled=True, decode_replicas=3
+        )
+        parsed = yaml.safe_load(result["contents"]["patch_decode"])
+
+        assert parsed["spec"]["replicas"] == 3
+        assert parsed["metadata"]["name"] == "decode"
+
+    def test_pd_kustomization_references_both_patches(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """Kustomization patches list has patch-prefill.yaml and patch-decode.yaml."""
+        result = llmd_generator.generate_all(sample_config, pd_enabled=True)
+        parsed = yaml.safe_load(result["contents"]["kustomization"])
+
+        patch_paths = [p["path"] for p in parsed["patches"]]
+        assert "patch-prefill.yaml" in patch_paths
+        assert "patch-decode.yaml" in patch_paths
+        assert "patch-vllm.yaml" not in patch_paths
+
+    def test_pd_all_outputs_valid_yaml(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """All contents parse as valid YAML when pd_enabled=True."""
+        result = llmd_generator.generate_all(sample_config, pd_enabled=True)
+
+        for key, content in result["contents"].items():
+            parsed = yaml.safe_load(content)
+            assert parsed is not None, f"{key} rendered as empty YAML"
+
+    def test_rejects_zero_prefill_replicas(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """generate_all() raises ValueError when prefill_replicas < 1."""
+        with pytest.raises(ValueError, match="must be between 1 and 32"):
+            llmd_generator.generate_all(sample_config, pd_enabled=True, prefill_replicas=0)
+
+    def test_rejects_zero_decode_replicas(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """generate_all() raises ValueError when decode_replicas < 1."""
+        with pytest.raises(ValueError, match="must be between 1 and 32"):
+            llmd_generator.generate_all(sample_config, pd_enabled=True, decode_replicas=0)
+
+    def test_rejects_prefill_replicas_above_max(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """generate_all() raises ValueError when prefill_replicas > 32."""
+        with pytest.raises(ValueError, match="must be between 1 and 32"):
+            llmd_generator.generate_all(sample_config, pd_enabled=True, prefill_replicas=33)
+
+    def test_rejects_decode_replicas_above_max(
+        self,
+        llmd_generator: LlmdDeploymentGenerator,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """generate_all() raises ValueError when decode_replicas > 32."""
+        with pytest.raises(ValueError, match="must be between 1 and 32"):
+            llmd_generator.generate_all(sample_config, pd_enabled=True, decode_replicas=33)
+
+
+@pytest.mark.unit
+class TestDeployAPINewParams:
+    """Tests for new parameters exposed in the deploy API endpoint."""
+
+    def test_generate_deployment_llmd_with_pd_enabled(
+        self, client: TestClient, sample_config: DeploymentConfiguration
+    ) -> None:
+        """POST with pd_enabled=True should return patch_prefill and patch_decode."""
+        response = client.post(
+            "/api/v1/generate-deployment",
+            json={
+                "configuration": sample_config.model_dump(),
+                "namespace": "test-ns",
+                "stack": "llm-d",
+                "pd_enabled": True,
+            },
+        )
+        assert response.status_code == 200
+        data = response.json()
+        assert "patch_prefill" in data["files"]
+        assert "patch_decode" in data["files"]
+
+    @pytest.mark.parametrize(
+        "field",
+        ["prefill_replicas", "decode_replicas"],
+    )
+    def test_generate_deployment_rejects_zero_replicas(
+        self,
+        client: TestClient,
+        sample_config: DeploymentConfiguration,
+        field: str,
+    ) -> None:
+        """API returns 422 when replica count is < 1."""
+        payload = {
+            "configuration": sample_config.model_dump(),
+            "namespace": "test-ns",
+            "stack": "llm-d",
+            "pd_enabled": True,
+            field: 0,
+        }
+        response = client.post("/api/v1/generate-deployment", json=payload)
+        assert response.status_code == 422
+
+    def test_generate_deployment_rejects_replicas_above_max(
+        self,
+        client: TestClient,
+        sample_config: DeploymentConfiguration,
+    ) -> None:
+        """API returns 422 when replica count exceeds 32."""
+        response = client.post(
+            "/api/v1/generate-deployment",
+            json={
+                "configuration": sample_config.model_dump(),
+                "namespace": "test-ns",
+                "stack": "llm-d",
+                "pd_enabled": True,
+                "prefill_replicas": 33,
+            },
+        )
+        assert response.status_code == 422