Add fields for working GPU count in rollup data and use working GPU count to compute total_throughput_per_gpu

src/srtctl/benchmarks/scripts/sa-bench/rollup.py

@@ -21,6 +21,9 @@
     "Config",
     "Total GPU Count",
     "Decode GPU Count",
+    "Total Working GPU Count",
+    "Decode Working GPU Count",
+    "Prefill Working GPU Count",
     "Concurrency",
     "Total Token Throughput",
     "Output Token Throughput",
@@ -73,7 +76,6 @@ def _as_int(value: Any) -> int:
     except (TypeError, ValueError):
         return 0
 
-
 def _compute_gpu_counts(resources: dict[str, Any]) -> tuple[int | None, int | None]:
     """Compute total and decode-serving GPU counts from resource settings."""
     gpus_per_node = _as_int(resources.get("gpus_per_node"))
@@ -151,6 +153,28 @@ def _extract_p90_decode_running_requests(log_dir: Path, metadata: dict[str, Any]
     return None
 
 
+def _compute_working_gpu_counts(resources: dict[str, Any]) -> tuple[int | None, int | None, int | None]:
+    """Compute prefill, decode, and total working GPU counts from worker-level resource settings."""
+    prefill_workers = _as_int(resources.get("prefill_workers"))
+    gpus_per_prefill = _as_int(resources.get("gpus_per_prefill"))
+    decode_workers = _as_int(resources.get("decode_workers"))
+    gpus_per_decode = _as_int(resources.get("gpus_per_decode"))
+
+    prefill_working = prefill_workers * gpus_per_prefill if prefill_workers > 0 and gpus_per_prefill > 0 else None
+    decode_working = decode_workers * gpus_per_decode if decode_workers > 0 and gpus_per_decode > 0 else None
+
+    if prefill_working is not None and decode_working is not None:
+        total_working = prefill_working + decode_working
+    elif prefill_working is not None:
+        total_working = prefill_working
+    elif decode_working is not None:
+        total_working = decode_working
+    else:
+        total_working = None
+
+    return prefill_working, decode_working, total_working
+
+
 def _safe_ratio(numerator: float | int | None, denominator: float | int | None) -> float | None:
     """Return numerator / denominator when both values are valid and denominator != 0."""
     if numerator is None or denominator in (None, 0):
@@ -175,14 +199,24 @@ def _build_csv_row(
     gpu_num: int | None,
     decode_gpu_count: int | None,
     p90_decode_running_requests: int | None,
+    prefill_working_gpu_count: int | None,
+    decode_working_gpu_count: int | None,
+    total_working_gpu_count: int | None,
 ) -> dict[str, object]:
     """Build one CSV row from a parsed sa-bench result."""
     total_token_throughput = data.get("total_token_throughput")
     median_tpot = data.get("median_tpot_ms")
+    # Fall back to node-based GPU counts when worker-level counts are unavailable (e.g. agg mode).
+    effective_total_working = total_working_gpu_count if total_working_gpu_count is not None else gpu_num
+    effective_decode_working = decode_working_gpu_count if decode_working_gpu_count is not None else decode_gpu_count
+    effective_prefill_working = prefill_working_gpu_count if prefill_working_gpu_count is not None else (0 if gpu_num is not None else None)
     row = {
         "Config": config_name,
         "Total GPU Count": gpu_num,
         "Decode GPU Count": decode_gpu_count,
+        "Total Working GPU Count": effective_total_working,
+        "Decode Working GPU Count": effective_decode_working,
+        "Prefill Working GPU Count": effective_prefill_working,
         "Concurrency": data.get("max_concurrency"),
         "Total Token Throughput": total_token_throughput,
         "Output Token Throughput": data.get("output_throughput"),
@@ -191,7 +225,7 @@ def _build_csv_row(
         "Median ITL": data.get("median_itl_ms"),
         "P90 Decode Running Requests": p90_decode_running_requests,
         "Output Token Throughput per User": _safe_ratio(1000.0, median_tpot),
-        "Total Token Throughput per GPU": _safe_ratio(total_token_throughput, gpu_num),
+        "Total Token Throughput per GPU": _safe_ratio(total_token_throughput, effective_total_working),
     }
     return {key: _format_csv_value(value) for key, value in row.items()}
 
@@ -210,6 +244,9 @@ def main(log_dir: Path) -> None:
     config_name = metadata.get("job_name") if metadata else None
     resources = metadata.get("resources") if metadata else None
     total_gpu_count, decode_gpu_count = _compute_gpu_counts(resources) if resources else (None, None)
+    prefill_working_gpu_count, decode_working_gpu_count, total_working_gpu_count = (
+        _compute_working_gpu_counts(resources) if resources else (None, None, None)
+    )
     p90_decode_running_requests = _extract_p90_decode_running_requests(log_dir, metadata)
 
     for result_file in result_files:
@@ -249,6 +286,9 @@ def main(log_dir: Path) -> None:
                 gpu_num=total_gpu_count,
                 decode_gpu_count=decode_gpu_count,
                 p90_decode_running_requests=p90_decode_running_requests,
+                prefill_working_gpu_count=prefill_working_gpu_count,
+                decode_working_gpu_count=decode_working_gpu_count,
+                total_working_gpu_count=total_working_gpu_count,
             )
         )
 

