Skip to content

GitLab

Unverified Commit 241b254d authored by Yan Ru Pei's avatar Yan Ru Pei Committed by GitHub
Browse files

chore(profiler): reshape replay_optimize into DGDR-aligned Pydantic specs (#8585) 


Signed-off-by: default avatarPeaBrane <yanrpei@gmail.com>
parent d401c1a4
Show whitespace changes
Inline Side-by-side
Showing with 1151 additions and 994 deletions
components/src/dynamo/profiler/tests/test_replay_optimize.py
View file @ 241b254d
......@@ -18,10 +18,13 @@ except ImportError:
from dynamo.profiler.utils import replay_optimize
from dynamo.profiler.utils.replay_optimize import (
DenseAggReplayState,
ReplayConstraints,
EngineSpec,
HardwareSpec,
ReplayObjective,
SyntheticReplayWorkload,
TraceReplayWorkload,
ReplayOptimizeSpec,
RouterSpec,
SLASpec,
WorkloadSpec,
compare_agg_and_disagg_with_replay,
optimize_dense_agg_with_replay,
optimize_dense_disagg_with_replay,
......@@ -100,6 +103,106 @@ def _write_trace(tmp_path: Path) -> Path:
return trace_path
def _synthetic_workload(
*,
isl: int = 64,
osl: int = 32,
request_count: int = 8,
concurrency: float = 4,
**extras: Any,
) -> WorkloadSpec:
return WorkloadSpec(
isl=isl,
osl=osl,
requestCount=request_count,
concurrency=concurrency,
**extras,
)
def _trace_workload(trace_file: Path, speedup: float = 100.0) -> WorkloadSpec:
return WorkloadSpec(traceFile=str(trace_file), arrivalSpeedupRatio=speedup)
def _sla(**bounds: float) -> SLASpec:
"""Build an SLASpec from keyword bounds.
Accepts the legacy `mean_e2e_latency_ms`, `mean_ttft_ms`, `mean_tpot_ms`,
etc. names for test readability; maps to camelCase DGDR field names.
"""
translate = {
"mean_ttft_ms": "ttft",
"mean_tpot_ms": "itl",
"mean_e2e_latency_ms": "e2eLatency",
"p95_ttft_ms": "p95Ttft",
"p95_tpot_ms": "p95Itl",
"p95_e2e_latency_ms": "p95E2eLatency",
}
return SLASpec(**{translate.get(k, k): v for k, v in bounds.items()})
def _disagg_spec(
*,
workload: WorkloadSpec | None = None,
total_gpus: int = 4,
sla: SLASpec | None = None,
objective: ReplayObjective = ReplayObjective.THROUGHPUT,
router_mode: str = "kv_router",
overlap_weights: list[float] | None = None,
base_router_config: KvRouterConfig | None = None,
max_parallel_evals: int = 1,
) -> ReplayOptimizeSpec:
return ReplayOptimizeSpec(
engine=EngineSpec(
model=_AIC_MODEL,
backend="vllm",
basePrefillEngineArgs=_base_prefill_args(),
baseDecodeEngineArgs=_base_decode_args(),
),
hardware=HardwareSpec(gpuSku=_AIC_SYSTEM, totalGpus=total_gpus),
workload=workload if workload is not None else _synthetic_workload(),
sla=sla if sla is not None else SLASpec(),
router=RouterSpec(
mode=router_mode,
overlapWeights=overlap_weights,
baseRouterConfig=base_router_config,
),
objective=objective,
maxParallelEvals=max_parallel_evals,
)
def _agg_spec(
*,
workload: WorkloadSpec | None = None,
total_gpus: int = 4,
sla: SLASpec | None = None,
router_mode: str = "kv_router",
overlap_weights: list[float] | None = None,
base_router_config: KvRouterConfig | None = None,
max_parallel_evals: int = 1,
) -> ReplayOptimizeSpec:
return ReplayOptimizeSpec(
engine=EngineSpec(
model=_AIC_MODEL,
backend="vllm",
baseEngineArgs=_base_agg_args(),
),
hardware=HardwareSpec(gpuSku=_AIC_SYSTEM, totalGpus=total_gpus),
workload=workload if workload is not None else _synthetic_workload(),
sla=sla if sla is not None else SLASpec(),
router=RouterSpec(
mode=router_mode,
overlapWeights=overlap_weights,
baseRouterConfig=base_router_config,
),
maxParallelEvals=max_parallel_evals,
)
# ---- internal-helper tests (unchanged from Phase 1 interface) ----
def test_enumerate_dense_tp_candidates_filters_to_tp_only(monkeypatch) -> None:
common = SimpleNamespace(BackendName=SimpleNamespace(vllm="vllm"))
task = SimpleNamespace(
......@@ -220,6 +323,9 @@ def test_iter_agg_worker_states_collapses_round_robin_overlap() -> None:
assert set(state.overlap_score_weight for state in states) == {0.0}
# ---- public-API tests (reshaped to ReplayOptimizeSpec) ----
def test_optimizer_finds_coordinate_optimum_and_reuses_cache(monkeypatch) -> None:
call_counter: Counter = Counter()
target_state = replay_optimize.DenseReplayState(2, 4, 2, 1, 2.0)
......@@ -253,21 +359,11 @@ def test_optimizer_finds_coordinate_optimum_and_reuses_cache(monkeypatch) -> Non
monkeypatch.setattr(replay_optimize.evaluate, "_run_replay_for_state", fake_run)
result = optimize_dense_disagg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=8,
constraints={"mean_e2e_latency_ms": 500.0},
overlap_score_weights=[0.0, 1.0, 2.0],
max_parallel_evals=1,
_disagg_spec(
total_gpus=8,
sla=_sla(mean_e2e_latency_ms=500.0),
overlap_weights=[0.0, 1.0, 2.0],
)
)
assert result.best_feasible is not None
......@@ -312,21 +408,12 @@ def test_agg_optimizer_finds_coordinate_optimum_and_reuses_cache(monkeypatch) ->
monkeypatch.setattr(replay_optimize.evaluate, "_run_agg_replay_for_state", fake_run)
result = optimize_dense_agg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
),
base_engine_args=_base_agg_args(),
max_total_gpus=8,
constraints={"mean_e2e_latency_ms": 500.0},
_agg_spec(
total_gpus=8,
sla=_sla(mean_e2e_latency_ms=500.0),
router_mode="both",
overlap_score_weights=[0.0, 1.0, 2.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0, 2.0],
)
)
assert result.best_feasible is not None
......@@ -371,21 +458,11 @@ def test_optimizer_uses_violation_penalty_when_no_state_is_feasible(
monkeypatch.setattr(replay_optimize.evaluate, "_run_replay_for_state", fake_run)
result = optimize_dense_disagg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=6,
constraints={"mean_e2e_latency_ms": 50.0},
overlap_score_weights=[0.0, 1.0],
max_parallel_evals=1,
_disagg_spec(
total_gpus=6,
sla=_sla(mean_e2e_latency_ms=50.0),
overlap_weights=[0.0, 1.0],
)
)
assert result.best_feasible is None
......@@ -429,21 +506,12 @@ def test_agg_optimizer_uses_violation_penalty_when_no_state_is_feasible(
monkeypatch.setattr(replay_optimize.evaluate, "_run_agg_replay_for_state", fake_run)
result = optimize_dense_agg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
),
base_engine_args=_base_agg_args(),
max_total_gpus=8,
constraints={"mean_e2e_latency_ms": 50.0},
_agg_spec(
total_gpus=8,
sla=_sla(mean_e2e_latency_ms=50.0),
router_mode="both",
overlap_score_weights=[0.0, 1.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0],
)
)
assert result.best_feasible is None
......@@ -479,26 +547,17 @@ def test_optimizer_supports_round_robin_router_mode(monkeypatch) -> None:
monkeypatch.setattr(replay_optimize.evaluate, "_run_replay_for_state", fake_run)
result = optimize_dense_disagg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=4,
constraints={"mean_e2e_latency_ms": 500.0},
_disagg_spec(
total_gpus=4,
sla=_sla(mean_e2e_latency_ms=500.0),
router_mode="round_robin",
overlap_score_weights=[0.0, 1.0, 2.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0, 2.0],
)
)
assert result.best_feasible is not None
assert set(seen_router_modes) == {"round_robin"}
# Guardrail #5: round_robin auto-collapses overlap weights to (0.0,)
assert set(seen_weights) == {0.0}
......@@ -533,22 +592,12 @@ def test_disagg_optimizer_supports_latency_objective(monkeypatch) -> None:
monkeypatch.setattr(replay_optimize.evaluate, "_run_replay_for_state", fake_run)
result = optimize_dense_disagg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=4,
constraints={"mean_e2e_latency_ms": 500.0},
objective="mean_e2e_latency",
overlap_score_weights=[0.0],
max_parallel_evals=1,
_disagg_spec(
total_gpus=4,
sla=_sla(mean_e2e_latency_ms=500.0),
objective=ReplayObjective.MEAN_E2E_LATENCY,
overlap_weights=[0.0],
)
)
assert result.best_feasible is not None
......@@ -561,22 +610,21 @@ def test_disagg_optimizer_supports_latency_objective(monkeypatch) -> None:
def test_disagg_optimizer_rejects_invalid_objective() -> None:
# Invalid objective strings raise at spec construction (Pydantic validates
# the enum field) and at direct ReplayObjective construction.
with pytest.raises(ValueError, match="not a valid ReplayObjective"):
optimize_dense_disagg_with_replay(
ReplayObjective("bad_objective")
with pytest.raises(ValueError):
ReplayOptimizeSpec(
engine=EngineSpec(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
basePrefillEngineArgs=_base_prefill_args(),
baseDecodeEngineArgs=_base_decode_args(),
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=4,
hardware=HardwareSpec(gpuSku=_AIC_SYSTEM, totalGpus=4),
workload=_synthetic_workload(),
objective="bad_objective",
max_parallel_evals=1,
)
......@@ -606,22 +654,12 @@ def test_disagg_optimizer_supports_router_mode_search(monkeypatch) -> None:
monkeypatch.setattr(replay_optimize.evaluate, "_run_replay_for_state", fake_run)
result = optimize_dense_disagg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=4,
constraints={"mean_e2e_latency_ms": 500.0},
_disagg_spec(
total_gpus=4,
sla=_sla(mean_e2e_latency_ms=500.0),
router_mode="both",
overlap_score_weights=[0.0, 1.0, 2.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0, 2.0],
)
)
assert result.best_feasible is not None
......@@ -658,21 +696,12 @@ def test_agg_optimizer_supports_router_mode_search(monkeypatch) -> None:
monkeypatch.setattr(replay_optimize.evaluate, "_run_agg_replay_for_state", fake_run)
result = optimize_dense_agg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
),
base_engine_args=_base_agg_args(),
max_total_gpus=4,
constraints={"mean_e2e_latency_ms": 500.0},
_agg_spec(
total_gpus=4,
sla=_sla(mean_e2e_latency_ms=500.0),
router_mode="both",
overlap_score_weights=[0.0, 1.0, 2.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0, 2.0],
)
)
assert result.best_feasible is not None
......@@ -721,31 +750,34 @@ def test_compare_agg_and_disagg_with_replay_picks_expected_mode(monkeypatch) ->
)
monkeypatch.setattr(
replay_optimize.bench, "optimize_dense_agg_with_replay", lambda **_: agg_result
replay_optimize.bench,
"optimize_dense_agg_with_replay",
lambda *_, **__: agg_result,
)
monkeypatch.setattr(
replay_optimize.bench,
"optimize_dense_disagg_with_replay",
lambda **_: disagg_result,
lambda *_, **__: disagg_result,
)
comparison = compare_agg_and_disagg_with_replay(
# Build a spec that populates both agg and disagg engine args so the
# comparison function accepts it (though the patched optimize_* bypass
# the actual engine-args assertions).
spec = ReplayOptimizeSpec(
engine=EngineSpec(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=64,
osl=32,
request_count=8,
replay_concurrency=4,
baseEngineArgs=_base_agg_args(),
basePrefillEngineArgs=_base_prefill_args(),
baseDecodeEngineArgs=_base_decode_args(),
),
base_engine_args=_base_agg_args(),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=8,
constraints={"mean_e2e_latency_ms": 500.0},
hardware=HardwareSpec(gpuSku=_AIC_SYSTEM, totalGpus=8),
workload=_synthetic_workload(),
sla=_sla(mean_e2e_latency_ms=500.0),
)
comparison = compare_agg_and_disagg_with_replay(spec)
assert comparison["chosen_mode"] == "agg"
assert comparison["chosen_best"] == agg_result.best_feasible
......@@ -761,6 +793,18 @@ def test_evaluate_state_prefers_normalized_metrics_over_report_payload() -> None
)
cache: dict[replay_optimize.DenseReplayState, dict[str, Any]] = {}
spec = ReplayOptimizeSpec(
engine=EngineSpec(
model="meta-llama/Llama-3.1-8B-Instruct",
backend="vllm",
basePrefillEngineArgs=_base_prefill_args(),
baseDecodeEngineArgs=_base_decode_args(),
),
hardware=HardwareSpec(gpuSku="h100_sxm", totalGpus=4),
workload=_synthetic_workload(isl=128, request_count=16),
sla=_sla(mean_e2e_latency_ms=1000.0),
)
with patch(
"dynamo.profiler.utils.replay_optimize.evaluate._run_replay_for_state",
return_value={
......@@ -773,20 +817,7 @@ def test_evaluate_state_prefers_normalized_metrics_over_report_payload() -> None
):
record = replay_optimize.evaluate._evaluate_state(
state=state,
workload=SyntheticReplayWorkload(
isl=128,
osl=32,
request_count=16,
replay_concurrency=4,
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
base_router_config=None,
model="meta-llama/Llama-3.1-8B-Instruct",
backend="vllm",
system="h100_sxm",
objective=ReplayObjective.THROUGHPUT,
constraints=ReplayConstraints(mean_e2e_latency_ms=1000.0),
spec=spec,
cache=cache,
)
......@@ -805,6 +836,17 @@ def test_evaluate_agg_state_prefers_normalized_metrics_over_report_payload() ->
)
cache: dict[DenseAggReplayState, dict[str, Any]] = {}
spec = ReplayOptimizeSpec(
engine=EngineSpec(
model="meta-llama/Llama-3.1-8B-Instruct",
backend="vllm",
baseEngineArgs=_base_agg_args(),
),
hardware=HardwareSpec(gpuSku="h100_sxm", totalGpus=4),
workload=_synthetic_workload(isl=128, request_count=16),
sla=_sla(mean_e2e_latency_ms=1000.0),
)
with patch(
"dynamo.profiler.utils.replay_optimize.evaluate._run_agg_replay_for_state",
return_value={
......@@ -817,19 +859,7 @@ def test_evaluate_agg_state_prefers_normalized_metrics_over_report_payload() ->
):
record = replay_optimize.evaluate._evaluate_agg_state(
state=state,
workload=SyntheticReplayWorkload(
isl=128,
osl=32,
request_count=16,
replay_concurrency=4,
),
base_engine_args=_base_agg_args(),
base_router_config=None,
model="meta-llama/Llama-3.1-8B-Instruct",
backend="vllm",
system="h100_sxm",
objective=ReplayObjective.THROUGHPUT,
constraints=ReplayConstraints(mean_e2e_latency_ms=1000.0),
spec=spec,
cache=cache,
)
......@@ -859,25 +889,17 @@ def test_agg_optimizer_synthetic_replay_smoke(monkeypatch) -> None:
)
result = optimize_dense_agg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=128,
osl=32,
request_count=8,
replay_concurrency=4,
_agg_spec(
workload=_synthetic_workload(isl=128, request_count=8),
total_gpus=4,
sla=_sla(
mean_ttft_ms=100000.0,
mean_tpot_ms=100000.0,
mean_e2e_latency_ms=100000.0,
),
base_engine_args=_base_agg_args(),
max_total_gpus=4,
constraints={
"mean_ttft_ms": 100000.0,
"mean_tpot_ms": 100000.0,
"mean_e2e_latency_ms": 100000.0,
},
router_mode="both",
overlap_score_weights=[0.0, 1.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0],
)
)
assert not result.evaluated_df.empty
......@@ -894,23 +916,17 @@ def test_agg_optimizer_timed_trace_smoke(tmp_path, monkeypatch) -> None:
)
result = optimize_dense_agg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=TraceReplayWorkload(
trace_file=_write_trace(tmp_path),
arrival_speedup_ratio=100.0,
_agg_spec(
workload=_trace_workload(_write_trace(tmp_path)),
total_gpus=4,
sla=_sla(
mean_ttft_ms=100000.0,
mean_tpot_ms=100000.0,
mean_e2e_latency_ms=100000.0,
),
base_engine_args=_base_agg_args(),
max_total_gpus=4,
constraints={
"mean_ttft_ms": 100000.0,
"mean_tpot_ms": 100000.0,
"mean_e2e_latency_ms": 100000.0,
},
router_mode="both",
overlap_score_weights=[0.0, 1.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0],
)
)
assert not result.evaluated_df.empty
......@@ -927,25 +943,16 @@ def test_optimizer_synthetic_replay_smoke(tmp_path, monkeypatch) -> None:
)
result = optimize_dense_disagg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=SyntheticReplayWorkload(
isl=128,
osl=32,
request_count=8,
replay_concurrency=4,
_disagg_spec(
workload=_synthetic_workload(isl=128, request_count=8),
total_gpus=4,
sla=_sla(
mean_ttft_ms=100000.0,
mean_tpot_ms=100000.0,
mean_e2e_latency_ms=100000.0,
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=4,
constraints={
"mean_ttft_ms": 100000.0,
"mean_tpot_ms": 100000.0,
"mean_e2e_latency_ms": 100000.0,
},
overlap_score_weights=[0.0, 1.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0],
)
)
assert not result.evaluated_df.empty
......@@ -962,24 +969,16 @@ def test_optimizer_timed_trace_smoke(tmp_path, monkeypatch) -> None:
)
result = optimize_dense_disagg_with_replay(
model=_AIC_MODEL,
backend="vllm",
system=_AIC_SYSTEM,
workload=TraceReplayWorkload(
trace_file=_write_trace(tmp_path),
arrival_speedup_ratio=100.0,
_disagg_spec(
workload=_trace_workload(_write_trace(tmp_path)),
total_gpus=4,
sla=_sla(
mean_ttft_ms=100000.0,
mean_tpot_ms=100000.0,
mean_e2e_latency_ms=100000.0,
),
base_prefill_engine_args=_base_prefill_args(),
base_decode_engine_args=_base_decode_args(),
max_total_gpus=4,
constraints={
"mean_ttft_ms": 100000.0,
"mean_tpot_ms": 100000.0,
"mean_e2e_latency_ms": 100000.0,
},
overlap_score_weights=[0.0, 1.0],
max_parallel_evals=1,
overlap_weights=[0.0, 1.0],
)
)
assert not result.evaluated_df.empty
assert result.best_feasible is not None
components/src/dynamo/profiler/utils/replay_optimize/README.md
View file @ 241b254d
......@@ -12,7 +12,7 @@ This experiment searches over disaggregated replay states to answer a concrete q
- given a fixed GPU budget
- for a workload with real prefix overlap
- and latency constraints that still permit meaningful throughput
- and latency SLAs that still permit meaningful throughput
which `(prefill_tp, decode_tp, prefill_workers, decode_workers, overlap_score_weight)` combination
produces the best offline replay result?
......@@ -20,6 +20,26 @@ produces the best offline replay result?
This is a heuristic search over replay states, not an exact optimizer over all feasible
configurations.
## Spec Shape (DGDR-aligned)
The public API takes a single [`ReplayOptimizeSpec`](specs.py) composed of:
- `EngineSpec` — model, backend, engine args (prefill + decode for disagg; single `baseEngineArgs`
for agg)
- `HardwareSpec` — GPU SKU + total GPU budget
- `WorkloadSpec` — synthetic workload knobs (isl/osl/concurrency/...) **or** a trace source
(`traceFile`/`arrivalSpeedupRatio`), discriminated by whether `traceFile` is set
- `SLASpec` — latency bounds (`ttft`, `itl`, `e2eLatency`, plus p95 variants); each is independent
and optional
- `RouterSpec` — router mode, overlap-score-weight sweep, KV-router base config
- `objective` — `ReplayObjective.THROUGHPUT` (default), `MEAN_TTFT`, or `MEAN_E2E_LATENCY`
- `maxParallelEvals` — how many replay evaluations to run concurrently
Field names use lowerCamelCase throughout to match the operator's auto-generated
[`dgdr_v1beta1_types.py`](../dgdr_v1beta1_types.py) so the eventual upstream merge into the Go
`DynamoGraphDeploymentRequestSpec` is mechanical. Method names (`.violation_penalty()`,
`.summarize()`, `.aic_task_kwargs()`) stay snake_case to match Pydantic convention.
## Prerequisites
Run from the repository root.
......@@ -72,7 +92,8 @@ export PYTHONPATH=lib/bindings/python/src:components/src
```
If the replay search uses multiple worker processes, prefer a real script file over a heredoc. This
matters on macOS because `ProcessPoolExecutor` child workers need a stable module path.
matters on macOS because `ProcessPoolExecutor` child workers need a stable module path, and the
driver module must guard its entry behind `if __name__ == "__main__":`.
For KV-router replay logs, this filter keeps the run readable without hiding useful `info` output:
......@@ -84,12 +105,12 @@ export DYN_LOG='info,dynamo_kv_router::scheduling::selector=warn'
This sweep uses:
- model: `Qwen/Qwen3-32B`
- backend: `vllm`
- system: `h200_sxm`
- router mode: `kv_router`
- workload type: `SyntheticReplayWorkload`
- GPU budget: `16`
- `EngineSpec.model`: `Qwen/Qwen3-32B`
- `EngineSpec.backend`: `vllm`
- `HardwareSpec.gpuSku`: `h200_sxm`
- `HardwareSpec.totalGpus`: `16`
- `RouterSpec.mode`: `kv_router`
- Synthetic `WorkloadSpec`
The GPU budget here is a simulated search constraint used by offline replay when it enumerates
candidate TP and worker configurations. You do not need 16 real GPUs locally to run this search.
......@@ -99,10 +120,10 @@ matter:
- `isl=32768`
- `osl=256`
- `request_count=5000`
- `replay_concurrency=200`
- `shared_prefix_ratio=0.5`
- `num_prefix_groups=50`
- `requestCount=5000`
- `concurrency=200`
- `sharedPrefixRatio=0.5`
- `numPrefixGroups=50`
The base engine args stay conservative:
......@@ -119,23 +140,37 @@ This setup does not force scheduler-specific bottlenecks such as:
Only add those when the experiment is specifically about scheduler limits.
## Search Strategy
`replay_optimize` runs a coordinate descent over three dimensions per round, iterating until the incumbent stops moving or `DEFAULT_SEARCH_ROUNDS` is reached:
```mermaid
flowchart LR
A["TP search<br/>choose TP shape<br/>(prefill_tp, decode_tp)<br/>under GPU budget"] --> B["Worker search<br/>choose worker split<br/>(prefill_workers, decode_workers)<br/>for the chosen TP"]
B --> C["Router search<br/>choose routing mode<br/>and overlap_score_weight"]
C --> A
```
Each step calls [`evaluate._evaluate_states`](evaluate.py), which replays the candidate state through `run_synthetic_trace_replay` or `run_trace_replay` (see [Mocker Trace Replay](../../../../../../docs/benchmarks/mocker-trace-replay.md) for the underlying harness) and ranks the resulting records with `scoring._pick_best_record`. The ranking key is `spec.objective` (throughput, mean_ttft, or mean_e2e_latency) subject to `spec.sla` bounds and `spec.hardware.totalGpus` as a feasibility gate.
The descent is budget-focused: each step prunes to near-budget-edge states so the sweep ends up at a TP/worker shape that actually consumes the available GPU budget, rather than at a throughput-per-GPU pareto point. Aggregated replay (`optimize_dense_agg_with_replay`) collapses dimensions 1 and 2 into `(tp, workers)` but is otherwise identical; see [`search.py`](search.py) for both entrypoints.
## Driver Script
The canonical starting point now lives in [example.py](example.py). Keeping it as a real module is
The canonical starting point lives in [example.py](example.py). Keeping it as a real module is
better than carrying a large inline snippet in the README, and it also satisfies the macOS
`ProcessPoolExecutor` requirement for a stable module path.
Treat [example.py](example.py) as a starting point, not a frozen harness. Modify it as needed for
your search:
- change the workload shape
- swap `SyntheticReplayWorkload` for `TraceReplayWorkload`
- change constraints
- change `overlap_score_weights`
- change the `WorkloadSpec` shape (or switch to a trace source with `traceFile=...`)
- add SLA bounds on `SLASpec` (`ttft`, `itl`, `e2eLatency`, or their p95 variants)
- change `RouterSpec.overlapWeights`
- print different columns from `result.evaluated_df` or `result.feasible_df`
- persist the tables to CSV or parquet if you want downstream analysis
If you need to understand which knobs are available, see [models.py](models.py), [search.py](search.py),
If you need to understand which knobs are available, see [specs.py](specs.py), [search.py](search.py),
and [evaluate.py](evaluate.py).
The default path in [example.py](example.py) is the synthetic disaggregated sweep documented in
......@@ -146,8 +181,8 @@ used for the Mooncake-style replay path below without rewriting the harness from
The returned object is a `DenseReplayOptimizationResult` with:
- `best_feasible`: best visited state that satisfies all constraints
- `best_infeasible`: best visited state that misses at least one constraint
- `best_feasible`: best visited state that satisfies all SLAs and the GPU budget
- `best_infeasible`: best visited state that misses at least one SLA bound or the budget
- `evaluated_df`: all visited states
- `feasible_df`: only the feasible visited states
......@@ -162,45 +197,41 @@ Useful columns to inspect:
- cache behavior: `prefix_cache_reused_ratio`
- latency: `mean_ttft_ms`, `mean_tpot_ms`, `mean_e2e_latency_ms`
In local testing, this setup produced a non-trivial mean-E2E winner around:
- `prefill_tp=2`
- `decode_tp=1`
- `prefill_workers=2`
- `decode_workers=4`
- `overlap_score_weight=0.5`
Note that the report DataFrame still uses the Rust replay runner's key names
(`mean_ttft_ms`, `mean_tpot_ms`, `mean_e2e_latency_ms`) even though the input `SLASpec` uses DGDR's
camelCase names (`ttft`, `itl`, `e2eLatency`). `SLASpec` carries an internal translation map;
renaming the Rust output keys is a follow-up in the Rust replay runner.
Ballpark metrics for that point were:
In local testing, this setup produced a non-trivial mean-E2E winner around:
- `prefix_cache_reused_ratio ~= 0.5`
- `output_throughput_tok_s ~= 4500`
- `mean_ttft_ms ~= 4500`
- `mean_tpot_ms ~= 26`
- `mean_e2e_latency_ms ~= 11150`
- `prefill_tp=4`, `decode_tp=1`, `prefill_workers=3`, `decode_workers=4`, `overlap_score_weight=0.5`
- `output_throughput_tok_s ~= 970`, `prefix_cache_reused_ratio ~= 0.5`,
`mean_ttft_ms ~= 42800`, `mean_tpot_ms ~= 35`, `mean_e2e_latency_ms ~= 51900`
Treat those as sanity-check ranges, not fixed assertions.
Treat those as sanity-check ranges, not fixed assertions. See the regression anchor table in the
`run-replay-bench` skill for the current reference frontier.
## Tuning This Sweep
To broaden or shift the search, vary one axis at a time:
- `max_total_gpus`
- `overlap_score_weights`
- `shared_prefix_ratio`
- `num_prefix_groups`
- `HardwareSpec.totalGpus`
- `RouterSpec.overlapWeights`
- `WorkloadSpec.sharedPrefixRatio`
- `WorkloadSpec.numPrefixGroups`
- base prefill/decode engine args
If you want to compare routing strategies directly, use `router_mode="both"` instead of the default
KV-router-only search.
If you want to compare routing strategies directly, use `RouterSpec(mode="both")` instead of the
default KV-router-only search.
## Real Traffic Replay
`replay_optimize` is wired up for trace-driven replay. In
[evaluate.py](evaluate.py), `TraceReplayWorkload` goes through `run_trace_replay(...)`, while
`SyntheticReplayWorkload` goes through `run_synthetic_trace_replay(...)`.
`replay_optimize` is wired up for trace-driven replay. `WorkloadSpec` takes a `traceFile` (plus
optional `arrivalSpeedupRatio`) discriminator; when set, [evaluate.py](evaluate.py) routes through
`run_trace_replay(...)` instead of `run_synthetic_trace_replay(...)`.
Use a separate trace-driven experiment when you want to evaluate the same search structure against a
real Mooncake-style workload instead of the synthetic shared-prefix workload above.
Use a separate trace-driven experiment when you want to evaluate the same search structure against
a real Mooncake-style workload instead of the synthetic shared-prefix workload above.
### Download a Mooncake Trace
......@@ -222,53 +253,44 @@ wget -O /tmp/toolagent_trace.jsonl \
If you use [example.py](example.py), pass `--trace-file /tmp/toolagent_trace.jsonl` and optionally
`--arrival-speedup-ratio 0.8`.
If you want to edit the driver directly, replace:
If you want to edit the driver directly, replace the synthetic `WorkloadSpec`:
```python
workload=SyntheticReplayWorkload(
workload=WorkloadSpec(
isl=32768,
osl=256,
request_count=5000,
replay_concurrency=200,
shared_prefix_ratio=0.5,
num_prefix_groups=50,
requestCount=5000,
concurrency=200,
sharedPrefixRatio=0.5,
numPrefixGroups=50,
),
```
with:
```python
from dynamo.profiler.utils.replay_optimize import TraceReplayWorkload
workload=TraceReplayWorkload(
trace_file="/tmp/toolagent_trace.jsonl",
arrival_speedup_ratio=1.0,
workload=WorkloadSpec(
traceFile="/tmp/toolagent_trace.jsonl",
arrivalSpeedupRatio=1.0,
),
```
If you want to replay the same trace at `0.80x` of its original arrival rate, keep the same file
and set:
```python
workload=TraceReplayWorkload(
trace_file="/tmp/toolagent_trace.jsonl",
arrival_speedup_ratio=0.8,
),
```
and set `arrivalSpeedupRatio=0.8`.
The main behavioral change is that the workload stops generating requests in memory and instead
replays request arrivals from the JSONL trace. In this path:
- `trace_file` points at the Mooncake-style JSONL input
- `arrival_speedup_ratio` compresses or stretches the trace arrival process
- synthetic-only knobs such as `isl`, `osl`, `request_count`, `replay_concurrency`,
`shared_prefix_ratio`, and `num_prefix_groups` no longer apply at the workload level
- `traceFile` points at the Mooncake-style JSONL input
- `arrivalSpeedupRatio` compresses or stretches the trace arrival process
- synthetic-only knobs such as `isl`, `osl`, `requestCount`, `concurrency`,
`sharedPrefixRatio`, `numPrefixGroups` are ignored by the trace replay path
Important notes for the public toolagent trace:
- the dataset uses Mooncake-style `hash_ids` with `512` tokens per block
- the underlying `run_trace_replay(...)` API defaults `trace_block_size` to `512`
- the current `TraceReplayWorkload` wrapper does not expose a separate `trace_block_size` field
- `WorkloadSpec` does not yet expose a separate `traceBlockSize` field
- the prefix-data-generator tools in
[Prefix Data Generator](../../../../../../benchmarks/prefix_data_generator/README.md)
are useful if you want to inspect the trace first or synthesize a larger derivative trace before
......
components/src/dynamo/profiler/utils/replay_optimize/__init__.py
View file @ 241b254d
......@@ -3,7 +3,7 @@
from __future__ import annotations
from . import aic, bench, engine_args, evaluate, scoring, search
from . import aic, bench, engine_args, evaluate, scoring, search, specs
from .aic import _enumerate_dense_tp_candidates, _load_aiconfigurator_modules
from .bench import compare_agg_and_disagg_with_replay, compare_aic_and_replay_disagg
from .engine_args import (
......@@ -11,15 +11,7 @@ from .engine_args import (
_build_candidate_engine_args,
_build_router_config,
)
from .models import (
DenseAggReplayState,
DenseReplayOptimizationResult,
DenseReplayState,
ReplayConstraints,
ReplayObjective,
SyntheticReplayWorkload,
TraceReplayWorkload,
)
from .models import DenseAggReplayState, DenseReplayOptimizationResult, DenseReplayState
from .scoring import _pick_best_record
from .search import (
_iter_agg_tp_states_with_max_workers,
......@@ -29,6 +21,16 @@ from .search import (
optimize_dense_agg_with_replay,
optimize_dense_disagg_with_replay,
)
from .specs import (
EngineSpec,
HardwareSpec,
ReplayObjective,
ReplayOptimizeSpec,
RouterMode,
RouterSpec,
SLASpec,
WorkloadSpec,
)
__all__ = [
"_build_agg_candidate_engine_args",
......@@ -46,10 +48,14 @@ __all__ = [
"DenseAggReplayState",
"DenseReplayOptimizationResult",
"DenseReplayState",
"ReplayConstraints",
"EngineSpec",
"HardwareSpec",
"ReplayObjective",
"SyntheticReplayWorkload",
"TraceReplayWorkload",
"ReplayOptimizeSpec",
"RouterMode",
"RouterSpec",
"SLASpec",
"WorkloadSpec",
"aic",
"bench",
"engine_args",
......@@ -58,4 +64,5 @@ __all__ = [
"optimize_dense_disagg_with_replay",
"scoring",
"search",
"specs",
]
components/src/dynamo/profiler/utils/replay_optimize/bench.py
View file @ 241b254d
......@@ -3,65 +3,49 @@
from __future__ import annotations
from collections.abc import Mapping
from typing import Any
import pandas as pd
from aiconfigurator.sdk.task import TaskConfig, TaskRunner
from dynamo.llm import MockEngineArgs
from .models import ReplayConstraints, SyntheticReplayWorkload, TraceReplayWorkload
from .scoring import _pick_best_record
from .search import optimize_dense_agg_with_replay, optimize_dense_disagg_with_replay
from .specs import ReplayOptimizeSpec
def compare_aic_and_replay_disagg(
*,
model: str,
backend: str,
system: str,
isl: int,
osl: int,
max_total_gpus: int,
replay_request_count: int,
replay_concurrency: int,
base_prefill_engine_args: MockEngineArgs,
base_decode_engine_args: MockEngineArgs,
constraints: Mapping[str, float] | None = None,
max_parallel_evals: int = 1,
spec: ReplayOptimizeSpec,
) -> dict[str, Any]:
aic_constraints = ReplayConstraints.from_mapping(constraints, max_total_gpus)
"""Run AIC pareto + replay optimization side-by-side for a disagg config.
Uses SLA bounds from `spec.sla` as AIC's latency targets; round_robin router
mode is forced for the replay run (AIC itself has no router sweep).
"""
if spec.workload.isTraceBased:
raise ValueError("compare_aic_and_replay_disagg requires a synthetic workload")
if spec.workload.requestCount is None or spec.workload.concurrency is None:
raise ValueError(
"compare_aic_and_replay_disagg requires synthetic WorkloadSpec with "
"requestCount and concurrency"
)
aic_task = TaskConfig(
serving_mode="disagg",
model_path=model,
system_name=system,
backend_name=backend,
total_gpus=max_total_gpus,
isl=isl,
osl=osl,
**aic_constraints.aic_task_kwargs(),
model_path=spec.engine.model,
system_name=str(spec.hardware.gpuSku),
backend_name=spec.engine.backend.value,
total_gpus=spec.hardware.totalGpus,
isl=spec.workload.isl,
osl=spec.workload.osl,
**spec.sla.aic_task_kwargs(),
)
aic_result = TaskRunner().run(aic_task)
aic_df = aic_result.get("pareto_df", pd.DataFrame())
replay_result = optimize_dense_disagg_with_replay(
model=model,
backend=backend,
system=system,
workload=SyntheticReplayWorkload(
isl=isl,
osl=osl,
request_count=replay_request_count,
replay_concurrency=replay_concurrency,
),
base_prefill_engine_args=base_prefill_engine_args,
base_decode_engine_args=base_decode_engine_args,
max_total_gpus=max_total_gpus,
constraints=constraints,
router_mode="round_robin",
max_parallel_evals=max_parallel_evals,
replay_spec = spec.model_copy(
update={"router": spec.router.model_copy(update={"mode": "round_robin"})}
)
replay_result = optimize_dense_disagg_with_replay(replay_spec)
aic_best = None
if not aic_df.empty:
......@@ -108,45 +92,15 @@ def compare_aic_and_replay_disagg(
def compare_agg_and_disagg_with_replay(
*,
model: str,
backend: str,
system: str,
workload: SyntheticReplayWorkload | TraceReplayWorkload,
base_engine_args: MockEngineArgs,
base_prefill_engine_args: MockEngineArgs,
base_decode_engine_args: MockEngineArgs,
max_total_gpus: int,
constraints: Mapping[str, float] | None = None,
router_mode: str = "kv_router",
overlap_score_weights: tuple[float, ...] | list[float] | None = None,
max_parallel_evals: int = 1,
spec: ReplayOptimizeSpec,
) -> dict[str, Any]:
agg_result = optimize_dense_agg_with_replay(
model=model,
backend=backend,
system=system,
workload=workload,
base_engine_args=base_engine_args,
max_total_gpus=max_total_gpus,
constraints=constraints,
router_mode=router_mode,
overlap_score_weights=overlap_score_weights,
max_parallel_evals=max_parallel_evals,
)
disagg_result = optimize_dense_disagg_with_replay(
model=model,
backend=backend,
system=system,
workload=workload,
base_prefill_engine_args=base_prefill_engine_args,
base_decode_engine_args=base_decode_engine_args,
max_total_gpus=max_total_gpus,
constraints=constraints,
router_mode=router_mode,
overlap_score_weights=overlap_score_weights,
max_parallel_evals=max_parallel_evals,
)
"""Run both agg and disagg replay optimizations on the same spec and pick the winner.
The spec must populate `spec.engine.baseEngineArgs` (agg path) plus
`basePrefillEngineArgs` / `baseDecodeEngineArgs` (disagg path).
"""
agg_result = optimize_dense_agg_with_replay(spec)
disagg_result = optimize_dense_disagg_with_replay(spec)
agg_best = agg_result.best_feasible
disagg_best = disagg_result.best_feasible
......
components/src/dynamo/profiler/utils/replay_optimize/constants.py
View file @ 241b254d
......@@ -13,14 +13,3 @@ AIC_BACKEND_VERSIONS = {
DEFAULT_OVERLAP_SCORE_WEIGHTS = (0.0, 0.25, 0.5, 1.0, 2.0, 4.0)
DEFAULT_MAX_PARALLEL_EVALS = min(8, os.cpu_count() or 1)
DEFAULT_SEARCH_ROUNDS = 3
SUPPORTED_CONSTRAINTS = frozenset(
{
"mean_ttft_ms",
"p95_ttft_ms",
"mean_tpot_ms",
"p95_tpot_ms",
"mean_e2e_latency_ms",
"p95_e2e_latency_ms",
"max_total_gpus",
}
)
components/src/dynamo/profiler/utils/replay_optimize/evaluate.py
View file @ 241b254d
......@@ -4,8 +4,8 @@
"""Replay evaluation helpers for the budget-focused dense search heuristic.
The search in `search.py` assumes we prefer to consume the available GPU budget
and therefore ranks visited states by raw output throughput, subject to replay
constraints, rather than by throughput normalized per GPU.
and therefore ranks visited states by the selected `spec.objective`, subject to
SLA and budget constraints, rather than by throughput normalized per GPU.
"""
from __future__ import annotations
......@@ -25,105 +25,119 @@ from .engine_args import (
_build_router_config,
)
from .logging import ensure_dynamo_logging, log_state_finish, log_state_start
from .models import (
DenseAggReplayState,
DenseReplayState,
ReplayConstraints,
from .models import DenseAggReplayState, DenseReplayState
from .specs import (
EngineSpec,
HardwareSpec,
ReplayObjective,
SyntheticReplayWorkload,
TraceReplayWorkload,
ReplayOptimizeSpec,
RouterSpec,
SLASpec,
WorkloadSpec,
)
def _run_replay_for_state(
*,
state: DenseReplayState,
workload: SyntheticReplayWorkload | TraceReplayWorkload,
workload: WorkloadSpec,
prefill_engine_args: MockEngineArgs,
decode_engine_args: MockEngineArgs,
router_config: KvRouterConfig | None,
) -> dict[str, Any]:
if isinstance(workload, SyntheticReplayWorkload):
return run_synthetic_trace_replay(
workload.isl,
workload.osl,
workload.request_count,
if workload.isTraceBased:
return run_trace_replay(
Path(workload.traceFile),
prefill_engine_args=prefill_engine_args,
decode_engine_args=decode_engine_args,
router_config=router_config,
num_prefill_workers=state.prefill_workers,
num_decode_workers=state.decode_workers,
replay_concurrency=workload.replay_concurrency,
replay_mode="offline",
router_mode=state.router_mode,
arrival_interval_ms=workload.arrival_interval_ms,
turns_per_session=workload.turns_per_session,
shared_prefix_ratio=workload.shared_prefix_ratio,
num_prefix_groups=workload.num_prefix_groups,
inter_turn_delay_ms=workload.inter_turn_delay_ms,
arrival_speedup_ratio=workload.arrivalSpeedupRatio,
)
return run_trace_replay(
Path(workload.trace_file),
return run_synthetic_trace_replay(
workload.isl,
workload.osl,
int(workload.requestCount),
prefill_engine_args=prefill_engine_args,
decode_engine_args=decode_engine_args,
router_config=router_config,
num_prefill_workers=state.prefill_workers,
num_decode_workers=state.decode_workers,
replay_concurrency=int(workload.concurrency),
replay_mode="offline",
router_mode=state.router_mode,
arrival_speedup_ratio=workload.arrival_speedup_ratio,
arrival_interval_ms=workload.arrivalIntervalMs,
turns_per_session=workload.turnsPerSession,
shared_prefix_ratio=workload.sharedPrefixRatio,
num_prefix_groups=workload.numPrefixGroups,
inter_turn_delay_ms=workload.interTurnDelayMs,
)
def _run_agg_replay_for_state(
*,
state: DenseAggReplayState,
workload: SyntheticReplayWorkload | TraceReplayWorkload,
workload: WorkloadSpec,
engine_args: MockEngineArgs,
router_config: KvRouterConfig | None,
) -> dict[str, Any]:
if isinstance(workload, SyntheticReplayWorkload):
return run_synthetic_trace_replay(
workload.isl,
workload.osl,
workload.request_count,
if workload.isTraceBased:
return run_trace_replay(
Path(workload.traceFile),
extra_engine_args=engine_args,
router_config=router_config,
num_workers=state.workers,
replay_concurrency=workload.replay_concurrency,
replay_mode="offline",
router_mode=state.router_mode,
arrival_interval_ms=workload.arrival_interval_ms,
turns_per_session=workload.turns_per_session,
shared_prefix_ratio=workload.shared_prefix_ratio,
num_prefix_groups=workload.num_prefix_groups,
inter_turn_delay_ms=workload.inter_turn_delay_ms,
arrival_speedup_ratio=workload.arrivalSpeedupRatio,
)
return run_trace_replay(
Path(workload.trace_file),
return run_synthetic_trace_replay(
workload.isl,
workload.osl,
int(workload.requestCount),
extra_engine_args=engine_args,
router_config=router_config,
num_workers=state.workers,
replay_concurrency=int(workload.concurrency),
replay_mode="offline",
router_mode=state.router_mode,
arrival_speedup_ratio=workload.arrival_speedup_ratio,
arrival_interval_ms=workload.arrivalIntervalMs,
turns_per_session=workload.turnsPerSession,
shared_prefix_ratio=workload.sharedPrefixRatio,
num_prefix_groups=workload.numPrefixGroups,
inter_turn_delay_ms=workload.interTurnDelayMs,
)
def _feasibility(
*,
report: Mapping[str, Any],
state: DenseReplayState | DenseAggReplayState,
sla: SLASpec,
hardware: HardwareSpec,
) -> tuple[float, bool]:
"""Split SLA violation and hardware-budget gate, combine into one penalty.
The record's `violation_penalty` field keeps its Phase 1 role (drives
infeasible-record sorting in `_pick_best_record`); over-budget adds a
constant 1.0 so it outranks small SLA misses.
"""
sla_penalty = sla.violation_penalty(report)
over_budget = state.total_gpus_used > hardware.totalGpus
penalty = sla_penalty + (1.0 if over_budget else 0.0)
feasible = sla_penalty == 0.0 and not over_budget
return penalty, feasible
def _evaluate_state(
*,
state: DenseReplayState,
workload: SyntheticReplayWorkload | TraceReplayWorkload,
base_prefill_engine_args: MockEngineArgs,
base_decode_engine_args: MockEngineArgs,
base_router_config: KvRouterConfig | None,
model: str,
backend: str,
system: str,
objective: ReplayObjective,
constraints: ReplayConstraints,
spec: ReplayOptimizeSpec,
cache: dict[DenseReplayState, dict[str, Any]],
) -> dict[str, Any]:
ensure_dynamo_logging()
......@@ -133,54 +147,57 @@ def _evaluate_state(
log_state_start(state)
backend = spec.engine.backend.value
system = str(spec.hardware.gpuSku)
prefill_args = _build_candidate_engine_args(
base_args=base_prefill_engine_args,
base_args=spec.engine.basePrefillEngineArgs,
tp_size=state.prefill_tp,
worker_type="prefill",
backend=backend,
system=system,
model=model,
model=spec.engine.model,
)
decode_args = _build_candidate_engine_args(
base_args=base_decode_engine_args,
base_args=spec.engine.baseDecodeEngineArgs,
tp_size=state.decode_tp,
worker_type="decode",
backend=backend,
system=system,
model=model,
model=spec.engine.model,
)
router_config = None
if state.router_mode == "kv_router":
router_config = _build_router_config(
base_router_config, state.overlap_score_weight
spec.router.baseRouterConfig, state.overlap_score_weight
)
report = _run_replay_for_state(
state=state,
workload=workload,
workload=spec.workload,
prefill_engine_args=prefill_args,
decode_engine_args=decode_args,
router_config=router_config,
)
total_gpus_used = state.total_gpus_used
throughput = float(report["output_throughput_tok_s"])
score = objective.score(report)
penalty = constraints.violation_penalty(report, total_gpus_used)
feasible = penalty == 0.0
score = spec.objective.score(report)
penalty, feasible = _feasibility(
report=report, state=state, sla=spec.sla, hardware=spec.hardware
)
record = {
**report,
**asdict(state),
"total_gpus_used": total_gpus_used,
"total_gpus_used": state.total_gpus_used,
"output_throughput_tok_s": throughput,
"score": score,
"objective": objective.value,
"objective": spec.objective.value,
"feasible": feasible,
"violation_penalty": penalty,
}
log_state_finish(
state=state,
report=report,
constraints=constraints,
sla=spec.sla,
hardware=spec.hardware,
score=score,
feasible=feasible,
violation_penalty=penalty,
......@@ -192,14 +209,7 @@ def _evaluate_state(
def _evaluate_agg_state(
*,
state: DenseAggReplayState,
workload: SyntheticReplayWorkload | TraceReplayWorkload,
base_engine_args: MockEngineArgs,
base_router_config: KvRouterConfig | None,
model: str,
backend: str,
system: str,
objective: ReplayObjective,
constraints: ReplayConstraints,
spec: ReplayOptimizeSpec,
cache: dict[DenseAggReplayState, dict[str, Any]],
) -> dict[str, Any]:
ensure_dynamo_logging()
......@@ -209,44 +219,47 @@ def _evaluate_agg_state(
log_state_start(state)
backend = spec.engine.backend.value
system = str(spec.hardware.gpuSku)
engine_args = _build_agg_candidate_engine_args(
base_args=base_engine_args,
base_args=spec.engine.baseEngineArgs,
tp_size=state.tp,
backend=backend,
system=system,
model=model,
model=spec.engine.model,
)
router_config = None
if state.router_mode == "kv_router":
router_config = _build_router_config(
base_router_config, state.overlap_score_weight
spec.router.baseRouterConfig, state.overlap_score_weight
)
report = _run_agg_replay_for_state(
state=state,
workload=workload,
workload=spec.workload,
engine_args=engine_args,
router_config=router_config,
)
total_gpus_used = state.total_gpus_used
throughput = float(report["output_throughput_tok_s"])
score = objective.score(report)
penalty = constraints.violation_penalty(report, total_gpus_used)
feasible = penalty == 0.0
score = spec.objective.score(report)
penalty, feasible = _feasibility(
report=report, state=state, sla=spec.sla, hardware=spec.hardware
)
record = {
**report,
**asdict(state),
"total_gpus_used": total_gpus_used,
"total_gpus_used": state.total_gpus_used,
"output_throughput_tok_s": throughput,
"score": score,
"objective": objective.value,
"objective": spec.objective.value,
"feasible": feasible,
"violation_penalty": penalty,
}
log_state_finish(
state=state,
report=report,
constraints=constraints,
sla=spec.sla,
hardware=spec.hardware,
score=score,
feasible=feasible,
violation_penalty=penalty,
......@@ -255,47 +268,102 @@ def _evaluate_agg_state(
return record
def _evaluate_state_from_json_payloads(payload: Mapping[str, Any]) -> dict[str, Any]:
return _evaluate_state(
state=payload["state"],
workload=payload["workload"],
base_prefill_engine_args=MockEngineArgs.from_json(
payload["base_prefill_engine_args_json"]
# ---- Cross-process payload bridge ----
# `MockEngineArgs` and `KvRouterConfig` are Rust-bound and don't pickle through
# `ProcessPoolExecutor`; we round-trip them via their own `dump_json()` /
# `from_json()` methods. Everything else on `ReplayOptimizeSpec` is a Pydantic
# model and pickles natively, but we serialize the whole payload dict to keep
# cross-process transport explicit.
def _spec_to_payload(spec: ReplayOptimizeSpec) -> dict[str, Any]:
engine = spec.engine
router = spec.router
return {
"engine_model": engine.model,
"engine_backend": engine.backend.value,
"engine_base_agg_json": (
engine.baseEngineArgs.dump_json()
if engine.baseEngineArgs is not None
else None
),
"engine_base_prefill_json": (
engine.basePrefillEngineArgs.dump_json()
if engine.basePrefillEngineArgs is not None
else None
),
"engine_base_decode_json": (
engine.baseDecodeEngineArgs.dump_json()
if engine.baseDecodeEngineArgs is not None
else None
),
"hardware_json": spec.hardware.model_dump_json(),
"workload_json": spec.workload.model_dump_json(),
"sla_json": spec.sla.model_dump_json(),
"router_mode": router.mode,
"router_overlap_weights": (
None if router.overlapWeights is None else list(router.overlapWeights)
),
"router_base_config_json": (
router.baseRouterConfig.dump_json()
if router.baseRouterConfig is not None
else None
),
"objective": spec.objective.value,
"max_parallel_evals": spec.maxParallelEvals,
}
def _spec_from_payload(payload: Mapping[str, Any]) -> ReplayOptimizeSpec:
return ReplayOptimizeSpec(
engine=EngineSpec(
model=payload["engine_model"],
backend=payload["engine_backend"],
baseEngineArgs=(
MockEngineArgs.from_json(payload["engine_base_agg_json"])
if payload["engine_base_agg_json"] is not None
else None
),
basePrefillEngineArgs=(
MockEngineArgs.from_json(payload["engine_base_prefill_json"])
if payload["engine_base_prefill_json"] is not None
else None
),
baseDecodeEngineArgs=(
MockEngineArgs.from_json(payload["engine_base_decode_json"])
if payload["engine_base_decode_json"] is not None
else None
),
base_decode_engine_args=MockEngineArgs.from_json(
payload["base_decode_engine_args_json"]
),
base_router_config=(
KvRouterConfig.from_json(payload["base_router_config_json"])
if payload["base_router_config_json"] is not None
hardware=HardwareSpec.model_validate_json(payload["hardware_json"]),
workload=WorkloadSpec.model_validate_json(payload["workload_json"]),
sla=SLASpec.model_validate_json(payload["sla_json"]),
router=RouterSpec(
mode=payload["router_mode"],
overlapWeights=payload["router_overlap_weights"],
baseRouterConfig=(
KvRouterConfig.from_json(payload["router_base_config_json"])
if payload["router_base_config_json"] is not None
else None
),
model=payload["model"],
backend=payload["backend"],
system=payload["system"],
objective=payload["objective"],
constraints=payload["constraints"],
),
objective=ReplayObjective(payload["objective"]),
maxParallelEvals=payload["max_parallel_evals"],
)
def _evaluate_state_from_payload(payload: Mapping[str, Any]) -> dict[str, Any]:
return _evaluate_state(
state=payload["state"],
spec=_spec_from_payload(payload["spec"]),
cache={},
)
def _evaluate_agg_state_from_json_payloads(
payload: Mapping[str, Any]
) -> dict[str, Any]:
def _evaluate_agg_state_from_payload(payload: Mapping[str, Any]) -> dict[str, Any]:
return _evaluate_agg_state(
state=payload["state"],
workload=payload["workload"],
base_engine_args=MockEngineArgs.from_json(payload["base_engine_args_json"]),
base_router_config=(
KvRouterConfig.from_json(payload["base_router_config_json"])
if payload["base_router_config_json"] is not None
else None
),
model=payload["model"],
backend=payload["backend"],
system=payload["system"],
objective=payload["objective"],
constraints=payload["constraints"],
spec=_spec_from_payload(payload["spec"]),
cache={},
)
......@@ -303,15 +371,7 @@ def _evaluate_agg_state_from_json_payloads(
def _evaluate_states(
*,
states: Sequence[DenseReplayState],
workload: SyntheticReplayWorkload | TraceReplayWorkload,
base_prefill_engine_args: MockEngineArgs,
base_decode_engine_args: MockEngineArgs,
base_router_config: KvRouterConfig | None,
model: str,
backend: str,
system: str,
objective: ReplayObjective,
constraints: ReplayConstraints,
spec: ReplayOptimizeSpec,
cache: dict[DenseReplayState, dict[str, Any]],
max_parallel_evals: int,
executor: Executor | None = None,
......@@ -333,43 +393,13 @@ def _evaluate_states(
if max_parallel_evals <= 1 or len(uncached_states) == 1 or executor is None:
for index, state in zip(uncached_indices, uncached_states, strict=True):
records[index] = _evaluate_state(
state=state,
workload=workload,
base_prefill_engine_args=base_prefill_engine_args,
base_decode_engine_args=base_decode_engine_args,
base_router_config=base_router_config,
model=model,
backend=backend,
system=system,
objective=objective,
constraints=constraints,
cache=cache,
)
records[index] = _evaluate_state(state=state, spec=spec, cache=cache)
return [record for record in records if record is not None]
base_prefill_engine_args_json = base_prefill_engine_args.dump_json()
base_decode_engine_args_json = base_decode_engine_args.dump_json()
base_router_config_json = (
None if base_router_config is None else base_router_config.dump_json()
)
payloads = [
{
"state": state,
"workload": workload,
"base_prefill_engine_args_json": base_prefill_engine_args_json,
"base_decode_engine_args_json": base_decode_engine_args_json,
"base_router_config_json": base_router_config_json,
"model": model,
"backend": backend,
"system": system,
"objective": objective,
"constraints": constraints,
}
for state in uncached_states
]
spec_payload = _spec_to_payload(spec)
payloads = [{"state": state, "spec": spec_payload} for state in uncached_states]
future_records = list(executor.map(_evaluate_state_from_json_payloads, payloads))
future_records = list(executor.map(_evaluate_state_from_payload, payloads))
for index, state, record in zip(
uncached_indices,
......@@ -386,14 +416,7 @@ def _evaluate_states(
def _evaluate_agg_states(
*,
states: Sequence[DenseAggReplayState],
workload: SyntheticReplayWorkload | TraceReplayWorkload,
base_engine_args: MockEngineArgs,
base_router_config: KvRouterConfig | None,
model: str,
backend: str,
system: str,
objective: ReplayObjective,
constraints: ReplayConstraints,
spec: ReplayOptimizeSpec,
cache: dict[DenseAggReplayState, dict[str, Any]],
max_parallel_evals: int,
executor: Executor | None = None,
......@@ -415,42 +438,13 @@ def _evaluate_agg_states(
if max_parallel_evals <= 1 or len(uncached_states) == 1 or executor is None:
for index, state in zip(uncached_indices, uncached_states, strict=True):
records[index] = _evaluate_agg_state(
state=state,
workload=workload,
base_engine_args=base_engine_args,
base_router_config=base_router_config,
model=model,
backend=backend,
system=system,
objective=objective,
constraints=constraints,
cache=cache,
)
records[index] = _evaluate_agg_state(state=state, spec=spec, cache=cache)
return [record for record in records if record is not None]
base_engine_args_json = base_engine_args.dump_json()
base_router_config_json = (
None if base_router_config is None else base_router_config.dump_json()
)
payloads = [
{
"state": state,
"workload": workload,
"base_engine_args_json": base_engine_args_json,
"base_router_config_json": base_router_config_json,
"model": model,
"backend": backend,
"system": system,
"objective": objective,
"constraints": constraints,
}
for state in uncached_states
]
spec_payload = _spec_to_payload(spec)
payloads = [{"state": state, "spec": spec_payload} for state in uncached_states]
future_records = list(
executor.map(_evaluate_agg_state_from_json_payloads, payloads)
)
future_records = list(executor.map(_evaluate_agg_state_from_payload, payloads))
for index, state, record in zip(
uncached_indices,
......
components/src/dynamo/profiler/utils/replay_optimize/example.py
View file @ 241b254d
......@@ -8,16 +8,20 @@ from collections.abc import Sequence
from dynamo.llm import KvRouterConfig, MockEngineArgs
from dynamo.profiler.utils.replay_optimize import (
SyntheticReplayWorkload,
TraceReplayWorkload,
EngineSpec,
HardwareSpec,
ReplayOptimizeSpec,
RouterSpec,
SLASpec,
WorkloadSpec,
optimize_dense_disagg_with_replay,
)
MODEL = "Qwen/Qwen3-32B"
BACKEND = "vllm"
SYSTEM = "h200_sxm"
MAX_TOTAL_GPUS = 16
OVERLAP_SCORE_WEIGHTS = (0.0, 0.5, 1.0, 2.0)
GPU_SKU = "h200_sxm"
TOTAL_GPUS = 16
OVERLAP_WEIGHTS = [0.0, 0.5, 1.0, 2.0]
RESULT_COLUMNS: Sequence[str] = (
"prefill_tp",
"decode_tp",
......@@ -37,24 +41,24 @@ def _build_workload(
*,
trace_file: str | None,
arrival_speedup_ratio: float,
) -> SyntheticReplayWorkload | TraceReplayWorkload:
) -> WorkloadSpec:
if trace_file is not None:
return TraceReplayWorkload(
trace_file=trace_file,
arrival_speedup_ratio=arrival_speedup_ratio,
return WorkloadSpec(
traceFile=trace_file,
arrivalSpeedupRatio=arrival_speedup_ratio,
)
return SyntheticReplayWorkload(
return WorkloadSpec(
isl=32768,
osl=256,
request_count=5000,
replay_concurrency=200,
shared_prefix_ratio=0.5,
num_prefix_groups=50,
requestCount=5000,
concurrency=200,
sharedPrefixRatio=0.5,
numPrefixGroups=50,
)
def _build_engine_args(*, worker_type: str) -> MockEngineArgs:
def _engine_args(worker_type: str) -> MockEngineArgs:
return MockEngineArgs(
block_size=512,
num_gpu_blocks=20000,
......@@ -69,27 +73,28 @@ def run_example(
arrival_speedup_ratio: float = 1.0,
max_parallel_evals: int = 1,
) -> None:
result = optimize_dense_disagg_with_replay(
spec = ReplayOptimizeSpec(
engine=EngineSpec(
model=MODEL,
backend=BACKEND,
system=SYSTEM,
basePrefillEngineArgs=_engine_args("prefill"),
baseDecodeEngineArgs=_engine_args("decode"),
),
hardware=HardwareSpec(gpuSku=GPU_SKU, totalGpus=TOTAL_GPUS),
workload=_build_workload(
trace_file=trace_file,
arrival_speedup_ratio=arrival_speedup_ratio,
),
base_prefill_engine_args=_build_engine_args(worker_type="prefill"),
base_decode_engine_args=_build_engine_args(worker_type="decode"),
base_router_config=KvRouterConfig(),
max_total_gpus=MAX_TOTAL_GPUS,
constraints={
"mean_ttft_ms": 50000.0,
"mean_tpot_ms": 100.0,
"mean_e2e_latency_ms": 60000.0,
},
overlap_score_weights=OVERLAP_SCORE_WEIGHTS,
max_parallel_evals=max_parallel_evals,
sla=SLASpec(ttft=50000.0, itl=100.0, e2eLatency=60000.0),
router=RouterSpec(
baseRouterConfig=KvRouterConfig(),
overlapWeights=OVERLAP_WEIGHTS,
),
maxParallelEvals=max_parallel_evals,
)
result = optimize_dense_disagg_with_replay(spec)
print("Best feasible:")
print(result.best_feasible)
print()
......
components/src/dynamo/profiler/utils/replay_optimize/logging.py
View file @ 241b254d
......@@ -9,7 +9,8 @@ from typing import Any
from dynamo.runtime.logging import configure_dynamo_logging
from .models import DenseAggReplayState, DenseReplayState, ReplayConstraints
from .models import DenseAggReplayState, DenseReplayState
from .specs import HardwareSpec, SLASpec
logger = logging.getLogger(__name__)
_LOGGING_CONFIGURED = False
......@@ -27,20 +28,31 @@ def log_state_start(state: DenseReplayState | DenseAggReplayState) -> None:
logger.info("Replay optimize evaluating %s", state.format_summary())
def _budget_summary(
state: DenseReplayState | DenseAggReplayState, hardware: HardwareSpec
) -> str:
used = state.total_gpus_used
budget = hardware.totalGpus
status = "satisfied" if used <= budget else "unsatisfied"
return f"totalGpus={used}<={budget} {status}"
def log_state_finish(
*,
state: DenseReplayState | DenseAggReplayState,
report: Mapping[str, Any],
constraints: ReplayConstraints,
sla: SLASpec,
hardware: HardwareSpec,
score: float,
feasible: bool,
violation_penalty: float,
) -> None:
logger.info(
"Replay optimize finished %s score=%.3f feasible=%s violation_penalty=%.6f %s",
"Replay optimize finished %s score=%.3f feasible=%s violation_penalty=%.6f %s %s",
state.format_summary(),
score,
feasible,
violation_penalty,
constraints.summarize(report, state.total_gpus_used),
sla.summarize(report),
_budget_summary(state, hardware),
)
components/src/dynamo/profiler/utils/replay_optimize/models.py
View file @ 241b254d
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Optimizer-internal dataclasses for replay_optimize search states.
User-facing configuration lives in `specs.py`. This module carries only the
hot-path state types (used as dict keys in the search cache, instantiated
once per visited candidate) and the result bundle returned from the
public API.
"""
from __future__ import annotations
import math
import os
from collections.abc import Iterator, Mapping
from dataclasses import dataclass, fields
from enum import Enum
from dataclasses import dataclass
from typing import Any
import pandas as pd
from .constants import SUPPORTED_CONSTRAINTS
@dataclass(frozen=True)
class ReplayConstraints:
mean_ttft_ms: float | None = None
p95_ttft_ms: float | None = None
mean_tpot_ms: float | None = None
p95_tpot_ms: float | None = None
mean_e2e_latency_ms: float | None = None
p95_e2e_latency_ms: float | None = None
max_total_gpus: int | None = None
@classmethod
def from_mapping(
cls,
mapping: Mapping[str, float] | None,
max_total_gpus: int,
) -> ReplayConstraints:
raw = dict(mapping or {})
unknown = sorted(set(raw) - SUPPORTED_CONSTRAINTS)
if unknown:
raise ValueError(
"unsupported constraints: "
+ ", ".join(unknown)
+ f"; supported constraints are {sorted(SUPPORTED_CONSTRAINTS)}"
)
raw_gpus = raw.get("max_total_gpus")
if raw_gpus is not None and int(raw_gpus) != max_total_gpus:
raise ValueError(
"constraints['max_total_gpus'] must match max_total_gpus when both are provided"
)
def _bound(key: str) -> float | None:
value = raw.get(key)
return None if value is None or value <= 0 else float(value)
return cls(
mean_ttft_ms=_bound("mean_ttft_ms"),
p95_ttft_ms=_bound("p95_ttft_ms"),
mean_tpot_ms=_bound("mean_tpot_ms"),
p95_tpot_ms=_bound("p95_tpot_ms"),
mean_e2e_latency_ms=_bound("mean_e2e_latency_ms"),
p95_e2e_latency_ms=_bound("p95_e2e_latency_ms"),
max_total_gpus=int(max_total_gpus),
)
def _active(
self, report: Mapping[str, Any], total_gpus_used: int
) -> Iterator[tuple[str, float | None, float]]:
for field in fields(self):
if field.name == "max_total_gpus":
continue
bound = getattr(self, field.name)
if bound is None:
continue
value = report.get(field.name)
yield field.name, None if value is None else float(value), bound
if self.max_total_gpus is not None:
yield (
"max_total_gpus",
float(total_gpus_used),
float(self.max_total_gpus),
)
def violation_penalty(
self, report: Mapping[str, Any], total_gpus_used: int
) -> float:
penalty = 0.0
for _, metric, bound in self._active(report, total_gpus_used):
if metric is None:
penalty += math.inf
continue
penalty += max(metric / bound - 1.0, 0.0)
return penalty
def summarize(self, report: Mapping[str, Any], total_gpus_used: int) -> str:
statuses: list[str] = []
for name, metric, bound in self._active(report, total_gpus_used):
if metric is None:
statuses.append(f"{name}=missing<={bound:g} unsatisfied")
continue
state = "satisfied" if metric <= bound else "unsatisfied"
statuses.append(f"{name}={metric:.3f}<={bound:g} {state}")
return "constraints=" + ", ".join(statuses) if statuses else "constraints=none"
def aic_task_kwargs(self) -> dict[str, float | None]:
return {
"ttft": self.mean_ttft_ms,
"tpot": self.mean_tpot_ms,
"request_latency": self.mean_e2e_latency_ms,
}
class ReplayObjective(str, Enum):
THROUGHPUT = "throughput"
MEAN_TTFT = "mean_ttft"
MEAN_E2E_LATENCY = "mean_e2e_latency"
def score(self, report: Mapping[str, Any]) -> float:
if self is ReplayObjective.THROUGHPUT:
return float(report["output_throughput_tok_s"])
if self is ReplayObjective.MEAN_TTFT:
return -float(report["mean_ttft_ms"])
return -float(report["mean_e2e_latency_ms"])
@dataclass(frozen=True)
class SyntheticReplayWorkload:
isl: int
osl: int
request_count: int
replay_concurrency: int
arrival_interval_ms: float = 0.0
turns_per_session: int = 1
shared_prefix_ratio: float = 0.0
num_prefix_groups: int = 0
inter_turn_delay_ms: float = 0.0
@dataclass(frozen=True)
class TraceReplayWorkload:
trace_file: str | os.PathLike[str]
arrival_speedup_ratio: float = 1.0
@dataclass(frozen=True)
class DenseReplayState:
......
components/src/dynamo/profiler/utils/replay_optimize/search.py
View file @ 241b254d
......@@ -30,66 +30,21 @@ from collections.abc import Mapping, Sequence
from concurrent.futures import ProcessPoolExecutor
from typing import Literal
from dynamo.llm import KvRouterConfig, MockEngineArgs
from . import aic, evaluate
from .constants import (
AIC_BACKEND_VERSIONS,
DEFAULT_MAX_PARALLEL_EVALS,
DEFAULT_OVERLAP_SCORE_WEIGHTS,
DEFAULT_SEARCH_ROUNDS,
)
from .models import (
DenseAggReplayState,
DenseReplayOptimizationResult,
DenseReplayState,
ReplayConstraints,
ReplayObjective,
SyntheticReplayWorkload,
TraceReplayWorkload,
)
from .constants import DEFAULT_SEARCH_ROUNDS
from .models import DenseAggReplayState, DenseReplayOptimizationResult, DenseReplayState
from .scoring import _finalize_result, _pick_best_record
def _validate_backend(backend: str) -> str:
if backend not in AIC_BACKEND_VERSIONS:
raise ValueError(
f"backend must be one of {sorted(AIC_BACKEND_VERSIONS)}, got {backend!r}"
)
return backend
def _normalize_overlap_score_weights(
overlap_score_weights: Sequence[float] | None,
) -> tuple[float, ...]:
if overlap_score_weights is None:
return DEFAULT_OVERLAP_SCORE_WEIGHTS
weights = tuple(float(weight) for weight in overlap_score_weights)
if not weights:
raise ValueError("overlap_score_weights must not be empty")
return weights
def _normalize_router_mode(
router_mode: str,
) -> Literal["kv_router", "round_robin", "both"]:
if router_mode not in {"kv_router", "round_robin", "both"}:
raise ValueError(
"router_mode must be one of ['kv_router', 'round_robin', 'both'], "
f"got {router_mode!r}"
)
return router_mode
from .specs import ReplayOptimizeSpec, RouterMode
def _router_states(
*,
router_mode: Literal["kv_router", "round_robin", "both"],
router_mode: RouterMode,
overlap_score_weights: Sequence[float],
) -> list[tuple[str, float]]:
if router_mode == "round_robin":
if router_mode is RouterMode.ROUND_ROBIN:
return [("round_robin", 0.0)]
if router_mode == "kv_router":
if router_mode is RouterMode.KV_ROUTER:
return [("kv_router", float(weight)) for weight in overlap_score_weights]
return [("round_robin", 0.0)] + [
("kv_router", float(weight)) for weight in overlap_score_weights
......@@ -267,27 +222,15 @@ def _record_to_agg_state(
def optimize_dense_disagg_with_replay(
*,
model: str,
backend: Literal["vllm", "sglang"],
system: str,
workload: SyntheticReplayWorkload | TraceReplayWorkload,
base_prefill_engine_args: MockEngineArgs,
base_decode_engine_args: MockEngineArgs,
base_router_config: KvRouterConfig | None = None,
max_total_gpus: int,
constraints: Mapping[str, float] | None = None,
objective: Literal["throughput", "mean_e2e_latency", "mean_ttft"] = "throughput",
router_mode: Literal["kv_router", "round_robin", "both"] = "kv_router",
overlap_score_weights: Sequence[float] | None = None,
max_parallel_evals: int = DEFAULT_MAX_PARALLEL_EVALS,
spec: ReplayOptimizeSpec,
) -> DenseReplayOptimizationResult:
"""Run a heuristic block search over dense disaggregated offline replay configs.
This routine assumes we want to use as much of `max_total_gpus` as possible,
then ranks visited states by the selected `objective` subject to replay
constraints. Supported objectives: `"throughput"` (default, maximize
`output_throughput_tok_s`), `"mean_e2e_latency"` and `"mean_ttft"` (minimize
Assumes the optimizer should consume as much of `spec.hardware.totalGpus` as
possible, then ranks visited states by `spec.objective` subject to the SLA
bounds in `spec.sla`. Supported objectives: `ReplayObjective.THROUGHPUT`
(default, maximize `output_throughput_tok_s`),
`ReplayObjective.MEAN_E2E_LATENCY` and `ReplayObjective.MEAN_TTFT` (minimize
the corresponding report metric). The descended dimensions are:
1. `(prefill_tp, decode_tp)` at equal worker counts that fit the budget.
2. `(prefill_workers, decode_workers)` on the budget edge for the incumbent TP
......@@ -296,17 +239,24 @@ def optimize_dense_disagg_with_replay(
Returned "best" records are best among visited states, not a global optimum.
"""
backend = _validate_backend(backend)
router_mode = _normalize_router_mode(router_mode)
typed_objective = ReplayObjective(objective)
if max_total_gpus < 2:
raise ValueError("max_total_gpus must be at least 2 for disaggregated replay")
typed_constraints = ReplayConstraints.from_mapping(constraints, max_total_gpus)
overlap_weights = _normalize_overlap_score_weights(overlap_score_weights)
if router_mode == "round_robin":
overlap_weights = (0.0,)
max_parallel_evals = max(1, int(max_parallel_evals))
# Guardrail #8: disagg path requires both prefill and decode engine args.
if (
spec.engine.basePrefillEngineArgs is None
or spec.engine.baseDecodeEngineArgs is None
):
raise ValueError(
"optimize_dense_disagg_with_replay requires both "
"EngineSpec.basePrefillEngineArgs and EngineSpec.baseDecodeEngineArgs"
)
if spec.hardware.totalGpus < 2:
raise ValueError(
"hardware.totalGpus must be at least 2 for disaggregated replay"
)
backend = spec.engine.backend.value
system = str(spec.hardware.gpuSku)
overlap_weights = spec.router.effectiveOverlapWeights
max_parallel_evals = max(1, int(spec.maxParallelEvals))
prefill_tps, decode_tps = aic._enumerate_dense_tp_candidates(backend, system)
if not prefill_tps or not decode_tps:
raise ValueError(
......@@ -318,7 +268,7 @@ def optimize_dense_disagg_with_replay(
prefill_tps=prefill_tps,
decode_tps=decode_tps,
overlap_score_weight=overlap_weights[0],
max_total_gpus=max_total_gpus,
max_total_gpus=spec.hardware.totalGpus,
)
executor = (
......@@ -335,19 +285,11 @@ def optimize_dense_disagg_with_replay(
decode_tps=decode_tps,
router_mode=incumbent.router_mode,
overlap_score_weight=incumbent.overlap_score_weight,
max_total_gpus=max_total_gpus,
max_total_gpus=spec.hardware.totalGpus,
)
tp_records = evaluate._evaluate_states(
states=tp_states,
workload=workload,
base_prefill_engine_args=base_prefill_engine_args,
base_decode_engine_args=base_decode_engine_args,
base_router_config=base_router_config,
model=model,
backend=backend,
system=system,
objective=typed_objective,
constraints=typed_constraints,
spec=spec,
cache=cache,
max_parallel_evals=max_parallel_evals,
executor=executor,
......@@ -359,19 +301,11 @@ def optimize_dense_disagg_with_replay(
decode_tp=incumbent.decode_tp,
router_mode=incumbent.router_mode,
overlap_score_weight=incumbent.overlap_score_weight,
max_total_gpus=max_total_gpus,
max_total_gpus=spec.hardware.totalGpus,
)
worker_records = evaluate._evaluate_states(
states=worker_states,
workload=workload,
base_prefill_engine_args=base_prefill_engine_args,
base_decode_engine_args=base_decode_engine_args,
base_router_config=base_router_config,
model=model,
backend=backend,
system=system,
objective=typed_objective,
constraints=typed_constraints,
spec=spec,
cache=cache,
max_parallel_evals=max_parallel_evals,
executor=executor,
......@@ -389,19 +323,11 @@ def optimize_dense_disagg_with_replay(
router_mode=mode,
)
for mode, weight in _router_states(
router_mode=router_mode,
router_mode=spec.router.mode,
overlap_score_weights=overlap_weights,
)
],
workload=workload,
base_prefill_engine_args=base_prefill_engine_args,
base_decode_engine_args=base_decode_engine_args,
base_router_config=base_router_config,
model=model,
backend=backend,
system=system,
objective=typed_objective,
constraints=typed_constraints,
spec=spec,
cache=cache,
max_parallel_evals=max_parallel_evals,
executor=executor,
......@@ -418,44 +344,31 @@ def optimize_dense_disagg_with_replay(
def optimize_dense_agg_with_replay(
*,
model: str,
backend: Literal["vllm", "sglang"],
system: str,
workload: SyntheticReplayWorkload | TraceReplayWorkload,
base_engine_args: MockEngineArgs,
base_router_config: KvRouterConfig | None = None,
max_total_gpus: int,
constraints: Mapping[str, float] | None = None,
objective: Literal["throughput", "mean_e2e_latency", "mean_ttft"] = "throughput",
router_mode: Literal["kv_router", "round_robin", "both"] = "kv_router",
overlap_score_weights: Sequence[float] | None = None,
max_parallel_evals: int = DEFAULT_MAX_PARALLEL_EVALS,
spec: ReplayOptimizeSpec,
) -> DenseReplayOptimizationResult:
"""Run a heuristic block search over dense aggregated offline replay configs.
This routine assumes we want to use as much of `max_total_gpus` as possible,
then ranks visited states by the selected `objective` subject to replay
constraints. Supported objectives: `"throughput"` (default, maximize
`output_throughput_tok_s`), `"mean_e2e_latency"` and `"mean_ttft"` (minimize
the corresponding report metric). The descended dimensions are:
Assumes the optimizer should consume as much of `spec.hardware.totalGpus` as
possible, then ranks visited states by `spec.objective` subject to the SLA
bounds in `spec.sla`. The descended dimensions are:
1. `tp` at the maximum worker count that fits the budget.
2. `workers` for the incumbent `tp`.
3. `(router_mode, overlap_score_weight)`.
Returned "best" records are best among visited states, not a global optimum.
"""
backend = _validate_backend(backend)
router_mode = _normalize_router_mode(router_mode)
typed_objective = ReplayObjective(objective)
if max_total_gpus < 1:
raise ValueError("max_total_gpus must be at least 1 for aggregated replay")
typed_constraints = ReplayConstraints.from_mapping(constraints, max_total_gpus)
overlap_weights = _normalize_overlap_score_weights(overlap_score_weights)
if router_mode == "round_robin":
overlap_weights = (0.0,)
max_parallel_evals = max(1, int(max_parallel_evals))
# Guardrail #8: agg path requires the single baseEngineArgs.
if spec.engine.baseEngineArgs is None:
raise ValueError(
"optimize_dense_agg_with_replay requires EngineSpec.baseEngineArgs"
)
if spec.hardware.totalGpus < 1:
raise ValueError("hardware.totalGpus must be at least 1 for aggregated replay")
backend = spec.engine.backend.value
system = str(spec.hardware.gpuSku)
overlap_weights = spec.router.effectiveOverlapWeights
max_parallel_evals = max(1, int(spec.maxParallelEvals))
tps, _ = aic._enumerate_dense_tp_candidates(backend, system)
if not tps:
raise ValueError(
......@@ -463,7 +376,9 @@ def optimize_dense_agg_with_replay(
)
cache: dict[DenseAggReplayState, dict[str, float | int | bool | str]] = {}
incumbent = _select_initial_agg_state(tps=tps, max_total_gpus=max_total_gpus)
incumbent = _select_initial_agg_state(
tps=tps, max_total_gpus=spec.hardware.totalGpus
)
executor = (
ProcessPoolExecutor(max_workers=max_parallel_evals)
......@@ -478,18 +393,11 @@ def optimize_dense_agg_with_replay(
tps=tps,
router_mode=incumbent.router_mode,
overlap_score_weight=incumbent.overlap_score_weight,
max_total_gpus=max_total_gpus,
max_total_gpus=spec.hardware.totalGpus,
)
tp_records = evaluate._evaluate_agg_states(
states=tp_states,
workload=workload,
base_engine_args=base_engine_args,
base_router_config=base_router_config,
model=model,
backend=backend,
system=system,
objective=typed_objective,
constraints=typed_constraints,
spec=spec,
cache=cache,
max_parallel_evals=max_parallel_evals,
executor=executor,
......@@ -500,18 +408,11 @@ def optimize_dense_agg_with_replay(
tp=incumbent.tp,
router_mode=incumbent.router_mode,
overlap_score_weight=incumbent.overlap_score_weight,
max_total_gpus=max_total_gpus,
max_total_gpus=spec.hardware.totalGpus,
)
worker_records = evaluate._evaluate_agg_states(
states=worker_states,
workload=workload,
base_engine_args=base_engine_args,
base_router_config=base_router_config,
model=model,
backend=backend,
system=system,
objective=typed_objective,
constraints=typed_constraints,
spec=spec,
cache=cache,
max_parallel_evals=max_parallel_evals,
executor=executor,
......@@ -527,7 +428,7 @@ def optimize_dense_agg_with_replay(
overlap_score_weight=weight,
)
for mode, weight in _router_states(
router_mode=router_mode,
router_mode=spec.router.mode,
overlap_score_weights=overlap_weights,
)
if _supports_agg_router_mode(
......@@ -535,14 +436,7 @@ def optimize_dense_agg_with_replay(
router_mode=mode,
)
],
workload=workload,
base_engine_args=base_engine_args,
base_router_config=base_router_config,
model=model,
backend=backend,
system=system,
objective=typed_objective,
constraints=typed_constraints,
spec=spec,
cache=cache,
max_parallel_evals=max_parallel_evals,
executor=executor,
......
components/src/dynamo/profiler/utils/replay_optimize/specs.py 0 → 100644
View file @ 241b254d
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
"""Pydantic specs for replay_optimize, shaped to mirror DGDR v1beta1.
Field names follow DGDR's lowerCamelCase convention so that the eventual
upstream-back into `deploy/operator/api/v1beta1/...` is a clean merge.
Method names follow Python / Pydantic convention (snake_case), matching how
`dgdr_v1beta1_types.py` itself defines its one internal method
(`_validate_sla_options`).
DGDR shapes we clone / extend:
- `EngineSpec` — local extension (DGDR has `model`/`backend` flat on the
outer; we need engine-args carriers)
- `HardwareSpec` — subset clone of DGDR.HardwareSpec (gpuSku + totalGpus only)
- `WorkloadSpec` — DGDR.WorkloadSpec + replay extensions, unified synthetic/
trace with a `traceFile` discriminator
- `SLASpec` — DGDR.SLASpec field names + p95 variants; we explicitly
do NOT clone DGDR's "ttft+itl XOR e2eLatency" validator
because our optimizer treats them as independent additive
penalties (all three may be set)
- `RouterSpec` — our sweep-oriented router config; analogous location to
DGDR.KVRouterSpec but different semantics (runtime flag
there vs. dev-time sweep here)
- `ReplayOptimizeSpec` — top-level bundle, analog to
DGDR.DynamoGraphDeploymentRequestSpec
"""
from __future__ import annotations
import math
from collections.abc import Iterator, Mapping
from enum import Enum
from typing import Any
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from dynamo.llm import KvRouterConfig, MockEngineArgs
from dynamo.profiler.utils.dgdr_v1beta1_types import BackendType, GPUSKUType
from .constants import (
AIC_BACKEND_VERSIONS,
DEFAULT_MAX_PARALLEL_EVALS,
DEFAULT_OVERLAP_SCORE_WEIGHTS,
)
class RouterMode(str, Enum):
"""Router mode for the replay search.
`BOTH` triggers a combined sweep across `KV_ROUTER` and `ROUND_ROBIN`;
`round_robin` collapses `overlapWeights` to `(0.0,)` (guardrail #5).
Subclasses `str` for Pydantic coercion and wire-compatibility with the
existing `router_mode` field on `DenseReplayState` / `DenseAggReplayState`.
"""
KV_ROUTER = "kv_router"
ROUND_ROBIN = "round_robin"
BOTH = "both"
class ReplayObjective(str, Enum):
"""Optimization objective driving state ranking in the search."""
THROUGHPUT = "throughput"
MEAN_TTFT = "mean_ttft"
MEAN_E2E_LATENCY = "mean_e2e_latency"
def score(self, report: Mapping[str, Any]) -> float:
if self is ReplayObjective.THROUGHPUT:
return float(report["output_throughput_tok_s"])
if self is ReplayObjective.MEAN_TTFT:
return -float(report["mean_ttft_ms"])
return -float(report["mean_e2e_latency_ms"])
class EngineSpec(BaseModel):
"""Model + backend + engine-arg carriers.
DGDR has `model: str` and `backend: BackendType` flat on the outer spec and
no engine-args equivalent, so this spec is a replay-local extension.
Carries engine args for both agg and disagg paths; the relevant
`optimize_dense_*` entry asserts the right fields are populated
(guardrail #8).
`backend` has no default — pre-Phase-2 `optimize_dense_*` required it; keep
the explicit contract so a forgotten backend fails at spec construction
instead of silently falling through to a vLLM run.
"""
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
model: str
backend: BackendType
baseEngineArgs: MockEngineArgs | None = None
basePrefillEngineArgs: MockEngineArgs | None = None
baseDecodeEngineArgs: MockEngineArgs | None = None
@field_validator("backend", mode="after")
@classmethod
def _validate_backend_supported_by_aic(cls, backend: BackendType) -> BackendType:
# Guardrail #6: DGDR's BackendType allows Auto / Trtllm which AIC
# doesn't support in replay_optimize; reject at spec-construction time
# instead of crashing mid-search.
if backend.value not in AIC_BACKEND_VERSIONS:
raise ValueError(
f"backend must be one of {sorted(AIC_BACKEND_VERSIONS)}, "
f"got {backend.value!r}"
)
return backend
class HardwareSpec(BaseModel):
"""GPU budget + SKU. Subset clone of DGDR.HardwareSpec."""
model_config = ConfigDict(extra="forbid")
# Guardrail #7: keep the union so exotic AIC systems outside the current
# GPUSKUType enumeration still work. Pydantic coerces matching strings to
# the enum; non-matching strings stay as str and reach AIC untouched.
gpuSku: GPUSKUType | str
totalGpus: int = Field(gt=0)
_SYNTHETIC_ONLY_FIELDS: tuple[str, ...] = (
"isl",
"osl",
"concurrency",
"requestRate",
"requestCount",
)
class WorkloadSpec(BaseModel):
"""Workload, unified across synthetic and trace replay.
Extends DGDR.WorkloadSpec with:
- replay-specific synthetic knobs (request_count, shared-prefix / multi-
turn controls, arrival interval)
- trace-source fields (`traceFile`, `arrivalSpeedupRatio`)
`traceFile` acts as a discriminator:
- when set, the workload is trace-based and the synthetic-only fields
(`isl`, `osl`, `concurrency`, `requestRate`, `requestCount`) must not
be populated — the validator rejects mixed mode to avoid silent data loss
- when unset, the synthetic fields `isl`, `osl`, `concurrency`, and
`requestCount` are all required
"""
model_config = ConfigDict(extra="forbid")
# DGDR base fields
isl: int | None = None
osl: int | None = None
concurrency: float | None = None
requestRate: float | None = None
# Replay synthetic extensions
requestCount: int | None = None
sharedPrefixRatio: float = 0.0
numPrefixGroups: int = 0
turnsPerSession: int = 1
interTurnDelayMs: float = 0.0
arrivalIntervalMs: float = 0.0
# Replay trace-source extensions (mutually exclusive with synthetic fields)
traceFile: str | None = None
arrivalSpeedupRatio: float = 1.0
@model_validator(mode="after")
def _validate_source(self) -> "WorkloadSpec":
if self.traceFile is not None:
mixed = [
name
for name in _SYNTHETIC_ONLY_FIELDS
if getattr(self, name) is not None
]
if mixed:
raise ValueError(
"trace workload (traceFile set) must not also set synthetic "
f"fields: {mixed}"
)
return self
missing = [
name
for name in ("isl", "osl", "concurrency", "requestCount")
if getattr(self, name) is None
]
if missing:
raise ValueError(
"synthetic workload requires "
+ ", ".join(missing)
+ "; or set traceFile for trace replay"
)
return self
@property
def isTraceBased(self) -> bool:
return self.traceFile is not None
# SLASpec field names → Rust replay-report keys. The Rust runner emits
# snake_case + `_ms` suffix (`mean_ttft_ms`, `mean_tpot_ms`, ...) while DGDR's
# convention is camelCase + no unit suffix. This table bridges the two; the
# Rust side renaming is out of scope here.
_SLA_REPORT_KEYS: dict[str, str] = {
"ttft": "mean_ttft_ms",
"itl": "mean_tpot_ms",
"e2eLatency": "mean_e2e_latency_ms",
"p95Ttft": "p95_ttft_ms",
"p95Itl": "p95_tpot_ms",
"p95E2eLatency": "p95_e2e_latency_ms",
}
class SLASpec(BaseModel):
"""Latency SLA bounds.
Guardrail #1 + #2: defaults are `None` (unconstrained); we do NOT clone
DGDR's `_validate_sla_options` — our optimizer treats ttft / itl /
e2eLatency as independent additive penalties and the example uses all
three simultaneously.
`itl` is the DGDR-native name for what the Rust runner reports as
`mean_tpot_ms` (inter-token latency == time-per-output-token).
"""
model_config = ConfigDict(extra="forbid")
ttft: float | None = None
itl: float | None = None
e2eLatency: float | None = None
p95Ttft: float | None = None
p95Itl: float | None = None
p95E2eLatency: float | None = None
def _active(
self, report: Mapping[str, Any]
) -> Iterator[tuple[str, float | None, float]]:
"""Yield `(field_name, metric_from_report_or_None, bound)` for each set bound."""
for field_name, report_key in _SLA_REPORT_KEYS.items():
bound = getattr(self, field_name)
if bound is None or bound <= 0:
continue
value = report.get(report_key)
yield field_name, None if value is None else float(value), float(bound)
def violation_penalty(self, report: Mapping[str, Any]) -> float:
"""Sum of positive (metric/bound - 1) across active SLA bounds.
Missing report key contributes `inf` (fails the feasibility gate)
rather than silently scoring as zero.
"""
penalty = 0.0
for _, metric, bound in self._active(report):
if metric is None:
penalty += math.inf
continue
penalty += max(metric / bound - 1.0, 0.0)
return penalty
def summarize(self, report: Mapping[str, Any]) -> str:
statuses: list[str] = []
for field_name, metric, bound in self._active(report):
if metric is None:
statuses.append(f"{field_name}=missing<={bound:g} unsatisfied")
continue
state = "satisfied" if metric <= bound else "unsatisfied"
statuses.append(f"{field_name}={metric:.3f}<={bound:g} {state}")
return "sla=" + ", ".join(statuses) if statuses else "sla=none"
def aic_task_kwargs(self) -> dict[str, float | None]:
"""Translate to `aiconfigurator.sdk.task.TaskConfig` kwargs.
AIC's external API still uses `tpot` and `request_latency`; we keep
those wire names untouched.
"""
return {
"ttft": self.ttft,
"tpot": self.itl,
"request_latency": self.e2eLatency,
}
class RouterSpec(BaseModel):
"""Router config for the search.
Analogous location to DGDR.KVRouterSpec but semantically different: DGDR
has a single runtime on/off flag, we have a dev-time sweep over overlap
score weights plus a mode selector.
"""
model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
mode: RouterMode = RouterMode.KV_ROUTER
# None → fallback to DEFAULT_OVERLAP_SCORE_WEIGHTS (guardrail #3). Empty
# list rejected (guardrail #4). Round-robin auto-collapse happens in
# `effectiveOverlapWeights` (guardrail #5).
overlapWeights: list[float] | None = None
baseRouterConfig: KvRouterConfig | None = None
@field_validator("overlapWeights", mode="after")
@classmethod
def _reject_empty_weights(cls, weights: list[float] | None) -> list[float] | None:
if weights is not None and len(weights) == 0:
raise ValueError("overlapWeights must not be empty")
return weights
@property
def effectiveOverlapWeights(self) -> tuple[float, ...]:
"""Resolve to the concrete weight sweep used by the search."""
if self.mode is RouterMode.ROUND_ROBIN:
return (0.0,)
if self.overlapWeights is None:
return DEFAULT_OVERLAP_SCORE_WEIGHTS
return tuple(float(w) for w in self.overlapWeights)
class ReplayOptimizeSpec(BaseModel):
"""Top-level spec; analog to DGDR.DynamoGraphDeploymentRequestSpec."""
model_config = ConfigDict(extra="forbid")
engine: EngineSpec
hardware: HardwareSpec
workload: WorkloadSpec
sla: SLASpec = Field(default_factory=SLASpec)
router: RouterSpec = Field(default_factory=RouterSpec)
objective: ReplayObjective = ReplayObjective.THROUGHPUT
maxParallelEvals: int = Field(default=DEFAULT_MAX_PARALLEL_EVALS, gt=0)
docs/benchmarks/mocker-trace-replay.md
View file @ 241b254d
......@@ -19,6 +19,52 @@ Use this when you want to:
- compare timing and cache behavior across mocker configurations
- validate replay logic in CI without bringing up a distributed stack
## Harness Overview
The replay harness wires a load driver (trace file or synthetic workload generator) into one or more mocker engine simulations and tees request/token timing into a trace collector.
```mermaid
flowchart LR
LD[Load Driver] --> H[Replay Harness]
H --> SES[Single Engine Simulation]
H --> MES[Multi Engine Simulation]
SES --> H
MES --> H
H --> TC[Trace Collector]
```
The load driver is either a Mooncake-style JSONL trace (timestamps, ISL/OSL, `hash_ids`) or a synthetic generator parameterized by `isl`/`osl`/`concurrency`. Single-engine simulation (`SES`) is the fast path for `num_workers == 1` with the vLLM engine; multi-engine simulation (`MES`) covers aggregated multi-worker replay, disaggregated prefill/decode replay, and KV-router replay. The trace collector produces the AIPerf-style summary table, the JSON report, and the per-request timing fields consumed by downstream analysis.
Each simulation composes a different set of components. SES drives the engine core directly (scheduler + forward-pass modeling). MES composes multiple engine cores with KV transfer/offloading, KV routing, and planner simulation layered on top:
```mermaid
flowchart TD
subgraph SEC[Single Engine Core]
subgraph SCH[Scheduler Modeling]
F[Fwd Pass Modeling]
end
end
KV[KV Transfer + Offloading Simulation]
KR[KV Router Simulation]
P[Planner Simulation]
SES[Single Engine Simulation]
MES[Multi Engine Simulation]
SES --> SEC
MES --> SEC
MES --> KV
MES --> KR
MES --> P
```
See [`lib/mocker/src/replay/offline/README.md`](../../lib/mocker/src/replay/offline/README.md) for offline-harness internals (logical clock, event queue, worker model) and [`docs/mocker/mocker.md`](../mocker/mocker.md) for engine-core details (scheduler, KV block manager).
## Quick Start
Run offline replay through the dedicated replay CLI:
......
docs/mocker/mocker.md
View file @ 241b254d
......@@ -90,7 +90,7 @@ python -m dynamo.mocker \
| `--num-workers` | 1 | Workers per process |
| `--reasoning` | None | JSON config for emitting reasoning token spans, with `start_thinking_token_id`, `end_thinking_token_id`, and `thinking_ratio` |
| `--engine-type` | `vllm` | Engine simulation type: `vllm` or `sglang` |
| `--sglang-schedule-policy` | `fifo` / `fcfs` | SGLang scheduling policy override |
| `--sglang-schedule-policy` | `fifo` / `fcfs` | SGLang scheduling policy: `fifo`/`fcfs` (default) or `lpm` (longest prefix match) |
| `--sglang-page-size` | 1 | SGLang radix-cache page size in tokens. Also becomes the effective block size when `--engine-type sglang` and `--block-size` is omitted |
| `--sglang-max-prefill-tokens` | 16384 | SGLang max prefill-token budget per batch |
| `--sglang-chunked-prefill-size` | 8192 | SGLang chunked-prefill chunk size |
......@@ -366,7 +366,7 @@ kubectl apply -f examples/backends/mocker/deploy/disagg.yaml
## Architecture
The mocker is organized into several cooperating components that mirror the internal architecture of production LLM inference engines.
The mocker is organized into several cooperating components that mirror the internal architecture of production LLM inference engines. The scheduler (vLLM-style and SGLang-style variants) and KV block manager live inside the engine core. Multi-engine behavior — KV transfer/offloading simulation, KV router simulation, planner simulation — is added by the replay harness on top of multiple engine cores; see [Mocker Trace Replay](../benchmarks/mocker-trace-replay.md) for the component-level diagram and for offline replay internals under [`lib/mocker/src/replay/offline/`](../../lib/mocker/src/replay/offline/README.md).
### Scheduler
......@@ -388,43 +388,46 @@ When resources become constrained, the mocker simulates the engine's real recove
### KV Block Manager
The block manager tracks KV cache blocks using reference counting and an LRU eviction policy. Blocks exist in one of two pools:
The mocker's KV block manager is now built on [`kvbm-logical::BlockManager<G1>`](../../lib/kvbm-logical/), the same logical block manager the real Dynamo runtime uses. The mocker wraps it in [`lib/mocker/src/kv_manager/kvbm_backend.rs`](../../lib/mocker/src/kv_manager/kvbm_backend.rs) and translates its own `MoveBlock` protocol onto kvbm-logical's RAII lifecycle (`allocate → stage → register → drop`).
- **Active Pool** - Blocks currently in use by one or more sequences, tracked with reference counts
- **Inactive Pool** - Blocks no longer actively referenced but kept for potential reuse (prefix caching)
Blocks still conceptually live in one of two pools:
When a sequence needs blocks, the manager first checks if they already exist (cache hit). If not, it allocates new blocks, potentially evicting the least-recently-used inactive blocks to make room. When a sequence completes or is preempted, its blocks are either moved to the inactive pool (for potential reuse) or freed entirely.
- **Active** — blocks currently held by at least one sequence. Partial (still-filling) blocks are held as `MutableBlock<G1>`; full blocks are held as `ImmutableBlock<G1>` clones (the clone vec length is the mocker's refcount, one per `Use`).
- **Inactive** — blocks no longer referenced by any sequence but kept for prefix-cache reuse. Handled entirely by kvbm-logical's inactive pool; the mocker never tracks them manually.
The following diagram illustrates the block lifecycle, based on vLLM's block manager design:
The lifecycle is RAII: dropping the last `ImmutableBlock` clone transitions the block from active to inactive (kvbm-logical's `reset` pool), with no explicit `deref`/`evict` bookkeeping on the mocker side. When a sequence completes or is preempted, the mocker simply drops its handles; kvbm-logical recovers the capacity.
```mermaid
stateDiagram-v2
[*] --> Active : alloc
Active --> Inactive : deref
Inactive --> Active : cache hit (reuse)
Inactive --> Freed : evict
Active --> Freed : destroy (preemption)
[*] --> Active : allocate + stage + register
Active --> Inactive : last handle dropped (RAII)
Inactive --> Active : match_blocks(PLH) reuse
Inactive --> Freed : evicted by backend
Active --> Freed : explicit Removed (Destroy)
Freed --> [*]
state Active {
[*] --> Tracked : ref_count tracked
}
state Inactive {
[*] --> Ordered : LRU order
[*] --> Partial : MutableBlock<G1>
Partial --> Full : promote (PLH / SequenceHash)
[*] --> Full : ImmutableBlock<G1> clones
}
```
### Evictor
Three `Use` outcomes are tracked for KV-event emission: `ActiveHit` (bump refcount on an already-pinned block), `InactiveHit` (reactivate via `match_blocks(plh)`), and `NewStore` (fresh allocation). Only `NewStore` emits a `Stored` KV event — the router radix tree already knows about the other two and only forgets on explicit `Removed`.
### Eviction Backends
The LRU evictor maintains blocks ordered by a monotonic counter, enabling O(log n) eviction of the lowest-priority block. Each `insert` assigns the next counter value, so blocks inserted later have higher counters and survive longer.
The kvbm-logical inactive pool selects eviction victims via one of three backends, exposed as `MockerEvictionBackend` in [`lib/mocker/src/common/protocols.rs`](../../lib/mocker/src/common/protocols.rs):
This produces a **depth-aware eviction policy**: when a sequence completes, `free_signal` releases its blocks in reverse order (tail first). Deeper suffix blocks therefore receive lower counters and are evicted before shallower prefix blocks. This keeps shared prefixes cached longer, improving cache hit rates across requests with common prefixes.
- **`Lineage`** (default) — parent-chain aware: evicts leaf blocks first, preserving shared prefix chains. Subsumes the preemption-priority behavior the old hand-rolled `LRUEvictor::push_front` used to provide.
- **`Lru`** — plain recency-based LRU.
- **`MultiLru`** — 4-tier frequency-aware LRU built on a TinyLFU tracker.
The evictor also supports front-insertion (negative counters) for marking blocks for immediate eviction, though this is not currently used in the scheduler.
All three give the same "suffix blocks evicted before shared prefixes" outcome that the old evictor was designed to produce; `Lineage` does it structurally (via the block parent chain) rather than via monotonic counters.
### Sequence Tracking
Each active request is tracked as a sequence, managing its token blocks and generation state. As tokens are generated, the sequence tracks which blocks are partial (still being filled) versus full (complete and hashable for prefix caching). When a partial block fills up, it gets "promoted" to a full block with a content-based hash, enabling future cache hits from requests with matching prefixes.
Each active request is tracked as a sequence, managing its token blocks and generation state. As tokens are generated, the sequence tracks which blocks are partial (`MutableBlock<G1>`, still being filled) versus full (`ImmutableBlock<G1>`, complete and hashable for prefix caching). When a partial block fills up, it gets "promoted" to a full block with a content-based `SequenceHash` (or collapses onto an existing registered handle if the PLH is already present), enabling future cache hits from requests with matching prefixes.
### Performance Model
......
lib/mocker/src/replay/offline/README.md
View file @ 241b254d
......@@ -4,6 +4,8 @@ This directory contains the in-process offline replay harness used by `dynamo_mo
The goal is to simulate trace execution without spinning up async runtimes, network planes, or real worker tasks. Instead, the harness advances a logical clock, steps mock engine cores directly, and records request/token timing into `TraceCollector` in `lib/mocker/src/replay/collector.rs`.
For the harness-level picture (load driver → harness → SES/MES → trace collector) and operator-facing CLI docs, see [`docs/benchmarks/mocker-trace-replay.md`](../../../../../docs/benchmarks/mocker-trace-replay.md). This README dives into the offline-specific internals: logical clock, event queue, per-worker state machine.
## Where It Sits
The public replay entrypoints live one level up in `lib/mocker/src/replay/entrypoints.rs`. They:
......@@ -34,9 +36,20 @@ Offline replay starts in `lib/mocker/src/replay/offline/mod.rs`.
- `lib/mocker/src/replay/offline/state.rs`
Per-worker wrapper around `EngineCore`, including optional KV event capture.
- `lib/mocker/src/replay/offline/events.rs`
Priority-queue event type used by the multi-worker harness.
`SimulationEvent` + `SimulationEventKind` priority-queue types used by the multi-worker harness.
- `lib/mocker/src/replay/offline/core.rs`
Small `ReplayWorkerCore` wrapper used by the single-worker path.
- `lib/mocker/src/replay/offline/runtime_utils.rs`
Shared helpers used by `agg.rs` and `disagg.rs`: `WorkerCompletionPayload`, event scheduling, `next_timestamp`.
- `lib/mocker/src/replay/offline/progress.rs`
`ReplayProgress`, the indicatif-based progress bar used by the harnesses.
- `lib/mocker/src/replay/offline/components/`
Shared abstractions split out from the runtimes:
- `router.rs` — `OfflineReplayRouter` (synchronous in-process router, KV + round-robin modes) and `OfflineRouterSnapshot`.
- `engine.rs` — `EngineComponent`, `EngineEffects`, `EnginePassMode` wrappers around `EngineCore`.
- `admission.rs` — admission queue and trace/workload request gating.
- `types.rs` — `WorkerAdmission`, `RouterEffects`, `ScheduledWorkerCompletion`, `TrafficAccumulator`, `TrafficStats`, `ReplayMode`.
- `mod.rs` — re-exports.
## Single-Worker Fast Path
......@@ -117,21 +130,25 @@ So offline replay is not a toy simulator. It reuses the real per-pass mocker sch
## Completion Event Queue
The multi-worker and disagg harnesses use `SimulationEvent` from `lib/mocker/src/replay/offline/events.rs` as a min-time priority queue implemented with `BinaryHeap`.
Right now the only scheduled event type is:
The multi-worker and disagg harnesses use `SimulationEvent` from `lib/mocker/src/replay/offline/events.rs` as a min-time priority queue implemented with `BinaryHeap`. The event itself is a small struct carrying the scheduled timestamp, a sequence number for tie-breaking, and a typed payload:
- `WorkerCompletion`
```rust
pub(crate) struct SimulationEvent {
pub(crate) at_ms: f64,
pub(crate) seq_no: u64,
pub(crate) kind: SimulationEventKind,
}
That event carries:
- worker `stage` (`aggregated`, `prefill`, or `decode`)
- `worker_idx`
- `completed_requests`
- `output_signals`
- router-visible `kv_events`
pub(crate) enum SimulationEventKind {
WorkerCompletion { stage, worker_idx, completed_requests, output_signals, kv_events },
DecodeHandoff { uuid },
WorkerReady { stage, worker_id },
}
```
Those are emitted after a worker pass is executed and then applied later when the harness clock reaches `pass.end_ms`.
- `WorkerCompletion` is emitted after a worker pass is executed and applied when the harness clock reaches `pass.end_ms`. It carries the `stage` (`Aggregated`, `Prefill`, or `Decode`), `worker_idx`, `completed_requests`, `output_signals`, and router-visible `kv_events`.
- `DecodeHandoff` is used by the disaggregated harness to move a request from prefill to decode at the same logical timestamp (see below).
- `WorkerReady` marks the point at which a worker returns to the admission pool after a pass completes.
## Router Integration
......@@ -140,7 +157,7 @@ Offline replay can run in:
- `round_robin`
- `kv_router`
The router implementation for offline mode lives in `lib/mocker/src/replay/router/offline.rs`.
The router implementation for offline mode lives in `lib/mocker/src/replay/offline/components/router.rs` (`OfflineReplayRouter`).
This router is synchronous and in-process:
......
tests/fault_tolerance/migration/utils.py
View file @ 241b254d
......@@ -8,7 +8,6 @@ import time
import pytest
import requests
from openai import APIError, OpenAI
from tests.utils.constants import FAULT_TOLERANCE_MODEL_NAME
from tests.utils.managed_process import (
......@@ -16,6 +15,12 @@ from tests.utils.managed_process import (
)
from tests.utils.managed_process import ManagedProcess, terminate_process_tree
openai = pytest.importorskip(
"openai", reason="openai package is required for fault tolerance migration tests"
)
APIError = openai.APIError
OpenAI = openai.OpenAI
logger = logging.getLogger(__name__)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment