diff --git a/src/aiconfigurator/sdk/picking.py b/src/aiconfigurator/sdk/picking.py index a6c2b2b1f..76bd29e80 100644 --- a/src/aiconfigurator/sdk/picking.py +++ b/src/aiconfigurator/sdk/picking.py @@ -70,10 +70,12 @@ def _build_disagg_summary_dict( Returns: Dict with keys matching ``common.ColumnsDisagg``. """ - seq_s = min( - prefill_summary_dict["seq/s"] * prefill_num_worker * prefill_degradation_factor, - decode_summary_dict["seq/s"] * decode_num_worker * decode_degradation_factor, - ) + prefill_capacity = prefill_summary_dict["seq/s"] * prefill_num_worker * prefill_degradation_factor + decode_capacity = decode_summary_dict["seq/s"] * decode_num_worker * decode_degradation_factor + prefill_ttft_capacity = prefill_summary_dict["bs"] * prefill_num_worker * 1000.0 / prefill_summary_dict["ttft"] + + seq_s = min(prefill_capacity, decode_capacity, prefill_ttft_capacity) + prefill_gpus = prefill_summary_dict["pp"] * prefill_summary_dict["tp"] * prefill_summary_dict["dp"] decode_gpus = decode_summary_dict["pp"] * decode_summary_dict["tp"] * decode_summary_dict["dp"] num_total_gpus = prefill_gpus * prefill_num_worker + decode_gpus * decode_num_worker diff --git a/src/aiconfigurator/sdk/sweep.py b/src/aiconfigurator/sdk/sweep.py index ccde79a0d..da5650f9a 100644 --- a/src/aiconfigurator/sdk/sweep.py +++ b/src/aiconfigurator/sdk/sweep.py @@ -112,10 +112,10 @@ def _rate_match_dict( d = decode_summary_dict osl = p["osl"] - seq_s = min( - p["seq/s"] * prefill_num_worker * prefill_degradation, - d["seq/s"] * decode_num_worker * decode_degradation, - ) + prefill_capacity = p["seq/s"] * prefill_num_worker * prefill_degradation + decode_capacity = d["seq/s"] * decode_num_worker * decode_degradation + prefill_ttft_capacity = p["bs"] * prefill_num_worker * 1000.0 / p["ttft"] + seq_s = min(prefill_capacity, decode_capacity, prefill_ttft_capacity) prefill_gpus = p["pp"] * p["tp"] * p["dp"] decode_gpus = d["pp"] * d["tp"] * d["dp"] num_total_gpus = prefill_gpus * prefill_num_worker + decode_gpus * decode_num_worker diff --git a/tests/unit/sdk/sweep/test_rate_match_parity.py b/tests/unit/sdk/sweep/test_rate_match_parity.py index 7192a8160..2aac8c82c 100644 --- a/tests/unit/sdk/sweep/test_rate_match_parity.py +++ b/tests/unit/sdk/sweep/test_rate_match_parity.py @@ -140,6 +140,18 @@ def test_rate_match_with_custom_degradation_factors(): ) +def test_rate_match_prefill_ttft_capacity_can_bound_throughput(): + p = _make_prefill_dict(bs=1, ttft=500.0, **{"seq/s": 100.0}) + d = _make_decode_dict(**{"seq/s": 100.0}) + + new_result = _rate_match_dict(p, 2, d, 2) + old_result = _build_disagg_summary_dict(p, 2, d, 2) + + assert new_result["seq/s"] == 4.0 + for key in new_result: + assert new_result[key] == old_result[key] + + def test_rate_match_zero_osl_does_not_divide_by_zero(): """request_latency uses max(osl - 1, 0); osl=1 keeps decode_time=0.""" p = _make_prefill_dict(osl=1)