tests/test_sa_bench_rollup.py

@@ -99,6 +99,9 @@ def test_sa_bench_rollup_generates_json_and_csv_without_metadata(tmp_path):
     assert first["Config"] == "GLM-5-FP8"
     assert first["Total GPU Count"] == ""
     assert first["Decode GPU Count"] == ""
+    assert first["Total Working GPU Count"] == ""
+    assert first["Decode Working GPU Count"] == ""
+    assert first["Prefill Working GPU Count"] == ""
     assert first["Total Token Throughput"] == "9876"
     assert first["Output Token Throughput"] == "1234.5"
     assert first["Median TTFT"] == "95"
@@ -125,8 +128,11 @@ def test_sa_bench_rollup_uses_metadata_for_name_gpu_counts_and_p90(tmp_path):
                 "resources": {
                     "gpus_per_node": 8,
                     "prefill_nodes": 1,
+                    "prefill_workers": 1,
+                    "gpus_per_prefill": 8,
                     "decode_nodes": 2,
                     "decode_workers": 2,
+                    "gpus_per_decode": 8,
                     "agg_workers": 0,
                 },
             }
@@ -164,6 +170,9 @@ def test_sa_bench_rollup_uses_metadata_for_name_gpu_counts_and_p90(tmp_path):
     assert row["Config"] == "job-metadata-name"
     assert row["Total GPU Count"] == "24"
     assert row["Decode GPU Count"] == "16"
+    assert row["Total Working GPU Count"] == "24"
+    assert row["Decode Working GPU Count"] == "16"
+    assert row["Prefill Working GPU Count"] == "8"
     assert row["P90 Decode Running Requests"] == "32"
     assert row["Output Token Throughput per User"] == "40"
     assert row["Total Token Throughput per GPU"] == "33.333"
@@ -215,6 +224,9 @@ def test_sa_bench_rollup_aggregated_deployment_reports_all_gpus(tmp_path):
     assert row["Config"] == "glm47flash-agg-tp4-baseline"
     assert row["Total GPU Count"] == "4"
     assert row["Decode GPU Count"] == "4"
+    assert row["Total Working GPU Count"] == "4"
+    assert row["Decode Working GPU Count"] == "4"
+    assert row["Prefill Working GPU Count"] == "0"
     assert row["Total Token Throughput per GPU"] == "6307.5"
 
 
@@ -266,6 +278,62 @@ def test_sa_bench_rollup_tolerates_null_resource_fields(tmp_path):
     assert row["Total GPU Count"] == "4"
 
 
+def test_sa_bench_rollup_sub_node_prefill_working_gpu_counts(tmp_path):
+    """When gpus_per_prefill < gpus_per_node, working GPU counts differ from node-based counts."""
+    rollup = _load_rollup_module()
+
+    logs_dir = tmp_path / "logs"
+    result_dir = logs_dir / "sa-bench_isl_4096_osl_512"
+    result_dir.mkdir(parents=True)
+
+    # 1 prefill node (8 GPUs) but only 4 GPUs used per prefill worker → 1 worker × 4 = 4 prefill working GPUs
+    # 2 decode nodes, 2 workers × 8 GPUs = 16 decode working GPUs
+    # Total GPU Count  = (1 + 2) × 8 = 24  (node-based)
+    # Total Working GPU Count = 4 + 16 = 20 (worker-based, diverges from node count)
+    (tmp_path / "5001.json").write_text(
+        json.dumps(
+            {
+                "job_name": "sub-node-prefill-job",
+                "backend_type": "sglang",
+                "resources": {
+                    "gpus_per_node": 8,
+                    "prefill_nodes": 1,
+                    "prefill_workers": 1,
+                    "gpus_per_prefill": 4,
+                    "decode_nodes": 2,
+                    "decode_workers": 2,
+                    "gpus_per_decode": 8,
+                    "agg_workers": 0,
+                },
+            }
+        )
+    )
+
+    result = {
+        "model_id": "test-model",
+        "max_concurrency": 64,
+        "output_throughput": 500.0,
+        "total_token_throughput": 1000.0,
+        "median_ttft_ms": 50.0,
+        "median_tpot_ms": 20.0,
+        "median_itl_ms": 10.0,
+    }
+    (result_dir / "results_concurrency_64_gpus_24_ctx_4_gen_16.json").write_text(json.dumps(result))
+
+    rollup.main(logs_dir)
+
+    rows = _read_csv_rows(logs_dir / "benchmark-rollup.csv")
+    assert len(rows) == 1
+    row = rows[0]
+    assert row["Total GPU Count"] == "24"
+    assert row["Decode GPU Count"] == "16"
+    assert row["Prefill Working GPU Count"] == "4"
+    assert row["Decode Working GPU Count"] == "16"
+    assert row["Total Working GPU Count"] == "20"
+    # 1000 / 20 = 50 (uses working count, not node count)
+    assert row["Total Token Throughput per GPU"] == "50"
+
+
 def test_compute_gpu_counts_handles_none_values():
     """_compute_gpu_counts should never raise on null-valued keys."""
     rollup = _load_rollup_module()