feat(matrix): budget decode GPUs for P/D-shared DRAM offload pools / feat(matrix): 为 P/D 共享的 DRAM 卸载池计入解码 GPU 预算

inferencex-e2e/configs/CONFIGS.md

@@ -117,6 +117,8 @@ The below list describes what each field is:
   - `agentic-coding`: Agentic trace replay benchmarks using real conversation traces. Each entry must have:
     - `trace-source`: Identifier for the trace dataset to use.
     - `search-space`: Same structure as `fixed-seq-len` search-space entries.
+    - `dram-utilization`: Required when a search-space entry sets `kv-offloading: dram`; a fraction in (0, 1]. `total-cpu-dram-gb` is this fraction of the node's available CPU DRAM, capped at the AgentX ceiling, times the share of the node's GPUs that claim it.
+    - (Optional) `dram-budget-scope`: Which GPUs of a disaggregated entry claim that share. `prefill` (default) counts one prefill replica's GPUs on its node, for a pool only prefill workers own. `all-workers` counts every prefill and decode GPU of the deployment, the GPUs that normalize per-GPU throughput, for one pool that prefill reaches on every worker's node, such as a Mooncake Store with owners beside decode. That total may span nodes; each node holds at most the share of the deployment's GPUs on it. `all-workers` requires `prefill` and `decode` workers without per-worker `hardware`.
 
 `router` and `kv-p2p-transfer` may be omitted independently. Router metadata
 requires non-empty `name` and `version` strings. Its `version` must be the

inferencex-e2e/infx/matrix/generate.py

@@ -434,6 +434,16 @@ def multinode_worker_pair(benchmark: dict, disagg: bool) -> tuple[dict, dict]:
     return prefill, decode
 
 
+def worker_replica_gpus(worker: dict) -> int:
+    """Return GPUs one worker replica allocates; EP and DCP reuse its TP GPUs."""
+    return (
+        worker[Fields.TP.value]
+        * worker.get(Fields.PP.value, 1)
+        * worker.get(Fields.PCP_SIZE.value, 1)
+        * worker.get(Fields.DWDP_SIZE.value, 1)
+    )
+
+
 def worker_gpus_per_node(worker: dict, gpus_per_node: int) -> int:
     """Return GPUs a single worker replica occupies on one node.
 
@@ -449,12 +459,7 @@ def worker_gpus_per_node(worker: dict, gpus_per_node: int) -> int:
     * A replica within one node must divide it evenly so co-located replicas of
       the same role tile the node without overlap.
     """
-    gpus_per_replica = (
-        worker[Fields.TP.value]
-        * worker.get(Fields.PP.value, 1)
-        * worker.get(Fields.PCP_SIZE.value, 1)
-        * worker.get(Fields.DWDP_SIZE.value, 1)
-    )
+    gpus_per_replica = worker_replica_gpus(worker)
     factors = f"{Fields.TP.value}*{Fields.PP.value}*{Fields.PCP_SIZE.value}"
     if worker.get(Fields.DWDP_SIZE.value, 1) > 1:
         factors += f"*{Fields.DWDP_SIZE.value}"
@@ -475,18 +480,27 @@ def worker_gpus_per_node(worker: dict, gpus_per_node: int) -> int:
     return gpus_per_replica
 
 
+def worker_pool_gpus(worker: dict, gpus_per_node: int) -> int:
+    """Return GPUs all of a worker pool's replicas allocate, across every node they span."""
+    worker_gpus_per_node(worker, gpus_per_node)  # same node-tiling rules as a replica
+    return worker[Fields.NUM_WORKER.value] * worker_replica_gpus(worker)
+
+
 def agentic_dram_offload_gb(
     agentic_config: dict, benchmark: dict, runner: str, runner_data: dict
 ) -> int:
-    """Return the DRAM offload budget (GB) for one agentic server node.
+    """Return the DRAM offload budget (GB) of one agentic offload pool.
 
     The budget scales the node's (MAX-capped) available CPU DRAM by the
-    utilization and by the fraction of the node's GPUs in use:
+    utilization and by the GPUs that claim it, as a fraction of a node's GPUs:
 
     * Single-node entries use the TP/PP/PCP topology, which must fit one node.
-    * Disaggregated multinode entries use the prefill worker's per-node GPU
-      footprint, since only prefill offloads KV to CPU DRAM today (decode can be
-      budgeted separately if it ever gains its own pool).
+    * Disaggregated multinode entries default to the prefill worker's per-node GPU
+      footprint, for pools only prefill workers own.
+    * dram-budget-scope 'all-workers' budgets one pool that prefill reaches on
+      every worker's node, such as a Mooncake Store with owners beside decode: each
+      prefill and decode GPU claims its node share, the GPUs that normalize
+      per-GPU throughput. The total may span nodes.
     """
     kv_offloading = benchmark.get(Fields.KV_OFFLOADING.value, "none")
     if kv_offloading != "dram":
@@ -501,6 +515,10 @@ def agentic_dram_offload_gb(
 
     if Fields.WORKER.value in benchmark:
         gpu_count = worker_gpus_per_node(benchmark[Fields.WORKER.value], gpus_per_node)
+    elif agentic_config.get(Fields.DRAM_BUDGET_SCOPE.value, "prefill") == "all-workers":
+        gpu_count = worker_pool_gpus(
+            benchmark[Fields.PREFILL.value], gpus_per_node
+        ) + worker_pool_gpus(benchmark[Fields.DECODE.value], gpus_per_node)
     elif Fields.PREFILL.value in benchmark:
         gpu_count = worker_gpus_per_node(benchmark[Fields.PREFILL.value], gpus_per_node)
     else:

inferencex-e2e/infx/matrix/validation.py

@@ -77,6 +77,7 @@ class Fields(Enum):
     TOTAL_CPU_DRAM_GB = "total-cpu-dram-gb"
     AVAILABLE_CPU_DRAM_MIB = "available-cpu-dram-mib"
     DRAM_UTILIZATION = "dram-utilization"
+    DRAM_BUDGET_SCOPE = "dram-budget-scope"
     GPUS_PER_NODE = "gpus-per-node"
     NUM_NODES = "num-nodes"
     NODE_COUNT = "node-count"
@@ -778,6 +779,9 @@ class AgenticCodingConfig(BaseModel):
     dram_utilization: float | None = Field(
         default=None, alias=Fields.DRAM_UTILIZATION.value, gt=0, le=1
     )
+    dram_budget_scope: Literal["prefill", "all-workers"] = Field(
+        default="prefill", alias=Fields.DRAM_BUDGET_SCOPE.value
+    )
 
     @model_validator(mode="after")
     def validate_dram_offload_capacity(self) -> Self:
@@ -789,6 +793,19 @@ def validate_dram_offload_capacity(self) -> Self:
                     f"{Fields.KV_OFFLOADING.value}='dram' requires "
                     f"{Fields.DRAM_UTILIZATION.value} with runner hardware metadata"
                 )
+            if self.dram_budget_scope != "all-workers":
+                continue
+            if entry.prefill is None or entry.decode is None:
+                raise ValueError(
+                    f"{Fields.DRAM_BUDGET_SCOPE.value}='all-workers' requires "
+                    f"'{Fields.PREFILL.value}' and '{Fields.DECODE.value}' workers"
+                )
+            # The budget reads one runner's node DRAM for both roles.
+            if entry.prefill.hardware is not None:
+                raise ValueError(
+                    f"{Fields.DRAM_BUDGET_SCOPE.value}='all-workers' does not support "
+                    f"per-worker '{Fields.HARDWARE.value}'"
+                )
         return self
 
 

inferencex-e2e/infx/tests/matrix/test_generate_sweep_configs.py

@@ -2867,6 +2867,63 @@ def test_multinode_agentic_budget_ignores_decode_topology(
         assert len(result) == 1
         assert result[0]["total-cpu-dram-gb"] == 2399
 
+    @pytest.mark.parametrize(
+        ("prefill", "decode", "budget"),
+        [
+            # DCP reuses decode's TP GPUs: 4 + 4 GPUs claim one whole node.
+            (
+                {"num-worker": 1, "tp": 1, "pp": 4, "ep": 1, "dp-attn": False},
+                {"num-worker": 1, "tp": 4, "dcp-size": 4, "ep": 1, "dp-attn": False},
+                2399,
+            ),
+            # Each decode replica claims its own share: 4 + 2 * 4 GPUs.
+            (
+                {"num-worker": 1, "tp": 4, "ep": 1, "dp-attn": False},
+                {"num-worker": 2, "tp": 4, "ep": 1, "dp-attn": False},
+                3599,
+            ),
+            # A prefill replica spanning two nodes claims both: 16 + 8 GPUs.
+            (
+                {"num-worker": 1, "tp": 8, "pp": 2, "ep": 1, "dp-attn": False},
+                {"num-worker": 1, "tp": 8, "ep": 1, "dp-attn": False},
+                7199,
+            ),
+        ],
+    )
+    def test_multinode_agentic_all_workers_budget_counts_decode_gpus(
+        self, sample_runner_config, generate_agentic_sweep, prefill, decode, budget
+    ):
+        config = {
+            "dsv4-agentic-shared-pool": {
+                "image": "sglang-rocm",
+                "model": "deepseek-ai/DeepSeek-V4-Pro",
+                "model-prefix": "dsv4",
+                "precision": "fp4",
+                "framework": "sglang-disagg",
+                "runner": "cluster:mi355x-amds",
+                "multinode": True,
+                "disagg": True,
+                "kv-p2p-transfer": "mooncake",
+                "scenarios": {
+                    "agentic-coding": [{
+                        "dram-utilization": 0.80,
+                        "dram-budget-scope": "all-workers",
+                        "search-space": [{
+                            "conc-list": [16],
+                            "kv-offloading": "dram",
+                            "kv-offload-backend": {"name": "mooncake", "version": "0.3.13"},
+                            "prefill": prefill,
+                            "decode": decode,
+                        }],
+                    }],
+                },
+            },
+        }
+
+        result = generate_agentic_sweep(config, sample_runner_config)
+
+        assert [entry["total-cpu-dram-gb"] for entry in result] == [budget]
+
     def test_multinode_agentic_rejects_node_misaligned_prefill(
         self, sample_runner_config, generate_agentic_sweep
     ):

inferencex-e2e/infx/tests/matrix/test_validation.py

@@ -428,6 +428,36 @@ def test_dram_kv_offload_requires_dram_utilization(self):
                 }],
             })
 
+    @pytest.mark.parametrize(
+        ("entry", "match"),
+        [
+            (
+                {"tp": 4, "kv-offloading": "dram", "kv-offload-backend": {"name": "native"}},
+                "requires 'prefill' and 'decode' workers",
+            ),
+            (
+                {
+                    "kv-offloading": "dram",
+                    "kv-offload-backend": {"name": "mooncake"},
+                    "prefill": {
+                        "hardware": "gb200", "num-worker": 1, "tp": 4, "ep": 1, "dp-attn": False,
+                    },
+                    "decode": {
+                        "hardware": "h100", "num-worker": 1, "tp": 4, "ep": 1, "dp-attn": False,
+                    },
+                },
+                "per-worker 'hardware'",
+            ),
+        ],
+    )
+    def test_all_workers_dram_budget_rejects_unsupported_entries(self, entry, match):
+        with pytest.raises(ValidationError, match=match):
+            AgenticCodingConfig(**{
+                "dram-utilization": 0.80,
+                "dram-budget-scope": "all-workers",
+                "search-space": [{**entry, "conc-list": [16]}],
+            })
+
     def test_agentic_search_space_rejects_total_cpu_dram_gb(self):
         with pytest.raises(ValidationError, match="total-cpu-dram-gb"):
             AgenticCodingSearchSpaceEntry(**{