From 6f23ce608f03e90590174c79a0470060944b1bc7 Mon Sep 17 00:00:00 2001 From: hongkuanz Date: Mon, 11 May 2026 09:55:28 -0700 Subject: [PATCH] feat(planner): rate-bound (Little's Law) closed form for decode scale-down Replaces the linear "scale KV by N/(N-1)" projection used in ``_decode_load_decision``'s SLA check with the closed-form fixed point of the steady-state Little's Law equation. The new projection captures catastrophic saturation as ``itl_curr`` approaches the rate-bound capacity ``N * intercept``, which the linear extrapolation misses. ## Math The decode regression decomposes ITL into a fixed cost plus a load-dependent variable cost: itl_curr = intercept + V_curr where ``V_curr = c_req * num_req + c_kv * kv`` is the variable portion. At steady state Little's Law ties per-worker concurrency to the product of arrival rate and time-in-system. Time-in-system for a decode token is ~ITL, so the variable load on the survivor after N -> N-1 scales by: 1. ``N/(N-1)``: arrival rate per worker after losing a worker. 2. ``itl_post / itl_curr``: longer ITL means each request lingers longer in the batch, further inflating concurrency. The variable cost on the survivor becomes: V_post = (itl_post / itl_curr) * (N / (N-1)) * V_curr Substituting ``itl_post = intercept + V_post`` and solving the linear fixed point in ``itl_post`` yields: itl_post = (N-1) * intercept * itl_curr / (N * intercept - itl_curr) The denominator goes to zero as ``itl_curr -> N * intercept``; past that, one fewer worker physically cannot sustain the offered load. ## Why this matters For the customer's logs (regression intercept ~17.9 ms, current ITL ~30.5 ms at 2-worker steady state with 177K KV each, N=2): - Linear "scale kv" extrapolation: predicts ``itl_post ~ 43 ms`` (barely above 32 ms threshold). - Rate-bound closed form: predicts ``itl_post ~ 104 ms``. Both refuse against the 32 ms threshold in this case, but for borderline scenarios near the rate-bound capacity the linear approximation under-predicts and lets unsafe scale-downs through. ## Changes - ``DecodeRegressionModel.intercept_seconds`` (new property) and ``estimate_post_consolidation_itl`` (closed-form helper) -- both carry the math derivation in their docstrings. - ``_decode_load_decision`` queries the closed form first; falls back to the previous linear projection when the closed form is unavailable (non-positive intercept from noisy fit, unfitted regression). - Hard cache feasibility check remains unchanged -- still independent of the SLA model. - Tests: * ``TestEstimatePostConsolidationItl`` unit tests for the closed form (saturation, sub-saturation, fallback paths, N-sensitivity). * ``TestDecodeConsolidationAwareScaleDown`` integration tests rewritten to exercise the rate-bound regime explicitly: below-saturation permit, SLA breach refusal, rate-bound saturation refusal, cache fail-safe. All 380 planner unit tests pass. Note: prefill / agg-prefill / agg-decode keep their existing projections. The rate-bound derivation is specific to decode's batched-iteration physics; prefill's queue-induced TTFT separation remains the right model there. Signed-off-by: hongkuanz --- .../src/dynamo/planner/core/load_scaling.py | 53 ++++- .../dynamo/planner/core/perf_model/decode.py | 98 +++++++++ .../planner/tests/unit/test_state_machine.py | 205 ++++++++++++++---- 3 files changed, 305 insertions(+), 51 deletions(-) diff --git a/components/src/dynamo/planner/core/load_scaling.py b/components/src/dynamo/planner/core/load_scaling.py index 4aa07b330c07..297da63db54d 100644 --- a/components/src/dynamo/planner/core/load_scaling.py +++ b/components/src/dynamo/planner/core/load_scaling.py @@ -503,13 +503,31 @@ def _decode_load_decision( estimates: list[float] = [] # Consolidation-aware scale-down. Two safety checks per worker: - # 1. Hard cache-feasibility: post-consolidation KV must fit within - # ``max_kv_tokens``. Exceeding the cache forces request queueing - # / block eviction, a non-linear regime the regression cannot - # model, so refuse outright when crossed. - # 2. SLA check: predicted ITL at the survivor's post-consolidation KV - # must stay within ``SLA * sensitivity``. Decouples from cache - # size -- engines often saturate latency well before cache. + # + # 1. **Hard cache feasibility**: post-consolidation KV must fit + # within ``max_kv_tokens``. Exceeding the cache forces block + # eviction / request queueing, a non-linear regime the + # regression cannot model, so we refuse outright when crossed. + # + # 2. **Rate-bound SLA check**: predict the survivor's *steady-state* + # ITL after losing a worker via a closed-form Little's-Law fixed + # point, then compare to ``SLA * sensitivity``. + # + # The closed form (derived in + # ``DecodeRegressionModel.estimate_post_consolidation_itl``) is:: + # + # itl_post = (N-1) * intercept * itl_curr + # / (N * intercept - itl_curr) + # + # and tends to +inf as ``itl_curr -> N * intercept`` -- the + # rate-bound capacity limit beyond which one fewer worker + # physically cannot sustain the offered load. This captures + # catastrophic saturation that a "scale KV by N/(N-1)" linear + # extrapolation would miss. + # + # If the regression's intercept is non-positive (noisy fit), we + # fall back to the linear projection ``estimate_next_itl`` at + # ``post_kv = (sched + queued) * N/(N-1)``. can_scale_down = num_workers > 1 consolidation_refused = False for (wid, dp), fpm in fpm_stats.items(): @@ -534,11 +552,22 @@ def _decode_load_decision( can_scale_down = False consolidation_refused = True continue - # (2) SLA check via regression at post-consolidation kv - post_itl = self._decode_regression.estimate_next_itl( - scheduled_decode_kv=post_sched_kv, - queued_decode_kv=0, - ) + # (2) rate-bound SLA check at the survivor's steady-state ITL + post_itl: Optional[float] = None + if est is not None: + post_itl = self._decode_regression.estimate_post_consolidation_itl( + itl_curr=est, + num_workers=num_workers, + ) + if post_itl is None: + # Closed-form unavailable (regression cold or unstable + # intercept) -- fall back to linear extrapolation at + # the scaled kv. Less accurate near saturation but + # safe. + post_itl = self._decode_regression.estimate_next_itl( + scheduled_decode_kv=post_sched_kv, + queued_decode_kv=0, + ) if post_itl is None: can_scale_down = False elif post_itl * 1000 >= self._config.itl * sensitivity: diff --git a/components/src/dynamo/planner/core/perf_model/decode.py b/components/src/dynamo/planner/core/perf_model/decode.py index e23be4a7cf92..5ab361230c63 100644 --- a/components/src/dynamo/planner/core/perf_model/decode.py +++ b/components/src/dynamo/planner/core/perf_model/decode.py @@ -64,6 +64,25 @@ def _update_moving_averages(self, fpm: ForwardPassMetrics) -> None: def avg_decode_length(self) -> float: return self._avg_decode_len.value + @property + def intercept_seconds(self) -> Optional[float]: + """Regression intercept (no-load wall_time per iteration, in seconds). + + Returns ``None`` if the model has not been fitted yet. The intercept + represents fixed per-iter overhead (kernel launches, framework + bookkeeping) that does not scale with batch composition. It is the + physical lower bound on ITL -- the wall_time prediction approaches + this value as ``num_req -> 0`` and ``kv -> 0``. + + Used by the rate-bound consolidation predictor: at steady state the + post-survival ITL diverges as the system approaches the rate-bound + capacity ``ITL = N * intercept``, beyond which one fewer worker + cannot sustain the offered request rate. + """ + if not self._is_fitted: + return None + return float(self._model.intercept_) + def _predict_2d(self, num_requests: float, kv_tokens: float) -> float: return max( 1e-6, float(self._model.predict(np.array([[num_requests, kv_tokens]]))[0]) @@ -81,6 +100,85 @@ def estimate_next_itl( num_req = self._avg_num_decode.value + 1 return self._predict_2d(num_req, total_kv) + def estimate_post_consolidation_itl( + self, + itl_curr: float, + num_workers: int, + ) -> Optional[float]: + """Estimate steady-state ITL on the survivor after scaling N -> N-1. + + Closed-form Little's-Law solution for the post-consolidation ITL. + Derivation: + + The regression decomposes ITL into a fixed cost (intercept) plus a + load-dependent variable cost:: + + itl_curr = intercept + V_curr + + where ``V_curr = c_req * num_req_curr + c_kv * kv_curr`` is the + variable portion (load * regression slopes). + + At steady state, Little's Law ties per-worker concurrency to the + product of arrival rate and time-in-system. Time-in-system for a + decode token is ~ITL, so the variable load scales with both: + + 1. ``N/(N-1)``: arrival rate per worker after losing a worker + (cluster offered load is invariant; per-worker share grows). + 2. ``itl_post / itl_curr``: longer ITL means each request lingers + longer in the batch, further inflating concurrency / kv. + + Hence the variable cost on the survivor at the new steady state is:: + + V_post = (itl_post / itl_curr) * (N / (N-1)) * V_curr + = (itl_post / itl_curr) * (N / (N-1)) * (itl_curr - intercept) + + and ``itl_post = intercept + V_post``. Substituting and solving the + linear fixed point in ``itl_post``:: + + itl_post = (N-1) * intercept * itl_curr / (N * intercept - itl_curr) + + Equivalently:: + + itl_post = itl_curr * (1 + (itl_curr - intercept) / + (N * intercept - itl_curr)) + + **Saturation**: the denominator ``N * intercept - itl_curr`` goes to + zero as ``itl_curr -> N * intercept``. Past this point one fewer + worker physically cannot sustain the offered request rate (the + survivor would need to process tokens faster than its intercept + permits). Returns ``+inf`` in that regime so callers refuse + scale-down. + + Args: + itl_curr: Current per-worker ITL in seconds. + num_workers: Current worker count (must be >= 2 to consolidate). + + Returns: + Predicted post-consolidation ITL in seconds; ``+inf`` if the + system is rate-bound (infeasible to lose a worker); + ``None`` if the regression is not fitted or the intercept is + non-positive (a noisy fit we can't trust for this projection). + """ + if not self._ensure_fitted() or num_workers < 2: + return None + intercept = self.intercept_seconds + if intercept is None or intercept <= 0: + # A non-positive intercept means the fit was dominated by noise + # or extrapolation past training data -- the rate-bound formula + # would divide by something meaningless. Caller should fall back + # to a direct ``estimate_next_itl`` at scaled inputs instead. + return None + if itl_curr <= intercept: + # Below the fixed cost floor (regression noise) -- treat survivor + # as also at the floor; nothing to amplify. + return intercept + denom = num_workers * intercept - itl_curr + if denom <= 0: + # Rate-bound: even with infinite cache, one fewer worker cannot + # keep up with the offered request rate. + return float("inf") + return (num_workers - 1) * intercept * itl_curr / denom + def find_best_engine_decode_rps( self, itl: float, diff --git a/components/src/dynamo/planner/tests/unit/test_state_machine.py b/components/src/dynamo/planner/tests/unit/test_state_machine.py index 2c289c1d8f13..0cadd0787c59 100644 --- a/components/src/dynamo/planner/tests/unit/test_state_machine.py +++ b/components/src/dynamo/planner/tests/unit/test_state_machine.py @@ -347,27 +347,54 @@ def _decode_caps_with_max_kv(max_kv_tokens: int) -> WorkerCapabilities: ) +def _train_decode_regression_with_intercept(core: PlannerStateMachine) -> None: + """Trains the decode regression with a sizable fixed-cost intercept (~30ms). + + The rate-bound consolidation formula (Little's Law fixed point) only + behaves non-trivially when the regression's intercept is large relative + to the variable load -- the "interesting" regime is + ``intercept < itl_curr < N * intercept``. With the default training + helper's ~1 ms intercept the formula always reports saturation, so + rate-bound tests use this larger fixed cost. + """ + fpms = [ + _make_fpm( + sum_decode_kv_tokens=kv, + num_decode_requests=n, + wall_time=1e-5 * kv + 0.03, # intercept = 30 ms + ) + for n, kv in [(5, 5000), (10, 10000), (20, 20000), (30, 30000), (40, 40000)] + ] + core.load_benchmark_fpms(decode_fpms=fpms) + + class TestDecodeConsolidationAwareScaleDown: - """Decode scale-down uses two checks at the survivor's post-consolidation - KV (current sched+queued scaled by N/(N-1)): + """Decode scale-down uses two checks per worker: + + 1. **Hard cache feasibility** -- post_kv must fit within + ``max_kv_tokens``; crossing it forces block eviction / queueing, + a non-linear regime outside the regression's domain. + 2. **Rate-bound SLA check** -- closed-form Little's-Law projection of + the survivor's steady-state ITL after losing a worker, compared to + ``SLA * sensitivity``. + + The closed form is - 1. **Cache feasibility** -- post_kv must fit within ``max_kv_tokens``; - crossing it forces block eviction / queueing, a non-linear regime - outside the regression's domain. - 2. **SLA check** -- regression-predicted ITL at post_kv must stay - within ``SLA * sensitivity``. + itl_post = (N-1) * intercept * itl_curr / (N * intercept - itl_curr) - Either failure refuses the scale-down. + and diverges as ``itl_curr -> N * intercept`` (rate-bound capacity). + See ``DecodeRegressionModel.estimate_post_consolidation_itl`` for the + full derivation. """ - def _setup(self, *, itl_sla: float = 100.0, max_kv_tokens: int = 100_000): + def _setup(self, *, itl_sla: float = 300.0, max_kv_tokens: int = 100_000): core = _make_core( mode="decode", itl=itl_sla, load_scaling_down_sensitivity=80, ) core._capabilities = _decode_caps_with_max_kv(max_kv_tokens) - _train_decode_regression(core) + _train_decode_regression_with_intercept(core) return core def _tick(self, *, num_workers: int, sched_kv_per_worker: int) -> TickInput: @@ -377,10 +404,9 @@ def _tick(self, *, num_workers: int, sched_kv_per_worker: int) -> TickInput: worker_id=f"w{i}", sum_decode_kv_tokens=sched_kv_per_worker, num_decode_requests=max(1, sched_kv_per_worker // 1000), - # Match _train_decode_regression's wall_time formula so the - # post-bootstrap refit on each tick stays monotone in kv; - # otherwise the regression rejects the fit and decisions skip. - wall_time=0.00001 * sched_kv_per_worker + 0.001, + # Match the training regression's slope+intercept so per-tick + # refits stay monotone (otherwise the fit can reject). + wall_time=1e-5 * sched_kv_per_worker + 0.03, ) return TickInput( now_s=5.0, @@ -388,27 +414,44 @@ def _tick(self, *, num_workers: int, sched_kv_per_worker: int) -> TickInput: worker_counts=WorkerCounts(ready_num_decode=num_workers), ) - def test_post_consolidation_within_sla_permits(self): - """Light load: post_kv well under cache and SLA -> ALLOW. + def test_rate_bound_below_saturation_permits(self): + """Light load: ``itl_curr`` well below ``N * intercept`` -> ALLOW. - N=2, sched_kv=1500. post_kv = 3000. Predicted ITL ~= 0.001 + - 0.00001 * 4000 ~= 41 ms (with internal avg_decode_len), under - the 80 ms threshold. No scale-up either (under 100 ms SLA). + intercept ~= 30 ms. With sched_kv=1000, predicted itl_curr ~= 50 ms + (intercept + 0.01ms*kv*kv + avg_decode_len adjustment). Rate-bound + post-consolidation itl_post = 1 * 30 * 50 / (60 - 50) = 150 ms, + below SLA * sensitivity = 240 ms. """ - core = self._setup(itl_sla=100.0) - tick = self._tick(num_workers=2, sched_kv_per_worker=1_500) + core = self._setup(itl_sla=300.0) + tick = self._tick(num_workers=2, sched_kv_per_worker=1_000) effects = core.on_tick(_tick_for(tick), tick) assert effects.scale_to is not None assert effects.scale_to.num_decode == 1 - def test_post_consolidation_breaches_sla_refuses(self): - """Cache fine but predicted ITL > SLA*sensitivity -> SLA check refuses. + def test_rate_bound_breaches_sla_refuses(self): + """Sub-saturation but predicted post-itl > SLA*sensitivity -> REFUSE. - N=2, sched_kv=8000. post_kv=16000 (well below 100K cache). Predicted - ITL ~= 0.001 + 0.00001 * 17000 ~= 171 ms, above the 80 ms threshold. + Same load (itl_post ~= 150 ms by rate-bound formula) but a tighter + SLA of 150 ms (threshold 120 ms). 150 > 120 -> SLA check refuses. """ - core = self._setup(itl_sla=100.0) - tick = self._tick(num_workers=2, sched_kv_per_worker=8_000) + core = self._setup(itl_sla=150.0) + tick = self._tick(num_workers=2, sched_kv_per_worker=1_000) + effects = core.on_tick(_tick_for(tick), tick) + assert effects.scale_to is None or effects.scale_to.num_decode == 2 + assert ( + effects.diagnostics.load_decision_reason + == "scale_down_refused_consolidation" + ) + + def test_rate_bound_saturation_refuses(self): + """At the rate-bound capacity: ``itl_curr >= N * intercept`` -> REFUSE. + + sched_kv=4000 gives itl_curr ~= 70 ms > 60 ms = 2 * intercept. The + closed form's denominator goes non-positive -> post_itl = +inf -> + scale-down is infeasible at any sensitivity. + """ + core = self._setup(itl_sla=10_000.0) # SLA effectively off + tick = self._tick(num_workers=2, sched_kv_per_worker=4_000) effects = core.on_tick(_tick_for(tick), tick) assert effects.scale_to is None or effects.scale_to.num_decode == 2 assert ( @@ -417,10 +460,10 @@ def test_post_consolidation_breaches_sla_refuses(self): ) def test_post_consolidation_exceeds_max_kv_refuses(self): - """Hard cache fail-safe: post_kv >= max_kv -> refuse outright. + """Hard cache fail-safe still applies independently of rate-bound. - N=2, sched_kv=60_000. post_kv=120_000 >= max_kv 100_000. SLA is - effectively off (10s) so only the cache check can refuse. + sched_kv=60_000 -> post_kv = 120_000 >= max_kv = 100_000. Refused + without consulting the rate-bound model (SLA is effectively off). """ core = self._setup(itl_sla=10_000.0, max_kv_tokens=100_000) tick = self._tick(num_workers=2, sched_kv_per_worker=60_000) @@ -431,23 +474,107 @@ def test_post_consolidation_exceeds_max_kv_refuses(self): == "scale_down_refused_consolidation" ) - def test_no_max_kv_falls_through_to_sla_check(self): - """Without max_kv_tokens, only the SLA check governs. - - Cache check is skipped (no denominator); the regression still gates - scale-down by predicted ITL. Light load passes -> ALLOW. - """ - core = self._setup(itl_sla=100.0) - # Erase max_kv: cache check becomes a no-op. + def test_no_max_kv_falls_through_to_rate_bound_check(self): + """Without max_kv_tokens, only the rate-bound SLA check governs.""" + core = self._setup(itl_sla=300.0) core._capabilities = WorkerCapabilities( decode=EngineCapabilities(num_gpu=1, max_num_batched_tokens=2048), ) - tick = self._tick(num_workers=2, sched_kv_per_worker=1_500) + tick = self._tick(num_workers=2, sched_kv_per_worker=1_000) effects = core.on_tick(_tick_for(tick), tick) assert effects.scale_to is not None assert effects.scale_to.num_decode == 1 +class TestEstimatePostConsolidationItl: + """Unit tests for the closed-form rate-bound projection. + + Math: itl_post = (N-1) * intercept * itl_curr / (N * intercept - itl_curr) + """ + + def _fitted_regression(self): + core = _make_core(mode="decode") + _train_decode_regression_with_intercept(core) + # Force a fit (regression lazy-fits on first prediction). + core.decode_regression.estimate_next_itl( + scheduled_decode_kv=0, queued_decode_kv=0 + ) + return core.decode_regression + + def test_returns_none_when_unfitted(self): + core = _make_core(mode="decode") + assert ( + core.decode_regression.estimate_post_consolidation_itl( + itl_curr=0.05, num_workers=2 + ) + is None + ) + + def test_returns_none_for_n_below_2(self): + reg = self._fitted_regression() + assert reg.estimate_post_consolidation_itl(itl_curr=0.05, num_workers=1) is None + + def test_below_intercept_returns_intercept(self): + reg = self._fitted_regression() + intercept = reg.intercept_seconds + assert intercept is not None and intercept > 0 + # itl_curr below the fixed cost floor: survivor also pegged at floor. + result = reg.estimate_post_consolidation_itl( + itl_curr=intercept * 0.5, num_workers=2 + ) + assert result == intercept + + def test_at_saturation_returns_infinity(self): + reg = self._fitted_regression() + intercept = reg.intercept_seconds + assert intercept is not None + # itl_curr at exactly N * intercept -> denominator 0 -> infeasible. + result = reg.estimate_post_consolidation_itl( + itl_curr=2 * intercept, num_workers=2 + ) + assert result == float("inf") + + def test_past_saturation_returns_infinity(self): + reg = self._fitted_regression() + intercept = reg.intercept_seconds + assert intercept is not None + result = reg.estimate_post_consolidation_itl( + itl_curr=3 * intercept, num_workers=2 + ) + assert result == float("inf") + + def test_closed_form_matches_formula(self): + """Direct verification of (N-1)*c*itl / (N*c - itl) at a finite point.""" + reg = self._fitted_regression() + intercept = reg.intercept_seconds + assert intercept is not None + N = 2 + # Pick itl_curr midway between intercept and N*intercept. + itl_curr = 1.5 * intercept + expected = (N - 1) * intercept * itl_curr / (N * intercept - itl_curr) + actual = reg.estimate_post_consolidation_itl(itl_curr=itl_curr, num_workers=N) + assert actual is not None + assert abs(actual - expected) < 1e-9 + + def test_higher_n_more_permissive(self): + """At N=10 the saturation point is 10*intercept (vs 2*intercept at N=2), + so the same itl_curr is sub-saturation at high N but saturated at low N. + """ + reg = self._fitted_regression() + intercept = reg.intercept_seconds + assert intercept is not None + itl_curr = 1.8 * intercept + + at_n2 = reg.estimate_post_consolidation_itl(itl_curr, num_workers=2) + at_n10 = reg.estimate_post_consolidation_itl(itl_curr, num_workers=10) + + # N=2: itl_curr 1.8*c is approaching saturation -> very large + # N=10: same itl_curr is well below saturation -> finite, modest + assert at_n2 is not None and at_n10 is not None + assert at_n2 > at_n10 + assert at_n10 < 5 * intercept # comfortably finite + + def _train_slow_prefill_regression(core: PlannerStateMachine) -> None: """Trains a regression with low slope so chunked TTFTs stay tractable.