diff --git a/inferencex-e2e/configs/CONFIGS.md b/inferencex-e2e/configs/CONFIGS.md index 2e63d9eee9..b097e3163b 100644 --- a/inferencex-e2e/configs/CONFIGS.md +++ b/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 diff --git a/inferencex-e2e/infx/matrix/generate.py b/inferencex-e2e/infx/matrix/generate.py index 5865039af3..7cf057e829 100644 --- a/inferencex-e2e/infx/matrix/generate.py +++ b/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: diff --git a/inferencex-e2e/infx/matrix/validation.py b/inferencex-e2e/infx/matrix/validation.py index 1ae17ab77f..fa5a22e7ea 100644 --- a/inferencex-e2e/infx/matrix/validation.py +++ b/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 diff --git a/inferencex-e2e/infx/tests/matrix/test_generate_sweep_configs.py b/inferencex-e2e/infx/tests/matrix/test_generate_sweep_configs.py index 66f5871307..5507f92b90 100644 --- a/inferencex-e2e/infx/tests/matrix/test_generate_sweep_configs.py +++ b/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 ): diff --git a/inferencex-e2e/infx/tests/matrix/test_validation.py b/inferencex-e2e/infx/tests/matrix/test_validation.py index ac78746d81..58a38dc1c6 100644 --- a/inferencex-e2e/infx/tests/matrix/test_validation.py +++ b/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(**{