[codex] support optional heterogeneous hardware metadata (#2100)

* feat: add heterogeneous disaggregated hardware metadata

* fix: make heterogeneous hardware metadata optional

* fix: restrict disaggregation to multinode configs

* refactor: colocate hardware with worker configs

* fix: surface heterogeneous hardware in agentic results

Author: cquil11

.github/workflows/benchmark-multinode-tmpl.yml

@@ -39,6 +39,16 @@ on:
       disagg:
         required: true
         type: string
+      prefill-hardware:
+        description: "GPU SKU used by prefill workers"
+        required: false
+        type: string
+        default: ""
+      decode-hardware:
+        description: "GPU SKU used by decode workers"
+        required: false
+        type: string
+        default: ""
 
       max-model-len:
         required: true
@@ -144,6 +154,8 @@ env:
   CONC_LIST: ${{ join(fromJson(inputs.conc-list), ' ') }}
   SPEC_DECODING: ${{ inputs.spec-decoding }}
   DISAGG: ${{ inputs.disagg }}
+  PREFILL_HARDWARE: ${{ inputs.prefill-hardware }}
+  DECODE_HARDWARE: ${{ inputs.decode-hardware }}
   RUN_EVAL: ${{ inputs.run-eval }}
   EVAL_ONLY: ${{ inputs.eval-only }}
   EVAL_CONC: ${{ inputs.eval-conc }}

.github/workflows/e2e-tests.yml

@@ -106,6 +106,8 @@ jobs:
             conc-list: ${{ toJson(matrix.config.conc) }}
             spec-decoding: ${{ matrix.config.spec-decoding }}
             disagg: ${{ matrix.config.disagg }}
+            prefill-hardware: ${{ matrix.config.prefill.hardware }}
+            decode-hardware: ${{ matrix.config.decode.hardware }}
 
             prefill-num-worker: ${{ matrix.config.prefill.num-worker }}
             prefill-tp: ${{ matrix.config.prefill.tp }}
@@ -145,6 +147,8 @@ jobs:
             conc-list: ${{ toJson(matrix.config.conc) }}
             spec-decoding: ${{ matrix.config.spec-decoding }}
             disagg: ${{ matrix.config.disagg }}
+            prefill-hardware: ${{ matrix.config.prefill.hardware }}
+            decode-hardware: ${{ matrix.config.decode.hardware }}
 
             prefill-num-worker: ${{ matrix.config.prefill.num-worker }}
             prefill-tp: ${{ matrix.config.prefill.tp }}
@@ -221,6 +225,8 @@ jobs:
             conc-list: ${{ toJson(matrix.config.conc) }}
             spec-decoding: ${{ matrix.config.spec-decoding }}
             disagg: ${{ matrix.config.disagg }}
+            prefill-hardware: ${{ matrix.config.prefill.hardware }}
+            decode-hardware: ${{ matrix.config.decode.hardware }}
             prefill-num-worker: ${{ matrix.config.prefill.num-worker }}
             prefill-tp: ${{ matrix.config.prefill.tp }}
             prefill-ep: ${{ matrix.config.prefill.ep }}

.github/workflows/run-sweep.yml

@@ -369,6 +369,8 @@ jobs:
             conc-list: ${{ toJson(matrix.config.conc) }}
             spec-decoding: ${{ matrix.config.spec-decoding }}
             disagg: ${{ matrix.config.disagg }}
+            prefill-hardware: ${{ matrix.config.prefill.hardware }}
+            decode-hardware: ${{ matrix.config.decode.hardware }}
 
             prefill-num-worker: ${{ matrix.config.prefill.num-worker }}
             prefill-tp: ${{ matrix.config.prefill.tp }}
@@ -531,6 +533,8 @@ jobs:
             conc-list: '[${{ matrix.config.conc }}]'
             spec-decoding: ${{ matrix.config.spec-decoding }}
             disagg: ${{ matrix.config.disagg }}
+            prefill-hardware: ${{ matrix.config.prefill.hardware }}
+            decode-hardware: ${{ matrix.config.decode.hardware }}
             prefill-num-worker: ${{ matrix.config.prefill.num-worker }}
             prefill-tp: ${{ matrix.config.prefill.tp }}
             prefill-ep: ${{ matrix.config.prefill.ep }}
@@ -618,6 +622,8 @@ jobs:
             conc-list: ${{ toJson(matrix.config.conc) }}
             spec-decoding: ${{ matrix.config.spec-decoding }}
             disagg: ${{ matrix.config.disagg }}
+            prefill-hardware: ${{ matrix.config.prefill.hardware }}
+            decode-hardware: ${{ matrix.config.decode.hardware }}
             prefill-num-worker: ${{ matrix.config.prefill.num-worker }}
             prefill-tp: ${{ matrix.config.prefill.tp }}
             prefill-ep: ${{ matrix.config.prefill.ep }}

configs/CONFIGS.md

@@ -28,6 +28,33 @@ entry-name:
       - { tp: int, conc-start: int, conc-end: int }
       - ...
 ```
+
+Heterogeneous disaggregated search-space entries declare hardware on each
+worker pool. Omit both `hardware` fields for homogeneous hardware:
+
+```yaml
+multinode: true
+disagg: true
+scenarios:
+  fixed-seq-len:
+  - isl: 1024
+    osl: 1024
+    search-space:
+    - conc-list: [64]
+      prefill:
+        hardware: b200
+        num-worker: 1
+        tp: 8
+        ep: 8
+        dp-attn: false
+      decode:
+        hardware: h100
+        num-worker: 2
+        tp: 8
+        ep: 8
+        dp-attn: false
+```
+
 Note: while not required, `entry-name` typically takes the format `<INFMAX_MODEL_PREFIX>-<PRECISION>-<GPU>-<FRAMEWORK>`.
 
 The below list describes what each field is:
@@ -41,6 +68,13 @@ The below list describes what each field is:
   fleet.
 - `precision`: The precision to run the benchmark. Again, this is used to find which script to run in `benchmarks/`.
 - `framework`: The framework (serving runtime) to serve the benchmark, e.g., `vllm`, `sglang`, `trt`.
+- `disagg`: Enables disaggregated serving and may only be `true` when
+  `multinode` is also `true`.
+- `hardware`: Optional metadata within each `prefill` and `decode` worker block
+  for heterogeneous disaggregated deployments. If one worker declares a GPU
+  SKU, the other must also declare one. Omit both fields for homogeneous
+  hardware. These values flow into aggregate results but do not affect runner
+  scheduling.
 - `scenarios`: A dictionary of benchmark scenario types. At least one must be specified. Currently supported:
   - `fixed-seq-len`: Fixed input/output sequence length benchmarks. Each entry must have:
     - `isl`: An integer representing the input sequence length, e.g., `1024`

utils/agentic/aggregation/process_agentic_result.py

@@ -67,6 +67,12 @@ def _gpu_shape() -> tuple[dict[str, Any], int, int, int, str]:
     decode_tp = env_int("DECODE_TP")
     decode_ep = env_int("DECODE_EP", 1)
     decode_dp_attention = os.environ.get("DECODE_DP_ATTN", "false")
+    prefill_hardware = os.environ.get("PREFILL_HARDWARE", "")
+    decode_hardware = os.environ.get("DECODE_HARDWARE", "")
+    if bool(prefill_hardware) != bool(decode_hardware):
+        raise SystemExit(
+            "PREFILL_HARDWARE and DECODE_HARDWARE must be specified together."
+        )
     num_prefill_gpu = prefill_num_workers * prefill_tp
     num_decode_gpu = decode_num_workers * decode_tp
     num_gpus = num_prefill_gpu + num_decode_gpu
@@ -91,6 +97,9 @@ def _gpu_shape() -> tuple[dict[str, Any], int, int, int, str]:
             "num_decode_gpu": num_decode_gpu,
         }
     )
+    if prefill_hardware:
+        fields["prefill_hw"] = prefill_hardware
+        fields["decode_hw"] = decode_hardware
     return fields, num_gpus, tp, ep, dp_attention
 
 

utils/agentic/aggregation/test_process_agentic_result.py

@@ -26,6 +26,7 @@
     load_aggregate,
     load_records,
 )
+from utils.agentic.aggregation.process_agentic_result import _gpu_shape
 from utils.agentic.aggregation.server_metrics import (
     compute_server_metrics,
     load_server_metrics,
@@ -308,6 +309,8 @@ def _run_processor(
     env_overrides: dict[str, str] | None = None,
 ) -> dict:
     env = os.environ.copy()
+    env.pop("PREFILL_HARDWARE", None)
+    env.pop("DECODE_HARDWARE", None)
     env.update(
         {
             "RESULT_DIR": str(result_dir),
@@ -453,6 +456,70 @@ def test_processor_surfaces_allocated_cpu_dram(tmp_path: Path):
     assert agg["allocated_cpu_dram_gb"] == 2400
 
 
+def test_multinode_processor_surfaces_heterogeneous_hardware(tmp_path: Path):
+    result_dir = _write_fixture(tmp_path)
+    agg = _run_processor(
+        result_dir,
+        tmp_path / "out",
+        env_overrides={
+            "IS_MULTINODE": "true",
+            "DISAGG": "true",
+            "PREFILL_NUM_WORKERS": "1",
+            "PREFILL_TP": "8",
+            "PREFILL_EP": "8",
+            "PREFILL_DP_ATTN": "false",
+            "PREFILL_HARDWARE": "b200",
+            "DECODE_NUM_WORKERS": "2",
+            "DECODE_TP": "8",
+            "DECODE_EP": "8",
+            "DECODE_DP_ATTN": "false",
+            "DECODE_HARDWARE": "h100",
+        },
+    )
+
+    assert agg["prefill_hw"] == "b200"
+    assert agg["decode_hw"] == "h100"
+
+
+def test_multinode_processor_omits_homogeneous_hardware(tmp_path: Path):
+    result_dir = _write_fixture(tmp_path)
+    agg = _run_processor(
+        result_dir,
+        tmp_path / "out",
+        env_overrides={
+            "IS_MULTINODE": "true",
+            "DISAGG": "true",
+            "PREFILL_NUM_WORKERS": "1",
+            "PREFILL_TP": "8",
+            "DECODE_NUM_WORKERS": "2",
+            "DECODE_TP": "8",
+        },
+    )
+
+    assert "prefill_hw" not in agg
+    assert "decode_hw" not in agg
+
+
+@pytest.mark.parametrize(
+    ("present_var", "missing_var"),
+    [
+        ("PREFILL_HARDWARE", "DECODE_HARDWARE"),
+        ("DECODE_HARDWARE", "PREFILL_HARDWARE"),
+    ],
+)
+def test_multinode_processor_rejects_one_sided_hardware(
+    monkeypatch: pytest.MonkeyPatch,
+    present_var: str,
+    missing_var: str,
+):
+    monkeypatch.setenv("IS_MULTINODE", "true")
+    monkeypatch.setenv(present_var, "b200")
+    monkeypatch.delenv(missing_var, raising=False)
+
+    with pytest.raises(SystemExit, match="must be specified together"):
+        _gpu_shape()
+
+
 def test_processor_surfaces_request_accounting(tmp_path: Path):
     result_dir = tmp_path / "results"
     artifact = result_dir / "aiperf_artifacts"

utils/matrix_logic/test_generate_sweep_configs.py

@@ -78,6 +78,7 @@ def sample_multinode_config():
                             {
                                 "conc-list": [2150],
                                 "prefill": {
+                                    "hardware": "gb200",
                                     "num-worker": 5,
                                     "tp": 4,
                                     "ep": 4,
@@ -88,6 +89,7 @@ def sample_multinode_config():
                                     ],
                                 },
                                 "decode": {
+                                    "hardware": "h100",
                                     "num-worker": 1,
                                     "tp": 8,
                                     "ep": 8,
@@ -981,6 +983,8 @@ def test_multinode_entry_structure(self, sample_multinode_config, sample_runner_
         assert entry["prefill"]["num-worker"] == 5
         assert entry["decode"]["num-worker"] == 1
         assert entry["disagg"] is True
+        assert entry["prefill"]["hardware"] == "gb200"
+        assert entry["decode"]["hardware"] == "h100"
 
     def test_multinode_conc_as_list(self, sample_multinode_config, sample_runner_config, full_sweep_args_multi_node):
         """Multinode conc should be passed as list."""
@@ -2004,8 +2008,8 @@ def test_multinode_agentic_groups_concurrencies_per_search_entry(self):
                             "search-space": [
                                 {
                                     "conc-list": [16, 32, 64, 128, 256],
-                                    "prefill": {"num-worker": 2, "tp": 8, "ep": 8, "dp-attn": False},
-                                    "decode": {"num-worker": 1, "tp": 8, "ep": 1, "dp-attn": False},
+                                    "prefill": {"hardware": "gb200", "num-worker": 2, "tp": 8, "ep": 8, "dp-attn": False},
+                                    "decode": {"hardware": "h100", "num-worker": 1, "tp": 8, "ep": 1, "dp-attn": False},
                                 }
                             ],
                         }
@@ -2151,8 +2155,8 @@ def test_node_type_filters_apply_to_agentic_configs(
                         "search-space": [
                             {
                                 "conc-list": [16, 32],
-                                "prefill": {"num-worker": 2, "tp": 8, "ep": 8, "dp-attn": False},
-                                "decode": {"num-worker": 1, "tp": 8, "ep": 1, "dp-attn": False},
+                                "prefill": {"hardware": "gb200", "num-worker": 2, "tp": 8, "ep": 8, "dp-attn": False},
+                                "decode": {"hardware": "h100", "num-worker": 1, "tp": 8, "ep": 1, "dp-attn": False},
                             },
                         ],
                     }],