fix(replay/planner): planner-in-the-loop replay diagnostics + planner scaling logics fixes

[tedzhouhk]

Summary

Adds interactive HTML diagnostics reports to offline planner-in-the-loop replay, bootstraps the throughput regression from the mocker's perf model (via AIC), and fixes several planner correctness issues surfaced by the new reports.

Before this PR, `dynamo.replay --planner-config ...` produced only a text summary of trace metrics. After this PR, every run emits a Plotly HTML report showing replica counts, observed vs predicted traffic, TTFT/ITL vs SLA, per-engine KV load, engine capacity, and color-coded decision timelines — giving a per-tick view of how the planner reacted to the simulated workload.

What's in the PR

Replay diagnostics

• HTML report — wire `DiagnosticsRecorder` into `ReplayPlannerAdapter` so offline replay produces the same interactive chart set as the live planner, plus:

  • `report_filename` config for a stable path (refresh-in-browser)
  • Throughput lower-bound dashed line on the Replica Counts chart (separate prefill / decode lines in disagg)
  • Combined queued + inflight decode KV with a horizontal `max_kv_tokens` capacity line
  • Per-component Load / Throughput decision timelines in disagg (prefill on y=2, decode on y=1)
  • Hidden per-engine legend entries to stop the legend covering charts during scaling
  • `connectgaps=True` on throughput-tick-only traces (engine RPS, predicted RPS/ISL/OSL)

• TTFT / ITL plumbed through the traffic accumulator (Rust) — the mocker's `TraceCollector` already tracks arrival / first-token / per-token timestamps; now `TrafficAccumulator::on_request` reads `request_latencies()` and exposes `avg_ttft_ms` / `avg_itl_ms` via `drain_traffic()`. The Python adapter feeds these into `Metrics` so the report's "Observed TTFT / ITL vs SLA" charts have data.

• Synthetic benchmark from AIC — instead of requiring the user to supply NPZ profiling data, replay now sweeps AIC's `predict_prefill` / `predict_decode` at a configurable `--benchmark-granularity` (default 8, matching the profiler) and bootstraps the planner's regression models with batch-size-aware data consistent with the mocker's timing. Prefill sweep caps at `max_num_batched_tokens`; agg-mode bootstrap combines prefill + decode points so both regression features have variance.

Planner correctness fixes (surfaced by replay diagnostics)

• `find_best_engine_agg_rps` max_bs cap — the previous formula `int(max_num_batched_tokens / avg_ctx) * 2` collapsed to 2 whenever `ISL > MBT`, so `engine_rps` was wildly underestimated for prefill-heavy workloads. Replaced with a binary search in `[1, min(kv_cap, seq_cap)]` that enforces the prefill / decode steady-state balance `isl / (max_num_batched_tokens - x) <= osl`. Took engine_rps from 1.15 to ~3.5 on a representative workload.

• `find_best_engine_prefill_rps` chunking — for `ISL > max_num_batched_tokens`, the wall-time prediction was being extrapolated beyond the regression's training domain. Now chunks `ISL` into `ceil(ISL / MBT)` forward passes and sums predictions per chunk.

• `find_best_engine_decode_rps` capability caps — was using `_max_observed_kv / context_length` as the sweep upper bound, which stays low until the engine runs hot. Now uses `max_kv_tokens / context_length` and `max_num_seqs` from engine capabilities, with the observed-KV fallback only when capabilities aren't provided.

• State-machine tick ordering — on a combined load + throughput tick, load scaling ran first with stale `throughput_lower_bound*`, then throughput updated the floor afterward, letting the load decision violate the brand-new floor. Swapped the order so throughput scaling updates the floor before load scaling reads it (precedence unchanged: load decision still wins over throughput decision when both produce one).

• Throughput floor enforced on no-change — `_advance_load_disagg` and the SLA-mode agg path both returned `None` for "no change" before applying the floor, so replicas could linger below a raised floor for many ticks. Now computes the final count as `max(load_or_current, min_endpoint, throughput_floor)` first, then checks no-change against the post-floor value.

• Relaxable negative coefficients — `DecodeRegressionModel` fits on `(num_decode_requests, sum_decode_kv_tokens)`; total KV dominates wall time, so `num_decode_requests` can flip slightly negative under multicollinearity + noisy FPMs. Added a `_relaxable_feature_indices` hook in the base class; decode marks index 0 as relaxable so minor negatives don't reject the whole fit. `sum_decode_kv_tokens` (and all agg / prefill features) still hard-reject on negatives.

• Per-component decision reasons — `TickDiagnostics` now carries `{load,throughput}decision_reason{prefill,decode}` in disagg mode. Aggregate reason is derived by priority (`scale_up > scale_down_capped_by_throughput > scale_down > no_change`) for back-compat.

How to run

Trace preparation

Any Mooncake-format trace works. Example:

```bash
mkdir -p traces/mooncake_fast25
curl -sL https://raw.githubusercontent.com/kvcache-ai/Mooncake/refs/heads/main/FAST25-release/traces/toolagent_trace.jsonl
-o traces/mooncake_fast25/toolagent_trace.jsonl
head -2361 traces/mooncake_fast25/toolagent_trace.jsonl > traces/mooncake_fast25/toolagent_trace_10pct.jsonl
```

Build Python bindings

```bash
.venv/bin/maturin develop --uv -m lib/bindings/python/Cargo.toml
export PYTHONPATH=lib/bindings/python/src:components/src
```

Install AIC (for throughput-based scaling)

```bash
uv pip install --python .venv/bin/python aiconfigurator
```

Agg mode with SLA-based scaling

```bash
.venv/bin/python -m dynamo.replay traces/mooncake_fast25/toolagent_trace_10pct.jsonl
--planner-config '{
"mode": "agg",
"optimization_target": "sla",
"ttft": 2500,
"itl": 50,
"enable_throughput_scaling": true,
"enable_load_scaling": true,
"pre_deployment_sweeping_mode": "rapid",
"throughput_adjustment_interval": 60,
"load_adjustment_interval": 5,
"report_filename": "replay_agg.html"
}'
--extra-engine-args '{"aic_backend": "vllm", "aic_system": "h200_sxm", "aic_model_path": "nvidia/Llama-3.1-8B-Instruct-FP8", "aic_tp_size": 1}'
--num-workers 2
--arrival-speedup-ratio 1.0
```

Disagg mode with SLA-based scaling

```bash
.venv/bin/python -m dynamo.replay traces/mooncake_fast25/toolagent_trace_10pct.jsonl
--planner-config '{
"mode": "disagg",
"optimization_target": "sla",
"ttft": 2500,
"itl": 50,
"enable_throughput_scaling": true,
"enable_load_scaling": true,
"pre_deployment_sweeping_mode": "rapid",
"throughput_adjustment_interval": 60,
"load_adjustment_interval": 5,
"report_filename": "replay_disagg.html"
}'
--prefill-engine-args '{"worker_type": "prefill", "aic_backend": "vllm", "aic_system": "h200_sxm", "aic_model_path": "nvidia/Llama-3.1-8B-Instruct-FP8", "aic_tp_size": 1}'
--decode-engine-args '{"worker_type": "decode", "aic_backend": "vllm", "aic_system": "h200_sxm", "aic_model_path": "nvidia/Llama-3.1-8B-Instruct-FP8", "aic_tp_size": 1}'
--num-prefill-workers 1
--num-decode-workers 1
--arrival-speedup-ratio 1.0
```

Results

Example run on 10% of the FAST25 toolagent trace (2,361 requests, ISL8.8k / OSL173).

Agg (90 planner ticks, SLA TTFT=2500ms / ITL=50ms):

Metric	avg	p90	p99
TTFT	861 ms	2,063 ms	4,765 ms
ITL	24 ms	29 ms	195 ms
E2E	5,165 ms	14,586 ms	31,729 ms

Disagg (83 planner ticks, same SLA):

Metric	avg	p90	p99
TTFT	952 ms	2,263 ms	4,612 ms
ITL	19 ms	26 ms	32 ms
E2E	4,443 ms	11,759 ms	19,147 ms

Disagg's ITL is tighter because decode is isolated from prefill chunking; agg takes a slightly lower avg TTFT because the same engine can start serving prefill without cross-worker KV transfer.

Full interactive reports

Both HTML reports (~4.7 MB each, with observed TTFT / ITL vs SLA, predicted vs observed traffic, per-engine KV queue + inflight with capacity line, engine RPS, load / throughput decision timelines):

• Agg: https://htmlpreview.github.io/?https://gist.githubusercontent.com/tedzhouhk/71e38c99b39fc370c92b8e7ff86c70b8/raw/replay_agg.html
• Disagg: https://htmlpreview.github.io/?https://gist.githubusercontent.com/tedzhouhk/71e38c99b39fc370c92b8e7ff86c70b8/raw/replay_disagg.html

(Hosted on a gist — htmlpreview renders the Plotly charts in-browser. Use "View raw" on this gist if the preview link stalls.)

Test plan

• All 86 planner unit tests pass (`pytest components/src/dynamo/planner/tests/unit/`)
• Agg SLA replay produces scaling events respecting throughput floor
• Disagg SLA replay produces separate prefill / decode decision timelines
• `connectgaps` renders sparse throughput traces as continuous lines
• Pre-commit clean on all touched files

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Enhanced diagnostics for load and throughput scaling with per-component decision tracking (prefill vs. decode).
  • Extended performance models to respect KV token capacity and sequence concurrency limits.
  • Improved traffic metrics capturing first-token latency and inter-token latency data.

• Improvements

  • Refined load scaling evaluation to properly apply post-bounds constraints.
  • Better throughput lower bound visualization in monitoring dashboards.
  • More accurate batch size calculations respecting token limits.

---

[coderabbitai]

Walkthrough

This PR enhances scaling decision diagnostics by segregating load and throughput reasons per component (prefill/decode), extends performance regression models with resource constraints (KV tokens, sequence counts, batch sizes), and enriches traffic tracking with latency metrics (TTFT, inter-token latency) across the Rust backend and Python planner.

Changes

Cohort / File(s)	Summary
Load Scaling Diagnostics components/src/dynamo/planner/core/load_scaling.py	Removed early-return for unchanged load decisions; refactored to apply bounds/budgets before concluding no scaling. Added per-component diagnostic computation (distinguishing scale_up, scale_down, scale_down_capped_by_throughput, no_change) and aggregate prioritization logic. Updated aggregated desired handling to allow throughput enforcement even on unchanged load decisions.
Performance Model Constraints components/src/dynamo/planner/core/perf_model/agg.py, decode.py, prefill.py	Extended regression model APIs with optional resource constraints: find_best_engine_agg_rps and find_best_engine_decode_rps now accept max_kv_tokens and max_num_seqs; find_best_engine_prefill_rps accepts max_num_batched_tokens. Prefill adds conditional chunking logic when input size exceeds batch limit. Decode and agg compute batch size upper bound as minimum of KV-capacity, sequence concurrency, and balance constraints via binary search.
Regression Model Foundation components/src/dynamo/planner/core/perf_model/base.py	Introduced _relaxable_feature_indices to allow selected regression coefficients to be negative without invalidating the fit; updated validation to reject only non-relaxable negative coefficients.
State Machine & Scaling Resolution components/src/dynamo/planner/core/state_machine.py, throughput_scaling.py	Added per-component diagnostic fields (_diag_load_reason_prefill, _diag_load_reason_decode, _diag_throughput_reason_prefill, _diag_throughput_reason_decode) and reordered tick resolution: throughput scaling executes first, load scaling follows, and load decision takes precedence when present. Extended regression queries to pass capability limits (KV tokens, sequence counts, batch sizes) to perf models.
Diagnostics Data & Reporting components/src/dynamo/planner/core/types.py, monitoring/diagnostics_recorder.py	Extended TickDiagnostics with throughput lower bounds and per-component decision reasons for prefill/decode. Updated DiagnosticsRecorder to plumb max_kv_tokens and enriched Plotly charts: conditional lower-bound overlay lines, disaggregated decision timelines with per-component markers, combined KV token visualizations, and replica transition counting instead of decision presence.
Traffic Latency Tracking (Python) components/src/dynamo/planner/offline/replay_adapter.py	Refined traffic tracking to conditionally record latency fields (TTFT, inter-token latency) only when requests are present; moved _last_traffic updates from diagnostics recording to traffic-metric-building phase.
Traffic Statistics Struct & Drain APIs (Rust) lib/mocker/src/replay/offline/components/types.rs, agg.rs, disagg.rs, planner_handle.rs, lib/bindings/python/rust/llm/replay.rs	Created new TrafficStats struct (adds avg_ttft_ms, avg_itl_ms to existing duration/count/throughput fields). Updated TrafficAccumulator::on_request to accept optional latency tuple and accumulate TTFT/ITL metrics. Changed drain APIs across agg, disagg, and planner_handle to return TrafficStats instead of tuple. Extended Python binding to expose new latency fields in JSON output.
Latency Collection & Benchmarking lib/mocker/src/replay/collector.rs, components/src/dynamo/planner/offline/replay_adapter.py, lib/bindings/python/src/dynamo/replay/main.py	Added request_latencies() method to compute per-request TTFT and inter-token latency from trace data. Integrated AIC-based synthetic benchmark generation into replay planner when model is not "easy", sweeping input parameters with engine KV budget constraints and merging points for aggregated mode.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 3

✅ Passed checks (3 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Title check	✅ Passed	The title clearly identifies the PR's main focus: planner-in-the-loop replay diagnostics and planner scaling logic fixes, which aligns with the primary changes across multiple files.
Description check	✅ Passed	The PR description comprehensively covers the changes, objectives, use cases, and test results, but does not follow the template structure with explicit sections for Overview, Details, Where to start, and Related Issues.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[jthomson04]

Looks good. One nit (can be for a follow up): Would it be possible to get the x-axes of all the charts to scale together? If I want to narrow down on a specific time period, I have to separately zoom into that area on each graph.

[tedzhouhk]

Looks good. One nit (can be for a follow up): Would it be possible to get the x-axes of all the charts to scale together? If I want to narrow down on a specific time period, I have to separately zoom into that area on each graph.

sure, added a TODO