diff --git a/components/src/dynamo/common/configuration/groups/kv_router_args.py b/components/src/dynamo/common/configuration/groups/kv_router_args.py index da5747dd8cca..8a97729ce860 100644 --- a/components/src/dynamo/common/configuration/groups/kv_router_args.py +++ b/components/src/dynamo/common/configuration/groups/kv_router_args.py @@ -26,6 +26,7 @@ "router_track_active_blocks", "router_track_output_blocks", "router_assume_kv_reuse", + "router_track_prefill_tokens", "router_snapshot_threshold", "router_reset_states", "router_ttl_secs", @@ -51,6 +52,7 @@ class KvRouterConfigBase(ConfigBase): router_track_active_blocks: bool router_track_output_blocks: bool router_assume_kv_reuse: bool + router_track_prefill_tokens: bool router_snapshot_threshold: int router_reset_states: bool router_ttl_secs: float @@ -173,6 +175,18 @@ def add_arguments(self, parser) -> None: ), obsolete_flag="--assume-kv-reuse", ) + add_negatable_bool_argument( + g, + flag_name="--router-track-prefill-tokens", + env_var="DYN_ROUTER_TRACK_PREFILL_TOKENS", + default=True, + dest="router_track_prefill_tokens", + help=( + "KV Router: Include prompt-side prefill tokens in active load accounting. " + "Use --no-router-track-prefill-tokens to ignore prompt tokens in router " + "prefill-token load, queue pressure, and active_prefill_tokens metrics." + ), + ) add_argument( g, flag_name="--router-snapshot-threshold", diff --git a/components/src/dynamo/mocker/aic_session.py b/components/src/dynamo/mocker/aic_session.py index eb715b067bb7..2c719b183f5b 100644 --- a/components/src/dynamo/mocker/aic_session.py +++ b/components/src/dynamo/mocker/aic_session.py @@ -14,7 +14,7 @@ from aiconfigurator.sdk.backends.factory import get_backend from aiconfigurator.sdk.inference_session import InferenceSession from aiconfigurator.sdk.models import get_model -from aiconfigurator.sdk.perf_database import get_database +from aiconfigurator.sdk.perf_database import get_database, get_supported_databases logger = logging.getLogger(__name__) @@ -40,6 +40,14 @@ def __init__( ) database = get_database(system=system, backend=backend_name, version=version) + if database is None: + supported = get_supported_databases().get(system, {}).get(backend_name, []) + supported_versions = ", ".join(supported) if supported else "" + raise RuntimeError( + "AIC perf database not found for " + f"system={system!r}, backend={backend_name!r}, version={version!r}. " + f"Supported versions for this system/backend: {supported_versions}" + ) model_config = config.ModelConfig(tp_size=tp_size) model = get_model( model_path=model_path, diff --git a/components/src/dynamo/mocker/config.py b/components/src/dynamo/mocker/config.py index 3d3e4414ba78..59716325dd82 100644 --- a/components/src/dynamo/mocker/config.py +++ b/components/src/dynamo/mocker/config.py @@ -42,6 +42,13 @@ def _build_sglang_args(args: argparse.Namespace) -> SglangArgs | None: def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: + worker_type = ( + "prefill" + if getattr(args, "is_prefill_worker", False) + else "decode" + if getattr(args, "is_decode_worker", False) + else "aggregated" + ) aic_backend = None aic_system = None aic_backend_version = None @@ -53,7 +60,6 @@ def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: aic_backend_version = getattr(args, "aic_backend_version", None) aic_tp_size = getattr(args, "aic_tp_size", None) aic_model_path = getattr(args, "model_path", None) - return MockEngineArgs( engine_type=getattr(args, "engine_type", None) or "vllm", num_gpu_blocks=getattr(args, "num_gpu_blocks", _DEFAULT_NUM_GPU_BLOCKS), @@ -64,18 +70,12 @@ def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: ), enable_prefix_caching=getattr(args, "enable_prefix_caching", True), enable_chunked_prefill=getattr(args, "enable_chunked_prefill", True), - preemption_mode=getattr(args, "preemption_mode", "lifo"), speedup_ratio=getattr(args, "speedup_ratio", 1.0), decode_speedup_ratio=getattr(args, "decode_speedup_ratio", 1.0), dp_size=getattr(args, "dp_size", 1), startup_time=getattr(args, "startup_time", None), - worker_type=( - "prefill" - if getattr(args, "is_prefill_worker", False) - else "decode" - if getattr(args, "is_decode_worker", False) - else "aggregated" - ), + worker_type=worker_type, + planner_profile_data=getattr(args, "planner_profile_data", None), aic_backend=aic_backend, aic_system=aic_system, aic_backend_version=aic_backend_version, @@ -85,6 +85,7 @@ def build_mocker_engine_args(args: argparse.Namespace) -> MockEngineArgs: kv_transfer_bandwidth=getattr(args, "kv_transfer_bandwidth", None), reasoning=_parse_reasoning_config(getattr(args, "reasoning", None)), sglang=_build_sglang_args(args), + preemption_mode=getattr(args, "preemption_mode", "lifo"), ) diff --git a/components/src/dynamo/router/README.md b/components/src/dynamo/router/README.md index a738380c28db..2ddacd3d49d8 100644 --- a/components/src/dynamo/router/README.md +++ b/components/src/dynamo/router/README.md @@ -29,7 +29,7 @@ python -m dynamo.router \ - `--endpoint`: Full endpoint path for workers in the format `namespace.component.endpoint` (e.g., `dynamo.prefill.generate`) **Router Configuration:** -All router options use the `--router-*` prefix (e.g., `--router-block-size`, `--router-kv-overlap-score-weight`, `--router-temperature`, `--router-kv-events` / `--no-router-kv-events`, `--router-replica-sync`, `--router-snapshot-threshold`, `--router-reset-states`, `--router-track-active-blocks` / `--no-router-track-active-blocks`). Legacy names without the prefix (e.g., `--block-size`, `--kv-events`) are still accepted but deprecated. For detailed descriptions, see the [Router Guide](/docs/components/router/router-guide.md). +All router options use the `--router-*` prefix (e.g., `--router-block-size`, `--router-kv-overlap-score-weight`, `--router-temperature`, `--router-kv-events` / `--no-router-kv-events`, `--router-replica-sync`, `--router-snapshot-threshold`, `--router-reset-states`, `--router-track-active-blocks` / `--no-router-track-active-blocks`, `--router-track-prefill-tokens` / `--no-router-track-prefill-tokens`). Legacy names without the prefix (e.g., `--block-size`, `--kv-events`) are still accepted but deprecated. For detailed descriptions, see the [Router Guide](/docs/components/router/router-guide.md). ## Architecture @@ -74,6 +74,9 @@ python -m dynamo.vllm --model MODEL_NAME --block-size 64 --disaggregation-mode p > **Why `--no-router-track-active-blocks` for prefill routing?** > Active block tracking is used for load balancing across decode (generation) phases. For prefill-only routing, decode load is not relevant, so disabling this reduces overhead and simplifies the router state. > +> **When should I use `--no-router-track-prefill-tokens`?** +> Use it on decode-only routers that should ignore already-completed prompt work. This keeps `active_prefill_tokens`, queue pressure, and load estimates focused on decode-side work after a prefill-to-decode handoff. +> > **Why `--router-block-size` is required for standalone routers:** > Unlike the frontend router which can infer block size from the ModelDeploymentCard (MDC) during worker registration, standalone routers cannot access the MDC and must have the block size explicitly specified. This is a work in progress to enable automatic inference. diff --git a/components/src/dynamo/router/__main__.py b/components/src/dynamo/router/__main__.py index 4adf13077347..c7fa95efcbe1 100644 --- a/components/src/dynamo/router/__main__.py +++ b/components/src/dynamo/router/__main__.py @@ -171,6 +171,7 @@ async def worker(runtime: DistributedRuntime): f"router_track_active_blocks={config.router_track_active_blocks}, " f"router_track_output_blocks={config.router_track_output_blocks}, " f"router_assume_kv_reuse={config.router_assume_kv_reuse}, " + f"router_track_prefill_tokens={config.router_track_prefill_tokens}, " f"router_ttl_secs={config.router_ttl_secs}, " f"router_max_tree_size={config.router_max_tree_size}, " f"router_prune_target_ratio={config.router_prune_target_ratio}" diff --git a/docs/benchmarks/mocker-trace-replay.md b/docs/benchmarks/mocker-trace-replay.md index b6ee91753c3d..e3c0ec8d0a6c 100644 --- a/docs/benchmarks/mocker-trace-replay.md +++ b/docs/benchmarks/mocker-trace-replay.md @@ -117,6 +117,8 @@ The dedicated replay CLI exposes: - `--replay-mode offline|online` - `--router-mode round_robin|kv_router` - `--num-workers` +- `--num-prefill-workers` +- `--num-decode-workers` - `--replay-concurrency` - `--arrival-interval-ms` - `--arrival-speedup-ratio` @@ -125,6 +127,8 @@ The dedicated replay CLI exposes: - `--num-prefix-groups` - `--inter-turn-delay-ms` - `--extra-engine-args` (JSON string) +- `--prefill-engine-args` (JSON string) +- `--decode-engine-args` (JSON string) - `--router-config` (JSON string) - `--report-json` @@ -164,6 +168,19 @@ as `block_size`, `engine_type`, `dp_size`, `speedup_ratio`, and `decode_speedup_ `--extra-engine-args`, not as top-level replay CLI flags. Unspecified fields fall back to the same defaults used by `MockEngineArgs::default()` and `KvRouterConfig::default()`. +Offline disagg replay uses staged engine args instead of `--extra-engine-args`: + +- `--prefill-engine-args` for the prefill worker config +- `--decode-engine-args` for the decode worker config +- `--num-prefill-workers` and `--num-decode-workers` for pool sizes + +For offline disagg replay, the staged JSON must set `worker_type` explicitly: + +- `--prefill-engine-args` must use `worker_type: "prefill"` +- `--decode-engine-args` must use `worker_type: "decode"` + +The staged configs must also use the same `block_size`. + ### Synthetic Replay Synthetic replay bypasses trace loading and generates in-memory requests with fixed input/output @@ -320,9 +337,12 @@ If `--report-json` is not provided, `python -m dynamo.replay` writes a timestamp Shared replay constraints: -- aggregated mode - `extra_engine_args.engine_type` must be `vllm` or `sglang` -- `extra_engine_args.dp_size` must be `1` +- aggregated replay requires the existing aggregated args path +- disagg replay requires both `prefill_engine_args` and `decode_engine_args` +- disagg replay requires `router_mode=kv_router` +- replay `dp_size` must be `1` +- disagg replay requires matching `block_size` in `prefill_engine_args` and `decode_engine_args` Additional offline constraints: @@ -330,6 +350,7 @@ Additional offline constraints: - single-worker offline replay is still a dedicated fast path for `vllm`, but it now supports both flat request replay and workload-driven multi-turn replay - `sglang` still goes through the shared multi-worker replay runtime even when `num_workers=1` +- offline disagg replay is a separate two-stage runtime with prefill and decode worker pools Additional online constraints: @@ -343,12 +364,13 @@ If you violate those constraints, replay fails immediately with a validation err either a trace file, or all of `--input-tokens`, `--output-tokens`, and `--request-count` - `--replay-concurrency` works with both trace replay and synthetic replay - mocker compute-speed knobs such as `speedup_ratio` still affect simulated timing when passed via - `--extra-engine-args` + the engine-args JSON for the chosen replay mode - `--arrival-speedup-ratio` affects trace timestamps, not worker compute speed - `--arrival-interval-ms` only applies to synthetic replay - `--turns-per-session`, `--shared-prefix-ratio`, `--num-prefix-groups`, and `--inter-turn-delay-ms` only apply to synthetic replay -- `--extra-engine-args` and `--router-config` are JSON strings on the standalone replay CLI +- `--extra-engine-args`, `--prefill-engine-args`, `--decode-engine-args`, and `--router-config` + are JSON strings on the standalone replay CLI - offline replay does not need planner runtime setup, router registration, or external event transport - the replay block size should match the trace block size, because token synthesis expands `hash_ids` using the configured block size diff --git a/docs/components/router/README.md b/docs/components/router/README.md index 79488ae77864..670b2a020760 100644 --- a/docs/components/router/README.md +++ b/docs/components/router/README.md @@ -23,6 +23,7 @@ For Kubernetes, set `DYN_ROUTER_MODE=kv` on the Frontend service. Workers automa | `--no-router-kv-events` | enabled | Fall back to approximate routing (no event consumption from workers) | | `--router-queue-threshold` | `4.0` | Backpressure queue threshold; enables priority scheduling via `nvext.agent_hints.priority` | | `--router-queue-policy` | `fcfs` | Queue scheduling policy: `fcfs` (tail TTFT), `wspt` (avg TTFT), or `lcfs` (comparison-only reverse ordering) | +| `--no-router-track-prefill-tokens` | disabled | Ignore prompt-side prefill tokens in router load accounting; useful for decode-only routing paths | ### Standalone Router diff --git a/docs/components/router/router-guide.md b/docs/components/router/router-guide.md index 95f0f713e8fd..3e2dfc292fc2 100644 --- a/docs/components/router/router-guide.md +++ b/docs/components/router/router-guide.md @@ -256,6 +256,8 @@ The main KV-aware routing arguments (frontend uses the same `--router-*` flag na - `--no-router-assume-kv-reuse`: When tracking active blocks, disables the assumption of KV cache reuse. By default (`router_assume_kv_reuse=true`), the router computes actual block hashes for sequence tracking to deduplicate blocks and optimize load balancing. When disabled via this flag, the router generates random hashes for sequence blocks, treating each request's blocks as unique. This is useful in disaggregated setups where prefill transfers blocks to decode workers that may already have those blocks cached, but the engine cannot coordinate transfers to avoid duplication. Without this flag, the router's load balancing heuristics would undercount decode blocks when duplicates exist. +- `--no-router-track-prefill-tokens`: Disables prompt-side prefill token accounting in the router's active load model. By default (`router_track_prefill_tokens=true`), the router counts uncached prompt tokens toward `active_prefill_tokens`, queue pressure, and potential prefill-token load. Disable this for decode-only routing paths where prompt processing has already happened elsewhere and the decode router should ignore transferred prompt load. In normal live disaggregated serving, the decode-stage override applies this behavior automatically. + - `--router-replica-sync`: Disabled by default. Enables NATS-based synchronization of local routing decisions between router replicas. When enabled, routers share their active sequence information and local predictions of block usage, improving routing consistency across instances. Note that this does not sync the radix tree or cached KV block states themselves - in JetStream mode those are synchronized through JetStream events; in local indexer mode (default) each router queries workers directly. ### KV Indexer / Approx KV Indexer @@ -280,6 +282,8 @@ Use `--no-router-kv-events` when you are not confident that your backend engine Use `--no-router-assume-kv-reuse` in disaggregated setups where the decode worker does not reuse transferred KV cache blocks. By default the router assumes KV blocks transferred from prefill to decode will be deduplicated on the decode side, but vLLM and SGLang decode workers currently do not support this — only TensorRT-LLM does. Without this flag, the router undercounts decode blocks when duplicates exist, leading to inaccurate load estimates. +Use `--no-router-track-prefill-tokens` when a router is serving decode-only traffic and prompt processing has already completed elsewhere. This keeps decode routing decisions focused on decode-side load instead of briefly charging prompt tokens to the decode worker after handoff. The built-in live disaggregated decode path applies the equivalent per-request override automatically. + Use `--router-track-output-blocks` **(experimental)** when your workload is output-heavy and you want the router to account for output-side KV cache growth in load balancing. This is useful in two scenarios: (1) workloads with long output sequences and little multi-turn reuse, where output blocks dominate the KV cache footprint; (2) agentic schedulers (e.g. NAT or other LLM routers) that can accurately predict the expected output sequence length per request. When enabled, the router adds placeholder blocks as tokens are generated. If you additionally pass `nvext.agent_hints.osl` (expected output sequence length in tokens) per request, the router applies fractional decay to output blocks — each output block's weight starts at 1.0 and decays linearly toward 0.0 as generation approaches the expected OSL. This lets the router predict that a request nearing completion will soon free its blocks, effectively modeling the future load trajectory rather than just the current snapshot. Without `osl`, output blocks are added at full weight with no decay. The flag requires `--router-track-active-blocks` (the default). The `--router-queue-threshold` (default: 4.0) controls when incoming requests are held in a priority queue. The router holds requests while all workers exceed the given fraction of `max_num_batched_tokens`, releasing them as capacity frees up. This defers the routing decision so it is made with the freshest load metrics, rather than dispatching into an already-saturated system. It also enables priority scheduling via `nvext.agent_hints.priority`. Set to None to disable queueing entirely. @@ -310,6 +314,11 @@ The prefill router is automatically created when: - **Seamlessly integrated** into the request pipeline between preprocessing and decode routing - **Falls back gracefully** to decode-only mode if prefill fails or no prefill workers are available +**Key characteristics of the decode routing stage in disaggregated mode:** +- **Disables overlap scoring** (`overlap_score_weight=0`) because decode routing should not chase prefix reuse +- **Disables KV reuse assumption** (`assume_kv_reuse=false`) unless the backend can truly deduplicate transferred blocks +- **Disables prefill-token tracking** (`track_prefill_tokens=false`) so decode-side load reflects decode work rather than already-completed prompt work + ### Setup Example When both workers are registered, requests are automatically routed. diff --git a/docs/mocker/mocker.md b/docs/mocker/mocker.md index 20f2822cbae5..11bd917ae4b2 100644 --- a/docs/mocker/mocker.md +++ b/docs/mocker/mocker.md @@ -96,7 +96,7 @@ python -m dynamo.mocker \ | `--sglang-chunked-prefill-size` | 8192 | SGLang chunked-prefill chunk size | | `--sglang-clip-max-new-tokens` | 4096 | SGLang admission-budget cap for max new tokens | | `--sglang-schedule-conservativeness` | 1.0 | SGLang schedule conservativeness factor | -| `--aic-perf-model` | False | Use AIC SDK for latency prediction instead of interpolated/polynomial models. Requires `aiconfigurator` SDK installed (install with `pip install ai-dynamo[mocker]`) | +| `--aic-perf-model` | False | Use AIC SDK for latency prediction instead of interpolated/polynomial models. Opt-in only: default mocker and replay paths do not use AIC. Requires `aiconfigurator` installed and usable AIC systems/perf data for the requested `system/backend/version` tuple | | `--aic-system` | `h200_sxm` | AIC system name (e.g., `h200_sxm`). Used with `--aic-perf-model` | | `--aic-backend-version` | Auto | AIC backend engine version (e.g., `0.12.0` for vLLM). If not set, uses the default version for the backend | | `--aic-tp-size` | 1 | Tensor parallel size for AIC latency prediction. Only affects AIC performance model lookups, not mocker scheduling | @@ -126,10 +126,12 @@ python -m dynamo.mocker \ The mocker supports replaying Mooncake-style traces through the dedicated replay CLI, which exposes `offline|online`, `round_robin|kv_router`, `arrival_speedup_ratio`, closed-loop concurrency -admission, and synthetic workload generation directly: +admission, synthetic workload generation, and offline disaggregated prefill/decode replay directly: -The replay CLI defaults to `--replay-mode offline` and `--router-mode round_robin`. Engine settings -such as `block_size`, `engine_type`, and compute speedups still belong in `--extra-engine-args`. +The replay CLI defaults to `--replay-mode offline` and `--router-mode round_robin`. Aggregated +replay uses `--extra-engine-args`. Offline disagg replay instead uses +`--prefill-engine-args` plus `--decode-engine-args`, together with +`--num-prefill-workers` and `--num-decode-workers`. ```bash python -m dynamo.replay /path/to/mooncake_trace.jsonl \ @@ -197,6 +199,31 @@ Replay supports aggregated `vllm` and `sglang` engine configs. Internally replay `block_size`; for `sglang`, `sglang.page_size` is still accepted as a compatibility alias as long as it matches `block_size` when both are provided. +Offline replay also supports disaggregated `kv_router` mode. In that mode: + +- `--prefill-engine-args` must describe a prefill worker +- `--decode-engine-args` must describe a decode worker +- `--router-mode` must be `kv_router` +- only offline replay is supported + +Example: + +```bash +python -m dynamo.replay \ + --input-tokens 4096 \ + --output-tokens 256 \ + --request-count 100 \ + --replay-mode offline \ + --router-mode kv_router \ + --replay-concurrency 32 \ + --num-prefill-workers 2 \ + --num-decode-workers 6 \ + --prefill-engine-args '{"worker_type":"prefill","block_size":512}' \ + --decode-engine-args '{"worker_type":"decode","block_size":512}' \ + --router-config '{"router_queue_policy":"wspt"}' \ + --report-json /tmp/replay-report.json +``` + ## Performance Modeling Setup By default, the mocker uses hardcoded polynomial formulas to estimate prefill and decode timing. For more realistic simulations, pass `--planner-profile-data` with either: @@ -223,7 +250,7 @@ python -m dynamo.mocker \ To use the AIC SDK for latency prediction: ```bash -pip install ai-dynamo[mocker] +uv pip install '.[mocker]' python -m dynamo.mocker \ --model-path nvidia/Llama-3.1-8B-Instruct-FP8 \ @@ -234,13 +261,33 @@ python -m dynamo.mocker \ The AIC model automatically uses `--model-path` and `--engine-type` to select the appropriate performance data. Available systems include `h200_sxm`, `h100_sxm`, etc. (see AIC SDK documentation for the full list). -When using `python -m dynamo.replay`, there are no dedicated AIC flags. Pass the equivalent fields directly via `--extra-engine-args`: +Important notes: + +- AIC is opt-in. If you do not pass `--aic-perf-model`, `python -m dynamo.mocker` does not use AIC. +- `python -m dynamo.replay` also does not use AIC unless you explicitly put AIC fields in the engine-args JSON. +- `aiconfigurator` must be able to load the requested performance database for the selected `system/backend/version`. If the SDK is installed but the backing systems data is missing or unreadable, mocker now fails fast at startup with a clear error instead of failing later on first request. +- In development environments, this may require pointing Python at a source checkout of `aiconfigurator` with real Git LFS payloads materialized in its `systems/` directory. + +When using `python -m dynamo.replay`, there are no dedicated AIC flags. For aggregated replay, +pass the equivalent fields via `--extra-engine-args`: ```bash python -m dynamo.replay /path/to/trace.jsonl \ --extra-engine-args '{"aic_backend":"vllm","aic_system":"h200_sxm","aic_model_path":"nvidia/Llama-3.1-8B-Instruct-FP8","aic_tp_size":1}' ``` +For offline disagg replay, pass the staged engine configs instead: + +```bash +python -m dynamo.replay /path/to/trace.jsonl \ + --replay-mode offline \ + --router-mode kv_router \ + --prefill-engine-args '{"worker_type":"prefill","aic_backend":"vllm","aic_system":"h200_sxm","aic_model_path":"nvidia/Llama-3.1-8B-Instruct-FP8","aic_tp_size":1,"block_size":512}' \ + --decode-engine-args '{"worker_type":"decode","aic_backend":"vllm","aic_system":"h200_sxm","aic_model_path":"nvidia/Llama-3.1-8B-Instruct-FP8","aic_tp_size":1,"block_size":512}' \ + --num-prefill-workers 2 \ + --num-decode-workers 6 +``` + The `aic_backend` field enables the AIC perf model and should match `engine_type` (`"vllm"` or `"sglang"`). The `aic_model_path` field is the equivalent of `--model-path` in `dynamo.mocker`. Example `--reasoning` configuration: @@ -355,7 +402,7 @@ The mocker supports three timing prediction modes: **Interpolated Model:** Loads actual profiling data from an NPZ file containing measured prefill and decode latencies. The mocker interpolates between data points to predict timing for any input size. This enables high-fidelity simulation matching a specific hardware configuration. -**AIC Model (`--aic-perf-model`):** Uses the NVIDIA AI Configurator (AIC) SDK for latency prediction. AIC provides calibrated performance models for specific GPU/model/engine combinations, predicting prefill and decode latency as a function of batch size, sequence length, and prefix cache hits. The model path is automatically derived from `--model-path`, and the engine type from `--engine-type`. This mode requires the `aiconfigurator` SDK, installable via `pip install ai-dynamo[mocker]`. +**AIC Model (`--aic-perf-model`):** Uses the NVIDIA AI Configurator (AIC) SDK for latency prediction. AIC provides calibrated performance models for specific GPU/model/engine combinations, predicting prefill and decode latency as a function of batch size, sequence length, and prefix cache hits. The model path is automatically derived from `--model-path`, and the engine type from `--engine-type`. This mode is opt-in and requires both the `aiconfigurator` SDK and loadable systems/perf data for the requested tuple. ### Bootstrap Rendezvous (Disaggregated Serving) diff --git a/lib/bench/kv_router/active_sequences_bench.rs b/lib/bench/kv_router/active_sequences_bench.rs index cc7ea7081ba9..75f69756f009 100644 --- a/lib/bench/kv_router/active_sequences_bench.rs +++ b/lib/bench/kv_router/active_sequences_bench.rs @@ -384,6 +384,7 @@ async fn apply_entry( token_sequence: Some(block_hashes), isl, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: Some(output_length as u32), worker, lora_name: None, diff --git a/lib/bindings/c/src/lib.rs b/lib/bindings/c/src/lib.rs index 06b57b8d437f..3a2cad51f400 100644 --- a/lib/bindings/c/src/lib.rs +++ b/lib/bindings/c/src/lib.rs @@ -497,6 +497,8 @@ impl RouterHandles { let config_override = if is_disaggregated { Some(RouterConfigOverride { overlap_score_weight: Some(0.0), + assume_kv_reuse: Some(false), + track_prefill_tokens: Some(false), ..Default::default() }) } else { @@ -573,6 +575,9 @@ fn kv_router_config_from_env() -> KvRouterConfig { if let Some(v) = env_bool("DYN_ROUTER_TRACK_OUTPUT_BLOCKS") { cfg.router_track_output_blocks = v; } + if let Some(v) = env_bool("DYN_ROUTER_TRACK_PREFILL_TOKENS") { + cfg.router_track_prefill_tokens = v; + } if let Some(v) = env_f64("DYN_ROUTER_QUEUE_THRESHOLD") { cfg.router_queue_threshold = Some(v); } @@ -584,6 +589,7 @@ fn kv_router_config_from_env() -> KvRouterConfig { router_replica_sync = cfg.router_replica_sync, router_track_active_blocks = cfg.router_track_active_blocks, router_track_output_blocks = cfg.router_track_output_blocks, + router_track_prefill_tokens = cfg.router_track_prefill_tokens, router_queue_threshold = ?cfg.router_queue_threshold, "KvRouterConfig initialized (DYN_* env overrides applied)" ); @@ -862,6 +868,12 @@ pub unsafe extern "C" fn add_request( tokio::time::timeout(timeout_duration, async { let worker = WorkerWithDpRank::new(worker_id, dp_rank); + let router_config_override = RouterConfigOverride { + overlap_score_weight: Some(0.0), + assume_kv_reuse: Some(false), + track_prefill_tokens: Some(false), + ..Default::default() + }; // Compute overlap_blocks using the public method let overlap_blocks = match decode_router @@ -884,7 +896,7 @@ pub unsafe extern "C" fn add_request( None, worker, None, // lora_name - None, // router_config_override + Some(&router_config_override), ) .await; diff --git a/lib/bindings/python/rust/llm/entrypoint.rs b/lib/bindings/python/rust/llm/entrypoint.rs index 5e3db4862723..5e31baec4f3f 100644 --- a/lib/bindings/python/rust/llm/entrypoint.rs +++ b/lib/bindings/python/rust/llm/entrypoint.rs @@ -58,7 +58,7 @@ impl KvRouterConfig { #[pymethods] impl KvRouterConfig { #[new] - #[pyo3(signature = (overlap_score_weight=1.0, router_temperature=0.0, use_kv_events=true, durable_kv_events=false, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_snapshot_threshold=1000000, router_reset_states=false, router_ttl_secs=120.0, router_max_tree_size=1048576, router_prune_target_ratio=0.8, router_queue_threshold=Some(4.0), router_event_threads=4, router_enable_cache_control=false, min_initial_workers=1, router_queue_policy="fcfs", remote_indexer_component=None))] + #[pyo3(signature = (overlap_score_weight=1.0, router_temperature=0.0, use_kv_events=true, durable_kv_events=false, router_replica_sync=false, router_track_active_blocks=true, router_track_output_blocks=false, router_assume_kv_reuse=true, router_track_prefill_tokens=true, router_snapshot_threshold=1000000, router_reset_states=false, router_ttl_secs=120.0, router_max_tree_size=1048576, router_prune_target_ratio=0.8, router_queue_threshold=Some(4.0), router_event_threads=4, router_enable_cache_control=false, min_initial_workers=1, router_queue_policy="fcfs", remote_indexer_component=None))] #[allow(clippy::too_many_arguments)] fn new( overlap_score_weight: f64, @@ -69,6 +69,7 @@ impl KvRouterConfig { router_track_active_blocks: bool, router_track_output_blocks: bool, router_assume_kv_reuse: bool, + router_track_prefill_tokens: bool, router_snapshot_threshold: Option, router_reset_states: bool, router_ttl_secs: f64, @@ -91,6 +92,7 @@ impl KvRouterConfig { router_track_active_blocks, router_track_output_blocks, router_assume_kv_reuse, + router_track_prefill_tokens, router_snapshot_threshold, router_reset_states, router_ttl_secs, diff --git a/lib/bindings/python/rust/llm/replay.rs b/lib/bindings/python/rust/llm/replay.rs index 636b055d85ec..d3f2cfbbe410 100644 --- a/lib/bindings/python/rust/llm/replay.rs +++ b/lib/bindings/python/rust/llm/replay.rs @@ -13,8 +13,10 @@ use dynamo_mocker::common::protocols::{ use dynamo_mocker::loadgen::{ ArrivalSpec, DelaySpec, LengthSpec, SyntheticTraceSpec, Trace as RsTrace, }; +use dynamo_mocker::replay::ReplayArgsMode; use pyo3::{exceptions::PyException, prelude::*}; use pythonize::pythonize; +use serde_json::json; use uuid::Uuid; use super::aic_callback::create_aic_callback; @@ -131,7 +133,7 @@ impl MockEngineArgs { #[pymethods] impl MockEngineArgs { #[new] - #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=16384, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, enable_local_indexer=false, bootstrap_port=None, kv_bytes_per_token=None, kv_transfer_bandwidth=None, reasoning=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None))] + #[pyo3(signature = (engine_type="vllm", num_gpu_blocks=16384, block_size=0, max_num_seqs=Some(256), max_num_batched_tokens=Some(8192), enable_prefix_caching=true, enable_chunked_prefill=true, speedup_ratio=1.0, decode_speedup_ratio=1.0, dp_size=1, startup_time=None, worker_type="aggregated", planner_profile_data=None, aic_backend=None, aic_system=None, aic_backend_version=None, aic_tp_size=None, aic_model_path=None, enable_local_indexer=false, bootstrap_port=None, kv_bytes_per_token=None, kv_transfer_bandwidth=None, reasoning=None, zmq_kv_events_port=None, zmq_replay_port=None, preemption_mode="lifo", router_queue_policy=None, sglang=None))] #[allow(clippy::too_many_arguments)] fn new( engine_type: &str, @@ -146,6 +148,7 @@ impl MockEngineArgs { dp_size: u32, startup_time: Option, worker_type: &str, + planner_profile_data: Option, aic_backend: Option, aic_system: Option, aic_backend_version: Option, @@ -173,7 +176,7 @@ impl MockEngineArgs { }) .transpose()?; - let inner = RsMockEngineArgs::builder() + let mut builder = RsMockEngineArgs::builder() .engine_type(engine_type) .num_gpu_blocks(num_gpu_blocks) .block_size(block_size) @@ -186,6 +189,7 @@ impl MockEngineArgs { .dp_size(dp_size) .startup_time(startup_time) .worker_type(worker_type) + .planner_profile_data(planner_profile_data.clone()) .aic_backend(aic_backend) .aic_system(aic_system) .aic_backend_version(aic_backend_version) @@ -200,7 +204,19 @@ impl MockEngineArgs { .zmq_replay_port(zmq_replay_port) .preemption_mode(preemption_mode) .router_queue_policy(router_queue_policy) - .sglang(sglang.map(|config| config.inner())) + .sglang(sglang.map(|config| config.inner())); + + if let Some(npz_path) = planner_profile_data { + let perf_model = PerfModel::from_npz(&npz_path).map_err(|e| { + PyException::new_err(format!( + "Failed to load planner_profile_data from {:?}: {e}", + npz_path + )) + })?; + builder = builder.perf_model(Arc::new(perf_model)); + } + + let inner = builder .build() .map_err(|e| PyException::new_err(format!("Failed to build MockEngineArgs: {e}")))? .normalized() @@ -218,6 +234,60 @@ impl MockEngineArgs { .map_err(|e| PyException::new_err(format!("Failed to parse MockEngineArgs JSON: {e}"))) } + fn dump_json(&self) -> PyResult { + let worker_type = match self.inner.worker_type { + RsWorkerType::Aggregated => "aggregated", + RsWorkerType::Prefill => "prefill", + RsWorkerType::Decode => "decode", + }; + let engine_type = match self.inner.engine_type { + RsMockerEngineType::Vllm => "vllm", + RsMockerEngineType::Sglang => "sglang", + }; + let preemption_mode = match self.inner.preemption_mode { + RsPreemptionMode::Lifo => "lifo", + RsPreemptionMode::Fifo => "fifo", + }; + let router_queue_policy = self + .inner + .router_queue_policy + .as_ref() + .map(|policy| policy.to_string()); + let payload = json!({ + "engine_type": engine_type, + "num_gpu_blocks": self.inner.num_gpu_blocks, + "block_size": self.inner.block_size, + "max_num_seqs": self.inner.max_num_seqs, + "max_num_batched_tokens": self.inner.max_num_batched_tokens, + "enable_prefix_caching": self.inner.enable_prefix_caching, + "enable_chunked_prefill": self.inner.enable_chunked_prefill, + "speedup_ratio": self.inner.speedup_ratio, + "decode_speedup_ratio": self.inner.decode_speedup_ratio, + "dp_size": self.inner.dp_size, + "startup_time": self.inner.startup_time, + "worker_type": worker_type, + "planner_profile_data": self.inner.planner_profile_data, + "aic_backend": self.inner.aic_backend, + "aic_system": self.inner.aic_system, + "aic_backend_version": self.inner.aic_backend_version, + "aic_tp_size": self.inner.aic_tp_size, + "aic_model_path": self.inner.aic_model_path, + "enable_local_indexer": self.inner.enable_local_indexer, + "bootstrap_port": self.inner.bootstrap_port, + "kv_bytes_per_token": self.inner.kv_bytes_per_token, + "kv_transfer_bandwidth": self.inner.kv_transfer_bandwidth, + "reasoning": self.inner.reasoning, + "zmq_kv_events_port": self.inner.zmq_kv_events_port, + "zmq_replay_port": self.inner.zmq_replay_port, + "preemption_mode": preemption_mode, + "router_queue_policy": router_queue_policy, + "sglang": self.inner.sglang, + "has_perf_model": true, + }); + serde_json::to_string_pretty(&payload) + .map_err(|e| PyException::new_err(format!("Failed to serialize MockEngineArgs: {e}"))) + } + #[getter] fn block_size(&self) -> usize { self.inner.block_size @@ -289,67 +359,113 @@ impl MockEngineArgs { } #[pyfunction] -#[pyo3(signature = (trace_file, extra_engine_args=None, router_config=None, num_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0))] +#[pyo3(signature = (trace_file, extra_engine_args=None, prefill_engine_args=None, decode_engine_args=None, router_config=None, num_workers=1, num_prefill_workers=1, num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0))] #[allow(clippy::too_many_arguments)] pub fn run_mocker_trace_replay( py: Python<'_>, trace_file: PathBuf, extra_engine_args: Option, + prefill_engine_args: Option, + decode_engine_args: Option, router_config: Option, num_workers: usize, + num_prefill_workers: usize, + num_decode_workers: usize, replay_concurrency: Option, replay_mode: &str, router_mode: &str, arrival_speedup_ratio: f64, ) -> PyResult { - let args = load_replay_mocker_args(py, extra_engine_args)?; + let args_selection = load_replay_args_selection( + py, + extra_engine_args, + prefill_engine_args, + decode_engine_args, + num_workers, + num_prefill_workers, + num_decode_workers, + )?; let router_config = load_replay_router_config(router_config); let replay_mode = replay_mode.to_owned(); let router_mode = parse_replay_router_mode(router_mode)?; let report = py.allow_threads(move || { let replay_concurrency = parse_replay_concurrency(replay_concurrency)?; - match (replay_mode.as_str(), replay_concurrency) { - ("offline", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_file_with_router_mode( - args, - router_config.clone(), - &trace_file, - max_in_flight, - num_workers, - router_mode, - ) - } - ("offline", None) => dynamo_mocker::replay::simulate_trace_file_with_router_mode( - args, - router_config.clone(), - &trace_file, - num_workers, - arrival_speedup_ratio, - router_mode, - ), - ("online", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_live_file_with_router_mode( - args, - router_config.clone(), - &trace_file, - max_in_flight, - num_workers, - router_mode, - ) + match args_selection { + ReplayArgsSelection::Aggregated(args) => { + match (replay_mode.as_str(), replay_concurrency) { + ("offline", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_file_with_router_mode( + *args, + router_config.clone(), + &trace_file, + max_in_flight, + num_workers, + router_mode, + ) + } + ("offline", None) => { + dynamo_mocker::replay::simulate_trace_file_with_router_mode( + *args, + router_config.clone(), + &trace_file, + num_workers, + arrival_speedup_ratio, + router_mode, + ) + } + ("online", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_live_file_with_router_mode( + *args, + router_config.clone(), + &trace_file, + max_in_flight, + num_workers, + router_mode, + ) + } + ("online", None) => { + dynamo_mocker::replay::simulate_trace_live_file_with_router_mode( + *args, + router_config.clone(), + &trace_file, + num_workers, + arrival_speedup_ratio, + router_mode, + ) + } + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + } } - ("online", None) => dynamo_mocker::replay::simulate_trace_live_file_with_router_mode( - args, - router_config.clone(), - &trace_file, - num_workers, - arrival_speedup_ratio, - router_mode, - ), - (other, _) => anyhow::bail!( - "replay_mode must be either 'offline' or 'online', got '{}'", - other - ), + ReplayArgsSelection::Disagg(config) => match (replay_mode.as_str(), replay_concurrency) + { + ("offline", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_file_disagg_with_router_mode( + *config, + router_config.clone(), + &trace_file, + max_in_flight, + router_mode, + ) + } + ("offline", None) => { + dynamo_mocker::replay::simulate_trace_file_disagg_with_router_mode( + *config, + router_config.clone(), + &trace_file, + arrival_speedup_ratio, + router_mode, + ) + } + ("online", _) => anyhow::bail!("disagg replay only supports replay_mode='offline'"), + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + }, } }); let report = report.map_err(to_pyerr)?; @@ -359,7 +475,7 @@ pub fn run_mocker_trace_replay( } #[pyfunction] -#[pyo3(signature = (input_tokens, output_tokens, request_count, extra_engine_args=None, router_config=None, num_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0, arrival_interval_ms=1.0, turns_per_session=1, shared_prefix_ratio=0.0, num_prefix_groups=0, inter_turn_delay_ms=0.0))] +#[pyo3(signature = (input_tokens, output_tokens, request_count, extra_engine_args=None, prefill_engine_args=None, decode_engine_args=None, router_config=None, num_workers=1, num_prefill_workers=1, num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", arrival_speedup_ratio=1.0, arrival_interval_ms=1.0, turns_per_session=1, shared_prefix_ratio=0.0, num_prefix_groups=0, inter_turn_delay_ms=0.0))] #[allow(clippy::too_many_arguments)] pub fn run_mocker_synthetic_trace_replay( py: Python<'_>, @@ -367,8 +483,12 @@ pub fn run_mocker_synthetic_trace_replay( output_tokens: usize, request_count: usize, extra_engine_args: Option, + prefill_engine_args: Option, + decode_engine_args: Option, router_config: Option, num_workers: usize, + num_prefill_workers: usize, + num_decode_workers: usize, replay_concurrency: Option, replay_mode: &str, router_mode: &str, @@ -379,10 +499,22 @@ pub fn run_mocker_synthetic_trace_replay( num_prefix_groups: usize, inter_turn_delay_ms: f64, ) -> PyResult { - let args = load_replay_mocker_args(py, extra_engine_args)?; + let args_selection = load_replay_args_selection( + py, + extra_engine_args, + prefill_engine_args, + decode_engine_args, + num_workers, + num_prefill_workers, + num_decode_workers, + )?; let router_config = load_replay_router_config(router_config); let replay_mode = replay_mode.to_owned(); let router_mode = parse_replay_router_mode(router_mode)?; + let block_size = match &args_selection { + ReplayArgsSelection::Aggregated(args) => args.block_size.max(1), + ReplayArgsSelection::Disagg(config) => config.prefill_args.block_size.max(1), + }; let report = py.allow_threads(move || { let replay_concurrency = parse_replay_concurrency(replay_concurrency)?; let use_workload = turns_per_session > 1 @@ -392,7 +524,7 @@ pub fn run_mocker_synthetic_trace_replay( if use_workload { let mut trace = build_synthetic_workload( - args.block_size.max(1), + block_size, input_tokens, output_tokens, request_count, @@ -406,103 +538,159 @@ pub fn run_mocker_synthetic_trace_replay( trace = trace.speed_up_timing(arrival_speedup_ratio)?; } - return match (replay_mode.as_str(), replay_concurrency) { + return match args_selection { + ReplayArgsSelection::Aggregated(args) => match (replay_mode.as_str(), replay_concurrency) + { + ("offline", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_workload_with_router_mode( + *args, + router_config.clone(), + trace, + max_in_flight, + num_workers, + router_mode, + ) + } + ("offline", None) => { + dynamo_mocker::replay::simulate_trace_workload_with_router_mode( + *args, + router_config.clone(), + trace, + num_workers, + router_mode, + ) + } + ("online", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_live_workload_with_router_mode( + *args, + router_config.clone(), + trace, + max_in_flight, + num_workers, + router_mode, + ) + } + ("online", None) => { + dynamo_mocker::replay::simulate_trace_live_workload_with_router_mode( + *args, + router_config.clone(), + trace, + num_workers, + router_mode, + ) + } + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + }, + ReplayArgsSelection::Disagg(config) => { + match (replay_mode.as_str(), replay_concurrency) { + ("offline", Some(max_in_flight)) => dynamo_mocker::replay::simulate_concurrency_workload_disagg_with_router_mode( + *config, + router_config.clone(), + trace, + max_in_flight, + router_mode, + ), + ("offline", None) => dynamo_mocker::replay::simulate_trace_workload_disagg_with_router_mode( + *config, + router_config.clone(), + trace, + router_mode, + ), + ("online", _) => anyhow::bail!( + "disagg replay only supports replay_mode='offline'" + ), + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + } + } + }; + } + + let requests = build_synthetic_requests( + input_tokens, + output_tokens, + request_count, + arrival_interval_ms, + replay_concurrency.is_none(), + )?; + + match args_selection { + ReplayArgsSelection::Aggregated(args) => match (replay_mode.as_str(), replay_concurrency) + { ("offline", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_workload_with_router_mode( - args, + dynamo_mocker::replay::simulate_concurrency_requests_with_router_mode( + *args, router_config.clone(), - trace, + requests, max_in_flight, num_workers, router_mode, ) } - ("offline", None) => { - dynamo_mocker::replay::simulate_trace_workload_with_router_mode( - args, + ("offline", None) => dynamo_mocker::replay::simulate_trace_requests_with_router_mode( + *args, + router_config.clone(), + requests, + num_workers, + arrival_speedup_ratio, + router_mode, + ), + ("online", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_live_requests_with_router_mode( + *args, router_config.clone(), - trace, + requests, + max_in_flight, num_workers, router_mode, ) } - ("online", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_live_workload_with_router_mode( - args, + ("online", None) => { + dynamo_mocker::replay::simulate_trace_live_requests_with_router_mode( + *args, router_config.clone(), - trace, - max_in_flight, + requests, num_workers, + arrival_speedup_ratio, router_mode, ) } - ("online", None) => { - dynamo_mocker::replay::simulate_trace_live_workload_with_router_mode( - args, + (other, _) => anyhow::bail!( + "replay_mode must be either 'offline' or 'online', got '{}'", + other + ), + }, + ReplayArgsSelection::Disagg(config) => match (replay_mode.as_str(), replay_concurrency) + { + ("offline", Some(max_in_flight)) => { + dynamo_mocker::replay::simulate_concurrency_requests_disagg_with_router_mode( + *config, router_config.clone(), - trace, - num_workers, + requests, + max_in_flight, + router_mode, + ) + } + ("offline", None) => { + dynamo_mocker::replay::simulate_trace_requests_disagg_with_router_mode( + *config, + router_config.clone(), + requests, + arrival_speedup_ratio, router_mode, ) } + ("online", _) => anyhow::bail!("disagg replay only supports replay_mode='offline'"), (other, _) => anyhow::bail!( "replay_mode must be either 'offline' or 'online', got '{}'", other ), - }; - } - - let requests = build_synthetic_requests( - input_tokens, - output_tokens, - request_count, - arrival_interval_ms, - replay_concurrency.is_none(), - )?; - - match (replay_mode.as_str(), replay_concurrency) { - ("offline", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_requests_with_router_mode( - args, - router_config.clone(), - requests, - max_in_flight, - num_workers, - router_mode, - ) - } - ("offline", None) => dynamo_mocker::replay::simulate_trace_requests_with_router_mode( - args, - router_config.clone(), - requests, - num_workers, - arrival_speedup_ratio, - router_mode, - ), - ("online", Some(max_in_flight)) => { - dynamo_mocker::replay::simulate_concurrency_live_requests_with_router_mode( - args, - router_config.clone(), - requests, - max_in_flight, - num_workers, - router_mode, - ) - } - ("online", None) => { - dynamo_mocker::replay::simulate_trace_live_requests_with_router_mode( - args, - router_config.clone(), - requests, - num_workers, - arrival_speedup_ratio, - router_mode, - ) - } - (other, _) => anyhow::bail!( - "replay_mode must be either 'offline' or 'online', got '{}'", - other - ), + }, } }); let report = report.map_err(to_pyerr)?; @@ -511,15 +699,62 @@ pub fn run_mocker_synthetic_trace_replay( .map(|obj| obj.unbind()) } -fn load_replay_mocker_args( +enum ReplayArgsSelection { + Aggregated(Box), + Disagg(Box), +} + +fn load_replay_args_selection( py: Python<'_>, extra_engine_args: Option, -) -> PyResult { - let mut args = match extra_engine_args { - Some(extra_args) => extra_args.inner(), - None => RsMockEngineArgs::default(), - }; + prefill_engine_args: Option, + decode_engine_args: Option, + num_workers: usize, + num_prefill_workers: usize, + num_decode_workers: usize, +) -> PyResult { + let aggregated_args = load_optional_replay_mocker_args(py, extra_engine_args)?; + let prefill_args = load_optional_replay_mocker_args(py, prefill_engine_args)?; + let decode_args = load_optional_replay_mocker_args(py, decode_engine_args)?; + + let replay_args_mode = dynamo_mocker::replay::validate_replay_args_mode( + aggregated_args.as_ref(), + prefill_args.as_ref(), + decode_args.as_ref(), + num_workers, + num_prefill_workers, + num_decode_workers, + ) + .map_err(to_pyerr)?; + + match replay_args_mode { + ReplayArgsMode::Aggregated => Ok(ReplayArgsSelection::Aggregated(Box::new( + aggregated_args.unwrap_or_default(), + ))), + ReplayArgsMode::Disagg => Ok(ReplayArgsSelection::Disagg(Box::new( + dynamo_mocker::replay::OfflineDisaggReplayConfig { + prefill_args: prefill_args.expect("validated disagg prefill args"), + decode_args: decode_args.expect("validated disagg decode args"), + num_prefill_workers, + num_decode_workers, + }, + ))), + } +} +fn load_optional_replay_mocker_args( + py: Python<'_>, + extra_engine_args: Option, +) -> PyResult> { + extra_engine_args + .map(|extra_args| materialize_replay_mocker_args(py, extra_args.inner())) + .transpose() +} + +fn materialize_replay_mocker_args( + py: Python<'_>, + mut args: RsMockEngineArgs, +) -> PyResult { if let Some(ref backend_name) = args.aic_backend.clone() { let backend = backend_name.clone(); let system = args.aic_system.as_deref().unwrap_or("h200_sxm").to_string(); diff --git a/lib/bindings/python/src/dynamo/_core.pyi b/lib/bindings/python/src/dynamo/_core.pyi index 9092eb20530e..8961f26845eb 100644 --- a/lib/bindings/python/src/dynamo/_core.pyi +++ b/lib/bindings/python/src/dynamo/_core.pyi @@ -1217,6 +1217,7 @@ class MockEngineArgs: dp_size: int = 1, startup_time: Optional[float] = None, worker_type: str = "aggregated", + planner_profile_data: Optional[str | os.PathLike[str]] = None, aic_backend: Optional[str] = None, aic_system: Optional[str] = None, aic_backend_version: Optional[str] = None, @@ -1239,6 +1240,8 @@ class MockEngineArgs: def from_json(config_json: str) -> "MockEngineArgs": ... + def dump_json(self) -> str: ... + @property def block_size(self) -> int: ... @@ -1376,8 +1379,12 @@ async def run_input(runtime: DistributedRuntime, input: str, engine_config: Engi def run_mocker_trace_replay( trace_file: str | os.PathLike[str], extra_engine_args: Optional[MockEngineArgs] = None, + prefill_engine_args: Optional[MockEngineArgs] = None, + decode_engine_args: Optional[MockEngineArgs] = None, router_config: Optional[KvRouterConfig] = None, num_workers: int = 1, + num_prefill_workers: int = 1, + num_decode_workers: int = 1, replay_concurrency: Optional[int] = None, replay_mode: Literal["offline", "online"] = "offline", router_mode: Literal["round_robin", "kv_router"] = "round_robin", @@ -1391,8 +1398,12 @@ def run_mocker_synthetic_trace_replay( output_tokens: int, request_count: int, extra_engine_args: Optional[MockEngineArgs] = None, + prefill_engine_args: Optional[MockEngineArgs] = None, + decode_engine_args: Optional[MockEngineArgs] = None, router_config: Optional[KvRouterConfig] = None, num_workers: int = 1, + num_prefill_workers: int = 1, + num_decode_workers: int = 1, replay_concurrency: Optional[int] = None, replay_mode: Literal["offline", "online"] = "offline", router_mode: Literal["round_robin", "kv_router"] = "round_robin", diff --git a/lib/bindings/python/src/dynamo/replay/api.py b/lib/bindings/python/src/dynamo/replay/api.py index 25d3a3207bb1..049a1afe6f8d 100644 --- a/lib/bindings/python/src/dynamo/replay/api.py +++ b/lib/bindings/python/src/dynamo/replay/api.py @@ -11,8 +11,12 @@ def run_trace_replay( trace_file, *, extra_engine_args=None, + prefill_engine_args=None, + decode_engine_args=None, router_config=None, num_workers=1, + num_prefill_workers=1, + num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", @@ -21,8 +25,12 @@ def run_trace_replay( return _run_mocker_trace_replay( trace_file, extra_engine_args=extra_engine_args, + prefill_engine_args=prefill_engine_args, + decode_engine_args=decode_engine_args, router_config=router_config, num_workers=num_workers, + num_prefill_workers=num_prefill_workers, + num_decode_workers=num_decode_workers, replay_concurrency=replay_concurrency, replay_mode=replay_mode, router_mode=router_mode, @@ -36,8 +44,12 @@ def run_synthetic_trace_replay( request_count, *, extra_engine_args=None, + prefill_engine_args=None, + decode_engine_args=None, router_config=None, num_workers=1, + num_prefill_workers=1, + num_decode_workers=1, replay_concurrency=None, replay_mode="offline", router_mode="round_robin", @@ -53,8 +65,12 @@ def run_synthetic_trace_replay( output_tokens, request_count, extra_engine_args=extra_engine_args, + prefill_engine_args=prefill_engine_args, + decode_engine_args=decode_engine_args, router_config=router_config, num_workers=num_workers, + num_prefill_workers=num_prefill_workers, + num_decode_workers=num_decode_workers, replay_concurrency=replay_concurrency, replay_mode=replay_mode, router_mode=router_mode, diff --git a/lib/bindings/python/src/dynamo/replay/main.py b/lib/bindings/python/src/dynamo/replay/main.py index 5cb254917e80..60448ccf270a 100644 --- a/lib/bindings/python/src/dynamo/replay/main.py +++ b/lib/bindings/python/src/dynamo/replay/main.py @@ -4,24 +4,80 @@ from __future__ import annotations import argparse +import importlib import json import os import sys from collections.abc import Sequence from pathlib import Path +from types import SimpleNamespace +from typing import Protocol os.environ.setdefault("DYNAMO_SKIP_PYTHON_LOG_INIT", "1") from dynamo.llm import KvRouterConfig, MockEngineArgs -from dynamo.mocker.args import resolve_planner_profile_data from dynamo.replay import run_synthetic_trace_replay, run_trace_replay from dynamo.replay.reporting import format_report_table, write_report_json +class PlannerProfileDataResult(Protocol): + npz_path: Path | None + + +def resolve_planner_profile_data( + planner_profile_data: Path | None, +) -> PlannerProfileDataResult: + try: + module = importlib.import_module("dynamo.mocker.args") + except ImportError: + if planner_profile_data is None: + return SimpleNamespace(npz_path=None) + return SimpleNamespace( + npz_path=planner_profile_data + if planner_profile_data.suffix == ".npz" + else None + ) + return module.resolve_planner_profile_data(planner_profile_data) + + +def _load_engine_args(raw_args: str | None): + if raw_args is None: + return None + + raw = json.loads(raw_args) + if not isinstance(raw, dict): + raise ValueError("engine-args must be a JSON object") + worker_type = raw.pop("worker_type", None) + if worker_type is not None: + if "is_prefill" in raw or "is_decode" in raw: + raise ValueError( + "worker_type cannot be combined with is_prefill or is_decode" + ) + if worker_type == "prefill": + raw["is_prefill"] = True + elif worker_type == "decode": + raw["is_decode"] = True + elif worker_type != "aggregated": + raise ValueError( + "worker_type must be one of 'aggregated', 'prefill', or 'decode'" + ) + if "planner_profile_data" in raw: + profile_data_result = resolve_planner_profile_data( + Path(raw["planner_profile_data"]) + ) + if profile_data_result.npz_path is not None: + raw["planner_profile_data"] = str(profile_data_result.npz_path) + else: + del raw["planner_profile_data"] + return MockEngineArgs.from_json(json.dumps(raw)) + + def main(argv: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(prog="python -m dynamo.replay") parser.add_argument("trace_file", nargs="?") parser.add_argument("--extra-engine-args") + parser.add_argument("--prefill-engine-args") + parser.add_argument("--decode-engine-args") parser.add_argument("--router-config") parser.add_argument("--input-tokens", type=int) parser.add_argument("--output-tokens", type=int) @@ -36,6 +92,8 @@ def main(argv: Sequence[str] | None = None) -> int: parser.add_argument("--num-prefix-groups", type=int, default=0) parser.add_argument("--inter-turn-delay-ms", type=float, default=0.0) parser.add_argument("--num-workers", type=int, default=1) + parser.add_argument("--num-prefill-workers", type=int, default=1) + parser.add_argument("--num-decode-workers", type=int, default=1) parser.add_argument("--replay-concurrency", type=int) parser.add_argument( "--replay-mode", @@ -74,24 +132,9 @@ def main(argv: Sequence[str] | None = None) -> int: "synthetic replay requires --input-tokens, --output-tokens, and --request-count" ) - # Resolve planner_profile_data directory -> NPZ before passing to Rust. - # Rust only accepts NPZ files; resolve_planner_profile_data handles conversion. - profile_data_result = None - if args.extra_engine_args is not None: - raw = json.loads(args.extra_engine_args) - if "planner_profile_data" in raw: - profile_data_result = resolve_planner_profile_data( - Path(raw["planner_profile_data"]) - ) - if profile_data_result.npz_path is not None: - raw["planner_profile_data"] = str(profile_data_result.npz_path) - else: - del raw["planner_profile_data"] - extra_engine_args = MockEngineArgs.from_json(json.dumps(raw)) - else: - extra_engine_args = MockEngineArgs.from_json(args.extra_engine_args) - else: - extra_engine_args = None + extra_engine_args = _load_engine_args(args.extra_engine_args) + prefill_engine_args = _load_engine_args(args.prefill_engine_args) + decode_engine_args = _load_engine_args(args.decode_engine_args) router_config = ( KvRouterConfig.from_json(args.router_config) if args.router_config is not None @@ -102,8 +145,12 @@ def main(argv: Sequence[str] | None = None) -> int: report = run_trace_replay( args.trace_file, extra_engine_args=extra_engine_args, + prefill_engine_args=prefill_engine_args, + decode_engine_args=decode_engine_args, router_config=router_config, num_workers=args.num_workers, + num_prefill_workers=args.num_prefill_workers, + num_decode_workers=args.num_decode_workers, replay_concurrency=args.replay_concurrency, replay_mode=args.replay_mode, router_mode=args.router_mode, @@ -115,8 +162,12 @@ def main(argv: Sequence[str] | None = None) -> int: args.output_tokens, args.request_count, extra_engine_args=extra_engine_args, + prefill_engine_args=prefill_engine_args, + decode_engine_args=decode_engine_args, router_config=router_config, num_workers=args.num_workers, + num_prefill_workers=args.num_prefill_workers, + num_decode_workers=args.num_decode_workers, replay_concurrency=args.replay_concurrency, replay_mode=args.replay_mode, router_mode=args.router_mode, diff --git a/lib/bindings/python/tests/test_replay.py b/lib/bindings/python/tests/test_replay.py index deeb6302ec35..b097ca883db4 100644 --- a/lib/bindings/python/tests/test_replay.py +++ b/lib/bindings/python/tests/test_replay.py @@ -5,12 +5,13 @@ import os import subprocess import sys +from pathlib import Path +import numpy as np import pytest from dynamo.llm import KvRouterConfig, MockEngineArgs from dynamo.replay import run_synthetic_trace_replay, run_trace_replay -from dynamo.replay.main import main from dynamo.replay.reporting import format_report_table, write_report_json pytestmark = [ @@ -75,6 +76,7 @@ def _router_config_payload(): "router_track_active_blocks": True, "router_track_output_blocks": False, "router_assume_kv_reuse": True, + "router_track_prefill_tokens": True, "router_snapshot_threshold": 1000000, "router_reset_states": False, "router_ttl_secs": 120.0, @@ -194,6 +196,14 @@ def _sglang_args(): return MockEngineArgs.from_json(json.dumps(_sglang_args_payload())) +def _prefill_args(): + return MockEngineArgs(block_size=64, speedup_ratio=1000.0, worker_type="prefill") + + +def _decode_args(): + return MockEngineArgs(block_size=64, speedup_ratio=1000.0, worker_type="decode") + + def _write_router_config(tmp_path): config_path = tmp_path / "router_config.json" config_path.write_text( @@ -229,8 +239,12 @@ def _assert_basic_report_metrics(report): def _replay_cli_env() -> dict[str, str]: + repo_root = Path(__file__).resolve().parents[4] env = os.environ.copy() - pythonpath_entries = ["lib/bindings/python/src", "components/src"] + pythonpath_entries = [ + str(repo_root / "lib/bindings/python/src"), + str(repo_root / "components/src"), + ] existing_pythonpath = env.get("PYTHONPATH") if existing_pythonpath: pythonpath_entries.append(existing_pythonpath) @@ -238,6 +252,33 @@ def _replay_cli_env() -> dict[str, str]: return env +def _planner_profile_data_npz_path() -> Path: + return ( + Path(__file__).resolve().parents[4] + / "benchmarks/results/H200_TP1P_TP1D_perf_data.npz" + ) + + +def _planner_profile_data_dir_path() -> Path: + return ( + Path(__file__).resolve().parents[4] + / "tests/planner/profiling_results/H200_TP1P_TP1D" + ) + + +def _write_planner_profile_data_npz(tmp_path: Path) -> Path: + planner_profile_data = tmp_path / "planner_profile_data.npz" + np.savez( + planner_profile_data, + prefill_isl=np.array([128.0, 256.0]), + prefill_ttft_ms=np.array([4.0, 8.0]), + decode_active_kv_tokens=np.array([1024.0, 2048.0]), + decode_context_length=np.array([128.0, 256.0]), + decode_itl=np.array([[1.0, 1.5], [2.0, 2.5]]), + ) + return planner_profile_data + + def _run_replay_cli(tmp_path, *args): return subprocess.run( [ @@ -264,18 +305,34 @@ def _assert_replay_cli_outputs(completed, report_path): @pytest.mark.parametrize("engine_type", ["vllm", "sglang"]) @pytest.mark.parametrize("replay_mode", ["offline", "online"]) @pytest.mark.parametrize("router_mode", ["round_robin", "kv_router"]) -def test_run_trace_replay_smoke_matrix(tmp_path, engine_type, replay_mode, router_mode): +@pytest.mark.parametrize("serving_mode", ["agg", "disagg"]) +def test_run_trace_replay_smoke_matrix( + tmp_path, engine_type, replay_mode, router_mode, serving_mode +): trace_path = _write_trace_and_args(tmp_path) - args_path = _vllm_args() if engine_type == "vllm" else _sglang_args() - num_workers = 1 if router_mode == "round_robin" else 2 - - report = run_trace_replay( - trace_path, - extra_engine_args=args_path, - num_workers=num_workers, - replay_mode=replay_mode, - router_mode=router_mode, - ) + if serving_mode == "disagg": + if replay_mode != "offline": + pytest.skip("disagg replay only supports offline mode") + report = run_trace_replay( + trace_path, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode=replay_mode, + router_mode=router_mode, + ) + else: + args_path = _vllm_args() if engine_type == "vllm" else _sglang_args() + num_workers = 1 if router_mode == "round_robin" else 2 + report = run_trace_replay( + trace_path, + extra_engine_args=args_path, + num_workers=num_workers, + replay_mode=replay_mode, + router_mode=router_mode, + ) _assert_basic_report_counts( report, @@ -345,22 +402,39 @@ def test_run_trace_replay_supports_multiturn_sessions(tmp_path, replay_mode): @pytest.mark.parametrize("engine_type", ["vllm", "sglang"]) @pytest.mark.parametrize("replay_mode", ["offline", "online"]) @pytest.mark.parametrize("router_mode", ["round_robin", "kv_router"]) +@pytest.mark.parametrize("serving_mode", ["agg", "disagg"]) def test_run_synthetic_trace_replay_smoke_matrix( - tmp_path, engine_type, replay_mode, router_mode + tmp_path, engine_type, replay_mode, router_mode, serving_mode ): - args_path = _vllm_args() if engine_type == "vllm" else _sglang_args() - num_workers = 1 if router_mode == "round_robin" else 2 - - report = run_synthetic_trace_replay( - 64, - 2, - 2, - extra_engine_args=args_path, - num_workers=num_workers, - replay_mode=replay_mode, - router_mode=router_mode, - arrival_interval_ms=5.0, - ) + if serving_mode == "disagg": + if replay_mode != "offline": + pytest.skip("disagg replay only supports offline mode") + report = run_synthetic_trace_replay( + 64, + 2, + 2, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode=replay_mode, + router_mode=router_mode, + arrival_interval_ms=5.0, + ) + else: + args_path = _vllm_args() if engine_type == "vllm" else _sglang_args() + num_workers = 1 if router_mode == "round_robin" else 2 + report = run_synthetic_trace_replay( + 64, + 2, + 2, + extra_engine_args=args_path, + num_workers=num_workers, + replay_mode=replay_mode, + router_mode=router_mode, + arrival_interval_ms=5.0, + ) _assert_basic_report_counts( report, @@ -553,6 +627,103 @@ def test_run_trace_replay_accepts_partial_extra_engine_args_json(tmp_path, repla ) +@pytest.mark.parametrize("router_mode", ["round_robin", "kv_router"]) +def test_run_trace_replay_supports_disagg_offline(tmp_path, router_mode): + trace_path = _write_trace_and_args(tmp_path) + + report = run_trace_replay( + trace_path, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode="offline", + router_mode=router_mode, + ) + + _assert_basic_report_counts( + report, + num_requests=2, + input_tokens=64, + output_tokens=2, + ) + _assert_basic_report_metrics(report) + + +@pytest.mark.parametrize("router_mode", ["round_robin", "kv_router"]) +def test_run_synthetic_trace_replay_disagg_preserves_expected_output_tokens( + router_mode, +): + report = run_synthetic_trace_replay( + 128, + 7, + 6, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode="offline", + router_mode=router_mode, + ) + + _assert_basic_report_counts( + report, + num_requests=6, + input_tokens=128, + output_tokens=7, + ) + _assert_basic_report_metrics(report) + + +def test_run_trace_replay_rejects_partial_disagg_args(tmp_path): + trace_path = _write_trace_and_args(tmp_path) + + with pytest.raises(Exception, match="must be provided together"): + run_trace_replay( + trace_path, + prefill_engine_args=_prefill_args(), + replay_mode="offline", + router_mode="kv_router", + ) + + +def test_run_trace_replay_rejects_online_disagg(tmp_path): + trace_path = _write_trace_and_args(tmp_path) + + with pytest.raises( + Exception, match="disagg replay only supports replay_mode='offline'" + ): + run_trace_replay( + trace_path, + prefill_engine_args=_prefill_args(), + decode_engine_args=_decode_args(), + router_config=_router_config(), + num_prefill_workers=2, + num_decode_workers=2, + replay_mode="online", + router_mode="kv_router", + ) + + +def test_run_trace_replay_rejects_disagg_worker_counts_for_aggregated_mode(tmp_path): + trace_path = _write_trace_and_args(tmp_path) + + with pytest.raises( + Exception, + match="num_prefill_workers and num_decode_workers are only used for disagg replay", + ): + run_trace_replay( + trace_path, + extra_engine_args=MockEngineArgs(block_size=64, speedup_ratio=1000.0), + num_workers=1, + num_prefill_workers=2, + num_decode_workers=2, + replay_mode="offline", + ) + + def test_format_report_table_matches_aiperf_shape(): report = { "mean_ttft_ms": 18.26, @@ -616,91 +787,49 @@ def test_write_report_json_creates_file(tmp_path): ) -def test_replay_cli_prints_table_and_saves_json(tmp_path, monkeypatch, capsys): - report = { - "mean_ttft_ms": 10.0, - "min_ttft_ms": 9.0, - "max_ttft_ms": 12.0, - "p99_ttft_ms": 12.0, - "p90_ttft_ms": 11.0, - "p75_ttft_ms": 10.5, - "std_ttft_ms": 1.0, - "output_throughput_tok_s": 123.0, - "request_throughput_rps": 4.0, - "completed_requests": 3, - } - - def fake_run(*args, **kwargs): - return report +@pytest.mark.timeout(30) +def test_replay_cli_subprocess_synthetic_smoke(tmp_path): + report_path = tmp_path / "synthetic_report.json" - monkeypatch.setattr("dynamo.replay.main.run_synthetic_trace_replay", fake_run) - report_path = tmp_path / "cli_report.json" + completed = _run_replay_cli( + tmp_path, + "--input-tokens", + "250", + "--output-tokens", + "25", + "--request-count", + "10", + "--num-workers", + "4", + "--replay-concurrency", + "4", + "--report-json", + str(report_path), + "--extra-engine-args", + '{"block_size":64,"speedup_ratio":1000.0}', + ) - exit_code = main( - [ - "--input-tokens", - "16", - "--output-tokens", - "8", - "--request-count", - "3", - "--report-json", - str(report_path), - ] - ) - - assert exit_code == 0 - stdout = capsys.readouterr().out - assert "NVIDIA AIPerf | LLM Metrics" in stdout - assert "Saved full report to:" in stdout - assert '"completed_requests"' not in stdout - assert json.loads(report_path.read_text(encoding="utf-8")) == report - - -def test_replay_cli_passes_multiturn_workload_kwargs(monkeypatch): - captured = {} - - def fake_run(*args, **kwargs): - captured["args"] = args - captured["kwargs"] = kwargs - return { - "completed_requests": 4, - "request_throughput_rps": 1.0, - "output_throughput_tok_s": 1.0, - } - - monkeypatch.setattr("dynamo.replay.main.run_synthetic_trace_replay", fake_run) - - exit_code = main( - [ - "--input-tokens", - "16", - "--output-tokens", - "8", - "--request-count", - "2", - "--turns-per-session", - "2", - "--shared-prefix-ratio", - "0.5", - "--num-prefix-groups", - "3", - "--inter-turn-delay-ms", - "7.0", - ] - ) - - assert exit_code == 0 - assert captured["args"] == (16, 8, 2) - assert captured["kwargs"]["turns_per_session"] == 2 - assert captured["kwargs"]["shared_prefix_ratio"] == 0.5 - assert captured["kwargs"]["num_prefix_groups"] == 3 - assert captured["kwargs"]["inter_turn_delay_ms"] == 7.0 + report = _assert_replay_cli_outputs(completed, report_path) + _assert_basic_report_counts( + report, + num_requests=10, + input_tokens=250, + output_tokens=25, + ) + _assert_basic_report_metrics(report) @pytest.mark.timeout(30) -def test_replay_cli_subprocess_synthetic_smoke(tmp_path): - report_path = tmp_path / "synthetic_report.json" +@pytest.mark.parametrize("planner_profile_data_kind", ["dir", "npz"]) +def test_replay_cli_subprocess_synthetic_smoke_accepts_planner_profile_data( + tmp_path, planner_profile_data_kind +): + report_path = tmp_path / f"synthetic_report_{planner_profile_data_kind}.json" + planner_profile_data = ( + _planner_profile_data_dir_path() + if planner_profile_data_kind == "dir" + else _write_planner_profile_data_npz(tmp_path) + ) completed = _run_replay_cli( tmp_path, @@ -717,7 +846,13 @@ def test_replay_cli_subprocess_synthetic_smoke(tmp_path): "--report-json", str(report_path), "--extra-engine-args", - '{"block_size":64,"speedup_ratio":1000.0}', + json.dumps( + { + "block_size": 64, + "speedup_ratio": 1000.0, + "planner_profile_data": str(planner_profile_data), + } + ), ) report = _assert_replay_cli_outputs(completed, report_path) @@ -798,6 +933,40 @@ def test_replay_cli_subprocess_trace_smoke(tmp_path): _assert_basic_report_metrics(report) +@pytest.mark.timeout(30) +def test_replay_cli_subprocess_trace_disagg_smoke(tmp_path): + trace_path = _write_cli_smoke_trace(tmp_path) + report_path = tmp_path / "trace_disagg_report.json" + + completed = _run_replay_cli( + tmp_path, + str(trace_path), + "--replay-mode", + "offline", + "--router-mode", + "kv_router", + "--num-prefill-workers", + "2", + "--num-decode-workers", + "2", + "--report-json", + str(report_path), + "--prefill-engine-args", + '{"block_size":64,"speedup_ratio":1000.0,"worker_type":"prefill"}', + "--decode-engine-args", + '{"block_size":64,"speedup_ratio":1000.0,"worker_type":"decode"}', + ) + + report = _assert_replay_cli_outputs(completed, report_path) + _assert_basic_report_counts( + report, + num_requests=10, + input_tokens=250, + output_tokens=25, + ) + _assert_basic_report_metrics(report) + + @pytest.mark.timeout(30) def test_replay_cli_subprocess_multiturn_trace_smoke(tmp_path): trace_path = _write_multiturn_trace(tmp_path) diff --git a/lib/kv-router/src/protocols.rs b/lib/kv-router/src/protocols.rs index daea6793aa43..a964f4b8c6ad 100644 --- a/lib/kv-router/src/protocols.rs +++ b/lib/kv-router/src/protocols.rs @@ -8,6 +8,10 @@ use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use xxhash_rust::xxh3; +const fn default_track_prefill_tokens() -> bool { + true +} + /// The event subject that workers publish KV cache events on. pub const KV_EVENT_SUBJECT: &str = "kv-events"; @@ -431,6 +435,8 @@ pub enum ActiveSequenceEventData { token_sequence: Option>, isl: usize, overlap: u32, + #[serde(default = "default_track_prefill_tokens")] + track_prefill_tokens: bool, expected_output_tokens: Option, }, Free, @@ -990,14 +996,6 @@ mod tests { assert_ne!(lora_a[0], lora_b[0]); } - #[test] - fn test_lora_name_none_matches_legacy() { - let tokens: Vec = (0..8).collect(); - let hashes_none = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); - let hashes_none2 = compute_block_hash_for_seq(&tokens, 4, BlockHashOptions::default()); - assert_eq!(hashes_none, hashes_none2); - } - #[test] fn test_lora_name_empty_string_normalized_to_none() { let tokens: Vec = (0..4).collect(); @@ -1172,16 +1170,6 @@ mod tests { assert_eq!(deserialized.block_hashes[1].0, 5); } - #[test] - fn test_router_request_mark_free_backwards_compatible_deserialization() { - let request: RouterRequest = serde_json::from_str(r#"{"method":"mark_free"}"#).unwrap(); - - assert!(matches!( - request, - RouterRequest::MarkFree { request_id: None } - )); - } - #[test] fn test_router_request_mark_free_serialization_with_request_id() { let request = RouterRequest::MarkFree { diff --git a/lib/kv-router/src/scheduling/config.rs b/lib/kv-router/src/scheduling/config.rs index bcd383b47f4b..f42a1a1ea297 100644 --- a/lib/kv-router/src/scheduling/config.rs +++ b/lib/kv-router/src/scheduling/config.rs @@ -17,6 +17,10 @@ const fn default_min_initial_workers() -> usize { 1 } +const fn default_track_prefill_tokens() -> bool { + true +} + #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "lowercase")] pub enum RouterQueuePolicy { @@ -63,6 +67,9 @@ pub struct RouterConfigOverride { #[builder(default)] pub assume_kv_reuse: Option, + + #[builder(default)] + pub track_prefill_tokens: Option, } /// KV Router configuration parameters @@ -98,6 +105,12 @@ pub struct KvRouterConfig { /// When false, generates random hashes (assuming no KV cache reuse). pub router_assume_kv_reuse: bool, + /// Whether to include prompt-side prefill tokens in active load accounting (default: true). + /// When false, prompt tokens are excluded from active prefill token tracking, queue pressure, + /// and potential prefill-token load calculations. + #[serde(default = "default_track_prefill_tokens")] + pub router_track_prefill_tokens: bool, + /// Threshold for triggering snapshots. If None, no snapshots will be performed. #[validate(range(min = 1))] pub router_snapshot_threshold: Option, @@ -171,6 +184,7 @@ impl Default for KvRouterConfig { router_track_active_blocks: true, router_track_output_blocks: false, router_assume_kv_reuse: true, + router_track_prefill_tokens: default_track_prefill_tokens(), router_snapshot_threshold: Some(1000000), router_reset_states: false, router_ttl_secs: 120.0, @@ -208,6 +222,18 @@ fn validate_kv_router_config(config: &KvRouterConfig) -> Result<(), ValidationEr } impl KvRouterConfig { + pub fn assume_kv_reuse(&self, config_override: Option<&RouterConfigOverride>) -> bool { + config_override + .and_then(|cfg| cfg.assume_kv_reuse) + .unwrap_or(self.router_assume_kv_reuse) + } + + pub fn track_prefill_tokens(&self, config_override: Option<&RouterConfigOverride>) -> bool { + config_override + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.router_track_prefill_tokens) + } + /// Compute sequence hashes for active block tracking based on configuration. /// /// Returns: @@ -231,9 +257,7 @@ impl KvRouterConfig { return Some(Vec::new()); } - let assume_kv_reuse = config_override - .and_then(|cfg| cfg.assume_kv_reuse) - .unwrap_or(self.router_assume_kv_reuse); + let assume_kv_reuse = self.assume_kv_reuse(config_override); if assume_kv_reuse { let block_hashes = match precomputed_block_hashes { @@ -290,6 +314,11 @@ mod tests { assert_eq!(KvRouterConfig::default().min_initial_workers, 1); } + #[test] + fn kv_router_config_defaults_to_tracking_prefill_tokens() { + assert!(KvRouterConfig::default().router_track_prefill_tokens); + } + #[test] fn kv_router_config_rejects_zero_initial_workers() { let cfg = KvRouterConfig { @@ -332,6 +361,17 @@ mod tests { assert_ne!(without_mm, with_mm); } + #[test] + fn router_config_override_serde_round_trip_preserves_track_prefill_tokens() { + let serialized = serde_json::to_string(&RouterConfigOverride { + track_prefill_tokens: Some(false), + ..Default::default() + }) + .unwrap(); + let deserialized: RouterConfigOverride = serde_json::from_str(&serialized).unwrap(); + assert_eq!(deserialized.track_prefill_tokens, Some(false)); + } + #[test] fn compute_seq_hashes_for_tracking_uses_precomputed_block_hashes() { let config = KvRouterConfig::default(); diff --git a/lib/kv-router/src/scheduling/local.rs b/lib/kv-router/src/scheduling/local.rs index 2304e5a7d27f..4da0291cb6db 100644 --- a/lib/kv-router/src/scheduling/local.rs +++ b/lib/kv-router/src/scheduling/local.rs @@ -30,6 +30,7 @@ where request_tx: mpsc::Sender, slots: Arc>, queue: Arc>, + track_prefill_tokens_default: bool, worker_type: &'static str, } @@ -48,6 +49,7 @@ where block_size: u32, selector: Sel, policy: S, + track_prefill_tokens_default: bool, cancellation_token: CancellationToken, worker_type: &'static str, monitor_worker_configs: bool, @@ -135,6 +137,7 @@ where request_tx, slots, queue, + track_prefill_tokens_default, worker_type, } } @@ -154,6 +157,9 @@ where allowed_worker_ids: Option>, ) -> Result { let (resp_tx, resp_rx) = tokio::sync::oneshot::channel(); + let track_prefill_tokens = router_config_override + .and_then(|cfg| cfg.track_prefill_tokens) + .unwrap_or(self.track_prefill_tokens_default); let request = SchedulingRequest { maybe_request_id, token_seq, @@ -161,6 +167,7 @@ where overlaps, decode_blocks: HashMap::new(), prefill_tokens: HashMap::new(), + track_prefill_tokens, router_config_override: router_config_override.cloned(), update_states, lora_name, @@ -224,10 +231,16 @@ where token_seq: Option>, isl_tokens: usize, overlaps: OverlapScores, + track_prefill_tokens: bool, ) -> Vec { - let (decode_blocks, prefill_tokens) = - self.slots - .potential_blocks_and_tokens(token_seq.as_deref(), isl_tokens, overlaps); + let (decode_blocks, prefill_tokens) = self + .slots + .potential_blocks_and_tokens_with_prefill_tracking( + token_seq.as_deref(), + isl_tokens, + overlaps, + track_prefill_tokens, + ); let mut workers: HashSet = HashSet::new(); workers.extend(decode_blocks.keys().copied()); @@ -300,6 +313,7 @@ mod tests { 64, DefaultWorkerSelector::new(None, "test"), FcfsPolicy, + true, cancel_token.clone(), "test", monitor_worker_configs, @@ -344,6 +358,48 @@ mod tests { cancel_token.cancel(); } + #[tokio::test] + async fn test_schedule_override_can_disable_prefill_tracking() { + let mut workers = HashMap::new(); + workers.insert( + 0, + SimpleWorkerConfig { + max_num_batched_tokens: Some(64), + ..Default::default() + }, + ); + let (scheduler, slots, _cfg_tx, cancel_token) = make_scheduler(workers, None, true); + + scheduler + .schedule( + Some("req-1".to_string()), + 64, + Some(vec![1, 2, 3, 4]), + OverlapScores::default(), + Some(&crate::config::RouterConfigOverride { + track_prefill_tokens: Some(false), + ..Default::default() + }), + true, + None, + 0.0, + None, + None, + ) + .await + .unwrap(); + + assert_eq!( + slots + .active_tokens() + .get(&WorkerWithDpRank::new(0, 0)) + .copied(), + Some(0) + ); + + cancel_token.cancel(); + } + #[tokio::test] async fn test_mark_prefill_completed_drains_pending_queue() { let mut workers = HashMap::new(); @@ -474,7 +530,7 @@ mod tests { .collect(); expected.sort_by_key(|load| (load.worker_id, load.dp_rank)); - let mut actual = scheduler.get_potential_loads(Some(token_seq), 128, overlaps); + let mut actual = scheduler.get_potential_loads(Some(token_seq), 128, overlaps, true); actual.sort_by_key(|load| (load.worker_id, load.dp_rank)); assert_eq!(actual.len(), expected.len()); @@ -500,7 +556,7 @@ mod tests { make_scheduler(HashMap::new(), None, false); scheduler.register_workers(&HashSet::from([42])); - let loads = scheduler.get_potential_loads(None, 64, OverlapScores::default()); + let loads = scheduler.get_potential_loads(None, 64, OverlapScores::default(), true); assert_eq!(loads.len(), 1); assert_eq!(loads[0].worker_id, 42); @@ -517,7 +573,7 @@ mod tests { assert_eq!( scheduler - .get_potential_loads(None, 64, OverlapScores::default()) + .get_potential_loads(None, 64, OverlapScores::default(), true) .len(), 1 ); @@ -536,7 +592,7 @@ mod tests { tokio::time::timeout(Duration::from_secs(1), async { loop { if scheduler - .get_potential_loads(None, 64, OverlapScores::default()) + .get_potential_loads(None, 64, OverlapScores::default(), true) .len() == 3 { @@ -550,4 +606,39 @@ mod tests { cancel_token.cancel(); } + + #[tokio::test] + async fn test_get_potential_loads_can_ignore_prefill_tokens() { + let mut workers = HashMap::new(); + workers.insert( + 0, + SimpleWorkerConfig { + max_num_batched_tokens: Some(256), + ..Default::default() + }, + ); + let (scheduler, _slots, _cfg_tx, cancel_token) = make_scheduler(workers, None, true); + + scheduler + .schedule( + Some("req-1".to_string()), + 64, + Some(vec![11, 22]), + OverlapScores::default(), + None, + true, + None, + 0.0, + None, + None, + ) + .await + .unwrap(); + + let loads = scheduler.get_potential_loads(None, 64, OverlapScores::default(), false); + assert_eq!(loads.len(), 1); + assert_eq!(loads[0].potential_prefill_tokens, 64); + + cancel_token.cancel(); + } } diff --git a/lib/kv-router/src/scheduling/policy.rs b/lib/kv-router/src/scheduling/policy.rs index 33548ee75c22..8a613da17469 100644 --- a/lib/kv-router/src/scheduling/policy.rs +++ b/lib/kv-router/src/scheduling/policy.rs @@ -135,6 +135,7 @@ mod tests { overlaps, decode_blocks: HashMap::new(), prefill_tokens: HashMap::new(), + track_prefill_tokens: true, router_config_override: None, update_states: false, lora_name: None, diff --git a/lib/kv-router/src/scheduling/queue.rs b/lib/kv-router/src/scheduling/queue.rs index 663a7842f921..93334f6e1061 100644 --- a/lib/kv-router/src/scheduling/queue.rs +++ b/lib/kv-router/src/scheduling/queue.rs @@ -191,11 +191,14 @@ impl< /// Run the full scheduling pipeline for a single request: /// compute potential load -> select worker -> respond -> book via add_request. async fn schedule(&self, mut request: SchedulingRequest) { - let (decode_blocks, prefill_tokens) = self.slots.potential_blocks_and_tokens( - request.token_seq.as_deref(), - request.isl_tokens, - request.overlaps.clone(), - ); + let (decode_blocks, prefill_tokens) = self + .slots + .potential_blocks_and_tokens_with_prefill_tracking( + request.token_seq.as_deref(), + request.isl_tokens, + request.overlaps.clone(), + request.track_prefill_tokens, + ); request.decode_blocks = decode_blocks; request.prefill_tokens = prefill_tokens; @@ -235,6 +238,7 @@ impl< token_sequence: request.token_seq, isl: request.isl_tokens, overlap: selection.overlap_blocks, + track_prefill_tokens: request.track_prefill_tokens, expected_output_tokens: request.expected_output_tokens, worker: selection.worker, lora_name: request.lora_name.clone(), @@ -376,6 +380,7 @@ mod tests { overlaps: OverlapScores::default(), decode_blocks: HashMap::new(), prefill_tokens: HashMap::new(), + track_prefill_tokens: true, router_config_override: None, update_states: true, lora_name: None, @@ -695,6 +700,7 @@ mod tests { overlaps: OverlapScores::default(), decode_blocks: HashMap::new(), prefill_tokens: HashMap::new(), + track_prefill_tokens: true, router_config_override: None, update_states: true, lora_name: None, @@ -719,4 +725,31 @@ mod tests { .unwrap(); slots.free(&"filter-0".to_string()).await.unwrap(); } + + #[tokio::test(flavor = "multi_thread")] + async fn test_queue_busy_check_ignores_untracked_prefill_tokens() { + let (queue, slots) = make_queue(1, 16, 256, Some(0.0)); + + let (mut req1, rx1) = make_request("req-1", 256); + req1.track_prefill_tokens = false; + queue.enqueue(req1).await; + let _resp1 = rx1.await.unwrap().unwrap(); + assert_eq!( + slots + .active_tokens() + .get(&WorkerWithDpRank::new(0, 0)) + .copied(), + Some(0) + ); + + let (req2, rx2) = make_request("req-2", 256); + queue.enqueue(req2).await; + let _resp2 = rx2.await.unwrap().unwrap(); + assert_eq!(queue.pending_count(), 0); + + let _ = slots.mark_prefill_completed(&"req-1".to_string()).await; + let _ = slots.free(&"req-1".to_string()).await; + let _ = slots.mark_prefill_completed(&"req-2".to_string()).await; + let _ = slots.free(&"req-2".to_string()).await; + } } diff --git a/lib/kv-router/src/scheduling/types.rs b/lib/kv-router/src/scheduling/types.rs index 9fe246601bc3..ce6bc2114fe7 100644 --- a/lib/kv-router/src/scheduling/types.rs +++ b/lib/kv-router/src/scheduling/types.rs @@ -42,6 +42,7 @@ pub struct SchedulingRequest { pub overlaps: OverlapScores, pub decode_blocks: HashMap, pub prefill_tokens: HashMap, + pub track_prefill_tokens: bool, pub router_config_override: Option, pub update_states: bool, pub lora_name: Option, diff --git a/lib/kv-router/src/sequences/multi_worker.rs b/lib/kv-router/src/sequences/multi_worker.rs index 50e496dbde73..3ef51298fe9e 100644 --- a/lib/kv-router/src/sequences/multi_worker.rs +++ b/lib/kv-router/src/sequences/multi_worker.rs @@ -97,6 +97,7 @@ pub struct SequenceRequest { pub token_sequence: Option>, pub isl: usize, pub overlap: u32, + pub track_prefill_tokens: bool, pub expected_output_tokens: Option, pub worker: WorkerWithDpRank, pub lora_name: Option, @@ -221,6 +222,7 @@ impl ActiveSequencesMultiWorker

{ token_sequence, isl, overlap, + track_prefill_tokens, expected_output_tokens, } => { self.request_to_worker @@ -233,12 +235,13 @@ impl ActiveSequencesMultiWorker

{ let table = self.workers.read(); if let Some(&idx) = table.index.get(&event.worker) { - table.slots[idx].1.write().add_request( + table.slots[idx].1.write().add_request_with_prefill_tracking( event.request_id.clone(), token_sequence.clone(), *isl, *overlap, *expected_output_tokens, + *track_prefill_tokens, ); } else { tracing::warn!( @@ -380,6 +383,7 @@ impl ActiveSequencesMultiWorker

{ token_sequence, isl, overlap, + track_prefill_tokens, expected_output_tokens, worker, lora_name, @@ -409,12 +413,13 @@ impl ActiveSequencesMultiWorker

{ .get(&worker) .ok_or(SequenceError::WorkerNotFound { worker })?; let mut seq = table.slots[idx].1.write(); - seq.add_request( + seq.add_request_with_prefill_tracking( request_id, token_sequence, isl, overlap, expected_output_tokens, + track_prefill_tokens, ) }; @@ -437,6 +442,7 @@ impl ActiveSequencesMultiWorker

{ token_sequence: req.token_sequence.clone(), isl: req.isl, overlap: req.overlap, + track_prefill_tokens: req.track_prefill_tokens, expected_output_tokens: req.expected_output_tokens, }, router_id: self.router_id, @@ -527,6 +533,10 @@ impl ActiveSequencesMultiWorker

{ /// /// Note: This operation is idempotent. Calling it multiple times for the same request /// will log a warning but not return an error (double free is allowed). + /// + /// This also performs the underlying prefill-complete cleanup via + /// [`ActiveSequences::free`], so callers do not need to call + /// [`Self::mark_prefill_completed`] before freeing a completed request. pub async fn free(&self, request_id: &RequestId) -> Result<(), SequenceError> { if !self.request_to_worker.contains_key(request_id) { tracing::debug!("Request {request_id} not found, already freed (idempotent)"); @@ -696,6 +706,19 @@ impl ActiveSequencesMultiWorker

{ ) -> ( HashMap, HashMap, + ) { + self.potential_blocks_and_tokens_with_prefill_tracking(token_sequence, isl, overlaps, true) + } + + pub fn potential_blocks_and_tokens_with_prefill_tracking( + &self, + token_sequence: Option<&[SequenceHash]>, + isl: usize, + overlaps: OverlapScores, + track_prefill_tokens: bool, + ) -> ( + HashMap, + HashMap, ) { #[cfg(feature = "bench")] let start = tokio::time::Instant::now(); @@ -711,9 +734,14 @@ impl ActiveSequencesMultiWorker

{ for (worker, lock) in &table.slots { let overlap = *overlaps.scores.get(worker).unwrap_or(&0); - let (blocks, tokens) = - lock.read() - .potential_blocks_and_tokens(token_sequence, isl, overlap); + let (blocks, tokens) = lock + .read() + .potential_blocks_and_tokens_with_prefill_tracking( + token_sequence, + isl, + overlap, + track_prefill_tokens, + ); potential_blocks.insert(*worker, blocks); potential_tokens.insert(*worker, tokens); } @@ -832,3 +860,44 @@ impl ActiveSequencesMultiWorker

{ }); } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use super::*; + use crate::test_utils::NoopSequencePublisher; + + fn make_sequences() -> ActiveSequencesMultiWorker { + ActiveSequencesMultiWorker::new( + NoopSequencePublisher, + 4, + HashMap::from([(1_u64, (0_u32, 1_u32))]), + false, + 0, + "test", + ) + } + + #[tokio::test] + async fn add_request_can_skip_prefill_token_tracking() { + let sequences = make_sequences(); + let worker = WorkerWithDpRank::new(1, 0); + + sequences + .add_request(SequenceRequest { + request_id: "req-1".to_string(), + token_sequence: Some(vec![1, 2, 3]), + isl: 12, + overlap: 0, + track_prefill_tokens: false, + expected_output_tokens: None, + worker, + lora_name: None, + }) + .await + .unwrap(); + + assert_eq!(sequences.active_tokens().get(&worker).copied(), Some(0)); + } +} diff --git a/lib/kv-router/src/sequences/single.rs b/lib/kv-router/src/sequences/single.rs index 99116eb298b0..f74116a1b192 100644 --- a/lib/kv-router/src/sequences/single.rs +++ b/lib/kv-router/src/sequences/single.rs @@ -143,6 +143,27 @@ impl ActiveSequences { isl: usize, overlap: u32, expected_output_tokens: Option, + ) -> HashSet { + self.add_request_with_prefill_tracking( + request_id, + token_sequence, + isl, + overlap, + expected_output_tokens, + true, + ) + } + + /// Add a new request with optional prompt-token load accounting. + /// Returns the set of expired request IDs that were removed during cleanup. + pub fn add_request_with_prefill_tracking( + &mut self, + request_id: RequestId, + token_sequence: Option>, + isl: usize, + overlap: u32, + expected_output_tokens: Option, + track_prefill_tokens: bool, ) -> HashSet { // Check for double-add and log error, returning early if self.active_seqs.contains_key(&request_id) { @@ -153,7 +174,11 @@ impl ActiveSequences { // Lazily check and clean up expired requests, capturing removed IDs let removed_requests = self.force_expiry(); - let prefill_tokens = self.new_tokens(isl, overlap); + let prefill_tokens = if track_prefill_tokens { + self.new_tokens(isl, overlap) + } else { + 0 + }; self.prefill_tokens .insert(request_id.clone(), prefill_tokens); self.active_tokens += prefill_tokens; @@ -208,13 +233,27 @@ impl ActiveSequences { token_sequence: Option<&[SequenceHash]>, isl: usize, overlap: u32, + ) -> (usize, usize) { + self.potential_blocks_and_tokens_with_prefill_tracking(token_sequence, isl, overlap, true) + } + + pub fn potential_blocks_and_tokens_with_prefill_tracking( + &self, + token_sequence: Option<&[SequenceHash]>, + isl: usize, + overlap: u32, + track_prefill_tokens: bool, ) -> (usize, usize) { let potential_blocks = if let Some(token_seq) = token_sequence { self.new_blocks(token_seq) + self.active_blocks() } else { self.active_blocks() }; - let potential_tokens = self.new_tokens(isl, overlap) + self.active_tokens; + let potential_tokens = if track_prefill_tokens { + self.new_tokens(isl, overlap) + self.active_tokens + } else { + self.active_tokens + }; (potential_blocks, potential_tokens) } @@ -232,7 +271,10 @@ impl ActiveSequences { self.new_blocks(token_sequence) + self.active_blocks() } - /// Free all blocks associated with a request + /// Free all blocks associated with a request. + /// + /// This implicitly calls [`Self::mark_prefill_completed`] first, so callers do not need + /// to invoke both when the request is finishing. pub fn free(&mut self, request_id: &RequestId) -> usize { self.mark_prefill_completed(request_id); @@ -424,6 +466,48 @@ mod tests { assert_eq!(seq_manager.active_tokens(), 0); } + #[test] + fn test_add_request_without_prefill_tracking_keeps_active_tokens_zero() { + let mut seq_manager = ActiveSequences::new(4); + + seq_manager.add_request_with_prefill_tracking( + "r1".to_string(), + Some(vec![1, 2, 3]), + 12, + 0, + None, + false, + ); + + assert_eq!(seq_manager.active_tokens(), 0); + seq_manager.mark_prefill_completed(&"r1".to_string()); + assert_eq!(seq_manager.active_tokens(), 0); + seq_manager.free(&"r1".to_string()); + assert_eq!(seq_manager.active_blocks(), 0); + } + + #[test] + fn test_potential_blocks_and_tokens_without_prefill_tracking_ignores_prompt_load() { + let mut seq_manager = ActiveSequences::new(4); + seq_manager.add_request_with_prefill_tracking( + "r1".to_string(), + Some(vec![1, 2, 3]), + 12, + 0, + None, + false, + ); + + let (blocks, tokens) = seq_manager.potential_blocks_and_tokens_with_prefill_tracking( + Some(&[1, 2, 3, 4]), + 16, + 0, + false, + ); + assert_eq!(blocks, 4); + assert_eq!(tokens, 0); + } + #[tokio::test(start_paused = true)] async fn test_force_expiry() { let block_size = 4; diff --git a/lib/llm/src/kv_router.rs b/lib/llm/src/kv_router.rs index 88c192856edc..b8ec84b6c632 100644 --- a/lib/llm/src/kv_router.rs +++ b/lib/llm/src/kv_router.rs @@ -1,20 +1,16 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use std::sync::Arc; -use std::time::{Duration, Instant}; +use std::time::Instant; use anyhow::Result; use dynamo_kv_router::{ - ConcurrentRadixTree, ThreadPoolIndexer, - approx::PruneConfig, config::{KvRouterConfig, RouterConfigOverride}, - indexer::{GetWorkersRequest, KvIndexer, KvIndexerInterface, KvIndexerMetrics, KvRouterError}, + indexer::KvRouterError, protocols::KV_EVENT_SUBJECT, protocols::{ - BlockExtraInfo, BlockHashOptions, DpRank, LocalBlockHash, OverlapScores, RouterEvent, - RouterRequest, RouterResponse, TokensWithHashes, WorkerId, WorkerWithDpRank, - compute_block_hash_for_seq, + BlockExtraInfo, BlockHashOptions, DpRank, RouterEvent, RouterRequest, RouterResponse, + TokensWithHashes, WorkerId, WorkerWithDpRank, compute_block_hash_for_seq, }, }; use dynamo_runtime::{ @@ -29,30 +25,29 @@ use dynamo_runtime::{ traits::DistributedRuntimeProvider, }; use futures::stream; -use tokio::sync::oneshot; use tracing::Instrument; use validator::Validate; pub mod cache_control; +pub mod indexer; mod jetstream; pub mod metrics; pub mod prefill_router; pub mod publisher; pub mod push_router; -pub mod remote_indexer; pub mod scheduler; pub mod sequence; pub mod subscriber; pub mod worker_query; pub use cache_control::{CacheControlClient, spawn_pin_prefix}; +pub use indexer::Indexer; pub use prefill_router::PrefillRouter; pub use push_router::{DirectRoutingRouter, KvPushRouter}; use crate::{ discovery::RuntimeConfigWatch, kv_router::{ - remote_indexer::RemoteIndexer, scheduler::{DefaultWorkerSelector, KvScheduler, PotentialLoad}, sequence::{SequenceError, SequenceRequest}, }, @@ -108,188 +103,6 @@ pub fn router_discovery_query(namespace: String, component: String) -> Discovery } } -#[derive(Clone)] -pub enum Indexer { - /// Single-threaded radix tree with channel-based event processing. - /// Supports TTL-based expiration and size-based pruning. - /// Has the ability to persist and snapshot states. - KvIndexer(KvIndexer), - - /// Concurrent radix tree with a thread pool for event processing. - /// Uses sticky worker routing for per-worker event serialization. - /// Does not support TTL/pruning. - Concurrent(Arc>), - - /// Forwards queries to a standalone KV indexer service via the request plane. - /// The standalone indexer manages its own radix tree and event subscription. - Remote(Arc), - - /// Used when we do not wish to use the indexer at all (e.g., when overlap_score_weight is 0). - /// Note: This will cause KV events to accumulate in JetStream as we do not regularly purge them. - None, -} - -impl Indexer { - pub async fn new( - component: &dynamo_runtime::component::Component, - kv_router_config: &KvRouterConfig, - block_size: u32, - model_name: Option, - ) -> Result { - if kv_router_config.overlap_score_weight == 0.0 { - return Ok(Indexer::None); - } - - // Remote indexer: forward queries to a standalone KV indexer service. - if let Some(ref indexer_component_name) = kv_router_config.remote_indexer_component { - let model_name = model_name.ok_or_else(|| { - anyhow::anyhow!( - "model_name is required when remote_indexer_component is configured" - ) - })?; - tracing::info!( - remote_indexer_component = %indexer_component_name, - model_name, - "Using remote KV indexer" - ); - let remote = RemoteIndexer::new(component, indexer_component_name, model_name).await?; - return Ok(Indexer::Remote(Arc::new(remote))); - } - - // Approximate mode (--no-kv-events): always use single-threaded KvIndexer - // with TTL/pruning regardless of event_threads, since updates come from - // routing decisions only, not live KV events from workers. - if !kv_router_config.use_kv_events { - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - let cancellation_token = component.drt().primary_token(); - let prune_config = Some(PruneConfig { - ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs), - max_tree_size: kv_router_config.router_max_tree_size, - prune_target_ratio: kv_router_config.router_prune_target_ratio, - }); - return Ok(Indexer::KvIndexer(KvIndexer::new_with_frequency( - cancellation_token, - None, - block_size, - kv_indexer_metrics, - prune_config, - ))); - } - - if kv_router_config.router_event_threads > 1 { - return Ok(Indexer::Concurrent(Arc::new(ThreadPoolIndexer::new( - ConcurrentRadixTree::new(), - kv_router_config.router_event_threads as usize, - block_size, - )))); - } - - let kv_indexer_metrics = KvIndexerMetrics::from_component(component); - let cancellation_token = component.drt().primary_token(); - - Ok(Indexer::KvIndexer(KvIndexer::new_with_frequency( - cancellation_token, - None, // expiration_duration for frequency tracking - block_size, - kv_indexer_metrics, - None, - ))) - } - - pub(crate) async fn find_matches( - &self, - sequence: Vec, - ) -> Result { - match self { - Indexer::KvIndexer(indexer) => indexer.find_matches(sequence).await, - Indexer::Concurrent(tpi) => tpi.find_matches(sequence).await, - Indexer::Remote(remote) => remote.find_matches(sequence).await.map_err(|e| { - tracing::warn!(error = %e, "Remote indexer query failed"); - KvRouterError::IndexerOffline - }), - Indexer::None => Ok(OverlapScores::new()), - } - } - - pub(crate) async fn dump_events(&self) -> Result, KvRouterError> { - match self { - Indexer::KvIndexer(indexer) => indexer.dump_events().await, - Indexer::Concurrent(tpi) => tpi.dump_events().await, - Indexer::Remote(_) => Ok(Vec::new()), - Indexer::None => { - panic!( - "Cannot dump events: indexer does not exist (is overlap_score_weight set to 0?)" - ); - } - } - } - - pub(crate) async fn process_routing_decision_for_request( - &self, - tokens_with_hashes: &mut TokensWithHashes, - worker: WorkerWithDpRank, - ) -> Result<(), KvRouterError> { - match self { - Indexer::KvIndexer(indexer) => { - indexer - .process_routing_decision_for_request(tokens_with_hashes, worker) - .await - } - Indexer::Concurrent(tpi) => { - tpi.process_routing_decision_for_request(tokens_with_hashes, worker) - .await - } - Indexer::Remote(_) => Ok(()), - Indexer::None => Ok(()), - } - } - - pub(crate) async fn apply_event(&self, event: RouterEvent) { - match self { - Indexer::KvIndexer(indexer) => { - if let Err(e) = indexer.event_sender().send(event).await { - tracing::warn!("Failed to send event to indexer: {e}"); - } - } - Indexer::Concurrent(tpi) => tpi.apply_event(event).await, - Indexer::Remote(_) => {} // standalone indexer gets events directly - Indexer::None => {} - } - } - - pub(crate) async fn remove_worker(&self, worker_id: WorkerId) { - match self { - Indexer::KvIndexer(indexer) => { - if let Err(e) = indexer.remove_worker_sender().send(worker_id).await { - tracing::warn!("Failed to send worker removal for {worker_id}: {e}"); - } - } - Indexer::Concurrent(tpi) => { - KvIndexerInterface::remove_worker(tpi.as_ref(), worker_id).await; - } - Indexer::Remote(_) => {} // standalone indexer manages its own workers - Indexer::None => {} - } - } - - pub(crate) async fn get_workers(&self) -> Vec { - match self { - Indexer::KvIndexer(indexer) => { - let (resp_tx, resp_rx) = oneshot::channel(); - let req = GetWorkersRequest { resp: resp_tx }; - if let Err(e) = indexer.get_workers_sender().send(req).await { - tracing::warn!("Failed to send get_workers request: {e}"); - return Vec::new(); - } - resp_rx.await.unwrap_or_default() - } - Indexer::Concurrent(tpi) => tpi.backend().get_workers(), - Indexer::Remote(_) => Vec::new(), - Indexer::None => Vec::new(), - } - } -} - /// A KvRouter only decides which worker you should use. It doesn't send you there. /// TODO: Rename this to indicate it only selects a worker, it does not route. pub struct KvRouter @@ -529,6 +342,9 @@ where hash_options, None, ); + let track_prefill_tokens = self + .kv_router_config + .track_prefill_tokens(router_config_override); if let Err(e) = self .scheduler @@ -537,6 +353,7 @@ where token_sequence: maybe_seq_hashes, isl: isl_tokens, overlap: overlap_blocks, + track_prefill_tokens, expected_output_tokens, worker, lora_name, @@ -623,12 +440,17 @@ where hash_options, Some(&block_hashes), ); - + let track_prefill_tokens = self + .kv_router_config + .track_prefill_tokens(router_config_override); let overlap_scores = self.indexer.find_matches(block_hashes).await?; - Ok(self - .scheduler - .get_potential_loads(maybe_seq_hashes, isl_tokens, overlap_scores)) + Ok(self.scheduler.get_potential_loads( + maybe_seq_hashes, + isl_tokens, + overlap_scores, + track_prefill_tokens, + )) } /// Dump all events from the indexer diff --git a/lib/llm/src/kv_router/indexer.rs b/lib/llm/src/kv_router/indexer.rs new file mode 100644 index 000000000000..84794e2cf8a5 --- /dev/null +++ b/lib/llm/src/kv_router/indexer.rs @@ -0,0 +1,233 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::time::Duration; + +use anyhow::Result; +use futures::StreamExt; + +use dynamo_kv_router::{ + ConcurrentRadixTree, ThreadPoolIndexer, + approx::PruneConfig, + config::KvRouterConfig, + indexer::{ + IndexerQueryRequest, IndexerQueryResponse, KV_INDEXER_QUERY_ENDPOINT, KvIndexer, + KvIndexerInterface, KvIndexerMetrics, KvRouterError, + }, + protocols::{ + LocalBlockHash, OverlapScores, RouterEvent, TokensWithHashes, WorkerId, WorkerWithDpRank, + }, +}; +use dynamo_runtime::{ + component::Component, + pipeline::{ManyOut, RouterMode, SingleIn, network::egress::push_router::PushRouter}, + traits::DistributedRuntimeProvider, +}; +use tokio::sync::oneshot; + +pub struct RemoteIndexer { + router: PushRouter, + model_name: String, + namespace: String, +} + +impl RemoteIndexer { + async fn new( + component: &Component, + indexer_component_name: &str, + model_name: String, + ) -> Result { + let namespace = component.namespace().name(); + let indexer_ns = component.namespace(); + let indexer_component = indexer_ns.component(indexer_component_name)?; + let endpoint = indexer_component.endpoint(KV_INDEXER_QUERY_ENDPOINT); + let client = endpoint.client().await?; + let router = + PushRouter::from_client_no_fault_detection(client, RouterMode::RoundRobin).await?; + Ok(Self { + router, + model_name, + namespace, + }) + } + + async fn find_matches(&self, block_hashes: Vec) -> Result { + let request = IndexerQueryRequest { + model_name: self.model_name.clone(), + namespace: self.namespace.clone(), + block_hashes, + }; + let mut stream: ManyOut = + self.router.round_robin(SingleIn::new(request)).await?; + + match stream.next().await { + Some(IndexerQueryResponse::Scores(scores)) => Ok(scores.into()), + Some(IndexerQueryResponse::Error(msg)) => { + Err(anyhow::anyhow!("Remote indexer error: {}", msg)) + } + None => Err(anyhow::anyhow!("Remote indexer returned empty response")), + } + } +} + +#[derive(Clone)] +pub enum Indexer { + KvIndexer(KvIndexer), + Concurrent(Arc>), + Remote(Arc), + None, +} + +impl Indexer { + pub async fn new( + component: &Component, + kv_router_config: &KvRouterConfig, + block_size: u32, + model_name: Option, + ) -> Result { + if kv_router_config.overlap_score_weight == 0.0 { + return Ok(Self::None); + } + + if let Some(ref indexer_component_name) = kv_router_config.remote_indexer_component { + let model_name = model_name.ok_or_else(|| { + anyhow::anyhow!( + "model_name is required when remote_indexer_component is configured" + ) + })?; + tracing::info!( + remote_indexer_component = %indexer_component_name, + model_name, + "Using remote KV indexer" + ); + let remote = RemoteIndexer::new(component, indexer_component_name, model_name).await?; + return Ok(Self::Remote(Arc::new(remote))); + } + + if !kv_router_config.use_kv_events { + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + let cancellation_token = component.drt().primary_token(); + let prune_config = Some(PruneConfig { + ttl: Duration::from_secs_f64(kv_router_config.router_ttl_secs), + max_tree_size: kv_router_config.router_max_tree_size, + prune_target_ratio: kv_router_config.router_prune_target_ratio, + }); + return Ok(Self::KvIndexer(KvIndexer::new_with_frequency( + cancellation_token, + None, + block_size, + kv_indexer_metrics, + prune_config, + ))); + } + + if kv_router_config.router_event_threads > 1 { + return Ok(Self::Concurrent(Arc::new(ThreadPoolIndexer::new( + ConcurrentRadixTree::new(), + kv_router_config.router_event_threads as usize, + block_size, + )))); + } + + let kv_indexer_metrics = KvIndexerMetrics::from_component(component); + let cancellation_token = component.drt().primary_token(); + + Ok(Self::KvIndexer(KvIndexer::new_with_frequency( + cancellation_token, + None, + block_size, + kv_indexer_metrics, + None, + ))) + } + + pub(crate) async fn find_matches( + &self, + sequence: Vec, + ) -> Result { + match self { + Self::KvIndexer(indexer) => indexer.find_matches(sequence).await, + Self::Concurrent(tpi) => tpi.find_matches(sequence).await, + Self::Remote(remote) => remote.find_matches(sequence).await.map_err(|e| { + tracing::warn!(error = %e, "Remote indexer query failed"); + KvRouterError::IndexerOffline + }), + Self::None => Ok(OverlapScores::new()), + } + } + + pub(crate) async fn dump_events(&self) -> Result, KvRouterError> { + match self { + Self::KvIndexer(indexer) => indexer.dump_events().await, + Self::Concurrent(tpi) => tpi.dump_events().await, + Self::Remote(_) => Ok(Vec::new()), + Self::None => { + panic!( + "Cannot dump events: indexer does not exist (is overlap_score_weight set to 0?)" + ); + } + } + } + + pub(crate) async fn process_routing_decision_for_request( + &self, + tokens_with_hashes: &mut TokensWithHashes, + worker: WorkerWithDpRank, + ) -> Result<(), KvRouterError> { + match self { + Self::KvIndexer(indexer) => { + indexer + .process_routing_decision_for_request(tokens_with_hashes, worker) + .await + } + Self::Concurrent(tpi) => { + tpi.process_routing_decision_for_request(tokens_with_hashes, worker) + .await + } + Self::Remote(_) | Self::None => Ok(()), + } + } + + pub(crate) async fn apply_event(&self, event: RouterEvent) { + match self { + Self::KvIndexer(indexer) => { + if let Err(e) = indexer.event_sender().send(event).await { + tracing::warn!("Failed to send event to indexer: {e}"); + } + } + Self::Concurrent(tpi) => tpi.apply_event(event).await, + Self::Remote(_) | Self::None => {} + } + } + + pub(crate) async fn remove_worker(&self, worker_id: WorkerId) { + match self { + Self::KvIndexer(indexer) => { + if let Err(e) = indexer.remove_worker_sender().send(worker_id).await { + tracing::warn!("Failed to send worker removal for {worker_id}: {e}"); + } + } + Self::Concurrent(tpi) => { + KvIndexerInterface::remove_worker(tpi.as_ref(), worker_id).await; + } + Self::Remote(_) | Self::None => {} + } + } + + pub(crate) async fn get_workers(&self) -> Vec { + match self { + Self::KvIndexer(indexer) => { + let (resp_tx, resp_rx) = oneshot::channel(); + let req = dynamo_kv_router::indexer::GetWorkersRequest { resp: resp_tx }; + if let Err(e) = indexer.get_workers_sender().send(req).await { + tracing::warn!("Failed to send get_workers request: {e}"); + return Vec::new(); + } + resp_rx.await.unwrap_or_default() + } + Self::Concurrent(tpi) => tpi.backend().get_workers(), + Self::Remote(_) | Self::None => Vec::new(), + } + } +} diff --git a/lib/llm/src/kv_router/prefill_router.rs b/lib/llm/src/kv_router/prefill_router.rs deleted file mode 100644 index 7f21775bd7b1..000000000000 --- a/lib/llm/src/kv_router/prefill_router.rs +++ /dev/null @@ -1,746 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use std::collections::HashSet; -use std::sync::{Arc, OnceLock}; - -use anyhow::Result; -use futures::StreamExt; -use tokio::sync::{OwnedSemaphorePermit, oneshot}; -use tokio_util::sync::CancellationToken; -use tracing::Instrument; - -use dynamo_kv_router::{ - config::{KvRouterConfig, RouterConfigOverride}, - protocols::{BlockExtraInfo, WorkerId}, -}; -use dynamo_runtime::{ - component::Endpoint, - pipeline::{ - AsyncEngine, AsyncEngineContextProvider, Context, ManyOut, Operator, PushRouter, - RouterMode, ServerStreamingEngine, SingleIn, async_trait, - }, - protocols::{EndpointId, annotated::Annotated, maybe_error::MaybeError}, -}; - -use crate::{ - discovery::ModelManager, - kv_router::KvPushRouter, - protocols::common::llm_backend::{LLMEngineOutput, PreprocessedRequest}, - protocols::common::preprocessor::{BootstrapInfo, PrefillResult}, - protocols::common::timing::{RequestPhase, RequestTracker, WORKER_TYPE_PREFILL}, -}; - -/// Errors that can occur during prefill routing -#[derive(Debug, thiserror::Error)] -pub enum PrefillError { - /// Prefill router has not been activated yet - #[error("Prefill router not yet activated")] - NotActivated, - - /// TODO: Separate prefill worker error from prefill router error - /// Error during prefill execution - #[error("Prefill execution failed: {0}")] - PrefillError( - String, - #[source] Option>, - ), - - /// Disaggregated params not found in prefill response - #[error("No disaggregated params in prefill response: {0}")] - NoDisaggregatedParams(String), -} - -/// Result of the prefill phase in `generate()`. -enum PrefillOutcome { - /// Bootstrap optimization: prefill spawned in background, bootstrap info ready - Bootstrap(BootstrapInfo), - /// Synchronous prefill completed with result - Completed(PrefillResult), -} - -/// The inner router used by PrefillRouter -#[derive(Clone)] -enum InnerPrefillRouter { - /// KV-aware routing using KvPushRouter - KvRouter(Arc), - /// Simple routing (RoundRobin, Random, Direct) - /// Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only - /// available in KV routing mode where the router has actual bookkeeping. - SimpleRouter(Arc>>), -} - -impl InnerPrefillRouter { - /// Generate with optional direct routing to specific worker. - /// For KvRouter, target_worker is ignored since prefill_worker_id is already set on the request. - /// For SimpleRouter, target_worker triggers direct routing via router.direct(). - async fn generate_to_worker( - &self, - request: SingleIn, - target_worker: Option, - ) -> Result>> { - match (self, target_worker) { - // KvRouter: prefill_worker_id already set on request, KvPushRouter::select_worker uses it - (InnerPrefillRouter::KvRouter(router), _) => router.generate(request).await, - (InnerPrefillRouter::SimpleRouter(router), Some(worker_id)) => { - router.direct(request, worker_id).await - } - (InnerPrefillRouter::SimpleRouter(router), None) => router.generate(request).await, - } - } - - /// Select next worker (for non-KV modes only) - fn select_next_worker(&self) -> Option { - match self { - InnerPrefillRouter::SimpleRouter(router) => router.select_next_worker(), - InnerPrefillRouter::KvRouter(_) => None, - } - } -} - -/// PrefillRouter is a forward-only operator that sits between Migration and the decode router. -/// It optionally calls a prefill worker before routing to decode, extracting disaggregated_params -/// from the prefill response and injecting them into the decode request. -/// -/// Modes: -/// - Query-only: `query_instance_id` annotation present → returns worker IDs without execution -/// - Pre-routed: `prefill_worker_id`/`decode_worker_id` set → routes to specified workers -/// - Normal: Worker IDs determined by router based on KV cache state -pub struct PrefillRouter { - prefill_router: OnceLock, - model_manager: Arc, - endpoint_id: OnceLock, - cancel_token: CancellationToken, - router_mode: RouterMode, - enforce_disagg: bool, - /// Model name used to look up the worker monitor for prefill client registration - model_name: String, - /// Namespace used to look up the correct WorkerSet's worker monitor - namespace: String, - is_eagle: bool, -} - -impl PrefillRouter { - /// Create a disabled prefill router that will never activate (passthrough only) - pub fn disabled( - model_manager: Arc, - router_mode: RouterMode, - enforce_disagg: bool, - ) -> Arc { - Arc::new(Self { - prefill_router: OnceLock::new(), - model_manager, - endpoint_id: OnceLock::new(), - cancel_token: CancellationToken::new(), - router_mode, - enforce_disagg, - model_name: String::new(), // Not used for disabled router - namespace: String::new(), // Not used for disabled router - is_eagle: false, - }) - } - - #[expect(clippy::too_many_arguments)] - pub fn new( - activation_rx: oneshot::Receiver, - model_manager: Arc, - router_mode: RouterMode, - kv_cache_block_size: u32, - kv_router_config: Option, - enforce_disagg: bool, - model_name: String, - namespace: String, - is_eagle: bool, - ) -> Arc { - let prefill_router = OnceLock::new(); - let cancel_token = CancellationToken::new(); - - let router = Arc::new(Self { - prefill_router, - model_manager: model_manager.clone(), - endpoint_id: OnceLock::new(), - cancel_token: cancel_token.clone(), - router_mode, - enforce_disagg, - model_name, - namespace, - is_eagle, - }); - - // Spawn background task to wait for activation - let router_clone = router.clone(); - tokio::spawn(async move { - tokio::select! { - result = activation_rx => { - let Ok(endpoint) = result else { - tracing::debug!("Prefill router activation channel closed without receiving endpoint"); - return; - }; - - if let Err(e) = router_clone.activate( - endpoint, - model_manager, - kv_cache_block_size, - kv_router_config, - ).await { - tracing::error!(error = %e, "Failed to activate prefill router"); - } - } - _ = cancel_token.cancelled() => { - tracing::debug!("Prefill router activation cancelled"); - } - } - }); - - router - } - - /// Activate the prefill router with the provided endpoint - async fn activate( - &self, - endpoint: Endpoint, - model_manager: Arc, - kv_cache_block_size: u32, - kv_router_config: Option, - ) -> Result<()> { - tracing::info!( - router_mode = ?self.router_mode, - "Activating prefill router" - ); - - // Store endpoint_id for later use in resolve_prefill_worker - let _ = self.endpoint_id.set(endpoint.id()); - - // Start runtime config watcher for this endpoint (needed for get_disaggregated_endpoint) - // This must be done before creating the router so bootstrap info is available - model_manager - .get_or_create_runtime_config_watcher(&endpoint) - .await?; - - let inner_router = if self.router_mode.is_kv_routing() { - // Create KV chooser using the endpoint (this is a prefill router) - let kv_chooser = model_manager - .kv_chooser_for( - &endpoint, - kv_cache_block_size, - kv_router_config, - WORKER_TYPE_PREFILL, - Some(self.model_name.clone()), - self.is_eagle, - ) - .await?; - - // Extract client from kv_chooser to ensure shared state - let client = kv_chooser.client().clone(); - - // Register prefill client with worker monitor for TTFT metric cleanup in disaggregated mode - if let Some(monitor) = - model_manager.get_worker_monitor_for_namespace(&self.model_name, &self.namespace) - { - monitor.set_prefill_client(client.clone()); - } - - // Build the PushRouter for prefill with KV mode using the shared client - let push_router = PushRouter::>::from_client_with_threshold( - client, - RouterMode::KV, - None, // busy_threshold - None, // worker_monitor - ) - .await?; - - // Wrap it in KvPushRouter - InnerPrefillRouter::KvRouter(Arc::new(KvPushRouter::new(push_router, kv_chooser))) - } else { - // Create client for simple router - let client = endpoint.client().await?; - - // Register prefill client with worker monitor for TTFT metric cleanup in disaggregated mode - if let Some(monitor) = - model_manager.get_worker_monitor_for_namespace(&self.model_name, &self.namespace) - { - monitor.set_prefill_client(client.clone()); - } - - // Create simple push router with the frontend's router mode - // Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only - // available in KV routing mode where the router has actual bookkeeping. - let push_router = PushRouter::>::from_client_with_threshold( - client, - self.router_mode, - None, // busy_threshold - None, // worker_monitor - ) - .await?; - - InnerPrefillRouter::SimpleRouter(Arc::new(push_router)) - }; - - // Set the router (ignore error if already set) - let _ = self.prefill_router.set(inner_router); - - tracing::info!( - router_mode = ?self.router_mode, - "Prefill router activated successfully" - ); - - Ok(()) - } - - /// Select a prefill worker and resolve its bootstrap connection info. - /// If preselected_worker is provided (GAIE Stage 2), use it directly. - /// Otherwise, query for the best worker (KV mode) or select next worker (non-KV modes). - async fn resolve_prefill_worker( - &self, - req: &PreprocessedRequest, - preselected_worker: Option, - ) -> Option<(u64, u32, BootstrapInfo)> { - let endpoint_id = self.endpoint_id.get()?; - self.prefill_router.get()?; - - // Worker selection - let (worker_id, dp_rank) = if let Some(id) = preselected_worker { - let dp_rank = req.routing.as_ref().and_then(|r| r.dp_rank).unwrap_or(0); - tracing::debug!( - worker_id = id, - dp_rank = dp_rank, - "Using pre-selected prefill worker for bootstrap" - ); - (id, dp_rank) - } else { - // Use shared worker selection logic (update_states=false for peek behavior) - // Extract LORA name and priority jump from routing hints - let lora_name = req.routing.as_ref().and_then(|r| r.lora_name.clone()); - let priority_jump = req - .routing - .as_ref() - .and_then(|r| r.priority_jump) - .unwrap_or(0.0); - let allowed_worker_ids = req - .routing - .as_ref() - .and_then(|r| r.allowed_worker_ids.clone()); - let (routing_token_ids, block_mm_infos) = req.block_mm_routing_info(); - match self - .query_prefill_worker( - routing_token_ids, - block_mm_infos, - false, - lora_name, - priority_jump, - allowed_worker_ids, - ) - .await - { - Ok((worker_id, dp_rank)) => (worker_id, dp_rank), - Err(_) => return None, - } - }; - - // Get bootstrap info from ModelManager (works for ANY mode) - let endpoint = self - .model_manager - .get_disaggregated_endpoint(endpoint_id, worker_id)?; - let host = endpoint.bootstrap_host?; - let port = endpoint.bootstrap_port?; - - let bootstrap_room: u64 = rand::random_range(0..=i64::MAX.cast_unsigned()); - - tracing::debug!( - worker_id = worker_id, - dp_rank = dp_rank, - bootstrap_host = %host, - bootstrap_port = port, - bootstrap_room = bootstrap_room, - router_mode = ?self.router_mode, - "Built bootstrap_info upfront before prefill" - ); - - Some(( - worker_id, - dp_rank, - BootstrapInfo { - bootstrap_host: host, - bootstrap_port: port, - bootstrap_room, - }, - )) - } - - /// Execute prefill with the given router and extract structured result. - /// - /// Uses direct routing to target_worker when specified (for non-KV modes with bootstrap optimization). - /// - /// If `phase_permit` is provided, it is dropped after the first output is received, - /// allowing subsequent `set_phase` calls to proceed. This is used in the bootstrap - /// optimization path to ensure `record_worker_full` completes before the phase changes. - /// - /// Returns (PrefillResult, Option<(worker_id, dp_rank)>). - async fn execute_prefill( - router: Option, - request: SingleIn, - target_worker: Option, - phase_permit: Option, - ) -> Result<(PrefillResult, Option<(u64, u32)>), PrefillError> { - let router = router.ok_or(PrefillError::NotActivated)?; - let mut prefill_response = router - .generate_to_worker(request, target_worker) - .await - .map_err(|e| { - PrefillError::PrefillError( - "failed to route to prefill worker".to_string(), - Some(e.into()), - ) - })?; - - // Drop phase permit now - routing is complete, record_worker_full was called in select_worker. - // This unblocks set_phase(Decode) in the main task without waiting for prefill output. - drop(phase_permit); - - let Some(first_output) = prefill_response.next().await else { - return Err(PrefillError::PrefillError( - "Prefill router returned no output (stream ended)".to_string(), - None, - )); - }; - - if let Some(err) = first_output.err() { - return Err(PrefillError::PrefillError( - "Prefill router returned error in output".to_string(), - Some(Box::new(err)), - )); - } - - let mut prompt_tokens_details = first_output - .data - .as_ref() - .and_then(|o| o.completion_usage.as_ref()) - .and_then(|u| u.prompt_tokens_details.clone()); - - while let Some(next) = prefill_response.next().await { - if let Some(o) = next.data.as_ref() - && prompt_tokens_details.is_none() - { - prompt_tokens_details = o - .completion_usage - .as_ref() - .and_then(|u| u.prompt_tokens_details.clone()); - } - } - - let Some(output) = &first_output.data else { - return Err(PrefillError::NoDisaggregatedParams( - "Prefill router output has no data field".to_string(), - )); - }; - - let Some(disaggregated_params) = output.disaggregated_params.clone() else { - return Err(PrefillError::NoDisaggregatedParams( - "Prefill router output missing disaggregated_params".to_string(), - )); - }; - - // Extract prefill worker ID and dp_rank from disaggregated_params - let prefill_worker_info = - disaggregated_params - .get("worker_id") - .and_then(|worker_id_json| { - let worker_id = worker_id_json - .get("prefill_worker_id") - .and_then(|v| v.as_u64())?; - let dp_rank = worker_id_json - .get("prefill_dp_rank") - .and_then(|v| v.as_u64()) - .map(|r| r as u32) - .unwrap_or(0); - Some((worker_id, dp_rank)) - }); - Ok(( - PrefillResult { - disaggregated_params, - prompt_tokens_details, - }, - prefill_worker_info, - )) - } - - /// Spawn prefill as a background task. - /// - /// Uses direct routing to target_worker when specified (for non-KV modes with bootstrap optimization). - /// - /// The `phase_permit` is passed to the spawned task and dropped after the first output, - /// allowing the main task's `set_phase(Decode)` to proceed. - fn spawn_prefill_task( - &self, - prefill_request: SingleIn, - target_worker: Option, - phase_permit: OwnedSemaphorePermit, - ) { - let router = self.prefill_router.get().cloned(); - // Capture current span to propagate trace context to the spawned task - let span = tracing::Span::current(); - - tokio::spawn( - async move { - match Self::execute_prefill( - router, - prefill_request, - target_worker, - Some(phase_permit), - ) - .await - { - Ok(_) => { - tracing::debug!("Prefill background task completed"); - } - Err(e) => { - tracing::warn!("Prefill background task error: {e:?}"); - } - } - } - .instrument(span), - ); - } - - /// Query the best prefill worker without executing a request. - /// Returns (worker_id, dp_rank). - /// - /// This is the shared worker selection logic used by both `resolve_prefill_worker` - /// and `query_route`. - /// Register externally-provided workers in the prefill router's slot tracker. - pub fn register_workers(&self, worker_ids: &HashSet) { - if let Some(InnerPrefillRouter::KvRouter(r)) = self.prefill_router.get() { - r.chooser.register_workers(worker_ids); - } - } - - pub async fn query_prefill_worker( - &self, - token_ids: &[u32], - block_mm_infos: Option<&[Option]>, - update_states: bool, - lora_name: Option, - priority_jump: f64, - allowed_worker_ids: Option>, - ) -> Result<(u64, u32)> { - let prefill_router = self - .prefill_router - .get() - .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; - - match prefill_router { - InnerPrefillRouter::KvRouter(r) => { - let (worker, _overlap) = r - .chooser - .find_best_match( - None, - token_ids, - block_mm_infos, - None, - update_states, - lora_name, - priority_jump, - None, - allowed_worker_ids, - ) - .await?; - Ok((worker.worker_id, worker.dp_rank)) - } - InnerPrefillRouter::SimpleRouter(r) => { - let worker_id = if update_states { - r.select_next_worker() - } else { - r.peek_next_worker() - } - .ok_or_else(|| anyhow::anyhow!("No workers available for prefill"))?; - Ok((worker_id, u32::MAX)) - } - } - } - - /// Check if disaggregated mode is currently active (prefill router activated) - pub fn is_activated(&self) -> bool { - self.prefill_router.get().is_some() - } -} - -impl Drop for PrefillRouter { - fn drop(&mut self) { - tracing::debug!("Dropping PrefillRouter, cancelling background activation task"); - self.cancel_token.cancel(); - } -} - -#[async_trait] -impl - Operator< - SingleIn, - ManyOut>, - SingleIn, - ManyOut>, - > for PrefillRouter -{ - async fn generate( - &self, - request: SingleIn, - next: ServerStreamingEngine>, - ) -> Result>> { - // Extract request data while preserving context - let (mut req, context) = request.into_parts(); - let request_id = context.id().to_string(); - let engine_ctx = context.context(); - - // Save original max_tokens for decode - let original_max_tokens = req.stop_conditions.max_tokens; - - // If prefill router is not activated (no prefill workers discovered), - // this is aggregated mode — route directly to decode. - // With --enforce-disagg, fail instead of falling back. - if self.prefill_router.get().is_none() { - if self.enforce_disagg { - return Err(anyhow::anyhow!(PrefillError::NotActivated)); - } - return next.generate(context.map(|_| req)).await; - } - - // Ensure tracker exists for routing decisions in disaggregated mode. - // Create one if not provided by the upstream DeltaGenerator. - if req.tracker.is_none() { - req.tracker = Some(Arc::new(RequestTracker::new())); - } - let tracker = req.tracker.as_ref().unwrap(); - let prefill_phase_permit = tracker.set_phase(RequestPhase::Prefill).await; - - // Prepare prefill request with max_tokens = 1 (clone after tracker is set) - let mut prefill_req = req.clone(); - prefill_req.stop_conditions.max_tokens = Some(1); - - // Try to resolve prefill worker upfront: if we can get bootstrap info early, - // spawn prefill in background and proceed to decode immediately. - let preselected_worker = prefill_req - .routing - .as_ref() - .and_then(|r| r.prefill_worker_id); - - if self.router_mode.is_direct_routing() && preselected_worker.is_none() { - return Err(anyhow::anyhow!( - "Prefill worker ID required in Direct routing mode but none found in request. \ - Expected prefill_worker_id to be set via x-prefill-instance-id header by external router (e.g., EPP)." - )); - } - - let prefill_result = async { - if let Some((worker_id, dp_rank, bootstrap_info)) = self - .resolve_prefill_worker(&prefill_req, preselected_worker) - .await - { - // Bootstrap optimization path: spawn prefill in background - // We successfully used the peeked worker, so we must now advance the router state - // to ensure the next request gets a different worker. - if !self.router_mode.is_kv_routing() - && let Some(router) = self.prefill_router.get() - { - router.select_next_worker(); - } - - let routing = prefill_req.routing_mut(); - routing.prefill_worker_id = Some(worker_id); - routing.dp_rank = Some(dp_rank); - prefill_req.bootstrap_info = Some(bootstrap_info.clone()); - - let prefill_context = Context::with_id(prefill_req, request_id.clone()); - engine_ctx.link_child(prefill_context.context()); - - // Pass phase permit to spawned task - it drops after first output (record_worker_full complete) - // This allows set_phase(Decode) below to proceed only after prefill routing is done - self.spawn_prefill_task(prefill_context, Some(worker_id), prefill_phase_permit); - - Ok(PrefillOutcome::Bootstrap(bootstrap_info)) - } else { - // Original prefill path: wait for prefill to complete - tracing::debug!("Using original prefill path"); - - // Drop the phase permit - we wait for completion - // so there's no race with set_phase(Decode) below - drop(prefill_phase_permit); - - let prefill_context = Context::with_id(prefill_req, request_id.clone()); - engine_ctx.link_child(prefill_context.context()); - - // In Direct mode, pass preselected_worker so execute_prefill uses - // router.direct() instead of router.generate() (which bails in Direct mode). - let (result, _worker_info) = Self::execute_prefill( - self.prefill_router.get().cloned(), - prefill_context, - preselected_worker, - None, - ) - .await?; - - Ok(PrefillOutcome::Completed(result)) - } - } - .await; - - // Abort if cancelled during prefill - if engine_ctx.is_stopped() || engine_ctx.is_killed() { - tracing::debug!("Abort entering decode after context is stopped or killed"); - return Err(anyhow::anyhow!( - "Context id {} is stopped or killed", - engine_ctx.id() - )); - } - - // Handle prefill result - match prefill_result { - Ok(outcome) => { - tracing::debug!("Prefill completed, proceeding to decode"); - - // Set phase to Decode for the decode request. - // In bootstrap path, this blocks until the spawned prefill task drops its permit - // (after first output / record_worker_full completes), ensuring correct phase for routing. - if let Some(ref tracker) = req.tracker { - let _decode_permit = tracker.set_phase(RequestPhase::Decode).await; - // Permit is dropped immediately - decode proceeds, no need to hold it - } - - let mut decode_req = req; - - match outcome { - PrefillOutcome::Bootstrap(info) => { - decode_req.bootstrap_info = Some(info); - } - PrefillOutcome::Completed(result) => { - decode_req.prefill_result = Some(result); - } - } - - // Restore original max_tokens for decode - decode_req.stop_conditions.max_tokens = original_max_tokens; - - // Set router_config_override for decode: - // - overlap_score_weight = 0 (no KV cache overlap scoring for decode) - // - assume_kv_reuse = false (generate random hashes since decode workers - // may already have blocks cached from prefill transfer) - let existing_override = decode_req.router_config_override.take(); - decode_req.router_config_override = Some(RouterConfigOverride { - overlap_score_weight: Some(0.0), - assume_kv_reuse: Some(false), - ..existing_override.unwrap_or_default() - }); - - // Map the modified request through with preserved context - let decode_request = context.map(|_| decode_req); - next.generate(decode_request).await - } - Err(PrefillError::NotActivated) => { - tracing::error!("Prefill router not activated, failing request"); - Err(anyhow::anyhow!(PrefillError::NotActivated)) - } - Err(e) => { - tracing::error!(error = %e, "Remote prefill failed, failing request"); - Err(anyhow::anyhow!(e)) - } - } - } -} diff --git a/lib/llm/src/kv_router/prefill_router/activation.rs b/lib/llm/src/kv_router/prefill_router/activation.rs new file mode 100644 index 000000000000..8f2a75c3c6bf --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/activation.rs @@ -0,0 +1,188 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; + +use anyhow::Result; +use tokio::sync::oneshot; + +use dynamo_kv_router::config::KvRouterConfig; +use dynamo_runtime::{ + component::{Client, Endpoint}, + pipeline::{PushRouter, RouterMode}, + protocols::annotated::Annotated, +}; + +use super::{InnerPrefillRouter, PrefillRouter}; +use crate::{ + discovery::ModelManager, + kv_router::KvPushRouter, + protocols::common::{ + llm_backend::{LLMEngineOutput, PreprocessedRequest}, + timing::WORKER_TYPE_PREFILL, + }, +}; + +impl PrefillRouter { + /// Create a disabled prefill router that will never activate (passthrough only) + pub fn disabled( + model_manager: Arc, + router_mode: RouterMode, + enforce_disagg: bool, + ) -> Arc { + Arc::new(Self { + prefill_router: std::sync::OnceLock::new(), + model_manager, + endpoint_id: std::sync::OnceLock::new(), + cancel_token: tokio_util::sync::CancellationToken::new(), + router_mode, + enforce_disagg, + model_name: String::new(), // Not used for disabled router + namespace: String::new(), // Not used for disabled router + is_eagle: false, + }) + } + + #[expect(clippy::too_many_arguments)] + pub fn new( + activation_rx: oneshot::Receiver, + model_manager: Arc, + router_mode: RouterMode, + kv_cache_block_size: u32, + kv_router_config: Option, + enforce_disagg: bool, + model_name: String, + namespace: String, + is_eagle: bool, + ) -> Arc { + let prefill_router = std::sync::OnceLock::new(); + let cancel_token = tokio_util::sync::CancellationToken::new(); + + let router = Arc::new(Self { + prefill_router, + model_manager: model_manager.clone(), + endpoint_id: std::sync::OnceLock::new(), + cancel_token: cancel_token.clone(), + router_mode, + enforce_disagg, + model_name, + namespace, + is_eagle, + }); + + // Spawn background task to wait for activation + let router_clone = router.clone(); + tokio::spawn(async move { + tokio::select! { + result = activation_rx => { + let Ok(endpoint) = result else { + tracing::debug!("Prefill router activation channel closed without receiving endpoint"); + return; + }; + + if let Err(e) = router_clone.activate( + endpoint, + model_manager, + kv_cache_block_size, + kv_router_config, + ).await { + tracing::error!(error = %e, "Failed to activate prefill router"); + } + } + _ = cancel_token.cancelled() => { + tracing::debug!("Prefill router activation cancelled"); + } + } + }); + + router + } + + /// Activate the prefill router with the provided endpoint + async fn activate( + &self, + endpoint: Endpoint, + model_manager: Arc, + kv_cache_block_size: u32, + kv_router_config: Option, + ) -> Result<()> { + tracing::info!( + router_mode = ?self.router_mode, + "Activating prefill router" + ); + + // Store endpoint_id for later use in resolve_prefill_worker + let _ = self.endpoint_id.set(endpoint.id()); + + // Start runtime config watcher for this endpoint (needed for get_disaggregated_endpoint) + // This must be done before creating the router so bootstrap info is available + model_manager + .get_or_create_runtime_config_watcher(&endpoint) + .await?; + + let inner_router = if self.router_mode.is_kv_routing() { + // Create KV chooser using the endpoint (this is a prefill router) + let kv_chooser = model_manager + .kv_chooser_for( + &endpoint, + kv_cache_block_size, + kv_router_config, + WORKER_TYPE_PREFILL, + Some(self.model_name.clone()), + self.is_eagle, + ) + .await?; + + // Extract client from kv_chooser to ensure shared state + let client = kv_chooser.client().clone(); + self.register_prefill_client(model_manager.as_ref(), &client); + + // Build the PushRouter for prefill with KV mode using the shared client + let push_router = PushRouter::>::from_client_with_threshold( + client, + RouterMode::KV, + None, // busy_threshold + None, // worker_monitor + ) + .await?; + + // Wrap it in KvPushRouter + InnerPrefillRouter::KvRouter(Arc::new(KvPushRouter::new(push_router, kv_chooser))) + } else { + // Create client for simple router + let client = endpoint.client().await?; + self.register_prefill_client(model_manager.as_ref(), &client); + + // Create simple push router with the frontend's router mode + // Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only + // available in KV routing mode where the router has actual bookkeeping. + let push_router = PushRouter::>::from_client_with_threshold( + client, + self.router_mode, + None, // busy_threshold + None, // worker_monitor + ) + .await?; + + InnerPrefillRouter::SimpleRouter(Arc::new(push_router)) + }; + + // Set the router (ignore error if already set) + let _ = self.prefill_router.set(inner_router); + + tracing::info!( + router_mode = ?self.router_mode, + "Prefill router activated successfully" + ); + + Ok(()) + } + + fn register_prefill_client(&self, model_manager: &ModelManager, client: &Client) { + if let Some(monitor) = + model_manager.get_worker_monitor_for_namespace(&self.model_name, &self.namespace) + { + monitor.set_prefill_client(client.clone()); + } + } +} diff --git a/lib/llm/src/kv_router/prefill_router/execution.rs b/lib/llm/src/kv_router/prefill_router/execution.rs new file mode 100644 index 000000000000..1b5f52a5aa66 --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/execution.rs @@ -0,0 +1,322 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::HashSet; +use std::sync::Arc; + +use anyhow::Result; +use futures::StreamExt; +use tokio::sync::OwnedSemaphorePermit; +use tracing::Instrument; + +use dynamo_kv_router::protocols::{BlockExtraInfo, WorkerId}; +use dynamo_runtime::{ + engine::AsyncEngineContext, + pipeline::{AsyncEngineContextProvider, Context, SingleIn}, + protocols::maybe_error::MaybeError, +}; + +use super::{InnerPrefillRouter, PrefillError, PrefillResolveDecision, PrefillRouter}; +use crate::protocols::common::{ + llm_backend::PreprocessedRequest, + preprocessor::{BootstrapInfo, PrefillResult}, +}; + +impl PrefillRouter { + /// Select a prefill worker and resolve its bootstrap connection info. + /// If preselected_worker is provided (GAIE Stage 2), use it directly. + /// Otherwise, query for the best worker (KV mode) or select next worker (non-KV modes). + pub(super) async fn resolve_prefill_worker( + &self, + req: &PreprocessedRequest, + preselected_worker: Option, + ) -> PrefillResolveDecision { + let Some(endpoint_id) = self.endpoint_id.get() else { + return PrefillResolveDecision::NotActivated; + }; + if self.prefill_router.get().is_none() { + return PrefillResolveDecision::NotActivated; + } + + // Worker selection + let (worker_id, dp_rank) = if let Some(id) = preselected_worker { + let dp_rank = req.routing.as_ref().and_then(|r| r.dp_rank).unwrap_or(0); + tracing::debug!( + worker_id = id, + dp_rank = dp_rank, + "Using pre-selected prefill worker for bootstrap" + ); + (id, dp_rank) + } else { + // Use shared worker selection logic (update_states=false for peek behavior) + // Extract LORA name and priority jump from routing hints + let lora_name = req.routing.as_ref().and_then(|r| r.lora_name.clone()); + let priority_jump = req + .routing + .as_ref() + .and_then(|r| r.priority_jump) + .unwrap_or(0.0); + let allowed_worker_ids = req + .routing + .as_ref() + .and_then(|r| r.allowed_worker_ids.clone()); + let (routing_token_ids, block_mm_infos) = req.block_mm_routing_info(); + match self + .query_prefill_worker( + routing_token_ids, + block_mm_infos, + false, + lora_name, + priority_jump, + allowed_worker_ids, + ) + .await + { + Ok((worker_id, dp_rank)) => (worker_id, dp_rank), + Err(_) => return PrefillResolveDecision::Unavailable, + } + }; + + // Get bootstrap info from ModelManager (works for ANY mode) + let Some(endpoint) = self + .model_manager + .get_disaggregated_endpoint(endpoint_id, worker_id) + else { + return PrefillResolveDecision::NoBootstrapEndpoint; + }; + let Some(host) = endpoint.bootstrap_host else { + return PrefillResolveDecision::NoBootstrapEndpoint; + }; + let Some(port) = endpoint.bootstrap_port else { + return PrefillResolveDecision::NoBootstrapEndpoint; + }; + + let bootstrap_room: u64 = rand::random_range(0..=i64::MAX.cast_unsigned()); + + tracing::debug!( + worker_id = worker_id, + dp_rank = dp_rank, + bootstrap_host = %host, + bootstrap_port = port, + bootstrap_room = bootstrap_room, + router_mode = ?self.router_mode, + "Built bootstrap_info upfront before prefill" + ); + + PrefillResolveDecision::Resolved { + worker_id, + dp_rank, + bootstrap_info: BootstrapInfo { + bootstrap_host: host, + bootstrap_port: port, + bootstrap_room, + }, + } + } + + /// Execute prefill with the given router and extract structured result. + /// + /// Uses direct routing to target_worker when specified (for non-KV modes with bootstrap optimization). + /// + /// If `phase_transition_permit` is provided, it is dropped immediately after routing completes, + /// allowing subsequent `set_phase` calls to proceed. This preserves the current synchronization: + /// the prefill route must finish `record_worker_full` before the phase can change to Decode. + /// + /// Returns (PrefillResult, Option<(worker_id, dp_rank)>). + pub(super) async fn execute_prefill( + router: Option, + request: SingleIn, + target_worker: Option, + phase_transition_permit: Option, + ) -> Result<(PrefillResult, Option<(u64, u32)>), PrefillError> { + let router = router.ok_or(PrefillError::NotActivated)?; + let mut prefill_response = router + .generate_to_worker(request, target_worker) + .await + .map_err(|e| { + PrefillError::PrefillError( + "failed to route to prefill worker".to_string(), + Some(e.into()), + ) + })?; + + // Release the phase barrier now that routing completed and record_worker_full already ran. + // Decode may proceed without waiting for prefill output streaming to finish. + drop(phase_transition_permit); + + let Some(first_output) = prefill_response.next().await else { + return Err(PrefillError::PrefillError( + "Prefill router returned no output (stream ended)".to_string(), + None, + )); + }; + + if let Some(err) = first_output.err() { + return Err(PrefillError::PrefillError( + "Prefill router returned error in output".to_string(), + Some(Box::new(err)), + )); + } + + let mut prompt_tokens_details = first_output + .data + .as_ref() + .and_then(|o| o.completion_usage.as_ref()) + .and_then(|u| u.prompt_tokens_details.clone()); + + while let Some(next) = prefill_response.next().await { + if let Some(o) = next.data.as_ref() + && prompt_tokens_details.is_none() + { + prompt_tokens_details = o + .completion_usage + .as_ref() + .and_then(|u| u.prompt_tokens_details.clone()); + } + } + + let Some(output) = &first_output.data else { + return Err(PrefillError::NoDisaggregatedParams( + "Prefill router output has no data field".to_string(), + )); + }; + + let Some(disaggregated_params) = output.disaggregated_params.clone() else { + return Err(PrefillError::NoDisaggregatedParams( + "Prefill router output missing disaggregated_params".to_string(), + )); + }; + + // Extract prefill worker ID and dp_rank from disaggregated_params + let prefill_worker_info = + disaggregated_params + .get("worker_id") + .and_then(|worker_id_json| { + let worker_id = worker_id_json + .get("prefill_worker_id") + .and_then(|v| v.as_u64())?; + let dp_rank = worker_id_json + .get("prefill_dp_rank") + .and_then(|v| v.as_u64()) + .map(|r| r as u32) + .unwrap_or(0); + Some((worker_id, dp_rank)) + }); + Ok(( + PrefillResult { + disaggregated_params, + prompt_tokens_details, + }, + prefill_worker_info, + )) + } + + /// Spawn prefill as a background task. + /// + /// Uses direct routing to target_worker when specified (for non-KV modes with bootstrap optimization). + /// + /// The `phase_transition_permit` is passed to the spawned task and released after routing + /// completes, allowing the main task's `set_phase(Decode)` to proceed. + pub(super) fn spawn_prefill_task( + &self, + prefill_request: SingleIn, + target_worker: Option, + phase_transition_permit: OwnedSemaphorePermit, + ) { + let router = self.prefill_router.get().cloned(); + // Capture current span to propagate trace context to the spawned task + let span = tracing::Span::current(); + + tokio::spawn( + async move { + match Self::execute_prefill( + router, + prefill_request, + target_worker, + Some(phase_transition_permit), + ) + .await + { + Ok(_) => { + tracing::debug!("Prefill background task completed"); + } + Err(e) => { + tracing::warn!("Prefill background task error: {e:?}"); + } + } + } + .instrument(span), + ); + } + + /// Query the best prefill worker without executing a request. + /// Returns (worker_id, dp_rank). + /// + /// This is the shared worker selection logic used by both `resolve_prefill_worker` + /// and `query_route`. + pub async fn query_prefill_worker( + &self, + token_ids: &[u32], + block_mm_infos: Option<&[Option]>, + update_states: bool, + lora_name: Option, + priority_jump: f64, + allowed_worker_ids: Option>, + ) -> Result<(u64, u32)> { + let prefill_router = self + .prefill_router + .get() + .ok_or_else(|| anyhow::anyhow!(PrefillError::NotActivated))?; + + match prefill_router { + InnerPrefillRouter::KvRouter(r) => { + let (worker, _overlap) = r + .chooser + .find_best_match( + None, + token_ids, + block_mm_infos, + None, + update_states, + lora_name, + priority_jump, + None, + allowed_worker_ids, + ) + .await?; + Ok((worker.worker_id, worker.dp_rank)) + } + InnerPrefillRouter::SimpleRouter(r) => { + let worker_id = if update_states { + r.select_next_worker() + } else { + r.peek_next_worker() + } + .ok_or_else(|| anyhow::anyhow!("No workers available for prefill"))?; + Ok((worker_id, 0)) + } + } + } + + /// Register externally-provided workers in the prefill router's slot tracker. + pub fn register_workers(&self, worker_ids: &HashSet) { + if let Some(InnerPrefillRouter::KvRouter(r)) = self.prefill_router.get() { + r.chooser.register_workers(worker_ids); + } + } + + /// Check if disaggregated mode is currently active (prefill router activated) + pub fn is_activated(&self) -> bool { + self.prefill_router.get().is_some() + } +} + +pub(super) fn link_child_context( + engine_ctx: &Arc, + request: T, + request_id: &str, +) -> Context { + let child_context = Context::with_id(request, request_id.to_string()); + engine_ctx.link_child(child_context.context()); + child_context +} diff --git a/lib/llm/src/kv_router/prefill_router/inner.rs b/lib/llm/src/kv_router/prefill_router/inner.rs new file mode 100644 index 000000000000..8b4f3508b67f --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/inner.rs @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; + +use anyhow::Result; + +use dynamo_runtime::{ + pipeline::{AsyncEngine, ManyOut, PushRouter, SingleIn}, + protocols::annotated::Annotated, +}; + +use crate::{ + kv_router::KvPushRouter, + protocols::common::llm_backend::{LLMEngineOutput, PreprocessedRequest}, +}; + +/// The inner router used by PrefillRouter +#[derive(Clone)] +pub(super) enum InnerPrefillRouter { + /// KV-aware routing using KvPushRouter + KvRouter(Arc), + /// Simple routing (RoundRobin, Random, Direct) + /// Note: Per-worker metrics (active_prefill_tokens, active_decode_blocks) are only + /// available in KV routing mode where the router has actual bookkeeping. + SimpleRouter(Arc>>), +} + +impl InnerPrefillRouter { + /// Generate with optional direct routing to specific worker. + /// For KvRouter, target_worker is ignored since prefill_worker_id is already set on the request. + /// For SimpleRouter, target_worker triggers direct routing via router.direct(). + pub(super) async fn generate_to_worker( + &self, + request: SingleIn, + target_worker: Option, + ) -> Result>> { + match (self, target_worker) { + // KvRouter: prefill_worker_id already set on request, KvPushRouter::select_worker uses it + (InnerPrefillRouter::KvRouter(router), _) => router.generate(request).await, + (InnerPrefillRouter::SimpleRouter(router), Some(worker_id)) => { + router.direct(request, worker_id).await + } + (InnerPrefillRouter::SimpleRouter(router), None) => router.generate(request).await, + } + } + + /// Select next worker (for non-KV modes only) + pub(super) fn select_next_worker(&self) -> Option { + match self { + InnerPrefillRouter::SimpleRouter(router) => router.select_next_worker(), + InnerPrefillRouter::KvRouter(_) => None, + } + } +} diff --git a/lib/llm/src/kv_router/prefill_router/mod.rs b/lib/llm/src/kv_router/prefill_router/mod.rs new file mode 100644 index 000000000000..c2a9aad4e14f --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/mod.rs @@ -0,0 +1,259 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::{Arc, OnceLock}; + +use anyhow::Result; +use tokio_util::sync::CancellationToken; + +use dynamo_runtime::{ + pipeline::{ + AsyncEngineContextProvider, ManyOut, Operator, RouterMode, ServerStreamingEngine, SingleIn, + async_trait, + }, + protocols::{EndpointId, annotated::Annotated}, +}; + +use crate::{ + discovery::ModelManager, + protocols::common::{ + llm_backend::{LLMEngineOutput, PreprocessedRequest}, + timing::{RequestPhase, RequestTracker}, + }, +}; + +mod activation; +mod execution; +mod inner; +mod types; + +use execution::link_child_context; +use inner::InnerPrefillRouter; +pub use types::PrefillError; +use types::{PrefillOutcome, PrefillResolveDecision, build_decode_router_override}; + +/// PrefillRouter is a forward-only operator that sits between Migration and the decode router. +/// It optionally calls a prefill worker before routing to decode, extracting disaggregated_params +/// from the prefill response and injecting them into the decode request. +/// +/// Modes: +/// - Query-only: `query_instance_id` annotation present → returns worker IDs without execution +/// - Pre-routed: `prefill_worker_id`/`decode_worker_id` set → routes to specified workers +/// - Normal: Worker IDs determined by router based on KV cache state +pub struct PrefillRouter { + prefill_router: OnceLock, + model_manager: Arc, + endpoint_id: OnceLock, + cancel_token: CancellationToken, + router_mode: RouterMode, + enforce_disagg: bool, + /// Model name used to look up the worker monitor for prefill client registration + model_name: String, + /// Namespace used to look up the correct WorkerSet's worker monitor + namespace: String, + is_eagle: bool, +} + +impl Drop for PrefillRouter { + fn drop(&mut self) { + tracing::debug!("Dropping PrefillRouter, cancelling background activation task"); + self.cancel_token.cancel(); + } +} + +#[async_trait] +impl + Operator< + SingleIn, + ManyOut>, + SingleIn, + ManyOut>, + > for PrefillRouter +{ + async fn generate( + &self, + request: SingleIn, + next: ServerStreamingEngine>, + ) -> Result>> { + // Extract request data while preserving context + let (mut req, context) = request.into_parts(); + let request_id = context.id().to_string(); + let engine_ctx = context.context(); + + // Save original max_tokens for decode + let original_max_tokens = req.stop_conditions.max_tokens; + + // If prefill router is not activated (no prefill workers discovered), + // this is aggregated mode — route directly to decode. + // With --enforce-disagg, fail instead of falling back. + if self.prefill_router.get().is_none() { + if self.enforce_disagg { + return Err(anyhow::anyhow!(PrefillError::NotActivated)); + } + return next.generate(context.map(|_| req)).await; + } + + // Ensure tracker exists for routing decisions in disaggregated mode. + // Create one if not provided by the upstream DeltaGenerator. + if req.tracker.is_none() { + req.tracker = Some(Arc::new(RequestTracker::new())); + } + let tracker = req.tracker.as_ref().unwrap(); + let prefill_phase_barrier = tracker.set_phase(RequestPhase::Prefill).await; + + // Prepare prefill request with max_tokens = 1 (clone after tracker is set) + let mut prefill_req = req.clone(); + prefill_req.stop_conditions.max_tokens = Some(1); + + // Try to resolve prefill worker upfront: if we can get bootstrap info early, + // spawn prefill in background and proceed to decode immediately. + let preselected_worker = prefill_req + .routing + .as_ref() + .and_then(|r| r.prefill_worker_id); + + if self.router_mode.is_direct_routing() && preselected_worker.is_none() { + return Err(anyhow::anyhow!( + "Prefill worker ID required in Direct routing mode but none found in request. \ + Expected prefill_worker_id to be set via x-prefill-instance-id header by external router (e.g., EPP)." + )); + } + + let prefill_result = match self + .resolve_prefill_worker(&prefill_req, preselected_worker) + .await + { + PrefillResolveDecision::Resolved { + worker_id, + dp_rank, + bootstrap_info, + } => { + // Bootstrap optimization path: spawn prefill in background + // We successfully used the peeked worker, so we must now advance the router state + // to ensure the next request gets a different worker. + if !self.router_mode.is_kv_routing() + && let Some(router) = self.prefill_router.get() + { + router.select_next_worker(); + } + + let routing = prefill_req.routing_mut(); + routing.prefill_worker_id = Some(worker_id); + routing.dp_rank = Some(dp_rank); + prefill_req.bootstrap_info = Some(bootstrap_info.clone()); + + let prefill_context = + link_child_context(&engine_ctx, prefill_req, request_id.as_str()); + + // Pass the phase barrier to the spawned task. It is released after routing + // completes so `record_worker_full` finishes before phase changes to Decode. + self.spawn_prefill_task(prefill_context, Some(worker_id), prefill_phase_barrier); + + Ok(PrefillOutcome::Bootstrap(bootstrap_info)) + } + PrefillResolveDecision::Unavailable + | PrefillResolveDecision::NotActivated + | PrefillResolveDecision::NoBootstrapEndpoint => { + // Original prefill path: wait for prefill to complete + tracing::debug!("Using original prefill path"); + + // Drop the phase barrier because we wait for prefill completion in this task, + // so there is no race with set_phase(Decode) below. + drop(prefill_phase_barrier); + + let prefill_context = + link_child_context(&engine_ctx, prefill_req, request_id.as_str()); + + // In Direct mode, pass preselected_worker so execute_prefill uses + // router.direct() instead of router.generate() (which bails in Direct mode). + let (result, _worker_info) = Self::execute_prefill( + self.prefill_router.get().cloned(), + prefill_context, + preselected_worker, + None, + ) + .await?; + + Ok(PrefillOutcome::Completed(result)) + } + }; + + // Abort if cancelled during prefill + if engine_ctx.is_stopped() || engine_ctx.is_killed() { + tracing::debug!("Abort entering decode after context is stopped or killed"); + return Err(anyhow::anyhow!( + "Context id {} is stopped or killed", + engine_ctx.id() + )); + } + + // Handle prefill result + match prefill_result { + Ok(outcome) => { + tracing::debug!("Prefill completed, proceeding to decode"); + + // Set phase to Decode for the decode request. + // In bootstrap path, this blocks until the spawned prefill task releases its + // phase barrier after routing completes, ensuring correct worker attribution. + if let Some(ref tracker) = req.tracker { + let _decode_permit = tracker.set_phase(RequestPhase::Decode).await; + // Permit is dropped immediately - decode proceeds, no need to hold it + } + + let mut decode_req = req; + + match outcome { + PrefillOutcome::Bootstrap(info) => { + decode_req.bootstrap_info = Some(info); + } + PrefillOutcome::Completed(result) => { + decode_req.prefill_result = Some(result); + } + } + + // Restore original max_tokens for decode + decode_req.stop_conditions.max_tokens = original_max_tokens; + + // Set router_config_override for decode: + // - overlap_score_weight = 0 (no KV cache overlap scoring for decode) + // - assume_kv_reuse = false (generate random hashes since decode workers + // may already have blocks cached from prefill transfer) + // - track_prefill_tokens = false (decode router should ignore prompt-side load) + let existing_override = decode_req.router_config_override.take(); + decode_req.router_config_override = + Some(build_decode_router_override(existing_override)); + + // Map the modified request through with preserved context + let decode_request = context.map(|_| decode_req); + next.generate(decode_request).await + } + Err(PrefillError::NotActivated) => { + tracing::error!("Prefill router not activated, failing request"); + Err(anyhow::anyhow!(PrefillError::NotActivated)) + } + Err(e) => { + tracing::error!(error = %e, "Remote prefill failed, failing request"); + Err(anyhow::anyhow!(e)) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use dynamo_kv_router::config::RouterConfigOverride; + + #[test] + fn decode_router_override_disables_overlap_and_prefill_tracking() { + let override_config = build_decode_router_override(Some(RouterConfigOverride { + router_temperature: Some(0.7), + ..Default::default() + })); + + assert_eq!(override_config.overlap_score_weight, Some(0.0)); + assert_eq!(override_config.assume_kv_reuse, Some(false)); + assert_eq!(override_config.track_prefill_tokens, Some(false)); + assert_eq!(override_config.router_temperature, Some(0.7)); + } +} diff --git a/lib/llm/src/kv_router/prefill_router/types.rs b/lib/llm/src/kv_router/prefill_router/types.rs new file mode 100644 index 000000000000..1acf1de53d86 --- /dev/null +++ b/lib/llm/src/kv_router/prefill_router/types.rs @@ -0,0 +1,56 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use dynamo_kv_router::config::RouterConfigOverride; + +use crate::protocols::common::preprocessor::{BootstrapInfo, PrefillResult}; + +/// Errors that can occur during prefill routing +#[derive(Debug, thiserror::Error)] +pub enum PrefillError { + /// Prefill router has not been activated yet + #[error("Prefill router not yet activated")] + NotActivated, + + /// TODO: Separate prefill worker error from prefill router error + /// Error during prefill execution + #[error("Prefill execution failed: {0}")] + PrefillError( + String, + #[source] Option>, + ), + + /// Disaggregated params not found in prefill response + #[error("No disaggregated params in prefill response: {0}")] + NoDisaggregatedParams(String), +} + +/// Result of the prefill phase in `generate()`. +pub(super) enum PrefillOutcome { + /// Bootstrap optimization: prefill spawned in background, bootstrap info ready + Bootstrap(BootstrapInfo), + /// Synchronous prefill completed with result + Completed(PrefillResult), +} + +pub(super) enum PrefillResolveDecision { + Resolved { + worker_id: u64, + dp_rank: u32, + bootstrap_info: BootstrapInfo, + }, + Unavailable, + NotActivated, + NoBootstrapEndpoint, +} + +pub(super) fn build_decode_router_override( + existing_override: Option, +) -> RouterConfigOverride { + RouterConfigOverride { + overlap_score_weight: Some(0.0), + assume_kv_reuse: Some(false), + track_prefill_tokens: Some(false), + ..existing_override.unwrap_or_default() + } +} diff --git a/lib/llm/src/kv_router/publisher/event_processor.rs b/lib/llm/src/kv_router/publisher/event_processor.rs new file mode 100644 index 000000000000..b8e9259f0809 --- /dev/null +++ b/lib/llm/src/kv_router/publisher/event_processor.rs @@ -0,0 +1,326 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::future::Future; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use anyhow::Result; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use dynamo_kv_router::RouterEventSink; +use dynamo_kv_router::indexer::LocalKvIndexer; +use dynamo_kv_router::protocols::*; +use dynamo_runtime::transports::event_plane::EventPublisher; +use dynamo_runtime::transports::nats::NatsQueue; + +use crate::kv_router::KV_EVENT_SUBJECT; + +use super::{DEFAULT_MAX_BATCH_BLOCKS, kv_publisher_metrics}; + +/// Accumulator for in-flight KV cache events that will be merged into a single +/// [`RouterEvent`] before being forwarded to the event sink. +#[derive(Debug)] +pub(super) struct BatchingState { + pub(super) pending_removed: Option, + pub(super) pending_stored: Option, + pub(super) next_publish_id: u64, + pub(super) last_dp_rank: u32, + pub(super) last_flush_time: Instant, +} + +impl BatchingState { + pub(super) fn new() -> Self { + Self { + pending_removed: None, + pending_stored: None, + next_publish_id: 1, + last_dp_rank: 0, + last_flush_time: Instant::now(), + } + } + + pub(super) fn has_pending(&self) -> bool { + self.pending_removed.is_some() || self.pending_stored.is_some() + } + + pub(super) fn pending_block_count(&self) -> usize { + self.pending_removed + .as_ref() + .map(|r| r.block_hashes.len()) + .unwrap_or(0) + + self + .pending_stored + .as_ref() + .map(|s| s.blocks.len()) + .unwrap_or(0) + } + + pub(super) fn record_flush_time(&mut self) { + self.last_flush_time = Instant::now(); + } + + pub(super) fn remaining_timeout(&self, timeout_ms: u64) -> Duration { + let timeout = Duration::from_millis(timeout_ms); + let elapsed = self.last_flush_time.elapsed(); + if elapsed >= timeout { + Duration::ZERO + } else { + timeout - elapsed + } + } + + pub(super) fn is_timeout_elapsed(&self, timeout_ms: u64) -> bool { + self.remaining_timeout(timeout_ms) == Duration::ZERO + } + + async fn flush( + &mut self, + publisher: &P, + local_indexer: &Option>, + worker_id: u64, + ) { + if !self.has_pending() { + return; + } + let id = self.next_publish_id; + let dp_rank = self.last_dp_rank; + if let Some(data) = self.pending_removed.take() { + emit( + publisher, + local_indexer, + worker_id, + KvCacheEvent { + event_id: id, + data: KvCacheEventData::Removed(data), + dp_rank, + }, + ) + .await; + } + if let Some(data) = self.pending_stored.take() { + emit( + publisher, + local_indexer, + worker_id, + KvCacheEvent { + event_id: id, + data: KvCacheEventData::Stored(data), + dp_rank, + }, + ) + .await; + } + self.next_publish_id += 1; + self.record_flush_time(); + } +} + +pub(super) struct EventPlanePublisher(pub(super) EventPublisher); + +impl RouterEventSink for EventPlanePublisher { + fn publish_event(&self, event: &RouterEvent) -> impl Future> + Send { + self.0.publish(event) + } +} + +pub(super) struct JetStreamPublisher(pub(super) NatsQueue); + +impl RouterEventSink for JetStreamPublisher { + fn publish_event(&self, event: &RouterEvent) -> impl Future> + Send { + NatsQueue::publish_event(&self.0, KV_EVENT_SUBJECT, event) + } +} + +async fn emit( + publisher: &P, + local_indexer: &Option>, + worker_id: u64, + event: KvCacheEvent, +) { + let router_event = RouterEvent::new(worker_id, event); + if let Some(indexer) = local_indexer + && let Err(e) = indexer.apply_event_with_buffer(router_event.clone()).await + { + tracing::warn!(worker_id, error = %e, "Failed to apply event to local indexer"); + } + if let Err(e) = publisher.publish_event(&router_event).await { + tracing::error!(worker_id, error = %e, "Failed to publish event"); + } +} + +pub(super) async fn run_event_processor_loop( + publisher: P, + worker_id: u64, + cancellation_token: CancellationToken, + mut rx: mpsc::UnboundedReceiver, + local_indexer: Option>, + timeout_ms: Option, + max_batch_blocks: usize, +) { + let mut batching_state = BatchingState::new(); + let mut last_raw_input_id: Option = None; + + loop { + tokio::select! { + _ = cancellation_token.cancelled() => { + tracing::info!("KV Event source received cancellation signal"); + batching_state.flush(&publisher, &local_indexer, worker_id).await; + break; + } + event = rx.recv() => { + let Some(placement_event) = event else { + tracing::debug!("Event processor channel closed."); + batching_state.flush(&publisher, &local_indexer, worker_id).await; + break; + }; + + let raw_event_id = placement_event.event.event_id; + if let Some(last_id) = last_raw_input_id + && raw_event_id > last_id + 1 + { + let gap = raw_event_id - last_id - 1; + tracing::warn!( + worker_id, + last_raw_input_id = last_id, + raw_event_id, + gap, + "Input event gap detected: raw events dropped before batching" + ); + if let Some(metrics) = kv_publisher_metrics() { + metrics.increment_engines_dropped_events(worker_id, gap); + } else { + tracing::warn!( + worker_id, + gap, + "Failed to record dropped events metric: metrics not initialized" + ); + } + } + last_raw_input_id = Some(raw_event_id); + + if !placement_event.placement.is_local_gpu() { + tracing::trace!( + worker_id, + ?placement_event.placement, + event_id = placement_event.event.event_id, + "Skipping non-local-GPU placement event" + ); + continue; + } + + let event = placement_event.event; + tracing::trace!( + "Event processor for worker_id {} processing event: {:?}", + worker_id, + event.data + ); + + let dp_rank_changed = + batching_state.has_pending() && event.dp_rank != batching_state.last_dp_rank; + + match event.data { + KvCacheEventData::Removed(data) => { + if batching_state.pending_stored.is_some() || dp_rank_changed { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + } + match &mut batching_state.pending_removed { + Some(pending) => pending.block_hashes.extend(data.block_hashes), + None => { + batching_state.pending_removed = Some(data); + } + } + } + KvCacheEventData::Stored(data) => { + let should_flush = dp_rank_changed + || batching_state.pending_removed.is_some() + || batching_state.pending_stored.as_ref().is_some_and(|p| { + data.parent_hash != p.blocks.last().map(|b| b.block_hash) + }); + if should_flush { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + } + match &mut batching_state.pending_stored { + Some(pending) => pending.blocks.extend(data.blocks), + None => { + batching_state.pending_stored = Some(data); + } + } + } + KvCacheEventData::Cleared => { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + emit( + &publisher, + &local_indexer, + worker_id, + KvCacheEvent { + event_id: batching_state.next_publish_id, + data: KvCacheEventData::Cleared, + dp_rank: event.dp_rank, + }, + ) + .await; + batching_state.next_publish_id += 1; + } + } + + batching_state.last_dp_rank = event.dp_rank; + + if batching_state.has_pending() + && (timeout_ms.is_none_or(|ms| batching_state.is_timeout_elapsed(ms)) + || batching_state.pending_block_count() > max_batch_blocks) + { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + } + } + _ = tokio::time::sleep( + timeout_ms + .map(|ms| batching_state.remaining_timeout(ms)) + .unwrap_or(Duration::from_secs(3600)) + ), if timeout_ms.is_some() && batching_state.has_pending() => { + batching_state.flush(&publisher, &local_indexer, worker_id).await; + } + } + } +} + +pub(super) async fn start_event_processor( + publisher: P, + worker_id: u64, + cancellation_token: CancellationToken, + rx: mpsc::UnboundedReceiver, + local_indexer: Option>, + batching_timeout_ms: Option, +) { + run_event_processor_loop( + publisher, + worker_id, + cancellation_token, + rx, + local_indexer, + batching_timeout_ms, + DEFAULT_MAX_BATCH_BLOCKS, + ) + .await +} + +pub(super) async fn start_event_processor_jetstream( + publisher: NatsQueue, + worker_id: u64, + cancellation_token: CancellationToken, + rx: mpsc::UnboundedReceiver, + local_indexer: Option>, + batching_timeout_ms: Option, +) { + run_event_processor_loop( + JetStreamPublisher(publisher), + worker_id, + cancellation_token, + rx, + local_indexer, + batching_timeout_ms, + DEFAULT_MAX_BATCH_BLOCKS, + ) + .await +} diff --git a/lib/llm/src/kv_router/publisher/mod.rs b/lib/llm/src/kv_router/publisher/mod.rs new file mode 100644 index 000000000000..357b7bf13abb --- /dev/null +++ b/lib/llm/src/kv_router/publisher/mod.rs @@ -0,0 +1,396 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::sync::OnceLock; +use std::sync::atomic::{AtomicU64, Ordering}; + +use anyhow::Result; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use dynamo_kv_router::indexer::{KvIndexerMetrics, LocalKvIndexer}; +use dynamo_kv_router::protocols::*; +pub use dynamo_kv_router::zmq_wire::create_stored_blocks; +#[cfg(test)] +use dynamo_kv_router::zmq_wire::*; +use dynamo_runtime::config::environment_names::nats as env_nats; +use dynamo_runtime::metrics::MetricsHierarchy; +use dynamo_runtime::metrics::prometheus_names::kv_publisher; +use dynamo_runtime::traits::DistributedRuntimeProvider; +use dynamo_runtime::{ + component::Component, + transports::nats::{NatsQueue, Slug}, +}; + +use crate::kv_router::{ + KV_EVENT_SUBJECT, WORKER_KV_INDEXER_BUFFER_SIZE, worker_query::start_worker_kv_query_endpoint, +}; + +mod event_processor; +#[cfg(test)] +mod tests; +mod worker_metrics; +mod zmq_listener; + +#[cfg(test)] +use event_processor::{BatchingState, run_event_processor_loop}; +use event_processor::{ + EventPlanePublisher, start_event_processor, start_event_processor_jetstream, +}; +pub use worker_metrics::WorkerMetricsPublisher; +use zmq_listener::start_zmq_listener; +#[cfg(test)] +use zmq_listener::{ + INITIAL_BACKOFF_MS, MAX_BACKOFF_EXPONENT, MAX_BACKOFF_MS, MAX_CONSECUTIVE_ERRORS, + calculate_backoff_ms, +}; + +const MAX_BATCHING_TIMEOUT_MS: u64 = 15_000; +pub const DEFAULT_BATCHING_TIMEOUT_MS: Option = None; +const DEFAULT_MAX_BATCH_BLOCKS: usize = 128; + +/// Helper function to create a KV stream name from a component and subject. +/// +/// Generates a slugified stream name in the format: +/// `namespace-{namespace}-component-{component}-{subject}` +fn create_kv_stream_name(component: &Component, subject: &str) -> String { + Slug::slugify(&format!( + "namespace.{}.component.{}.{}", + component.namespace().name(), + component.name(), + subject + )) + .to_string() + .replace("_", "-") +} + +/// Metrics for the KV publisher, created via the MetricsHierarchy API. +/// This provides automatic `dynamo_namespace`, `dynamo_component`, and other +/// hierarchy labels for free. +pub(super) struct KvPublisherMetrics { + /// Total number of raw events dropped by engines before reaching publisher + pub engines_dropped_events_total: prometheus::IntCounterVec, +} + +static KV_PUBLISHER_METRICS: OnceLock> = OnceLock::new(); + +impl KvPublisherMetrics { + /// Create from a Component, memoized in a static OnceLock. + /// Uses the MetricsHierarchy API which auto-prepends `dynamo_component_`, + /// injects hierarchy labels, and registers with the DRT `MetricsRegistry`. + pub fn from_component(component: &Component) -> Arc { + KV_PUBLISHER_METRICS + .get_or_init(|| { + let metrics = component.metrics(); + match metrics.create_intcountervec( + kv_publisher::ENGINES_DROPPED_EVENTS_TOTAL, + "Total number of raw events dropped by engines before reaching publisher (detected via event_id gaps)", + &["worker_id"], + &[], + ) { + Ok(engines_dropped_events_total) => { + Arc::new(Self { engines_dropped_events_total }) + } + Err(e) => { + tracing::warn!("Failed to create kv_publisher metrics from component: {}. Using unregistered metrics as fallback.", e); + Arc::new(Self::new_unregistered()) + } + } + }) + .clone() + } + + /// Creates unregistered metrics for use when the MetricsRegistry is not available. + /// This is used as a fallback when metric creation fails. + pub fn new_unregistered() -> Self { + Self { + engines_dropped_events_total: prometheus::IntCounterVec::new( + prometheus::Opts::new( + kv_publisher::ENGINES_DROPPED_EVENTS_TOTAL, + "Total number of raw events dropped by engines before reaching publisher (detected via event_id gaps)", + ), + &["worker_id"], + ) + .expect("failed to create engines_dropped_events_total counter"), + } + } + + /// Increment the engines dropped events counter by the given amount. + pub fn increment_engines_dropped_events(&self, worker_id: u64, count: u64) { + self.engines_dropped_events_total + .with_label_values(&[&worker_id.to_string()]) + .inc_by(count); + } +} + +fn kv_publisher_metrics() -> Option> { + KV_PUBLISHER_METRICS.get().cloned() +} + +/// Configure the source of KV events. +/// Currently, only ZMQ is supported. +pub enum KvEventSourceConfig { + Zmq { endpoint: String, topic: String }, +} + +enum KvEventSource { + Zmq { + zmq_handle: tokio::task::JoinHandle<()>, + }, +} + +impl KvEventSource { + fn start( + component: Component, + worker_id: WorkerId, + kv_block_size: u32, + source_config: KvEventSourceConfig, + cancellation_token: CancellationToken, + tx: mpsc::UnboundedSender, + next_event_id: Arc, + ) -> Result { + match source_config { + KvEventSourceConfig::Zmq { endpoint, topic } => { + let zmq_handle = component + .drt() + .runtime() + .secondary() + .spawn(start_zmq_listener( + endpoint, + topic, + worker_id, + tx, + cancellation_token.clone(), + kv_block_size, + next_event_id, + )); + + Ok(KvEventSource::Zmq { zmq_handle }) + } + } + } + + fn shutdown(&self) { + match self { + KvEventSource::Zmq { zmq_handle } => { + zmq_handle.abort(); + } + } + } +} + +/// A publisher of KV events. +pub struct KvEventPublisher { + /// The size of the KV block. + kv_block_size: u32, + /// The source of KV events. + /// Can be `None` if all events provided through [`KvEventPublisher::publish`]. + source: Option, + /// The cancellation token. + cancellation_token: CancellationToken, + /// The ID of the local worker emitting placement events. + worker_id: WorkerId, + /// The channel to send events to. + tx: mpsc::UnboundedSender, + /// Internal monotonic event ID counter. Shared with the ZMQ listener if present. + next_event_id: Arc, +} + +impl KvEventPublisher { + pub fn new( + component: Component, + kv_block_size: u32, + source_config: Option, + ) -> Result { + Self::new_with_local_indexer( + component, + kv_block_size, + source_config, + false, + 0, + DEFAULT_BATCHING_TIMEOUT_MS, + ) + } + + pub fn new_with_local_indexer( + component: Component, + kv_block_size: u32, + source_config: Option, + enable_local_indexer: bool, + dp_rank: DpRank, + batching_timeout_ms: Option, + ) -> Result { + let cancellation_token = CancellationToken::new(); + let batching_timeout_ms = batching_timeout_ms + .filter(|&ms| { + if ms > MAX_BATCHING_TIMEOUT_MS { + tracing::warn!( + requested_ms = ms, + max_ms = MAX_BATCHING_TIMEOUT_MS, + "batching_timeout_ms too high, capping to 15s" + ); + } + ms > 0 + }) + .map(|ms| ms.min(MAX_BATCHING_TIMEOUT_MS)); + + let (tx, rx) = mpsc::unbounded_channel::(); + let worker_id = component.drt().connection_id(); + + KvPublisherMetrics::from_component(&component); + + let component_name = component.name(); + tracing::info!( + "Initializing KvEventPublisher for worker {worker_id} in component {component_name}" + ); + + if enable_local_indexer { + tracing::info!( + "LocalKvIndexer enabled for worker {worker_id} in component {component_name}" + ); + } + + let next_event_id = Arc::new(AtomicU64::new(0)); + + let mut source = None; + if let Some(config) = source_config { + source = Some(KvEventSource::start( + component.clone(), + worker_id, + kv_block_size, + config, + cancellation_token.clone(), + tx.clone(), + next_event_id.clone(), + )?); + } + + let local_indexer = if enable_local_indexer { + let metrics = Arc::new(KvIndexerMetrics::new_unregistered()); + Some(Arc::new(LocalKvIndexer::new( + cancellation_token.clone(), + kv_block_size, + metrics, + WORKER_KV_INDEXER_BUFFER_SIZE, + ))) + } else { + None + }; + + let _local_indexer_query_handle = local_indexer.as_ref().map(|local_indexer_ref| { + let component = component.clone(); + let local_indexer = local_indexer_ref.clone(); + + component + .drt() + .runtime() + .secondary() + .spawn(start_worker_kv_query_endpoint( + component, + worker_id, + dp_rank, + local_indexer, + )) + }); + + let cancellation_token_clone = cancellation_token.clone(); + let local_indexer_clone = local_indexer.clone(); + + if enable_local_indexer { + tracing::info!("Using event plane for KV event publishing (local_indexer mode)"); + let component_clone = component.clone(); + component.drt().runtime().secondary().spawn(async move { + let event_publisher = + match dynamo_runtime::transports::event_plane::EventPublisher::for_component( + &component_clone, + KV_EVENT_SUBJECT, + ) + .await + { + Ok(publisher) => publisher, + Err(e) => { + tracing::error!("Failed to create event publisher: {}", e); + return; + } + }; + + start_event_processor( + EventPlanePublisher(event_publisher), + worker_id, + cancellation_token_clone, + rx, + local_indexer_clone, + batching_timeout_ms, + ) + .await + }); + } else { + let stream_name = create_kv_stream_name(&component, KV_EVENT_SUBJECT); + let nats_server = std::env::var(env_nats::NATS_SERVER) + .unwrap_or_else(|_| "nats://localhost:4222".to_string()); + let mut nats_queue = NatsQueue::new_without_consumer( + stream_name, + nats_server, + std::time::Duration::from_secs(60), + ); + + component.drt().runtime().secondary().spawn(async move { + if let Err(e) = nats_queue.connect().await { + tracing::error!("Failed to connect NatsQueue: {e}"); + return; + } + start_event_processor_jetstream( + nats_queue, + worker_id, + cancellation_token_clone, + rx, + local_indexer_clone, + batching_timeout_ms, + ) + .await + }); + } + + Ok(Self { + kv_block_size, + source, + cancellation_token, + worker_id, + tx, + next_event_id, + }) + } + + pub fn publish(&self, event: KvCacheEvent) -> Result<(), mpsc::error::SendError> { + let placement_event = PlacementEvent::local_gpu(self.worker_id, event); + match self.tx.send(placement_event) { + Ok(()) => Ok(()), + Err(err) => Err(mpsc::error::SendError(err.0.event)), + } + } + + pub fn next_event_id(&self) -> u64 { + self.next_event_id.fetch_add(1, Ordering::SeqCst) + } + + pub fn kv_block_size(&self) -> u32 { + self.kv_block_size + } + + pub fn shutdown(&mut self) { + if !self.cancellation_token.is_cancelled() { + self.cancellation_token.cancel(); + } + + if let Some(source) = self.source.take() { + source.shutdown(); + } + } +} + +impl Drop for KvEventPublisher { + fn drop(&mut self) { + self.shutdown(); + } +} diff --git a/lib/llm/src/kv_router/publisher.rs b/lib/llm/src/kv_router/publisher/tests.rs similarity index 67% rename from lib/llm/src/kv_router/publisher.rs rename to lib/llm/src/kv_router/publisher/tests.rs index 5cc7b9015afe..c4c07e67ffbf 100644 --- a/lib/llm/src/kv_router/publisher.rs +++ b/lib/llm/src/kv_router/publisher/tests.rs @@ -1,1049 +1,20 @@ // SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use super::*; +#[allow(unused_imports)] +use bytes::Bytes; +#[allow(unused_imports)] +use dynamo_kv_router::RouterEventSink; +#[allow(unused_imports)] +use rmp_serde as rmps; +#[allow(unused_imports)] use std::future::Future; -use std::sync::Arc; +#[allow(unused_imports)] use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; -use std::time::{Duration, Instant}; - -use anyhow::Result; -use rmp_serde as rmps; -use tokio::sync::mpsc; -use tokio_util::sync::CancellationToken; -use zeromq::{Socket, SocketRecv, SubSocket}; - -use dynamo_runtime::metrics::MetricsHierarchy; -use dynamo_runtime::traits::DistributedRuntimeProvider; -use dynamo_runtime::transports::event_plane::EventPublisher; -use dynamo_runtime::{ - component::{Component, Namespace}, - transports::nats::{NatsQueue, Slug}, -}; - -/// Helper function to create a KV stream name from a component and subject. -/// -/// Generates a slugified stream name in the format: -/// `namespace-{namespace}-component-{component}-{subject}` -fn create_kv_stream_name(component: &Component, subject: &str) -> String { - Slug::slugify(&format!( - "namespace.{}.component.{}.{}", - component.namespace().name(), - component.name(), - subject - )) - .to_string() - .replace("_", "-") -} - -use dynamo_kv_router::indexer::{KvIndexerMetrics, LocalKvIndexer}; -use dynamo_kv_router::protocols::*; -pub use dynamo_kv_router::zmq_wire::create_stored_blocks; -use dynamo_kv_router::zmq_wire::*; - -use crate::kv_router::{ - KV_EVENT_SUBJECT, KV_METRICS_SUBJECT, WORKER_KV_INDEXER_BUFFER_SIZE, - worker_query::start_worker_kv_query_endpoint, -}; -use dynamo_runtime::config::environment_names::nats as env_nats; - -// Error handling configuration for ZMQ operations -const INITIAL_BACKOFF_MS: u64 = 10; -const MAX_BACKOFF_MS: u64 = 5000; -const MAX_CONSECUTIVE_ERRORS: u32 = 10; -const MAX_BACKOFF_EXPONENT: u32 = 8; // Cap at 2^8 = 256x multiplier to prevent overflow - -// Batching configuration -const MAX_BATCHING_TIMEOUT_MS: u64 = 15_000; // 15 seconds, prevents misconfiguration -pub const DEFAULT_BATCHING_TIMEOUT_MS: Option = None; // disabled by default -const DEFAULT_MAX_BATCH_BLOCKS: usize = 128; // Max blocks to batch before flushing - -// --------------------------------------------------------------------------- -// Engines dropped events metric -// --------------------------------------------------------------------------- - -use std::sync::OnceLock; - -use dynamo_runtime::metrics::prometheus_names::kv_publisher; - -/// Metrics for the KV publisher, created via the MetricsHierarchy API. -/// This provides automatic `dynamo_namespace`, `dynamo_component`, and other -/// hierarchy labels for free. -pub struct KvPublisherMetrics { - /// Total number of raw events dropped by engines before reaching publisher - pub engines_dropped_events_total: prometheus::IntCounterVec, -} - -static KV_PUBLISHER_METRICS: OnceLock> = OnceLock::new(); - -impl KvPublisherMetrics { - /// Create from a Component, memoized in a static OnceLock. - /// Uses the MetricsHierarchy API which auto-prepends `dynamo_component_`, - /// injects hierarchy labels, and registers with the DRT `MetricsRegistry`. - pub fn from_component(component: &Component) -> Arc { - KV_PUBLISHER_METRICS - .get_or_init(|| { - let metrics = component.metrics(); - match metrics.create_intcountervec( - kv_publisher::ENGINES_DROPPED_EVENTS_TOTAL, - "Total number of raw events dropped by engines before reaching publisher (detected via event_id gaps)", - &["worker_id"], - &[], - ) { - Ok(engines_dropped_events_total) => { - Arc::new(Self { engines_dropped_events_total }) - } - Err(e) => { - tracing::warn!("Failed to create kv_publisher metrics from component: {}. Using unregistered metrics as fallback.", e); - Arc::new(Self::new_unregistered()) - } - } - }) - .clone() - } - - /// Creates unregistered metrics for use when the MetricsRegistry is not available. - /// This is used as a fallback when metric creation fails. - pub fn new_unregistered() -> Self { - Self { - engines_dropped_events_total: prometheus::IntCounterVec::new( - prometheus::Opts::new( - kv_publisher::ENGINES_DROPPED_EVENTS_TOTAL, - "Total number of raw events dropped by engines before reaching publisher (detected via event_id gaps)", - ), - &["worker_id"], - ) - .expect("failed to create engines_dropped_events_total counter"), - } - } - - /// Increment the engines dropped events counter by the given amount. - pub fn increment_engines_dropped_events(&self, worker_id: u64, count: u64) { - self.engines_dropped_events_total - .with_label_values(&[&worker_id.to_string()]) - .inc_by(count); - } -} - -/// Get the KV publisher metrics if initialized. -fn kv_publisher_metrics() -> Option> { - KV_PUBLISHER_METRICS.get().cloned() -} - -// ------------------------------------------------------------------------- -// Batching State ----------------------------------------------------------- -// ------------------------------------------------------------------------- - -/// Accumulator for in-flight KV cache events that will be merged into a single -/// [`RouterEvent`] before being forwarded to the event sink. -#[derive(Debug)] -struct BatchingState { - /// Block hashes accumulating for the next Removed event. - pending_removed: Option, - /// Blocks accumulating for the next Stored event. - pending_stored: Option, - /// Monotonic published-batch counter. Increments by 1 per flush so downstream - /// consumers always see consecutive event IDs, regardless of how many raw source - /// events were merged into the batch. - next_publish_id: u64, - /// dp_rank of the events in the current pending batch. - /// A change signals that the batch must be flushed before accumulating further. - last_dp_rank: u32, - /// When we last flushed (or initialized). Used to detect stale pending data: - /// if a new event arrives after a long idle period (exceeding timeout), - /// we flush immediately for lower latency on sparse important events. - last_flush_time: Instant, -} - -impl BatchingState { - fn new() -> Self { - Self { - pending_removed: None, - pending_stored: None, - next_publish_id: 1, - last_dp_rank: 0, - last_flush_time: Instant::now(), - } - } - - fn has_pending(&self) -> bool { - self.pending_removed.is_some() || self.pending_stored.is_some() - } - - fn pending_block_count(&self) -> usize { - self.pending_removed - .as_ref() - .map(|r| r.block_hashes.len()) - .unwrap_or(0) - + self - .pending_stored - .as_ref() - .map(|s| s.blocks.len()) - .unwrap_or(0) - } - - /// Records that a flush just happened. Called after every flush to track - /// idle periods for stale-data detection. - fn record_flush_time(&mut self) { - self.last_flush_time = Instant::now(); - } - - /// Returns the time remaining in the current batch window (zero if already elapsed). - fn remaining_timeout(&self, timeout_ms: u64) -> Duration { - let timeout = Duration::from_millis(timeout_ms); - let elapsed = self.last_flush_time.elapsed(); - if elapsed >= timeout { - Duration::ZERO - } else { - timeout - elapsed - } - } - - /// Returns `true` when the batch window has elapsed (or `timeout_ms` is zero). - fn is_timeout_elapsed(&self, timeout_ms: u64) -> bool { - self.remaining_timeout(timeout_ms) == Duration::ZERO - } -} - -// ------------------------------------------------------------------------- -// KV Event Publishers ----------------------------------------------------- -// ------------------------------------------------------------------------- - -/// Configure the source of KV events. -/// Currently, only ZMQ is supported. -pub enum KvEventSourceConfig { - Zmq { endpoint: String, topic: String }, -} - -/// The source of KV events. -enum KvEventSource { - Zmq { - zmq_handle: tokio::task::JoinHandle<()>, - }, -} - -impl KvEventSource { - /// Start the event source from a [`KvEventSourceConfig`]. - fn start( - component: Component, - worker_id: WorkerId, - kv_block_size: u32, - source_config: KvEventSourceConfig, - cancellation_token: CancellationToken, - tx: mpsc::UnboundedSender, - next_event_id: Arc, - ) -> Result { - match source_config { - KvEventSourceConfig::Zmq { endpoint, topic } => { - let zmq_handle = component - .drt() - .runtime() - .secondary() - .spawn(start_zmq_listener( - endpoint, - topic, - worker_id, - tx, - cancellation_token.clone(), - kv_block_size, - next_event_id, - )); - - Ok(KvEventSource::Zmq { zmq_handle }) - } - } - } - - fn shutdown(&self) { - match self { - KvEventSource::Zmq { zmq_handle } => { - zmq_handle.abort(); - } - } - } -} - -/// A publisher of KV events. -pub struct KvEventPublisher { - /// The size of the KV block. - kv_block_size: u32, - /// The source of KV events. - /// Can be `None` if all events provided through [`KvEventPublisher::publish`]. - source: Option, - /// The cancellation token. - cancellation_token: CancellationToken, - /// The ID of the local worker emitting placement events. - worker_id: WorkerId, - /// The channel to send events to. - tx: mpsc::UnboundedSender, - /// Internal monotonic event ID counter - ensures each event gets a unique, incrementing ID. - /// Shared with the ZMQ listener (if any) to maintain consistency. - next_event_id: Arc, -} - -impl KvEventPublisher { - pub fn new( - component: Component, - kv_block_size: u32, - source_config: Option, - ) -> Result { - Self::new_with_local_indexer( - component, - kv_block_size, - source_config, - false, - 0, - DEFAULT_BATCHING_TIMEOUT_MS, - ) - } - - pub fn new_with_local_indexer( - component: Component, - kv_block_size: u32, - source_config: Option, - enable_local_indexer: bool, - dp_rank: DpRank, - batching_timeout_ms: Option, - ) -> Result { - let cancellation_token = CancellationToken::new(); - // None = disabled (flush every event); Some(0) normalised to None; Some(ms) = opt-in. - // Cap at MAX_BATCHING_TIMEOUT_MS to prevent misconfiguration. - let batching_timeout_ms = batching_timeout_ms - .filter(|&ms| { - if ms > MAX_BATCHING_TIMEOUT_MS { - tracing::warn!( - requested_ms = ms, - max_ms = MAX_BATCHING_TIMEOUT_MS, - "batching_timeout_ms too high, capping to 15s" - ); - } - // if ms is 0, treat as disabled (None) - ms > 0 - }) - .map(|ms| ms.min(MAX_BATCHING_TIMEOUT_MS)); - - let (tx, rx) = mpsc::unbounded_channel::(); - - // Infer worker_id from component's connection - let worker_id = component.drt().connection_id(); - - // Initialize the KV publisher metrics via MetricsHierarchy API - // This provides automatic hierarchy labels (dynamo_namespace, dynamo_component, etc.) - KvPublisherMetrics::from_component(&component); - - let component_name = component.name(); - tracing::info!( - "Initializing KvEventPublisher for worker {worker_id} in component {component_name}" - ); - - if enable_local_indexer { - tracing::info!( - "LocalKvIndexer enabled for worker {worker_id} in component {component_name}" - ); - } - - // Internal monotonic event ID counter - shared with ZMQ listener if any - let next_event_id = Arc::new(AtomicU64::new(0)); - - // Create our event source (if any) - let mut source = None; - if let Some(config) = source_config { - source = Some(KvEventSource::start( - component.clone(), - worker_id, - kv_block_size, - config, - cancellation_token.clone(), - tx.clone(), - next_event_id.clone(), - )?); - } - - // Create local indexer if requested - let local_indexer = if enable_local_indexer { - let metrics = Arc::new(KvIndexerMetrics::new_unregistered()); - Some(Arc::new(LocalKvIndexer::new( - cancellation_token.clone(), - kv_block_size, - metrics, - WORKER_KV_INDEXER_BUFFER_SIZE, - ))) - } else { - None - }; - - // Spawn runtime for router->local indexer comm if requested - let _local_indexer_query_handle = local_indexer.as_ref().map(|local_indexer_ref| { - let component = component.clone(); - let local_indexer = local_indexer_ref.clone(); - - component - .drt() - .runtime() - .secondary() - .spawn(start_worker_kv_query_endpoint( - component, - worker_id, - dp_rank, - local_indexer, - )) - }); - - let cancellation_token_clone = cancellation_token.clone(); - let local_indexer_clone = local_indexer.clone(); - - if enable_local_indexer { - // When local indexer is enabled, use the event plane directly. - // EventPublisher handles transport selection (ZMQ or NATS) based on environment. - // Durability is provided by the local indexer's event buffer. - tracing::info!("Using event plane for KV event publishing (local_indexer mode)"); - let component_clone = component.clone(); - component.drt().runtime().secondary().spawn(async move { - let event_publisher = - match EventPublisher::for_component(&component_clone, KV_EVENT_SUBJECT).await { - Ok(publisher) => publisher, - Err(e) => { - tracing::error!("Failed to create event publisher: {}", e); - return; - } - }; - - start_event_processor( - EventPlanePublisher(event_publisher), - worker_id, - cancellation_token_clone, - rx, - local_indexer_clone, - batching_timeout_ms, - ) - .await - }); - } else { - // When local indexer is disabled, use JetStream (NatsQueue) for durability. - let stream_name = create_kv_stream_name(&component, KV_EVENT_SUBJECT); - let nats_server = std::env::var(env_nats::NATS_SERVER) - .unwrap_or_else(|_| "nats://localhost:4222".to_string()); - let mut nats_queue = NatsQueue::new_without_consumer( - stream_name, - nats_server, - std::time::Duration::from_secs(60), // 1 minute timeout - ); - - component.drt().runtime().secondary().spawn(async move { - if let Err(e) = nats_queue.connect().await { - tracing::error!("Failed to connect NatsQueue: {e}"); - return; - } - start_event_processor_jetstream( - JetStreamPublisher(nats_queue), - worker_id, - cancellation_token_clone, - rx, - local_indexer_clone, - batching_timeout_ms, - ) - .await - }); - } - - Ok(Self { - kv_block_size, - source, - cancellation_token, - worker_id, - tx, - next_event_id, - }) - } - - pub fn publish(&self, event: KvCacheEvent) -> Result<(), mpsc::error::SendError> { - let placement_event = PlacementEvent::local_gpu(self.worker_id, event); - match self.tx.send(placement_event) { - Ok(()) => Ok(()), - Err(err) => Err(mpsc::error::SendError(err.0.event)), - } - } - - /// Get and increment the next event ID atomically. - /// Use this to assign monotonically increasing event IDs to events before publishing. - pub fn next_event_id(&self) -> u64 { - self.next_event_id.fetch_add(1, Ordering::SeqCst) - } - - pub fn kv_block_size(&self) -> u32 { - self.kv_block_size - } - - pub fn shutdown(&mut self) { - if !self.cancellation_token.is_cancelled() { - self.cancellation_token.cancel(); - } - - if let Some(source) = self.source.take() { - source.shutdown(); - } - } -} - -impl Drop for KvEventPublisher { - fn drop(&mut self) { - self.shutdown(); - } -} - -use dynamo_kv_router::RouterEventSink; - -struct EventPlanePublisher(EventPublisher); - -impl RouterEventSink for EventPlanePublisher { - fn publish_event(&self, event: &RouterEvent) -> impl Future> + Send { - self.0.publish(event) - } -} - -struct JetStreamPublisher(NatsQueue); - -impl RouterEventSink for JetStreamPublisher { - fn publish_event(&self, event: &RouterEvent) -> impl Future> + Send { - NatsQueue::publish_event(&self.0, KV_EVENT_SUBJECT, event) - } -} - -/// Publishes a single [`KvCacheEvent`] to the event sink and, when present, the local indexer. -/// Errors are logged and swallowed so the caller loop can continue uninterrupted. -async fn emit( - publisher: &P, - local_indexer: &Option>, - worker_id: u64, - event: KvCacheEvent, -) { - let router_event = RouterEvent::new(worker_id, event); - if let Some(indexer) = local_indexer - && let Err(e) = indexer.apply_event_with_buffer(router_event.clone()).await - { - tracing::warn!(worker_id, error = %e, "Failed to apply event to local indexer"); - } - if let Err(e) = publisher.publish_event(&router_event).await { - tracing::error!(worker_id, error = %e, "Failed to publish event"); - } -} - -impl BatchingState { - /// Publishes any pending batch as a single [`RouterEvent`] and advances the monotonic - /// batch ID. No-ops when nothing is pending, so callers may call unconditionally. - async fn flush( - &mut self, - publisher: &P, - local_indexer: &Option>, - worker_id: u64, - ) { - if !self.has_pending() { - return; - } - let id = self.next_publish_id; - let dp_rank = self.last_dp_rank; - if let Some(data) = self.pending_removed.take() { - emit( - publisher, - local_indexer, - worker_id, - KvCacheEvent { - event_id: id, - data: KvCacheEventData::Removed(data), - dp_rank, - }, - ) - .await; - } - if let Some(data) = self.pending_stored.take() { - emit( - publisher, - local_indexer, - worker_id, - KvCacheEvent { - event_id: id, - data: KvCacheEventData::Stored(data), - dp_rank, - }, - ) - .await; - } - // Consecutive batch IDs (1, 2, 3, …) keep downstream gap-detection happy. - self.next_publish_id += 1; - // Record when we flushed for stale-data detection on next event. - self.record_flush_time(); - } -} - -/// Batching loop: accumulates Removed/Stored events and flushes them as a single -/// [`RouterEvent`] when any of the following conditions are met: -/// - Event type switches (Removed ↔ Stored) -/// - `dp_rank` changes between consecutive events -/// - A `Stored` event's `parent_hash` breaks the sequential chain -/// - The batch window expires (`Some(timeout_ms)`; `None` = disabled, flush every event) -/// - Channel is closed or a cancellation signal is received -async fn run_event_processor_loop( - publisher: P, - worker_id: u64, - cancellation_token: CancellationToken, - mut rx: mpsc::UnboundedReceiver, - local_indexer: Option>, - timeout_ms: Option, - max_batch_blocks: usize, -) { - let mut batching_state = BatchingState::new(); - // Track last raw input event_id for gap detection (dropped events before batching). - // The raw event_id is a globally monotonic counter assigned by the ZMQ listener, - // so any gap here means events were silently dropped (e.g. send error on the channel). - let mut last_raw_input_id: Option = None; - - loop { - tokio::select! { - _ = cancellation_token.cancelled() => { - tracing::info!("KV Event source received cancellation signal"); - batching_state.flush(&publisher, &local_indexer, worker_id).await; - break; - } - event = rx.recv() => { - let Some(placement_event) = event else { - tracing::debug!("Event processor channel closed."); - batching_state.flush(&publisher, &local_indexer, worker_id).await; - break; - }; - - // Warn if the raw input event_id is not consecutive — events were dropped - // (e.g. channel send error) before they reached the batching layer. - let raw_event_id = placement_event.event.event_id; - if let Some(last_id) = last_raw_input_id - && raw_event_id > last_id + 1 - { - let gap = raw_event_id - last_id - 1; - tracing::warn!( - worker_id, - last_raw_input_id = last_id, - raw_event_id, - gap, - "Input event gap detected: raw events dropped before batching" - ); - // Increment Prometheus counter for dropped events (if initialized) - if let Some(metrics) = kv_publisher_metrics() { - metrics.increment_engines_dropped_events(worker_id, gap); - } else { - tracing::warn!( - worker_id, - gap, - "Failed to record dropped events metric: metrics not initialized" - ); - } - } - last_raw_input_id = Some(raw_event_id); - - if !placement_event.placement.is_local_gpu() { - tracing::trace!( - worker_id, - ?placement_event.placement, - event_id = placement_event.event.event_id, - "Skipping non-local-GPU placement event" - ); - continue; - } - - let event = placement_event.event; - tracing::trace!("Event processor for worker_id {} processing event: {:?}", worker_id, event.data); - - let dp_rank_changed = batching_state.has_pending() - && event.dp_rank != batching_state.last_dp_rank; - - match event.data { - KvCacheEventData::Removed(data) => { - if batching_state.pending_stored.is_some() || dp_rank_changed { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - } - match &mut batching_state.pending_removed { - Some(pending) => pending.block_hashes.extend(data.block_hashes), - None => { - batching_state.pending_removed = Some(data); - } - } - } - KvCacheEventData::Stored(data) => { - // Flush if: type switch, dp_rank change, or the chain is broken - // (new event's parent_hash doesn't continue from the last stored block). - let should_flush = dp_rank_changed - || batching_state.pending_removed.is_some() - || batching_state.pending_stored.as_ref().is_some_and(|p| { - data.parent_hash != p.blocks.last().map(|b| b.block_hash) - }); - if should_flush { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - } - match &mut batching_state.pending_stored { - // Only extend blocks; parent_hash stays fixed from the first event. - Some(pending) => pending.blocks.extend(data.blocks), - None => { - batching_state.pending_stored = Some(data); - } - } - } - KvCacheEventData::Cleared => { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - emit(&publisher, &local_indexer, worker_id, KvCacheEvent { - event_id: batching_state.next_publish_id, - data: KvCacheEventData::Cleared, - dp_rank: event.dp_rank, - }).await; - batching_state.next_publish_id += 1; - } - } - - // Track dp_rank after the match so in-flight flushes use the old value. - batching_state.last_dp_rank = event.dp_rank; - - // Flush after every event when disabled (None), or when the window has elapsed, - // or when the batch exceeds the max block count. - // The sleep arm only arms when batching is enabled; this covers the disabled path. - if batching_state.has_pending() - && (timeout_ms.is_none_or(|ms| batching_state.is_timeout_elapsed(ms)) - || batching_state.pending_block_count() > max_batch_blocks) - { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - } - } - // if has some pending and has timeout, and no new events come in, then flush when timeout elapsed to prevent stale events - _ = tokio::time::sleep( - timeout_ms.map(|ms| batching_state.remaining_timeout(ms)).unwrap_or(Duration::from_secs(3600)) - ), if timeout_ms.is_some() && batching_state.has_pending() => { - batching_state.flush(&publisher, &local_indexer, worker_id).await; - } - } - } -} - -/// Batched event processor for ephemeral transports (NATS Core / ZMQ). -async fn start_event_processor( - publisher: P, - worker_id: u64, - cancellation_token: CancellationToken, - rx: mpsc::UnboundedReceiver, - local_indexer: Option>, - batching_timeout_ms: Option, -) { - run_event_processor_loop( - publisher, - worker_id, - cancellation_token, - rx, - local_indexer, - batching_timeout_ms, - DEFAULT_MAX_BATCH_BLOCKS, - ) - .await -} - -/// Batched event processor using JetStream (durable). -async fn start_event_processor_jetstream( - publisher: P, - worker_id: u64, - cancellation_token: CancellationToken, - rx: mpsc::UnboundedReceiver, - local_indexer: Option>, - batching_timeout_ms: Option, -) { - run_event_processor_loop( - publisher, - worker_id, - cancellation_token, - rx, - local_indexer, - batching_timeout_ms, - DEFAULT_MAX_BATCH_BLOCKS, - ) - .await -} - -/// Calculate exponential backoff duration based on consecutive error count -fn calculate_backoff_ms(consecutive_errors: u32) -> u64 { - std::cmp::min( - INITIAL_BACKOFF_MS * 2_u64.pow(consecutive_errors.min(MAX_BACKOFF_EXPONENT)), - MAX_BACKOFF_MS, - ) -} - -pub async fn start_zmq_listener( - zmq_endpoint: String, - zmq_topic: String, - worker_id: WorkerId, - tx: mpsc::UnboundedSender, - cancellation_token: CancellationToken, - kv_block_size: u32, - next_event_id: Arc, -) { - tracing::debug!( - "KVEventPublisher connecting to ZMQ endpoint {} (topic '{}')", - zmq_endpoint, - zmq_topic - ); - - let warning_count = Arc::new(AtomicU32::new(0)); - - let mut socket = SubSocket::new(); - - // Subscribe to the requested topic (empty string == all topics) - if let Err(e) = socket.subscribe(&zmq_topic).await { - tracing::error!("Failed to subscribe on ZMQ socket: {}", e); - return; - } - - // Connect to the ZMQ endpoint. SGLang binds locally, Dynamo connects. - // In multi-node setups, each node runs dynamo.sglang alongside local SGLang ranks, - // so ZMQ connections are always local. NATS handles cross-node event distribution. - if let Err(e) = socket.connect(&zmq_endpoint).await { - tracing::error!("Failed to connect ZMQ SUB socket to {zmq_endpoint}: {e}"); - return; - } - - let mut consecutive_errors = 0u32; - #[expect(unused_assignments)] - let mut exit_reason = "unknown"; - let mut messages_processed = 0u64; - - 'main: loop { - tokio::select! { - biased; - - // Check for cancellation - _ = cancellation_token.cancelled() => { - tracing::debug!("ZMQ listener received cancellation signal"); - exit_reason = "cancellation token cancelled"; - break 'main; - } - - // Receive message - msg_result = socket.recv() => { - let Ok(msg) = msg_result else { - let e = msg_result.unwrap_err(); - consecutive_errors += 1; - - if consecutive_errors >= MAX_CONSECUTIVE_ERRORS { - tracing::error!( - error=%e, - consecutive_errors=%consecutive_errors, - "Too many consecutive ZMQ errors, terminating listener" - ); - exit_reason = "too many consecutive errors"; - break 'main; - } - - // Simple exponential backoff with max exponent to prevent overflow - let backoff_ms = calculate_backoff_ms(consecutive_errors); - - tracing::warn!( - error=%e, - consecutive_errors=%consecutive_errors, - backoff_ms=%backoff_ms, - "Error reading from ZMQ socket, applying exponential backoff" - ); - - tokio::time::sleep(Duration::from_millis(backoff_ms)).await; - continue; - }; - // Reset error count on successful message - consecutive_errors = 0; - - // We expect multipart frames: [topic, seq, payload] - let mut frames: Vec> = msg.into_vec().into_iter().map(|frame| frame.to_vec()).collect(); - - if frames.len() != 3 { - tracing::warn!("Received unexpected ZMQ frame count: expected 3, actual {}", frames.len()); - continue; - } - - // Extract the payload and sequence number. - let payload = frames.pop().unwrap(); - let seq_bytes = frames.pop().unwrap(); - - if seq_bytes.len() != 8 { - tracing::warn!("Invalid sequence number byte length: expected 8, actual {}", seq_bytes.len()); - continue; - } - - // Note: We extract the engine's sequence number for logging but use our own - // internal monotonic counter for event_id to ensure per-dp_rank monotonicity - let engine_seq = u64::from_be_bytes(seq_bytes.try_into().unwrap()); - - // Decode our batch of events. - let batch_result = rmps::from_slice::(&payload); - let Ok(batch) = batch_result else { - let e = batch_result.unwrap_err(); - tracing::warn!("Failed to decode KVEventBatch msgpack: {e}"); - continue; - }; - - tracing::trace!( - "ZMQ listener on {} received batch with {} events (engine_seq={}, dp_rank={})", - zmq_endpoint, - batch.events.len(), - engine_seq, - batch.data_parallel_rank.unwrap_or(0) - ); - - let dp_rank = batch.data_parallel_rank.unwrap_or(0).cast_unsigned(); - for raw_event in batch.events.into_iter() { - // Use shared monotonic event_id counter instead of engine's sequence number - let event_id = next_event_id.fetch_add(1, Ordering::SeqCst); - let worker = WorkerWithDpRank::new(worker_id, dp_rank); - let event = convert_event( - raw_event, - event_id, - kv_block_size, - worker, - &warning_count, - ); - if tx.send(event).is_err() { - tracing::warn!("Failed to send message to channel - receiver dropped"); - exit_reason = "channel receiver dropped"; - break 'main; - } - messages_processed += 1; - } - } - } - } - tracing::debug!( - "ZMQ listener exiting, reason: {}, messages processed: {}", - exit_reason, - messages_processed - ); -} - -// ------------------------------------------------------------------------- -// Metrics Publishers ------------------------------------------------------ -// ------------------------------------------------------------------------- - -/// Metrics data passed through the channel for NATS publishing -#[derive(Debug, Clone, Default, PartialEq)] -struct WorkerMetrics { - dp_rank: DpRank, - active_decode_blocks: u64, -} - -pub struct WorkerMetricsPublisher { - tx: tokio::sync::watch::Sender, - rx: tokio::sync::watch::Receiver, -} - -impl WorkerMetricsPublisher { - pub fn new() -> Result { - let (tx, rx) = tokio::sync::watch::channel(WorkerMetrics::default()); - Ok(WorkerMetricsPublisher { tx, rx }) - } - - /// Publish worker metrics for load monitoring. - /// - /// # Arguments - /// * `dp_rank` - Data parallel rank of the worker (None defaults to 0) - /// * `active_decode_blocks` - Number of active KV cache blocks - pub fn publish(&self, dp_rank: Option, active_decode_blocks: u64) -> Result<()> { - let metrics = WorkerMetrics { - dp_rank: dp_rank.unwrap_or(0), - active_decode_blocks, - }; - tracing::trace!( - "Publish metrics: dp_rank={}, active_decode_blocks={}", - metrics.dp_rank, - metrics.active_decode_blocks - ); - self.tx - .send(metrics) - .map_err(|_| anyhow::anyhow!("metrics channel closed")) - } - - pub async fn create_endpoint(&self, component: Component) -> Result<()> { - let worker_id = component.drt().connection_id(); - self.start_nats_metrics_publishing(component.namespace().clone(), worker_id); - Ok(()) - } - - /// Starts a background task to publish metrics over NATS - /// - /// This task monitors metric changes (specifically active_decode_blocks) - /// and publishes stable metrics to NATS after they've been unchanged for 1ms. - fn start_nats_metrics_publishing(&self, namespace: Namespace, worker_id: u64) { - let nats_rx = self.rx.clone(); - - tokio::spawn(async move { - let event_publisher = - match EventPublisher::for_namespace(&namespace, KV_METRICS_SUBJECT).await { - Ok(publisher) => publisher, - Err(e) => { - tracing::error!("Failed to create metrics publisher: {}", e); - return; - } - }; - - let mut rx = nats_rx; - let mut last_metrics: Option = None; - let mut pending_publish: Option = None; - let mut publish_timer = - Box::pin(tokio::time::sleep(tokio::time::Duration::from_secs(0))); - publish_timer.as_mut().reset(tokio::time::Instant::now()); // Complete immediately - - loop { - tokio::select! { - // Handle metrics changes - result = rx.changed() => { - if result.is_err() { - tracing::debug!( - "Metrics publisher sender dropped, stopping NATS background task" - ); - break; - } - - let metrics = rx.borrow_and_update().clone(); - - // Check if metrics have changed - let has_changed = last_metrics.as_ref() != Some(&metrics); - - // If metrics changed, schedule a publish - if has_changed { - pending_publish = Some(metrics.clone()); - last_metrics = Some(metrics); - - // Start the 1ms timer - publish_timer.as_mut().reset( - tokio::time::Instant::now() + tokio::time::Duration::from_millis(1) - ); - } - } - // Timer expired - publish if we have pending metrics - _ = &mut publish_timer => { - if let Some(metrics) = pending_publish.take() { - let active_load = ActiveLoad { - worker_id, - dp_rank: metrics.dp_rank, - active_decode_blocks: Some(metrics.active_decode_blocks), - active_prefill_tokens: None, - }; - - if let Err(e) = event_publisher.publish(&active_load).await { - tracing::warn!("Failed to publish metrics: {}", e); - } - } - - // Reset timer to pending state to avoid tight loop - // It will be reset to 1ms when metrics actually change - publish_timer.as_mut().reset( - tokio::time::Instant::now() + tokio::time::Duration::from_secs(3600) - ); - } - } - } - }); - } -} - -// ------------------------------------------------------------------------- -// Testing ----------------------------------------------------------------- -// ------------------------------------------------------------------------- +use std::time::Duration; +#[allow(unused_imports)] +use zeromq::{PubSocket, Socket, SocketSend, ZmqMessage}; #[cfg(test)] mod test_event_processing { @@ -1459,9 +430,8 @@ mod test_event_processing { #[cfg(test)] mod tests_startup_helpers { use super::*; - use crate::kv_router::KvIndexer; use bytes::Bytes; - use dynamo_kv_router::indexer::{GetWorkersRequest, KvIndexerInterface}; + use dynamo_kv_router::indexer::{GetWorkersRequest, KvIndexer, KvIndexerInterface}; use dynamo_kv_router::protocols::{ExternalSequenceBlockHash, LocalBlockHash}; use std::sync::{Arc, Mutex}; use zeromq::{PubSocket, Socket, SocketSend, ZmqMessage}; @@ -2202,6 +1172,7 @@ mod test_exponential_backoff { #[cfg(all(test, feature = "integration"))] mod test_integration_publisher { use super::*; + use crate::kv_router::KV_METRICS_SUBJECT; use dynamo_kv_router::protocols::ActiveLoad; use dynamo_runtime::distributed_test_utils::create_test_drt_async; use dynamo_runtime::transports::event_plane::EventSubscriber; diff --git a/lib/llm/src/kv_router/publisher/worker_metrics.rs b/lib/llm/src/kv_router/publisher/worker_metrics.rs new file mode 100644 index 000000000000..d392e8e5cf2e --- /dev/null +++ b/lib/llm/src/kv_router/publisher/worker_metrics.rs @@ -0,0 +1,116 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use anyhow::Result; + +use dynamo_kv_router::protocols::{ActiveLoad, DpRank}; +use dynamo_runtime::component::{Component, Namespace}; +use dynamo_runtime::traits::DistributedRuntimeProvider; +use dynamo_runtime::transports::event_plane::EventPublisher; + +use crate::kv_router::KV_METRICS_SUBJECT; + +#[derive(Debug, Clone, Default, PartialEq)] +struct WorkerMetrics { + dp_rank: DpRank, + active_decode_blocks: u64, +} + +pub struct WorkerMetricsPublisher { + tx: tokio::sync::watch::Sender, + rx: tokio::sync::watch::Receiver, +} + +impl WorkerMetricsPublisher { + pub fn new() -> Result { + let (tx, rx) = tokio::sync::watch::channel(WorkerMetrics::default()); + Ok(Self { tx, rx }) + } + + pub fn publish(&self, dp_rank: Option, active_decode_blocks: u64) -> Result<()> { + let metrics = WorkerMetrics { + dp_rank: dp_rank.unwrap_or(0), + active_decode_blocks, + }; + tracing::trace!( + "Publish metrics: dp_rank={}, active_decode_blocks={}", + metrics.dp_rank, + metrics.active_decode_blocks + ); + self.tx + .send(metrics) + .map_err(|_| anyhow::anyhow!("metrics channel closed")) + } + + pub async fn create_endpoint(&self, component: Component) -> Result<()> { + let worker_id = component.drt().connection_id(); + self.start_nats_metrics_publishing(component.namespace().clone(), worker_id); + Ok(()) + } + + pub(super) fn start_nats_metrics_publishing(&self, namespace: Namespace, worker_id: u64) { + let nats_rx = self.rx.clone(); + + tokio::spawn(async move { + let event_publisher = + match EventPublisher::for_namespace(&namespace, KV_METRICS_SUBJECT).await { + Ok(publisher) => publisher, + Err(e) => { + tracing::error!("Failed to create metrics publisher: {}", e); + return; + } + }; + + let mut rx = nats_rx; + let mut last_metrics: Option = None; + let mut pending_publish: Option = None; + let mut publish_timer = + Box::pin(tokio::time::sleep(tokio::time::Duration::from_secs(0))); + publish_timer.as_mut().reset(tokio::time::Instant::now()); + + loop { + tokio::select! { + result = rx.changed() => { + if result.is_err() { + tracing::debug!( + "Metrics publisher sender dropped, stopping NATS background task" + ); + break; + } + + let metrics = rx.borrow_and_update().clone(); + let has_changed = last_metrics.as_ref() != Some(&metrics); + + if has_changed { + pending_publish = Some(metrics.clone()); + last_metrics = Some(metrics); + publish_timer.as_mut().reset( + tokio::time::Instant::now() + + tokio::time::Duration::from_millis(1) + ); + } + } + _ = &mut publish_timer => { + if let Some(metrics) = pending_publish.take() { + let active_load = ActiveLoad { + worker_id, + dp_rank: metrics.dp_rank, + active_decode_blocks: Some(metrics.active_decode_blocks), + active_prefill_tokens: None, + }; + + if let Err(e) = event_publisher.publish(&active_load).await { + tracing::warn!("Failed to publish metrics: {}", e); + } + } + + publish_timer.as_mut().reset( + tokio::time::Instant::now() + + tokio::time::Duration::from_secs(3600) + ); + } + } + } + }); + } +} diff --git a/lib/llm/src/kv_router/publisher/zmq_listener.rs b/lib/llm/src/kv_router/publisher/zmq_listener.rs new file mode 100644 index 000000000000..1a589604ad80 --- /dev/null +++ b/lib/llm/src/kv_router/publisher/zmq_listener.rs @@ -0,0 +1,161 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::sync::Arc; +use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use std::time::Duration; + +use rmp_serde as rmps; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; +use zeromq::{Socket, SocketRecv, SubSocket}; + +use dynamo_kv_router::protocols::*; +use dynamo_kv_router::zmq_wire::*; + +pub(super) const INITIAL_BACKOFF_MS: u64 = 10; +pub(super) const MAX_BACKOFF_MS: u64 = 5000; +pub(super) const MAX_CONSECUTIVE_ERRORS: u32 = 10; +pub(super) const MAX_BACKOFF_EXPONENT: u32 = 8; + +pub(super) fn calculate_backoff_ms(consecutive_errors: u32) -> u64 { + std::cmp::min( + INITIAL_BACKOFF_MS * 2_u64.pow(consecutive_errors.min(MAX_BACKOFF_EXPONENT)), + MAX_BACKOFF_MS, + ) +} + +pub(super) async fn start_zmq_listener( + zmq_endpoint: String, + zmq_topic: String, + worker_id: WorkerId, + tx: mpsc::UnboundedSender, + cancellation_token: CancellationToken, + kv_block_size: u32, + next_event_id: Arc, +) { + tracing::debug!( + "KVEventPublisher connecting to ZMQ endpoint {} (topic '{}')", + zmq_endpoint, + zmq_topic + ); + + let warning_count = Arc::new(AtomicU32::new(0)); + let mut socket = SubSocket::new(); + + if let Err(e) = socket.subscribe(&zmq_topic).await { + tracing::error!("Failed to subscribe on ZMQ socket: {}", e); + return; + } + + if let Err(e) = socket.connect(&zmq_endpoint).await { + tracing::error!("Failed to connect ZMQ SUB socket to {zmq_endpoint}: {e}"); + return; + } + + let mut consecutive_errors = 0u32; + #[expect(unused_assignments)] + let mut exit_reason = "unknown"; + let mut messages_processed = 0u64; + + 'main: loop { + tokio::select! { + biased; + + _ = cancellation_token.cancelled() => { + tracing::debug!("ZMQ listener received cancellation signal"); + exit_reason = "cancellation token cancelled"; + break 'main; + } + + msg_result = socket.recv() => { + let Ok(msg) = msg_result else { + let e = msg_result.unwrap_err(); + consecutive_errors += 1; + + if consecutive_errors >= MAX_CONSECUTIVE_ERRORS { + tracing::error!( + error=%e, + consecutive_errors=%consecutive_errors, + "Too many consecutive ZMQ errors, terminating listener" + ); + exit_reason = "too many consecutive errors"; + break 'main; + } + + let backoff_ms = calculate_backoff_ms(consecutive_errors); + + tracing::warn!( + error=%e, + consecutive_errors=%consecutive_errors, + backoff_ms=%backoff_ms, + "Error reading from ZMQ socket, applying exponential backoff" + ); + + tokio::time::sleep(Duration::from_millis(backoff_ms)).await; + continue; + }; + consecutive_errors = 0; + + let mut frames: Vec> = + msg.into_vec().into_iter().map(|frame| frame.to_vec()).collect(); + + if frames.len() != 3 { + tracing::warn!( + "Received unexpected ZMQ frame count: expected 3, actual {}", + frames.len() + ); + continue; + } + + let payload = frames.pop().unwrap(); + let seq_bytes = frames.pop().unwrap(); + + if seq_bytes.len() != 8 { + tracing::warn!( + "Invalid sequence number byte length: expected 8, actual {}", + seq_bytes.len() + ); + continue; + } + + let engine_seq = u64::from_be_bytes(seq_bytes.try_into().unwrap()); + + let batch_result = rmps::from_slice::(&payload); + let Ok(batch) = batch_result else { + let e = batch_result.unwrap_err(); + tracing::warn!("Failed to decode KVEventBatch msgpack: {e}"); + continue; + }; + + tracing::trace!( + "ZMQ listener on {} received batch with {} events (engine_seq={}, dp_rank={})", + zmq_endpoint, + batch.events.len(), + engine_seq, + batch.data_parallel_rank.unwrap_or(0) + ); + + let dp_rank = batch.data_parallel_rank.unwrap_or(0).cast_unsigned(); + for raw_event in batch.events { + let event_id = next_event_id.fetch_add(1, Ordering::SeqCst); + let worker = WorkerWithDpRank::new(worker_id, dp_rank); + let event = + convert_event(raw_event, event_id, kv_block_size, worker, &warning_count); + if tx.send(event).is_err() { + tracing::warn!("Failed to send message to channel - receiver dropped"); + exit_reason = "channel receiver dropped"; + break 'main; + } + messages_processed += 1; + } + } + } + } + + tracing::debug!( + "ZMQ listener exiting, reason: {}, messages processed: {}", + exit_reason, + messages_processed + ); +} diff --git a/lib/llm/src/kv_router/remote_indexer.rs b/lib/llm/src/kv_router/remote_indexer.rs deleted file mode 100644 index 1d0f161f0f73..000000000000 --- a/lib/llm/src/kv_router/remote_indexer.rs +++ /dev/null @@ -1,65 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -use anyhow::Result; -use futures::StreamExt; - -use dynamo_runtime::{ - component::Component, - pipeline::{ManyOut, RouterMode, SingleIn, network::egress::push_router::PushRouter}, -}; - -use dynamo_kv_router::{ - indexer::{IndexerQueryRequest, IndexerQueryResponse, KV_INDEXER_QUERY_ENDPOINT}, - protocols::{LocalBlockHash, OverlapScores}, -}; - -/// A remote indexer that queries a standalone KV indexer via the request plane. -/// -/// Used by the frontend when `remote_indexer_component` is configured. Instead of -/// maintaining a local radix tree, this forwards `find_matches` queries to the -/// standalone indexer service over the Dynamo request plane. -pub struct RemoteIndexer { - router: PushRouter, - model_name: String, - namespace: String, -} - -impl RemoteIndexer { - pub async fn new( - component: &Component, - indexer_component_name: &str, - model_name: String, - ) -> Result { - let namespace = component.namespace().name(); - let indexer_ns = component.namespace(); - let indexer_component = indexer_ns.component(indexer_component_name)?; - let endpoint = indexer_component.endpoint(KV_INDEXER_QUERY_ENDPOINT); - let client = endpoint.client().await?; - let router = - PushRouter::from_client_no_fault_detection(client, RouterMode::RoundRobin).await?; - Ok(Self { - router, - model_name, - namespace, - }) - } - - pub async fn find_matches(&self, block_hashes: Vec) -> Result { - let request = IndexerQueryRequest { - model_name: self.model_name.clone(), - namespace: self.namespace.clone(), - block_hashes, - }; - let mut stream: ManyOut = - self.router.round_robin(SingleIn::new(request)).await?; - - match stream.next().await { - Some(IndexerQueryResponse::Scores(scores)) => Ok(scores.into()), - Some(IndexerQueryResponse::Error(msg)) => { - Err(anyhow::anyhow!("Remote indexer error: {}", msg)) - } - None => Err(anyhow::anyhow!("Remote indexer returned empty response")), - } - } -} diff --git a/lib/llm/src/kv_router/scheduler.rs b/lib/llm/src/kv_router/scheduler.rs index 6b72f53e20e8..f38413b0a5dd 100644 --- a/lib/llm/src/kv_router/scheduler.rs +++ b/lib/llm/src/kv_router/scheduler.rs @@ -81,6 +81,7 @@ where block_size, selector, policy, + kv_router_config.router_track_prefill_tokens, component.drt().child_token(), worker_type, watch_worker_configs, @@ -180,9 +181,10 @@ where token_seq: Option>, isl_tokens: usize, overlaps: OverlapScores, + track_prefill_tokens: bool, ) -> Vec { self.inner - .get_potential_loads(token_seq, isl_tokens, overlaps) + .get_potential_loads(token_seq, isl_tokens, overlaps, track_prefill_tokens) } pub fn get_active_lora_counts(&self) -> HashMap { diff --git a/lib/llm/src/kv_router/sequence.rs b/lib/llm/src/kv_router/sequence.rs index ebf540645933..1620a01e497f 100644 --- a/lib/llm/src/kv_router/sequence.rs +++ b/lib/llm/src/kv_router/sequence.rs @@ -223,6 +223,7 @@ mod tests { token_sequence: Some(vec![0, 1, 2]), isl: 12, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::new(0, 0), lora_name: None, @@ -235,6 +236,7 @@ mod tests { token_sequence: Some(vec![3, 4]), isl: 8, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::new(0, 1), lora_name: None, @@ -247,6 +249,7 @@ mod tests { token_sequence: Some(vec![0, 1, 2, 3]), isl: 16, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::new(1, 0), lora_name: None, @@ -373,6 +376,7 @@ mod tests { token_sequence: None, isl: 12, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::from_worker_id(0), lora_name: None, @@ -385,6 +389,7 @@ mod tests { token_sequence: None, isl: 8, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::from_worker_id(1), lora_name: None, @@ -397,6 +402,7 @@ mod tests { token_sequence: None, isl: 16, overlap: 0, + track_prefill_tokens: true, expected_output_tokens: None, worker: WorkerWithDpRank::from_worker_id(2), lora_name: None, diff --git a/lib/llm/src/mocker.rs b/lib/llm/src/mocker.rs index 96ed4a718551..c418429463ab 100644 --- a/lib/llm/src/mocker.rs +++ b/lib/llm/src/mocker.rs @@ -23,7 +23,7 @@ use dynamo_mocker::common::protocols::{ DirectRequest, KvCacheEventSink, KvEventPublishers, MockEngineArgs, OutputSignal, RawKvEvent, RawKvEventSink, }; -use dynamo_mocker::common::utils::{compute_kv_transfer_delay, sleep_precise}; +use dynamo_mocker::common::utils::sleep_precise; use dynamo_mocker::engine::create_engine; use dynamo_mocker::scheduler::SchedulerHandle; use dynamo_runtime::DistributedRuntime; @@ -645,14 +645,6 @@ impl AsyncEngine, ManyOut, Error> let bootstrap_server = self.bootstrap_server.clone(); let reasoning = self.engine_args.reasoning.clone(); - // Compute KV transfer delay for prefill workers. - // Simulates the time to transfer KV cache from prefill to decode worker. - let kv_transfer_delay = if is_prefill { - compute_kv_transfer_delay(&self.engine_args, request.token_ids.len()) - } else { - None - }; - // Spawn a task to handle the complex async logic tokio::spawn(async move { let mut token_count = 0; @@ -693,17 +685,15 @@ impl AsyncEngine, ManyOut, Error> if signal.completed { let _ = stream_tx.send(output); - // Simulate KV transfer delay before prefill's first (and only) token. - // This models the time to transfer KV cache to the decode worker. - if token_count == 1 - && let Some(delay) = kv_transfer_delay + // Prefill-to-decode handoff delay is emitted by the shared mocker core. + if is_prefill + && let Some(delay_ms) = signal.handoff_delay_ms { - sleep_precise(delay).await; + sleep_precise(Duration::from_secs_f64(delay_ms / 1000.0)).await; } // Prefill: after first token, mark room complete (unblocks decode) if is_prefill - && token_count == 1 && let (Some(server), Some(room_id)) = (bootstrap_server.get(), bootstrap_room) { server.complete_room(room_id); diff --git a/lib/mocker/src/common/protocols.rs b/lib/mocker/src/common/protocols.rs index 08183abbb3f4..3aafa5e55898 100644 --- a/lib/mocker/src/common/protocols.rs +++ b/lib/mocker/src/common/protocols.rs @@ -139,6 +139,8 @@ impl PrefillCost { pub struct OutputSignal { pub uuid: Uuid, pub completed: bool, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub handoff_delay_ms: Option, } /// Preemption policy for evicting decode requests under memory pressure @@ -286,6 +288,10 @@ pub struct MockEngineArgs { #[builder(default = "WorkerType::Aggregated")] pub worker_type: WorkerType, + /// Original planner profile NPZ path used to materialize `perf_model`. + #[builder(default = "None")] + pub planner_profile_data: Option, + /// Performance model for timing predictions (not serialized, loaded from planner_profile_data) #[serde(skip)] #[builder(default = "Arc::new(PerfModel::default())")] @@ -691,6 +697,7 @@ impl MockEngineArgs { && let Some(path_str) = path_str.as_str() { let npz_path = PathBuf::from(path_str); + builder = builder.planner_profile_data(Some(npz_path.clone())); match PerfModel::from_npz(&npz_path) { Ok(model) => { tracing::info!("Successfully loaded performance model from: {:?}", npz_path); diff --git a/lib/mocker/src/common/utils.rs b/lib/mocker/src/common/utils.rs index 04ecc2c4be45..959d2565d786 100644 --- a/lib/mocker/src/common/utils.rs +++ b/lib/mocker/src/common/utils.rs @@ -3,32 +3,59 @@ use std::time::{Duration, Instant}; -use crate::common::protocols::MockEngineArgs; +use crate::common::protocols::{MockEngineArgs, WorkerType}; -/// Compute the KV transfer delay duration for a given number of input tokens. +/// Compute the modeled handoff delay after a prefill worker emits its terminal token. /// -/// Returns `None` if KV transfer simulation is disabled (bandwidth is 0 or not configured). -pub fn compute_kv_transfer_delay( - args: &MockEngineArgs, +/// NOTE: this intentionally does not model the internal prefill TTFT itself accurately, and the +/// exact prefill/decode boundary is backend dependent. For now we only care about decode-visible +/// TTFT, which is what the client observes, so modeling the delay as prefill-to-decode handoff is +/// good enough. +pub fn compute_prefill_handoff_delay_ms( + worker_type: WorkerType, + completed: bool, num_input_tokens: usize, -) -> Option { - match (args.kv_transfer_bandwidth, args.kv_bytes_per_token) { + kv_transfer_bandwidth: Option, + kv_bytes_per_token: Option, +) -> Option { + if worker_type != WorkerType::Prefill || !completed { + return None; + } + + match (kv_transfer_bandwidth, kv_bytes_per_token) { (Some(bw), Some(bpt)) if bw > 0.0 => { let kv_bytes = num_input_tokens as f64 * bpt as f64; - let delay = Duration::from_secs_f64(kv_bytes / (bw * 1e9)); + let delay_ms = kv_bytes / (bw * 1e9) * 1000.0; tracing::debug!( num_input_tokens, kv_bytes, bandwidth_gb_s = bw, - delay_ms = format!("{:.2}", delay.as_secs_f64() * 1000.0), - "KV transfer delay for prefill" + delay_ms = format!("{delay_ms:.2}"), + "KV handoff delay for prefill completion" ); - Some(delay) + Some(delay_ms) } _ => None, } } +/// Compute the KV transfer delay duration for a given number of input tokens. +/// +/// Returns `None` if KV transfer simulation is disabled (bandwidth is 0 or not configured). +pub fn compute_kv_transfer_delay( + args: &MockEngineArgs, + num_input_tokens: usize, +) -> Option { + compute_prefill_handoff_delay_ms( + args.worker_type, + true, + num_input_tokens, + args.kv_transfer_bandwidth, + args.kv_bytes_per_token, + ) + .map(|delay_ms| Duration::from_secs_f64(delay_ms / 1000.0)) +} + /// Sleep for the specified duration using timerfd on Linux for precision. pub async fn sleep_precise(duration: Duration) { sleep_until_precise(Instant::now() + duration).await; @@ -53,3 +80,42 @@ pub async fn sleep_until_precise(deadline: Instant) { tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)).await; } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_prefill_handoff_delay_only_applies_to_completed_prefill() { + let delay_ms = compute_prefill_handoff_delay_ms( + WorkerType::Prefill, + true, + 128, + Some(1.0), + Some(1_000_000), + ) + .expect("prefill completion should produce a handoff delay"); + assert!((delay_ms - 128.0).abs() < 1e-9); + + assert!( + compute_prefill_handoff_delay_ms( + WorkerType::Prefill, + false, + 128, + Some(1.0), + Some(1_000_000), + ) + .is_none() + ); + assert!( + compute_prefill_handoff_delay_ms( + WorkerType::Decode, + true, + 128, + Some(1.0), + Some(1_000_000), + ) + .is_none() + ); + } +} diff --git a/lib/mocker/src/loadgen/driver.rs b/lib/mocker/src/loadgen/driver.rs index ce04d46ff309..6a0c6be8e7db 100644 --- a/lib/mocker/src/loadgen/driver.rs +++ b/lib/mocker/src/loadgen/driver.rs @@ -7,7 +7,8 @@ use std::collections::{BinaryHeap, HashMap}; use anyhow::{Result, anyhow, bail}; use uuid::Uuid; -use super::types::{ReadyTurn, Trace, TurnTrace}; +use super::types::{ReadyTurn, ReplayRequestHashes, Trace}; +use crate::common::protocols::DirectRequest; #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum DriverMode { @@ -18,12 +19,20 @@ enum DriverMode { #[derive(Debug)] struct SessionRuntime { session_id: String, - turns: Vec, + turns: Vec, next_turn_index: usize, next_ready_at_ms: Option, in_flight: Option, } +#[derive(Debug)] +struct TurnRuntime { + tokens: Vec, + max_output_tokens: usize, + delay_after_previous_ms: f64, + replay_hashes: ReplayRequestHashes, +} + #[derive(Debug)] struct InFlightTurn { session_index: usize, @@ -66,7 +75,6 @@ impl PartialOrd for ReadySession { #[derive(Debug)] pub struct WorkloadDriver { mode: DriverMode, - block_size: usize, sessions: Vec, in_flight: HashMap, ready_sessions: BinaryHeap, @@ -82,20 +90,36 @@ impl WorkloadDriver { } fn new(trace: Trace, mode: DriverMode) -> Result { + let block_size = trace.block_size; let sessions: Vec = trace .sessions .into_iter() - .map(|session| SessionRuntime { - session_id: session.session_id, - turns: session.turns, - next_turn_index: 0, - next_ready_at_ms: Some(match mode { + .map(|session| -> Result { + let next_ready_at_ms = Some(match mode { DriverMode::Trace => session.first_arrival_timestamp_ms.unwrap_or(0.0), DriverMode::Concurrency => 0.0, - }), - in_flight: None, + }); + let turns = session + .turns + .into_iter() + .map(|turn| -> Result { + Ok(TurnRuntime { + tokens: turn.synthesize_tokens(block_size)?, + max_output_tokens: turn.max_output_tokens, + delay_after_previous_ms: turn.delay_after_previous_ms, + replay_hashes: turn.to_replay_hashes(block_size)?, + }) + }) + .collect::>>()?; + Ok(SessionRuntime { + session_id: session.session_id, + turns, + next_turn_index: 0, + next_ready_at_ms, + in_flight: None, + }) }) - .collect(); + .collect::>>()?; let ready_sessions = sessions .iter() @@ -111,7 +135,6 @@ impl WorkloadDriver { Ok(Self { mode, - block_size: trace.block_size, sessions, in_flight: HashMap::new(), ready_sessions, @@ -146,16 +169,18 @@ impl WorkloadDriver { .next_ready_at_ms .expect("ready session must have a timestamp"); let request_uuid = Uuid::new_v4(); - let replay_hashes = session.turns[turn_index] - .to_replay_hashes(self.block_size) - .expect("validated trace should always synthesize replay hashes"); + let turn = &session.turns[turn_index]; let arrival_timestamp_ms = match self.mode { DriverMode::Trace => Some(scheduled_ready_at_ms), DriverMode::Concurrency => None, }; - let request = session.turns[turn_index] - .to_direct_request(self.block_size, request_uuid, arrival_timestamp_ms) - .expect("validated trace should always synthesize into a direct request"); + let request = DirectRequest { + tokens: turn.tokens.clone(), + max_output_tokens: turn.max_output_tokens, + uuid: Some(request_uuid), + dp_rank: 0, + arrival_timestamp_ms, + }; session.in_flight = Some(request_uuid); session.next_ready_at_ms = None; self.in_flight.insert( @@ -170,7 +195,7 @@ impl WorkloadDriver { session_id: session.session_id.clone(), turn_index, scheduled_ready_at_ms, - replay_hashes: Some(replay_hashes), + replay_hashes: Some(turn.replay_hashes.clone()), request, }); } diff --git a/lib/mocker/src/loadgen/trace.rs b/lib/mocker/src/loadgen/trace.rs index a9931dc18a06..40818daa7c95 100644 --- a/lib/mocker/src/loadgen/trace.rs +++ b/lib/mocker/src/loadgen/trace.rs @@ -59,12 +59,7 @@ impl TurnTrace { Ok(()) } - pub fn to_direct_request( - &self, - block_size: usize, - request_uuid: Uuid, - arrival_timestamp_ms: Option, - ) -> Result { + pub(crate) fn synthesize_tokens(&self, block_size: usize) -> Result> { self.validate_block_size_and_capacity(block_size)?; let mut tokens = Vec::with_capacity(self.input_length); @@ -85,6 +80,16 @@ impl TurnTrace { ); } + Ok(tokens) + } + + pub fn to_direct_request( + &self, + block_size: usize, + request_uuid: Uuid, + arrival_timestamp_ms: Option, + ) -> Result { + let tokens = self.synthesize_tokens(block_size)?; Ok(DirectRequest { tokens, max_output_tokens: self.max_output_tokens, diff --git a/lib/mocker/src/replay/collector.rs b/lib/mocker/src/replay/collector.rs index 39c1ca880a6f..79ed61f71be3 100644 --- a/lib/mocker/src/replay/collector.rs +++ b/lib/mocker/src/replay/collector.rs @@ -309,6 +309,7 @@ impl TraceCollector { } let duration_s = (duration_ms / 1000.0).max(1e-9); + let itl_distribution = build_distribution_stats(&itls); TraceSimulationReport { request_counts: TraceRequestCounts { num_requests: requests.len(), @@ -335,8 +336,8 @@ impl TraceCollector { ttst: build_distribution_stats(&ttsts), tpot: build_distribution_stats(&tpots), itl: TraceInterTokenLatencyStats { - distribution: build_distribution_stats(&itls), - max_ms: max_value(&itls), + max_ms: itl_distribution.max_ms, + distribution: itl_distribution, }, e2e: build_distribution_stats(&e2e_latencies), output_token_throughput_per_user: build_distribution_stats( @@ -386,39 +387,41 @@ fn mean(values: &[f64]) -> f64 { } } -fn max_value(values: &[f64]) -> f64 { - values.iter().copied().reduce(f64::max).unwrap_or(0.0) -} - fn build_distribution_stats(values: &[f64]) -> TraceDistributionStats { + if values.is_empty() { + return TraceDistributionStats { + mean_ms: 0.0, + min_ms: 0.0, + max_ms: 0.0, + median_ms: 0.0, + p75_ms: 0.0, + p90_ms: 0.0, + p95_ms: 0.0, + p99_ms: 0.0, + std_ms: 0.0, + }; + } + + let mut sorted = values.to_vec(); + sorted.sort_by(|left, right| left.total_cmp(right)); TraceDistributionStats { mean_ms: mean(values), - min_ms: min_value(values), - max_ms: max_value(values), - median_ms: percentile(values, 50.0), - p75_ms: percentile(values, 75.0), - p90_ms: percentile(values, 90.0), - p95_ms: percentile(values, 95.0), - p99_ms: percentile(values, 99.0), + min_ms: sorted[0], + max_ms: *sorted.last().expect("sorted values must be non-empty"), + median_ms: percentile_sorted(&sorted, 50.0), + p75_ms: percentile_sorted(&sorted, 75.0), + p90_ms: percentile_sorted(&sorted, 90.0), + p95_ms: percentile_sorted(&sorted, 95.0), + p99_ms: percentile_sorted(&sorted, 99.0), std_ms: std_dev(values), } } -fn percentile(values: &[f64], percentile: f64) -> f64 { - if values.is_empty() { - return 0.0; - } - - let mut sorted = values.to_vec(); - sorted.sort_by(|left, right| left.total_cmp(right)); +fn percentile_sorted(sorted: &[f64], percentile: f64) -> f64 { let rank = ((sorted.len() - 1) as f64 * percentile / 100.0).round() as usize; sorted[rank.min(sorted.len() - 1)] } -fn min_value(values: &[f64]) -> f64 { - values.iter().copied().reduce(f64::min).unwrap_or(0.0) -} - fn std_dev(values: &[f64]) -> f64 { if values.is_empty() { return 0.0; diff --git a/lib/mocker/src/replay/entrypoints.rs b/lib/mocker/src/replay/entrypoints.rs index a25befff3349..13f80cd6cdc5 100644 --- a/lib/mocker/src/replay/entrypoints.rs +++ b/lib/mocker/src/replay/entrypoints.rs @@ -9,10 +9,11 @@ use dynamo_kv_router::config::KvRouterConfig; use super::online; use super::validate::{ - validate_offline_concurrency_args, validate_offline_replay_args, + validate_offline_concurrency_args, validate_offline_disagg_concurrency_args, + validate_offline_disagg_replay_args, validate_offline_replay_args, validate_online_concurrency_args, validate_online_replay_args, }; -use super::{ReplayRouterMode, TraceSimulationReport}; +use super::{OfflineDisaggReplayConfig, ReplayRouterMode, TraceSimulationReport}; use crate::common::protocols::{DirectRequest, MockEngineArgs}; use crate::loadgen::Trace; @@ -56,6 +57,28 @@ pub fn simulate_trace_file_with_router_mode( Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } +pub fn simulate_trace_file_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace_path: &Path, + arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_replay_args(&config, router_mode)?; + let trace = Trace::from_mooncake(trace_path, config.prefill_args.block_size)? + .normalize_session_starts()? + .speed_up_timing(arrival_speedup_ratio)?; + let started_at = Instant::now(); + let report = crate::replay::offline::simulate_trace_workload_disagg( + config, + router_config, + trace, + router_mode, + )?; + Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) +} + pub fn simulate_trace_live_file( args: MockEngineArgs, trace_path: &Path, @@ -130,6 +153,30 @@ pub fn simulate_trace_requests_with_router_mode( Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } +pub fn simulate_trace_requests_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_replay_args(&config, router_mode)?; + if requests.is_empty() { + bail!("trace replay requires at least one request"); + } + + let started_at = Instant::now(); + let report = crate::replay::offline::simulate_trace_disagg( + config, + router_config, + requests, + arrival_speedup_ratio, + router_mode, + )?; + Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) +} + pub fn simulate_trace_live_requests( args: MockEngineArgs, requests: Vec, @@ -209,6 +256,27 @@ pub fn simulate_concurrency_file_with_router_mode( Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } +pub fn simulate_concurrency_file_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace_path: &Path, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_concurrency_args(&config, max_in_flight, router_mode)?; + let trace = Trace::from_mooncake(trace_path, config.prefill_args.block_size)?; + let started_at = Instant::now(); + let report = simulate_concurrency_workload_disagg_with_router_mode( + config, + router_config, + trace, + max_in_flight, + router_mode, + )?; + Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) +} + pub fn simulate_concurrency_live_file( args: MockEngineArgs, trace_path: &Path, @@ -326,6 +394,28 @@ pub fn simulate_concurrency_requests_with_router_mode( ) } +pub fn simulate_concurrency_requests_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_concurrency_args(&config, max_in_flight, router_mode)?; + if requests.is_empty() { + bail!("concurrency replay requires at least one request"); + } + + crate::replay::offline::simulate_concurrency_disagg( + config, + router_config, + requests, + max_in_flight, + router_mode, + ) +} + pub fn simulate_trace_workload( args: MockEngineArgs, trace: Trace, @@ -360,6 +450,24 @@ pub fn simulate_trace_workload_with_router_mode( Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) } +pub fn simulate_trace_workload_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_replay_args(&config, router_mode)?; + let started_at = Instant::now(); + let report = crate::replay::offline::simulate_trace_workload_disagg( + config, + router_config, + trace, + router_mode, + )?; + Ok(report.with_wall_time_ms(started_at.elapsed().as_secs_f64() * 1000.0)) +} + pub fn simulate_trace_live_workload( args: MockEngineArgs, trace: Trace, @@ -422,6 +530,24 @@ pub fn simulate_concurrency_workload_with_router_mode( ) } +pub fn simulate_concurrency_workload_disagg_with_router_mode( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result { + let config = config.normalized()?; + validate_offline_disagg_concurrency_args(&config, max_in_flight, router_mode)?; + crate::replay::offline::simulate_concurrency_workload_disagg( + config, + router_config, + trace, + max_in_flight, + router_mode, + ) +} + pub fn simulate_concurrency_live_workload( args: MockEngineArgs, trace: Trace, diff --git a/lib/mocker/src/replay/mod.rs b/lib/mocker/src/replay/mod.rs index 30d4fef6dedf..9076bedf7b86 100644 --- a/lib/mocker/src/replay/mod.rs +++ b/lib/mocker/src/replay/mod.rs @@ -10,7 +10,7 @@ mod validate; use std::collections::VecDeque; -use crate::common::protocols::DirectRequest; +use crate::common::protocols::{DirectRequest, MockEngineArgs}; pub(crate) use collector::TraceCollector; #[cfg(test)] @@ -25,20 +25,50 @@ pub enum ReplayRouterMode { KvRouter, } +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ReplayArgsMode { + Aggregated, + Disagg, +} + +#[derive(Clone, Debug)] +pub struct OfflineDisaggReplayConfig { + pub prefill_args: MockEngineArgs, + pub decode_args: MockEngineArgs, + pub num_prefill_workers: usize, + pub num_decode_workers: usize, +} + +impl OfflineDisaggReplayConfig { + pub fn normalized(self) -> anyhow::Result { + Ok(Self { + prefill_args: self.prefill_args.normalized()?, + decode_args: self.decode_args.normalized()?, + num_prefill_workers: self.num_prefill_workers, + num_decode_workers: self.num_decode_workers, + }) + } +} + pub use entrypoints::{ - simulate_concurrency_file, simulate_concurrency_file_with_router_mode, - simulate_concurrency_live_file, simulate_concurrency_live_file_with_router_mode, - simulate_concurrency_live_requests, simulate_concurrency_live_requests_with_router_mode, - simulate_concurrency_live_workload, simulate_concurrency_live_workload_with_router_mode, - simulate_concurrency_requests, simulate_concurrency_requests_with_router_mode, - simulate_concurrency_workload, simulate_concurrency_workload_with_router_mode, - simulate_trace_file, simulate_trace_file_with_router_mode, simulate_trace_live_file, - simulate_trace_live_file_with_router_mode, simulate_trace_live_requests, - simulate_trace_live_requests_with_router_mode, simulate_trace_live_workload, - simulate_trace_live_workload_with_router_mode, simulate_trace_requests, + simulate_concurrency_file, simulate_concurrency_file_disagg_with_router_mode, + simulate_concurrency_file_with_router_mode, simulate_concurrency_live_file, + simulate_concurrency_live_file_with_router_mode, simulate_concurrency_live_requests, + simulate_concurrency_live_requests_with_router_mode, simulate_concurrency_live_workload, + simulate_concurrency_live_workload_with_router_mode, simulate_concurrency_requests, + simulate_concurrency_requests_disagg_with_router_mode, + simulate_concurrency_requests_with_router_mode, simulate_concurrency_workload, + simulate_concurrency_workload_disagg_with_router_mode, + simulate_concurrency_workload_with_router_mode, simulate_trace_file, + simulate_trace_file_disagg_with_router_mode, simulate_trace_file_with_router_mode, + simulate_trace_live_file, simulate_trace_live_file_with_router_mode, + simulate_trace_live_requests, simulate_trace_live_requests_with_router_mode, + simulate_trace_live_workload, simulate_trace_live_workload_with_router_mode, + simulate_trace_requests, simulate_trace_requests_disagg_with_router_mode, simulate_trace_requests_with_router_mode, simulate_trace_workload, - simulate_trace_workload_with_router_mode, + simulate_trace_workload_disagg_with_router_mode, simulate_trace_workload_with_router_mode, }; +pub use validate::validate_replay_args_mode; pub(crate) fn normalize_trace_requests( mut requests: Vec, diff --git a/lib/mocker/src/replay/offline/README.md b/lib/mocker/src/replay/offline/README.md index 2a9c2a73090c..d446c491b452 100644 --- a/lib/mocker/src/replay/offline/README.md +++ b/lib/mocker/src/replay/offline/README.md @@ -15,10 +15,11 @@ The public replay entrypoints live one level up in `lib/mocker/src/replay/entryp Offline replay starts in `lib/mocker/src/replay/offline/mod.rs`. -`offline/mod.rs` chooses between two implementations: +`offline/mod.rs` chooses between three implementations: - `lib/mocker/src/replay/offline/single.rs` for the special case `num_workers == 1` with the vLLM engine - `lib/mocker/src/replay/offline/multi.rs` for everything else, including multi-worker replay and `kv_router` replay +- `lib/mocker/src/replay/offline/disagg.rs` for offline disaggregated prefill/decode replay ## File Map @@ -28,6 +29,8 @@ Offline replay starts in `lib/mocker/src/replay/offline/mod.rs`. Minimal replay loop for one vLLM worker. - `lib/mocker/src/replay/offline/multi.rs` General offline cluster simulator for multi-worker replay and KV-router replay. +- `lib/mocker/src/replay/offline/disagg.rs` + Offline two-stage replay harness with separate prefill and decode pools. - `lib/mocker/src/replay/offline/state.rs` Per-worker wrapper around `EngineCore`, including optional KV event capture. - `lib/mocker/src/replay/offline/events.rs` @@ -114,7 +117,7 @@ So offline replay is not a toy simulator. It reuses the real per-pass mocker sch ## Completion Event Queue -The multi-worker harness uses `SimulationEvent` from `lib/mocker/src/replay/offline/events.rs` as a min-time priority queue implemented with `BinaryHeap`. +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: @@ -122,6 +125,7 @@ Right now the only scheduled event type is: That event carries: +- worker `stage` (`aggregated`, `prefill`, or `decode`) - `worker_idx` - `completed_requests` - `output_signals` @@ -167,7 +171,7 @@ flowchart LR M --> G ``` -### Why KV events are captured only here +### Why KV events are captured only where needed When offline replay uses `kv_router`, workers are created with KV event capture enabled via: @@ -177,6 +181,37 @@ When offline replay uses `kv_router`, workers are created with KV event capture That causes each pass to return router-visible `kv_events`, which the harness applies synchronously to the offline router indexer after the pass completes. In round-robin mode, this capture is skipped because nothing consumes those events. +In offline disagg replay, only the prefill workers capture and publish KV events; the decode workers +run with capture disabled because the decode router is overlap-blind and does not consume router +events. + +## Disaggregated Harness + +The disaggregated runtime in `lib/mocker/src/replay/offline/disagg.rs` models two distinct stages: + +- a prefill router and prefill worker pool +- a decode router and decode worker pool + +It keeps one logical clock and one completion-event heap, but request ownership moves through a +two-stage state machine instead of the aggregated single-pool lifecycle. + +The prefill router is derived from the main router config with `router_track_active_blocks = false`. +The decode router is derived with: + +- overlap disabled +- `assume_kv_reuse = false` +- `track_prefill_tokens = false` + +The prefill stage runs a hidden synthetic one-token bootstrap request. When prefill completes, the +harness: + +1. applies any prefill KV events +2. marks prefill complete in the prefill router +3. frees prefill router state +4. enqueues the original request into decode at the same logical timestamp + +Decode then runs with normal collector visibility. The public replay report remains decode-only, so +TTFT includes prefill queueing and prefill compute. ## Trace vs Concurrency Modes diff --git a/lib/mocker/src/replay/offline/disagg.rs b/lib/mocker/src/replay/offline/disagg.rs new file mode 100644 index 000000000000..3875dd5ca4a1 --- /dev/null +++ b/lib/mocker/src/replay/offline/disagg.rs @@ -0,0 +1,1129 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{BinaryHeap, HashMap, VecDeque}; + +use anyhow::{Result, anyhow, bail}; +use dynamo_kv_router::config::KvRouterConfig; +use dynamo_kv_router::protocols::RouterEvent; +use uuid::Uuid; + +use super::events::{SimulationEvent, SimulationWorkerStage}; +use super::normalize_trace_requests; +use super::runtime_utils::{ + WorkerCompletionPayload, next_timestamp as choose_next_timestamp, pop_next_concurrency_ready, + pop_next_trace_ready, pop_ready_decode_handoff, pop_ready_worker_completion, + push_decode_handoff, push_worker_completion, +}; +#[cfg(test)] +use super::state::DisaggPhase; +#[cfg(test)] +use super::state::DisaggRequestSnapshot; +use super::state::{DisaggRequestState, OfflineWorkerState}; +use crate::common::protocols::{DirectRequest, MockEngineArgs, OutputSignal}; +use crate::loadgen::{ReplayRequestHashes, Trace, WorkloadDriver}; +use crate::replay::router::OfflineReplayRouter; +use crate::replay::{ + OfflineDisaggReplayConfig, ReplayRouterMode, TraceCollector, TraceSimulationReport, +}; +use crate::scheduler::RouterEventVisibility; + +#[derive(Debug, Clone, Copy)] +enum ReplayMode { + Trace, + Concurrency { max_in_flight: usize }, +} + +enum AdmissionSource { + Requests(VecDeque), + Workload(WorkloadDriver), +} + +#[cfg(test)] +#[derive(Debug, Default, Clone, PartialEq)] +struct DisaggRuntimeStats { + request_snapshots: HashMap, + prefill_assignments: HashMap, + decode_assignments: HashMap, + handoff_ms: HashMap, + prefill_marked_count: usize, + prefill_freed_count: usize, + decode_freed_count: usize, + max_prefill_router_pending: usize, + max_decode_router_pending: usize, +} + +#[cfg(not(test))] +#[derive(Debug, Default, Clone, PartialEq, Eq)] +struct DisaggRuntimeStats; + +struct DisaggRuntime { + now_ms: f64, + next_prefill_worker_idx: usize, + next_decode_worker_idx: usize, + next_event_seq: u64, + admission: AdmissionSource, + prefill_workers: Vec, + decode_workers: Vec, + prefill_router: Option, + decode_router: Option, + requests: HashMap, + collector: TraceCollector, + events: BinaryHeap, + mode: ReplayMode, + stats: DisaggRuntimeStats, +} + +impl DisaggRuntime { + fn new( + config: &OfflineDisaggReplayConfig, + router_config: Option, + pending: VecDeque, + mode: ReplayMode, + router_mode: ReplayRouterMode, + ) -> Result { + Self::new_with_source( + config, + router_config, + AdmissionSource::Requests(pending), + mode, + router_mode, + ) + } + + fn new_workload( + config: &OfflineDisaggReplayConfig, + router_config: Option, + driver: WorkloadDriver, + mode: ReplayMode, + router_mode: ReplayRouterMode, + ) -> Result { + Self::new_with_source( + config, + router_config, + AdmissionSource::Workload(driver), + mode, + router_mode, + ) + } + + fn new_with_source( + config: &OfflineDisaggReplayConfig, + router_config: Option, + admission: AdmissionSource, + mode: ReplayMode, + router_mode: ReplayRouterMode, + ) -> Result { + let (prefill_router, decode_router) = match router_mode { + ReplayRouterMode::RoundRobin => (None, None), + ReplayRouterMode::KvRouter => { + let prefill_router_config = + derive_prefill_router_config(&config.prefill_args, router_config.clone()); + let decode_router_config = + derive_decode_router_config(&config.decode_args, router_config); + ( + Some(OfflineReplayRouter::new( + &config.prefill_args, + Some(prefill_router_config), + config.num_prefill_workers, + )?), + Some(OfflineReplayRouter::new( + &config.decode_args, + Some(decode_router_config), + config.num_decode_workers, + )?), + ) + } + }; + + Ok(Self { + now_ms: 0.0, + next_prefill_worker_idx: 0, + next_decode_worker_idx: 0, + next_event_seq: 0, + admission, + prefill_workers: (0..config.num_prefill_workers) + .map(|worker_idx| { + OfflineWorkerState::new( + worker_idx, + config.prefill_args.clone(), + prefill_router.is_some(), + ) + }) + .collect(), + decode_workers: (0..config.num_decode_workers) + .map(|worker_idx| { + OfflineWorkerState::new(worker_idx, config.decode_args.clone(), false) + }) + .collect(), + prefill_router, + decode_router, + requests: HashMap::new(), + collector: TraceCollector::default(), + events: BinaryHeap::new(), + mode, + #[cfg(test)] + stats: DisaggRuntimeStats::default(), + #[cfg(not(test))] + stats: DisaggRuntimeStats, + }) + } + + fn cluster_in_flight(&self) -> usize { + self.prefill_workers + .iter() + .map(OfflineWorkerState::in_flight) + .sum::() + + self + .decode_workers + .iter() + .map(OfflineWorkerState::in_flight) + .sum::() + + self + .prefill_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count) + + self + .decode_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count) + } + + fn next_prefill_worker(&mut self) -> usize { + let worker_idx = self.next_prefill_worker_idx; + self.next_prefill_worker_idx = + (self.next_prefill_worker_idx + 1) % self.prefill_workers.len(); + worker_idx + } + + fn next_decode_worker(&mut self) -> usize { + let worker_idx = self.next_decode_worker_idx; + self.next_decode_worker_idx = (self.next_decode_worker_idx + 1) % self.decode_workers.len(); + worker_idx + } + + fn record_router_pending(&mut self) { + #[cfg(test)] + { + self.stats.max_prefill_router_pending = self.stats.max_prefill_router_pending.max( + self.prefill_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count), + ); + self.stats.max_decode_router_pending = self.stats.max_decode_router_pending.max( + self.decode_router + .as_ref() + .map_or(0, OfflineReplayRouter::pending_count), + ); + } + } + + fn validate_worker_idx(&self, stage: SimulationWorkerStage, worker_idx: usize) -> Result<()> { + let worker_count = match stage { + SimulationWorkerStage::Prefill => self.prefill_workers.len(), + SimulationWorkerStage::Decode => self.decode_workers.len(), + SimulationWorkerStage::Aggregated => unreachable!("aggregated stage is not used"), + }; + if worker_idx >= worker_count { + bail!("offline disagg replay selected unknown {stage:?} worker index {worker_idx}"); + } + Ok(()) + } + + fn state(&self, uuid: Uuid) -> Result<&DisaggRequestState> { + self.requests + .get(&uuid) + .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}")) + } + + fn state_mut(&mut self, uuid: Uuid) -> Result<&mut DisaggRequestState> { + self.requests + .get_mut(&uuid) + .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}")) + } + + fn dispatch_prefill(&mut self, uuid: Uuid, worker_idx: usize) -> Result<()> { + self.validate_worker_idx(SimulationWorkerStage::Prefill, worker_idx)?; + let request = self.state(uuid)?.build_prefill_request()?; + self.prefill_workers[worker_idx].receive_request(request); + self.state_mut(uuid)?.start_prefill(worker_idx); + #[cfg(test)] + { + self.stats.prefill_assignments.insert(uuid, worker_idx); + } + Ok(()) + } + + fn dispatch_decode(&mut self, uuid: Uuid, worker_idx: usize) -> Result<()> { + self.validate_worker_idx(SimulationWorkerStage::Decode, worker_idx)?; + let request = self.state(uuid)?.build_decode_request()?; + self.decode_workers[worker_idx].receive_request(request); + self.state_mut(uuid)?.start_decode(worker_idx); + #[cfg(test)] + { + self.stats.decode_assignments.insert(uuid, worker_idx); + } + Ok(()) + } + + fn dispatch_prefill_admissions(&mut self, admissions: Vec<(Uuid, usize)>) -> Result<()> { + for (uuid, worker_idx) in admissions { + if !self.state(uuid)?.is_queued_prefill() { + bail!("offline disagg replay expected queued prefill request for {uuid}"); + } + self.dispatch_prefill(uuid, worker_idx)?; + } + Ok(()) + } + + fn dispatch_decode_admissions(&mut self, admissions: Vec<(Uuid, usize)>) -> Result<()> { + for (uuid, worker_idx) in admissions { + if !self.state(uuid)?.is_queued_decode() { + bail!("offline disagg replay expected queued decode request for {uuid}"); + } + self.dispatch_decode(uuid, worker_idx)?; + } + Ok(()) + } + + fn enqueue_decode(&mut self, uuid: Uuid) -> Result<()> { + let Some(decode_router) = self.decode_router.as_mut() else { + let worker_idx = self.next_decode_worker(); + self.dispatch_decode(uuid, worker_idx)?; + return Ok(()); + }; + let maybe_worker_idx = { + let requests = &self.requests; + let request = requests + .get(&uuid) + .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}"))? + .original_request()?; + decode_router.submit_request_with_hashes(request, None, self.now_ms)? + }; + self.record_router_pending(); + #[cfg(test)] + { + self.stats.handoff_ms.insert(uuid, self.now_ms); + } + if let Some(worker_idx) = maybe_worker_idx { + self.dispatch_decode(uuid, worker_idx)?; + return Ok(()); + } + + self.state_mut(uuid)?.queue_decode(); + Ok(()) + } + + fn on_external_arrival( + &mut self, + mut request: DirectRequest, + arrival_time_ms: f64, + replay_hashes: Option, + ) -> Result { + let uuid = request.uuid.unwrap_or_else(Uuid::new_v4); + request.uuid = Some(uuid); + request.arrival_timestamp_ms = Some(arrival_time_ms); + + self.collector.on_arrival( + uuid, + arrival_time_ms, + request.tokens.len(), + request.max_output_tokens, + ); + + self.requests + .insert(uuid, DisaggRequestState::new(request, arrival_time_ms)); + let Some(prefill_router) = self.prefill_router.as_mut() else { + let worker_idx = self.next_prefill_worker(); + self.dispatch_prefill(uuid, worker_idx)?; + return Ok(uuid); + }; + let maybe_worker_idx = { + let requests = &self.requests; + let request = requests + .get(&uuid) + .ok_or_else(|| anyhow!("offline disagg replay missing request state for {uuid}"))? + .original_request()?; + prefill_router.submit_request_with_hashes(request, replay_hashes, self.now_ms)? + }; + self.record_router_pending(); + if let Some(worker_idx) = maybe_worker_idx { + self.dispatch_prefill(uuid, worker_idx)?; + } + Ok(uuid) + } + + fn is_done(&self) -> bool { + self.events.is_empty() + && self.cluster_in_flight() == 0 + && match &self.admission { + AdmissionSource::Requests(pending) => pending.is_empty(), + AdmissionSource::Workload(driver) => driver.is_drained(), + } + && self + .prefill_workers + .iter() + .all(OfflineWorkerState::is_drained) + && self + .decode_workers + .iter() + .all(OfflineWorkerState::is_drained) + } + + fn next_timestamp(&mut self) -> Option { + let next_event_ms = self.events.peek().map(|event| event.at_ms); + let cluster_in_flight = self.cluster_in_flight(); + let next_arrival_ms = match (&self.mode, &mut self.admission) { + (ReplayMode::Trace, AdmissionSource::Requests(pending)) => pending + .front() + .and_then(|request| request.arrival_timestamp_ms), + (ReplayMode::Trace, AdmissionSource::Workload(driver)) => driver.next_ready_time_ms(), + (ReplayMode::Concurrency { max_in_flight }, AdmissionSource::Workload(driver)) => { + if cluster_in_flight < *max_in_flight { + driver.next_ready_time_ms() + } else { + None + } + } + (ReplayMode::Concurrency { .. }, AdmissionSource::Requests(_)) => None, + }; + + choose_next_timestamp(next_arrival_ms, next_event_ms) + } + + fn apply_prefill_router_events(&mut self, events: Vec) -> Result<()> { + let Some(prefill_router) = self.prefill_router.as_mut() else { + return Ok(()); + }; + for event in events { + prefill_router.apply_event(event)?; + } + Ok(()) + } + + fn process_prefill_signal(&mut self, signal: OutputSignal) -> Result<()> { + if !signal.completed { + return Ok(()); + } + + if self.prefill_router.is_some() { + let prefill_complete_admissions = { + let prefill_router = self.prefill_router.as_mut().expect("router checked above"); + prefill_router.mark_prefill_completed(signal.uuid)? + }; + #[cfg(test)] + { + self.stats.prefill_marked_count += 1; + } + self.record_router_pending(); + self.dispatch_prefill_admissions(prefill_complete_admissions)?; + + let admissions = { + let prefill_router = self.prefill_router.as_mut().expect("router checked above"); + prefill_router.free(signal.uuid)? + }; + #[cfg(test)] + { + self.stats.prefill_freed_count += 1; + } + self.record_router_pending(); + self.dispatch_prefill_admissions(admissions)?; + } + + self.enqueue_decode_after_handoff(signal.uuid, signal.handoff_delay_ms) + } + + fn process_decode_signal(&mut self, signal: OutputSignal) -> Result<()> { + if !signal.completed { + return Ok(()); + } + + let admissions = if let Some(decode_router) = self.decode_router.as_mut() { + let admissions = decode_router.free(signal.uuid)?; + #[cfg(test)] + { + self.stats.decode_freed_count += 1; + } + admissions + } else { + Vec::new() + }; + self.record_router_pending(); + if let AdmissionSource::Workload(driver) = &mut self.admission { + driver.on_complete(signal.uuid, self.now_ms)?; + } + self.state_mut(signal.uuid)?.mark_done(); + self.dispatch_decode_admissions(admissions)?; + Ok(()) + } + + fn process_prefill_pass( + &mut self, + worker_idx: usize, + completed_requests: usize, + output_signals: Vec, + kv_events: Vec, + ) -> Result<()> { + self.prefill_workers[worker_idx].mark_completed(completed_requests); + self.apply_prefill_router_events(kv_events)?; + for signal in output_signals { + self.process_prefill_signal(signal)?; + } + Ok(()) + } + + fn process_decode_pass( + &mut self, + worker_idx: usize, + completed_requests: usize, + output_signals: Vec, + ) -> Result<()> { + self.decode_workers[worker_idx].mark_completed(completed_requests); + for signal in output_signals { + self.process_decode_signal(signal)?; + } + Ok(()) + } + + fn apply_worker_completions(&mut self) -> Result { + let mut changed = false; + while let Some(WorkerCompletionPayload { + stage, + worker_idx, + completed_requests, + output_signals, + kv_events, + }) = pop_ready_worker_completion(&mut self.events, self.now_ms) + { + match stage { + SimulationWorkerStage::Prefill => { + self.prefill_workers[worker_idx].mark_idle(); + self.process_prefill_pass( + worker_idx, + completed_requests, + output_signals, + kv_events, + )?; + } + SimulationWorkerStage::Decode => { + self.decode_workers[worker_idx].mark_idle(); + self.process_decode_pass(worker_idx, completed_requests, output_signals)?; + } + SimulationWorkerStage::Aggregated => { + bail!("offline disagg replay received an aggregated completion event") + } + } + changed = true; + } + Ok(changed) + } + + fn apply_decode_handoffs(&mut self) -> Result { + let mut changed = false; + while let Some(uuid) = pop_ready_decode_handoff(&mut self.events, self.now_ms) { + self.enqueue_decode(uuid)?; + changed = true; + } + Ok(changed) + } + + fn enqueue_decode_after_handoff( + &mut self, + uuid: Uuid, + handoff_delay_ms: Option, + ) -> Result<()> { + if let Some(delay_ms) = handoff_delay_ms + && delay_ms > 0.0 + { + push_decode_handoff( + &mut self.events, + &mut self.next_event_seq, + self.now_ms + delay_ms, + uuid, + ); + return Ok(()); + } + + self.enqueue_decode(uuid) + } + + fn release_trace_arrivals(&mut self) -> Result { + let mut released_any = false; + if matches!(self.admission, AdmissionSource::Requests(_)) { + loop { + let next_ready = match &mut self.admission { + AdmissionSource::Requests(pending) => { + pop_next_trace_ready(pending, self.now_ms) + } + AdmissionSource::Workload(_) => unreachable!(), + }; + let Some((request, arrival_ms)) = next_ready else { + break; + }; + self.on_external_arrival(request, arrival_ms, None)?; + released_any = true; + } + return Ok(released_any); + } + + let ready_requests = match &mut self.admission { + AdmissionSource::Requests(_) => unreachable!(), + AdmissionSource::Workload(driver) => driver.pop_ready(self.now_ms, usize::MAX), + }; + for ready in ready_requests { + self.on_external_arrival( + ready.request, + ready.scheduled_ready_at_ms, + ready.replay_hashes, + )?; + released_any = true; + } + Ok(released_any) + } + + fn top_off_concurrency(&mut self, max_in_flight: usize) -> Result { + let mut released_any = false; + if matches!(self.admission, AdmissionSource::Requests(_)) { + loop { + let cluster_in_flight = self.cluster_in_flight(); + let next_ready = match &mut self.admission { + AdmissionSource::Requests(pending) => pop_next_concurrency_ready( + pending, + self.now_ms, + cluster_in_flight, + max_in_flight, + ), + AdmissionSource::Workload(_) => unreachable!(), + }; + let Some((request, arrival_ms)) = next_ready else { + break; + }; + self.on_external_arrival(request, arrival_ms, None)?; + released_any = true; + } + return Ok(released_any); + } + + let available = max_in_flight.saturating_sub(self.cluster_in_flight()); + if available == 0 { + return Ok(false); + } + + let ready_requests = match &mut self.admission { + AdmissionSource::Requests(_) => unreachable!(), + AdmissionSource::Workload(driver) => driver.pop_ready(self.now_ms, available), + }; + for ready in ready_requests { + self.on_external_arrival(ready.request, self.now_ms, ready.replay_hashes)?; + released_any = true; + } + Ok(released_any) + } + + fn drive_prefill_workers(&mut self) -> Result { + let mut changed = false; + for worker_idx in 0..self.prefill_workers.len() { + loop { + if !self.prefill_workers[worker_idx].is_ready() { + break; + } + + let executed = self.prefill_workers[worker_idx].execute_hidden_pass(self.now_ms); + changed = true; + + let completion_kv_events = + if executed.router_event_visibility == RouterEventVisibility::PassStart { + self.apply_prefill_router_events(executed.kv_events)?; + Vec::new() + } else { + executed.kv_events + }; + + if executed.end_ms == self.now_ms { + self.process_prefill_pass( + worker_idx, + executed.completed_requests, + executed.output_signals, + completion_kv_events, + )?; + continue; + } + + self.prefill_workers[worker_idx].mark_busy(); + push_worker_completion( + &mut self.events, + &mut self.next_event_seq, + executed.end_ms, + WorkerCompletionPayload { + stage: SimulationWorkerStage::Prefill, + worker_idx, + completed_requests: executed.completed_requests, + output_signals: executed.output_signals, + kv_events: completion_kv_events, + }, + ); + break; + } + } + Ok(changed) + } + + fn drive_decode_workers(&mut self) -> Result { + let mut changed = false; + for worker_idx in 0..self.decode_workers.len() { + loop { + if !self.decode_workers[worker_idx].is_ready() { + break; + } + + let executed = { + let (workers, collector) = (&mut self.decode_workers, &mut self.collector); + workers[worker_idx].execute_pass(collector, self.now_ms) + }; + changed = true; + + if executed.end_ms == self.now_ms { + self.process_decode_pass( + worker_idx, + executed.completed_requests, + executed.output_signals, + )?; + continue; + } + + self.decode_workers[worker_idx].mark_busy(); + push_worker_completion( + &mut self.events, + &mut self.next_event_seq, + executed.end_ms, + WorkerCompletionPayload { + stage: SimulationWorkerStage::Decode, + worker_idx, + completed_requests: executed.completed_requests, + output_signals: executed.output_signals, + kv_events: Vec::new(), + }, + ); + break; + } + } + Ok(changed) + } + + fn drain_current_timestamp(&mut self) -> Result<()> { + loop { + let mut changed = self.apply_worker_completions()?; + changed |= self.apply_decode_handoffs()?; + + changed |= match self.mode { + ReplayMode::Trace => self.release_trace_arrivals()?, + ReplayMode::Concurrency { max_in_flight } => { + self.top_off_concurrency(max_in_flight)? + } + }; + + changed |= self.drive_prefill_workers()?; + changed |= self.drive_decode_workers()?; + + if !changed { + break; + } + } + Ok(()) + } + + fn finish_test_stats(&mut self) { + #[cfg(test)] + { + self.stats.request_snapshots = self + .requests + .iter() + .map(|(uuid, state)| (*uuid, state.debug_snapshot())) + .collect(); + } + } + + fn run(mut self) -> Result<(TraceCollector, DisaggRuntimeStats)> { + self.drain_current_timestamp()?; + + while !self.is_done() { + let Some(next_timestamp_ms) = self.next_timestamp() else { + bail!( + "offline disagg replay reached a dead end with {} in-flight requests remaining", + self.cluster_in_flight() + ); + }; + self.now_ms = next_timestamp_ms; + self.drain_current_timestamp()?; + } + + self.finish_test_stats(); + Ok((self.collector, self.stats)) + } +} + +fn base_router_config( + args: &MockEngineArgs, + router_config: Option, +) -> KvRouterConfig { + let mut config = router_config.unwrap_or_default(); + if let Some(policy) = args.router_queue_policy { + config.router_queue_policy = policy; + } + config +} + +fn derive_prefill_router_config( + args: &MockEngineArgs, + router_config: Option, +) -> KvRouterConfig { + let mut config = base_router_config(args, router_config); + config.router_track_active_blocks = false; + config +} + +fn derive_decode_router_config( + args: &MockEngineArgs, + router_config: Option, +) -> KvRouterConfig { + let mut config = base_router_config(args, router_config); + config.overlap_score_weight = 0.0; + config.router_assume_kv_reuse = false; + config.router_track_prefill_tokens = false; + config +} + +pub(crate) fn simulate_trace_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, +) -> Result { + let pending = normalize_trace_requests(requests, arrival_speedup_ratio)?; + let (collector, _) = DisaggRuntime::new( + &config, + router_config, + pending, + ReplayMode::Trace, + router_mode, + )? + .run()?; + Ok(collector.finish()) +} + +pub(crate) fn simulate_concurrency_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result { + let pending = VecDeque::from(requests); + let (collector, _) = DisaggRuntime::new( + &config, + router_config, + pending, + ReplayMode::Concurrency { max_in_flight }, + router_mode, + )? + .run()?; + Ok(collector.finish()) +} + +pub(crate) fn simulate_trace_workload_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + router_mode: ReplayRouterMode, +) -> Result { + let driver = WorkloadDriver::new_trace(trace)?; + let (collector, _) = DisaggRuntime::new_workload( + &config, + router_config, + driver, + ReplayMode::Trace, + router_mode, + )? + .run()?; + Ok(collector.finish()) +} + +pub(crate) fn simulate_concurrency_workload_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> Result { + let driver = WorkloadDriver::new_concurrency(trace)?; + let (collector, _) = DisaggRuntime::new_workload( + &config, + router_config, + driver, + ReplayMode::Concurrency { max_in_flight }, + router_mode, + )? + .run()?; + Ok(collector.finish()) +} + +#[cfg(test)] +fn run_trace_collect( + config: &OfflineDisaggReplayConfig, + requests: Vec, + router_config: Option, + arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, +) -> (TraceCollector, DisaggRuntimeStats) { + let pending = normalize_trace_requests(requests, arrival_speedup_ratio).unwrap(); + DisaggRuntime::new( + config, + router_config, + pending, + ReplayMode::Trace, + router_mode, + ) + .unwrap() + .run() + .unwrap() +} + +#[cfg(test)] +fn run_concurrency_collect( + config: &OfflineDisaggReplayConfig, + requests: Vec, + router_config: Option, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> (TraceCollector, DisaggRuntimeStats) { + DisaggRuntime::new( + config, + router_config, + VecDeque::from(requests), + ReplayMode::Concurrency { max_in_flight }, + router_mode, + ) + .unwrap() + .run() + .unwrap() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::common::protocols::{MockEngineArgs, WorkerType}; + + fn staged_args(worker_type: WorkerType, speedup_ratio: f64) -> MockEngineArgs { + MockEngineArgs::builder() + .block_size(64) + .num_gpu_blocks(256) + .max_num_batched_tokens(Some(8192)) + .max_num_seqs(Some(8)) + .enable_prefix_caching(true) + .enable_chunked_prefill(true) + .speedup_ratio(speedup_ratio) + .decode_speedup_ratio(speedup_ratio) + .worker_type(worker_type) + .build() + .unwrap() + } + + fn disagg_config() -> OfflineDisaggReplayConfig { + OfflineDisaggReplayConfig { + prefill_args: staged_args(WorkerType::Prefill, 1000.0), + decode_args: staged_args(WorkerType::Decode, 1000.0), + num_prefill_workers: 2, + num_decode_workers: 2, + } + } + + fn disagg_config_with_handoff_delay() -> OfflineDisaggReplayConfig { + let mut config = disagg_config(); + config.prefill_args.kv_transfer_bandwidth = Some(1.0); + config.prefill_args.kv_bytes_per_token = Some(1_000_000); + config + } + + fn router_config() -> KvRouterConfig { + KvRouterConfig { + router_queue_threshold: Some(1.25), + ..KvRouterConfig::default() + } + } + + fn request( + uuid: u128, + prompt_tokens: usize, + output_tokens: usize, + arrival_ms: f64, + ) -> DirectRequest { + DirectRequest { + tokens: vec![1; prompt_tokens], + max_output_tokens: output_tokens, + uuid: Some(Uuid::from_u128(uuid)), + dp_rank: 0, + arrival_timestamp_ms: Some(arrival_ms), + } + } + + #[test] + fn test_derive_stage_router_configs_force_required_overrides() { + let config = KvRouterConfig { + overlap_score_weight: 2.0, + router_track_active_blocks: true, + router_assume_kv_reuse: true, + router_track_prefill_tokens: true, + ..KvRouterConfig::default() + }; + let args = staged_args(WorkerType::Prefill, 1.0); + let prefill = derive_prefill_router_config(&args, Some(config.clone())); + let decode = derive_decode_router_config(&args, Some(config)); + + assert!(!prefill.router_track_active_blocks); + assert_eq!(decode.overlap_score_weight, 0.0); + assert!(!decode.router_assume_kv_reuse); + assert!(!decode.router_track_prefill_tokens); + } + + #[rstest::rstest] + #[case(ReplayRouterMode::RoundRobin)] + #[case(ReplayRouterMode::KvRouter)] + fn test_trace_smoke_reports_decode_only_tokens(#[case] router_mode: ReplayRouterMode) { + let config = disagg_config(); + let requests = vec![request(1, 128, 3, 5.0)]; + + let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); + let (collector, stats) = + run_trace_collect(&config, requests, router_config, 1.0, router_mode); + let snapshot = collector.snapshot(Uuid::from_u128(1)).unwrap(); + let report = collector.finish(); + + assert_eq!(snapshot.arrival_time_ms, 0.0); + assert!(snapshot.first_admit_ms.is_some()); + assert!(snapshot.first_token_ms.is_some()); + assert_eq!(snapshot.output_length, 3); + assert_eq!(report.request_counts.completed_requests, 1); + assert_eq!( + stats.request_snapshots[&Uuid::from_u128(1)].phase, + DisaggPhase::Done + ); + } + + #[rstest::rstest] + #[case(ReplayRouterMode::RoundRobin)] + #[case(ReplayRouterMode::KvRouter)] + fn test_prefill_and_decode_use_separate_worker_pools(#[case] router_mode: ReplayRouterMode) { + let config = disagg_config(); + let requests = vec![request(1, 128, 2, 0.0), request(2, 128, 2, 10.0)]; + + let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); + let (_, stats) = run_trace_collect(&config, requests, router_config, 1.0, router_mode); + + for uuid in [Uuid::from_u128(1), Uuid::from_u128(2)] { + assert!(stats.prefill_assignments.contains_key(&uuid)); + assert!(stats.decode_assignments.contains_key(&uuid)); + } + } + + #[test] + fn test_prefill_overlap_prefers_same_worker_after_handoff_delay() { + let config = disagg_config(); + let requests = vec![request(1, 128, 2, 0.0), request(2, 128, 2, 100.0)]; + + let (_, stats) = run_trace_collect( + &config, + requests, + Some(router_config()), + 1.0, + ReplayRouterMode::KvRouter, + ); + + assert_eq!( + stats.prefill_assignments[&Uuid::from_u128(1)], + stats.prefill_assignments[&Uuid::from_u128(2)], + ); + } + + #[rstest::rstest] + #[case(ReplayRouterMode::RoundRobin)] + #[case(ReplayRouterMode::KvRouter)] + fn test_concurrency_backfill_waits_for_decode_completion( + #[case] router_mode: ReplayRouterMode, + ) { + let config = disagg_config(); + let requests = vec![ + DirectRequest { + tokens: vec![1; 128], + max_output_tokens: 3, + uuid: Some(Uuid::from_u128(1)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + DirectRequest { + tokens: vec![2; 128], + max_output_tokens: 3, + uuid: Some(Uuid::from_u128(2)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + ]; + + let router_config = (router_mode == ReplayRouterMode::KvRouter).then(router_config); + let (collector, _) = + run_concurrency_collect(&config, requests, router_config, 1, router_mode); + let first = collector.snapshot(Uuid::from_u128(1)).unwrap(); + let second = collector.snapshot(Uuid::from_u128(2)).unwrap(); + + assert_eq!(first.arrival_time_ms, 0.0); + assert_eq!(second.arrival_time_ms, first.last_token_ms.unwrap()); + } + + #[test] + fn test_prefill_completion_marks_and_frees_before_decode_handoff() { + let config = disagg_config(); + let requests = vec![request(1, 128, 2, 0.0)]; + + let (_, stats) = run_trace_collect( + &config, + requests, + Some(router_config()), + 1.0, + ReplayRouterMode::KvRouter, + ); + + assert_eq!(stats.prefill_marked_count, 1); + assert_eq!(stats.prefill_freed_count, 1); + assert_eq!(stats.decode_freed_count, 1); + } + + #[test] + fn test_handoff_delay_increases_decode_visible_ttft() { + let requests = vec![request(1, 128, 2, 0.0)]; + + let (baseline_collector, _) = run_trace_collect( + &disagg_config(), + requests.clone(), + None, + 1.0, + ReplayRouterMode::RoundRobin, + ); + let (delayed_collector, _) = run_trace_collect( + &disagg_config_with_handoff_delay(), + requests, + None, + 1.0, + ReplayRouterMode::RoundRobin, + ); + + let baseline = baseline_collector.snapshot(Uuid::from_u128(1)).unwrap(); + let delayed = delayed_collector.snapshot(Uuid::from_u128(1)).unwrap(); + let baseline_ttft = baseline.first_token_ms.unwrap() - baseline.arrival_time_ms; + let delayed_ttft = delayed.first_token_ms.unwrap() - delayed.arrival_time_ms; + + assert!( + delayed_ttft >= baseline_ttft + 120.0, + "expected delayed TTFT to include roughly 128ms of handoff delay, baseline={baseline_ttft}, delayed={delayed_ttft}" + ); + } +} diff --git a/lib/mocker/src/replay/offline/events.rs b/lib/mocker/src/replay/offline/events.rs index d7383a487551..d855f24f4ff1 100644 --- a/lib/mocker/src/replay/offline/events.rs +++ b/lib/mocker/src/replay/offline/events.rs @@ -4,15 +4,27 @@ use std::cmp::Ordering; use crate::common::protocols::OutputSignal; +use uuid::Uuid; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum SimulationWorkerStage { + Aggregated, + Prefill, + Decode, +} #[derive(Debug)] pub(crate) enum SimulationEventKind { WorkerCompletion { + stage: SimulationWorkerStage, worker_idx: usize, completed_requests: usize, output_signals: Vec, kv_events: Vec, }, + DecodeHandoff { + uuid: Uuid, + }, } #[derive(Debug)] diff --git a/lib/mocker/src/replay/offline/mod.rs b/lib/mocker/src/replay/offline/mod.rs index 225780a0af03..5bb23f1d6850 100644 --- a/lib/mocker/src/replay/offline/mod.rs +++ b/lib/mocker/src/replay/offline/mod.rs @@ -3,13 +3,16 @@ use crate::common::protocols::{DirectRequest, MockEngineArgs}; use crate::loadgen::Trace; +use crate::replay::OfflineDisaggReplayConfig; pub(crate) use crate::replay::normalize_trace_requests; use crate::replay::{ReplayRouterMode, TraceSimulationReport}; use dynamo_kv_router::config::KvRouterConfig; pub(crate) mod core; +pub(crate) mod disagg; pub(crate) mod events; pub(crate) mod multi; +pub(crate) mod runtime_utils; pub(crate) mod single; pub(crate) mod state; @@ -92,3 +95,54 @@ pub(crate) fn simulate_concurrency_workload( ) } } + +pub(crate) fn simulate_trace_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + arrival_speedup_ratio: f64, + router_mode: ReplayRouterMode, +) -> anyhow::Result { + disagg::simulate_trace_disagg( + config, + router_config, + requests, + arrival_speedup_ratio, + router_mode, + ) +} + +pub(crate) fn simulate_concurrency_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + requests: Vec, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> anyhow::Result { + disagg::simulate_concurrency_disagg(config, router_config, requests, max_in_flight, router_mode) +} + +pub(crate) fn simulate_trace_workload_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + router_mode: ReplayRouterMode, +) -> anyhow::Result { + disagg::simulate_trace_workload_disagg(config, router_config, trace, router_mode) +} + +pub(crate) fn simulate_concurrency_workload_disagg( + config: OfflineDisaggReplayConfig, + router_config: Option, + trace: Trace, + max_in_flight: usize, + router_mode: ReplayRouterMode, +) -> anyhow::Result { + disagg::simulate_concurrency_workload_disagg( + config, + router_config, + trace, + max_in_flight, + router_mode, + ) +} diff --git a/lib/mocker/src/replay/offline/multi.rs b/lib/mocker/src/replay/offline/multi.rs index 847133d556db..f973f7521d97 100644 --- a/lib/mocker/src/replay/offline/multi.rs +++ b/lib/mocker/src/replay/offline/multi.rs @@ -1,11 +1,15 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -use super::events::{SimulationEvent, SimulationEventKind}; +use super::events::{SimulationEvent, SimulationWorkerStage}; use super::normalize_trace_requests; +use super::runtime_utils::{ + WorkerCompletionPayload, next_timestamp as choose_next_timestamp, pop_next_concurrency_ready, + pop_next_trace_ready, pop_ready_worker_completion, push_worker_completion, +}; #[cfg(test)] use super::state::OfflineWorkerSnapshot; -use super::state::OfflineWorkerState; +use super::state::{AggRequestState, OfflineWorkerState}; use crate::common::protocols::{DirectRequest, MockEngineArgs, OutputSignal}; use crate::loadgen::{ReplayRequestHashes, Trace, WorkloadDriver}; use crate::replay::router::OfflineReplayRouter; @@ -16,7 +20,7 @@ use crate::scheduler::RouterEventVisibility; use anyhow::bail; use dynamo_kv_router::config::KvRouterConfig; use dynamo_kv_router::protocols::RouterEvent; -use std::collections::{BinaryHeap, HashMap, HashSet, VecDeque}; +use std::collections::{BinaryHeap, HashMap, VecDeque}; use uuid::Uuid; #[derive(Debug, Clone, Copy)] @@ -62,13 +66,13 @@ struct OfflineRuntime { next_worker_idx: usize, next_event_seq: u64, admission: AdmissionSource, - router_pending: HashMap, + requests: HashMap, + queued_requests: usize, workers: Vec, collector: TraceCollector, events: BinaryHeap, mode: ReplayMode, router: Option, - prefill_completed: HashSet, stats: OfflineRuntimeStats, #[cfg(test)] worker_active_requests: Vec>, @@ -135,7 +139,8 @@ impl OfflineRuntime { next_worker_idx: 0, next_event_seq: 0, admission, - router_pending: HashMap::new(), + requests: HashMap::new(), + queued_requests: 0, workers: (0..num_workers) .map(|worker_idx| { OfflineWorkerState::new(worker_idx, args.clone(), capture_kv_events) @@ -145,7 +150,6 @@ impl OfflineRuntime { events: BinaryHeap::new(), mode, router, - prefill_completed: HashSet::new(), #[cfg(test)] stats: OfflineRuntimeStats::default(), #[cfg(not(test))] @@ -162,7 +166,7 @@ impl OfflineRuntime { .iter() .map(OfflineWorkerState::in_flight) .sum::() - + self.router_pending.len() + + self.queued_requests } fn record_in_flight_peak(&mut self) { @@ -220,9 +224,14 @@ impl OfflineRuntime { fn dispatch_router_admissions(&mut self, admissions: Vec<(Uuid, usize)>) -> anyhow::Result<()> { for (uuid, worker_idx) in admissions { - let request = self.router_pending.remove(&uuid).ok_or_else(|| { - anyhow::anyhow!("offline replay missing queued request state for {uuid}") - })?; + let request = self + .requests + .get_mut(&uuid) + .ok_or_else(|| { + anyhow::anyhow!("offline replay missing queued request state for {uuid}") + })? + .take_queued_request(uuid)?; + self.queued_requests = self.queued_requests.saturating_sub(1); self.dispatch_to_worker(request, uuid, worker_idx)?; } Ok(()) @@ -248,6 +257,7 @@ impl OfflineRuntime { ); let Some(router) = self.router.as_mut() else { + self.requests.insert(uuid, AggRequestState::new_running()); let worker_idx = self.next_worker_idx; self.next_worker_idx = (self.next_worker_idx + 1) % self.workers.len(); self.dispatch_to_worker(request, uuid, worker_idx)?; @@ -258,11 +268,14 @@ impl OfflineRuntime { router.submit_request_with_hashes(&request, replay_hashes, self.now_ms)?; self.record_router_pending(); if let Some(worker_idx) = maybe_worker_idx { + self.requests.insert(uuid, AggRequestState::new_running()); self.dispatch_to_worker(request, uuid, worker_idx)?; return Ok(uuid); } - self.router_pending.insert(uuid, request); + self.requests + .insert(uuid, AggRequestState::new_queued(request)); + self.queued_requests += 1; self.record_in_flight_peak(); Ok(uuid) } @@ -295,21 +308,7 @@ impl OfflineRuntime { (ReplayMode::Concurrency { .. }, AdmissionSource::Requests(_)) => None, }; - match (next_arrival_ms, next_event_ms) { - (Some(arrival_ms), Some(event_ms)) => Some(arrival_ms.min(event_ms)), - (Some(arrival_ms), None) => Some(arrival_ms), - (None, Some(event_ms)) => Some(event_ms), - (None, None) => None, - } - } - - fn push_event(&mut self, at_ms: f64, kind: SimulationEventKind) { - self.events.push(SimulationEvent { - at_ms, - seq_no: self.next_event_seq, - kind, - }); - self.next_event_seq += 1; + choose_next_timestamp(next_arrival_ms, next_event_ms) } fn apply_completed_requests(&mut self, worker_idx: usize, completed_requests: usize) { @@ -339,7 +338,9 @@ impl OfflineRuntime { } self.record_router_pending(); } - self.prefill_completed.remove(&signal.uuid); + self.requests.remove(&signal.uuid).ok_or_else(|| { + anyhow::anyhow!("offline replay missing request state for {}", signal.uuid) + })?; if let AdmissionSource::Workload(driver) = &mut self.admission { driver.on_complete(signal.uuid, self.now_ms)?; } @@ -347,10 +348,23 @@ impl OfflineRuntime { return Ok(()); } - if !self.prefill_completed.insert(signal.uuid) { + let already_marked = self + .requests + .get(&signal.uuid) + .ok_or_else(|| { + anyhow::anyhow!("offline replay missing request state for {}", signal.uuid) + })? + .prefill_completed(); + if already_marked { return Ok(()); } + self.requests + .get_mut(&signal.uuid) + .ok_or_else(|| { + anyhow::anyhow!("offline replay missing request state for {}", signal.uuid) + })? + .mark_prefill_completed(); if let Some(router) = self.router.as_mut() { admissions = router.mark_prefill_completed(signal.uuid)?; #[cfg(test)] @@ -395,25 +409,15 @@ impl OfflineRuntime { fn apply_worker_completions(&mut self) -> anyhow::Result { let mut changed = false; - loop { - let Some(event) = self.events.peek() else { - break; - }; - if event.at_ms != self.now_ms { - break; - } - if !matches!(event.kind, SimulationEventKind::WorkerCompletion { .. }) { - break; - } - - let event = self.events.pop().expect("event must exist after peek"); - let SimulationEventKind::WorkerCompletion { - worker_idx, - completed_requests, - output_signals, - kv_events, - } = event.kind; - + while let Some(WorkerCompletionPayload { + stage, + worker_idx, + completed_requests, + output_signals, + kv_events, + }) = pop_ready_worker_completion(&mut self.events, self.now_ms) + { + debug_assert_eq!(stage, SimulationWorkerStage::Aggregated); self.workers[worker_idx].mark_idle(); self.process_completed_pass(worker_idx, completed_requests, output_signals, kv_events)?; changed = true; @@ -424,39 +428,33 @@ impl OfflineRuntime { fn release_trace_arrivals(&mut self) -> anyhow::Result { let mut released_any = false; - let mut ready_requests = Vec::new(); - - match &mut self.admission { - AdmissionSource::Requests(pending) => { - while pending - .front() - .and_then(|request| request.arrival_timestamp_ms) - .is_some_and(|arrival_ms| arrival_ms <= self.now_ms) - { - let request = pending - .pop_front() - .expect("front request must exist when arrival is ready"); - let arrival_ms = request - .arrival_timestamp_ms - .expect("trace replay requests must have an arrival timestamp"); - ready_requests.push((request, arrival_ms, None)); - } - } - AdmissionSource::Workload(driver) => { - ready_requests.extend(driver.pop_ready(self.now_ms, usize::MAX).into_iter().map( - |ready| { - ( - ready.request, - ready.scheduled_ready_at_ms, - ready.replay_hashes, - ) - }, - )); + if matches!(self.admission, AdmissionSource::Requests(_)) { + loop { + let next_ready = match &mut self.admission { + AdmissionSource::Requests(pending) => { + pop_next_trace_ready(pending, self.now_ms) + } + AdmissionSource::Workload(_) => unreachable!(), + }; + let Some((request, arrival_ms)) = next_ready else { + break; + }; + self.assign_request(request, arrival_ms, None)?; + released_any = true; } + return Ok(released_any); } - for (request, arrival_ms, replay_hashes) in ready_requests { - self.assign_request(request, arrival_ms, replay_hashes)?; + let ready_requests = match &mut self.admission { + AdmissionSource::Requests(_) => unreachable!(), + AdmissionSource::Workload(driver) => driver.pop_ready(self.now_ms, usize::MAX), + }; + for ready in ready_requests { + self.assign_request( + ready.request, + ready.scheduled_ready_at_ms, + ready.replay_hashes, + )?; released_any = true; } @@ -465,33 +463,38 @@ impl OfflineRuntime { fn top_off_concurrency(&mut self, max_in_flight: usize) -> anyhow::Result { let mut released_any = false; + if matches!(self.admission, AdmissionSource::Requests(_)) { + loop { + let cluster_in_flight = self.cluster_in_flight(); + let next_ready = match &mut self.admission { + AdmissionSource::Requests(pending) => pop_next_concurrency_ready( + pending, + self.now_ms, + cluster_in_flight, + max_in_flight, + ), + AdmissionSource::Workload(_) => unreachable!(), + }; + let Some((request, arrival_ms)) = next_ready else { + break; + }; + self.assign_request(request, arrival_ms, None)?; + released_any = true; + } + return Ok(released_any); + } + let available = max_in_flight.saturating_sub(self.cluster_in_flight()); if available == 0 { return Ok(false); } - let mut ready_requests = Vec::new(); - match &mut self.admission { - AdmissionSource::Requests(pending) => { - for _ in 0..available { - let Some(request) = pending.pop_front() else { - break; - }; - ready_requests.push((request, None)); - } - } - AdmissionSource::Workload(driver) => { - ready_requests.extend( - driver - .pop_ready(self.now_ms, available) - .into_iter() - .map(|ready| (ready.request, ready.replay_hashes)), - ); - } - } - - for (request, replay_hashes) in ready_requests { - self.assign_request(request, self.now_ms, replay_hashes)?; + let ready_requests = match &mut self.admission { + AdmissionSource::Requests(_) => unreachable!(), + AdmissionSource::Workload(driver) => driver.pop_ready(self.now_ms, available), + }; + for ready in ready_requests { + self.assign_request(ready.request, self.now_ms, ready.replay_hashes)?; released_any = true; } @@ -531,9 +534,12 @@ impl OfflineRuntime { } self.workers[worker_idx].mark_busy(); - self.push_event( + push_worker_completion( + &mut self.events, + &mut self.next_event_seq, executed.end_ms, - SimulationEventKind::WorkerCompletion { + WorkerCompletionPayload { + stage: SimulationWorkerStage::Aggregated, worker_idx, completed_requests: executed.completed_requests, output_signals: executed.output_signals, @@ -616,10 +622,19 @@ impl OfflineRuntime { #[cfg(test)] fn debug_snapshot(&self) -> OfflineRuntimeSnapshot { - let mut router_pending_request_ids = - self.router_pending.keys().copied().collect::>(); + let mut router_pending_request_ids = self + .requests + .iter() + .filter(|(_, state)| state.is_queued_at_router()) + .map(|(uuid, _)| *uuid) + .collect::>(); router_pending_request_ids.sort_unstable(); - let mut prefill_completed = self.prefill_completed.iter().copied().collect::>(); + let mut prefill_completed = self + .requests + .iter() + .filter(|(_, state)| state.prefill_completed()) + .map(|(uuid, _)| *uuid) + .collect::>(); prefill_completed.sort_unstable(); OfflineRuntimeSnapshot { @@ -1483,16 +1498,16 @@ mod tests { let request_1 = collector.snapshot(Uuid::from_u128(1)).unwrap(); let request_2 = collector.snapshot(Uuid::from_u128(2)).unwrap(); let request_3 = collector.snapshot(Uuid::from_u128(3)).unwrap(); + let first_unblock_ms = request_1 + .first_token_ms + .unwrap() + .min(request_2.first_token_ms.unwrap()); assert!(stats.max_router_pending > 0); assert!(request_3.first_admit_ms.unwrap() > request_3.arrival_time_ms); - assert!( - request_3.first_admit_ms.unwrap() - < request_1 - .last_token_ms - .unwrap() - .min(request_2.last_token_ms.unwrap()) - ); + assert_eq!(request_3.first_admit_ms.unwrap(), first_unblock_ms); + assert!(request_3.first_admit_ms.unwrap() < request_1.last_token_ms.unwrap()); + assert!(request_3.first_admit_ms.unwrap() < request_2.last_token_ms.unwrap()); } #[test] @@ -1561,6 +1576,72 @@ mod tests { ); } + #[test] + fn test_multi_worker_trace_kv_router_fcfs_and_lcfs_admit_queued_requests_in_opposite_timestamp_order() + { + let requests = vec![ + DirectRequest { + tokens: vec![10; 64], + max_output_tokens: 8, + uuid: Some(Uuid::from_u128(10)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.0), + }, + DirectRequest { + tokens: vec![20; 128], + max_output_tokens: 8, + uuid: Some(Uuid::from_u128(20)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.0), + }, + DirectRequest { + tokens: vec![30; 64], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(30)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.1), + }, + DirectRequest { + tokens: vec![40; 64], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(40)), + dp_rank: 0, + arrival_timestamp_ms: Some(0.2), + }, + ]; + + let (fcfs_collector, fcfs_stats) = run_trace_multi_collect_with_stats( + &queueing_router_args(RouterQueuePolicy::Fcfs), + requests.clone(), + 2, + ReplayRouterMode::KvRouter, + ); + let (lcfs_collector, lcfs_stats) = run_trace_multi_collect_with_stats( + &queueing_router_args(RouterQueuePolicy::Lcfs), + requests, + 2, + ReplayRouterMode::KvRouter, + ); + + let fcfs_request_30 = fcfs_collector.snapshot(Uuid::from_u128(30)).unwrap(); + let fcfs_request_40 = fcfs_collector.snapshot(Uuid::from_u128(40)).unwrap(); + let lcfs_request_30 = lcfs_collector.snapshot(Uuid::from_u128(30)).unwrap(); + let lcfs_request_40 = lcfs_collector.snapshot(Uuid::from_u128(40)).unwrap(); + + assert!(fcfs_stats.max_router_pending > 0); + assert!(lcfs_stats.max_router_pending > 0); + assert_eq!( + &fcfs_stats.dispatch_order[2..4], + &[Uuid::from_u128(30), Uuid::from_u128(40)] + ); + assert_eq!( + &lcfs_stats.dispatch_order[2..4], + &[Uuid::from_u128(40), Uuid::from_u128(30)] + ); + assert!(fcfs_request_30.first_admit_ms.unwrap() < fcfs_request_40.first_admit_ms.unwrap()); + assert!(lcfs_request_40.first_admit_ms.unwrap() < lcfs_request_30.first_admit_ms.unwrap()); + } + #[test] fn test_multi_worker_concurrency_kv_router_respects_max_in_flight() { let args = queueing_router_args(RouterQueuePolicy::Fcfs); @@ -1605,6 +1686,51 @@ mod tests { assert!(stats.max_router_pending > 0); } + #[test] + fn test_multi_worker_concurrency_kv_router_records_backfill_timing() { + let args = queueing_router_args(RouterQueuePolicy::Fcfs); + let (collector, stats) = run_concurrency_multi_collect_with_stats( + &args, + vec![ + DirectRequest { + tokens: vec![1; 64], + max_output_tokens: 2, + uuid: Some(Uuid::from_u128(11)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + DirectRequest { + tokens: vec![2; 64], + max_output_tokens: 4, + uuid: Some(Uuid::from_u128(22)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + DirectRequest { + tokens: vec![3; 64], + max_output_tokens: 2, + uuid: Some(Uuid::from_u128(33)), + dp_rank: 0, + arrival_timestamp_ms: None, + }, + ], + 2, + 2, + ReplayRouterMode::KvRouter, + ); + + let request_1 = collector.snapshot(Uuid::from_u128(11)).unwrap(); + let request_2 = collector.snapshot(Uuid::from_u128(22)).unwrap(); + let request_3 = collector.snapshot(Uuid::from_u128(33)).unwrap(); + + assert_eq!(request_1.arrival_time_ms, 0.0); + assert_eq!(request_2.arrival_time_ms, 0.0); + assert_eq!(request_3.arrival_time_ms, request_1.last_token_ms.unwrap()); + assert!(request_3.arrival_time_ms < request_2.last_token_ms.unwrap()); + assert_eq!(request_3.first_admit_ms.unwrap(), request_3.arrival_time_ms); + assert_eq!(stats.max_in_flight_seen, 2); + } + #[test] fn test_multi_worker_trace_single_worker_round_robin_matches_single_runtime() { let args = replay_args(true, true); diff --git a/lib/mocker/src/replay/offline/runtime_utils.rs b/lib/mocker/src/replay/offline/runtime_utils.rs new file mode 100644 index 000000000000..2f2e29a0165d --- /dev/null +++ b/lib/mocker/src/replay/offline/runtime_utils.rs @@ -0,0 +1,257 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use std::collections::{BinaryHeap, VecDeque}; + +use dynamo_kv_router::protocols::RouterEvent; + +use super::events::{SimulationEvent, SimulationEventKind, SimulationWorkerStage}; +use crate::common::protocols::{DirectRequest, OutputSignal}; + +#[derive(Debug)] +pub(super) struct WorkerCompletionPayload { + pub stage: SimulationWorkerStage, + pub worker_idx: usize, + pub completed_requests: usize, + pub output_signals: Vec, + pub kv_events: Vec, +} + +pub(super) fn next_timestamp( + next_arrival_ms: Option, + next_event_ms: Option, +) -> Option { + match (next_arrival_ms, next_event_ms) { + (Some(arrival_ms), Some(event_ms)) => Some(arrival_ms.min(event_ms)), + (Some(arrival_ms), None) => Some(arrival_ms), + (None, Some(event_ms)) => Some(event_ms), + (None, None) => None, + } +} + +pub(super) fn pop_next_trace_ready( + pending: &mut VecDeque, + now_ms: f64, +) -> Option<(DirectRequest, f64)> { + let arrival_ms = pending + .front() + .and_then(|request| request.arrival_timestamp_ms) + .filter(|arrival_ms| *arrival_ms <= now_ms)?; + let request = pending + .pop_front() + .expect("front request must exist when arrival is ready"); + Some((request, arrival_ms)) +} + +pub(super) fn pop_next_concurrency_ready( + pending: &mut VecDeque, + now_ms: f64, + cluster_in_flight: usize, + max_in_flight: usize, +) -> Option<(DirectRequest, f64)> { + if cluster_in_flight >= max_in_flight { + return None; + } + let request = pending.pop_front()?; + Some((request, now_ms)) +} + +pub(super) fn push_worker_completion( + events: &mut BinaryHeap, + next_event_seq: &mut u64, + at_ms: f64, + payload: WorkerCompletionPayload, +) { + events.push(SimulationEvent { + at_ms, + seq_no: *next_event_seq, + kind: SimulationEventKind::WorkerCompletion { + stage: payload.stage, + worker_idx: payload.worker_idx, + completed_requests: payload.completed_requests, + output_signals: payload.output_signals, + kv_events: payload.kv_events, + }, + }); + *next_event_seq += 1; +} + +pub(super) fn pop_ready_worker_completion( + events: &mut BinaryHeap, + now_ms: f64, +) -> Option { + let event = events.peek()?; + if event.at_ms != now_ms { + return None; + } + let SimulationEventKind::WorkerCompletion { .. } = &event.kind else { + return None; + }; + let event = events.pop().expect("event must exist after peek"); + let (stage, worker_idx, completed_requests, output_signals, kv_events) = match event.kind { + SimulationEventKind::WorkerCompletion { + stage, + worker_idx, + completed_requests, + output_signals, + kv_events, + } => ( + stage, + worker_idx, + completed_requests, + output_signals, + kv_events, + ), + SimulationEventKind::DecodeHandoff { .. } => { + unreachable!("peeked worker completion event must match popped event") + } + }; + Some(WorkerCompletionPayload { + stage, + worker_idx, + completed_requests, + output_signals, + kv_events, + }) +} + +pub(super) fn push_decode_handoff( + events: &mut BinaryHeap, + next_event_seq: &mut u64, + at_ms: f64, + uuid: uuid::Uuid, +) { + events.push(SimulationEvent { + at_ms, + seq_no: *next_event_seq, + kind: SimulationEventKind::DecodeHandoff { uuid }, + }); + *next_event_seq += 1; +} + +pub(super) fn pop_ready_decode_handoff( + events: &mut BinaryHeap, + now_ms: f64, +) -> Option { + let event = events.peek()?; + if event.at_ms != now_ms { + return None; + } + let SimulationEventKind::DecodeHandoff { .. } = &event.kind else { + return None; + }; + let event = events.pop().expect("event must exist after peek"); + let SimulationEventKind::DecodeHandoff { uuid } = event.kind else { + unreachable!("peeked decode handoff event must match popped event"); + }; + Some(uuid) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::replay::offline::events::SimulationWorkerStage; + use uuid::Uuid; + + fn direct_request(uuid: u128, arrival_timestamp_ms: Option) -> DirectRequest { + DirectRequest { + tokens: vec![1; 8], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(uuid)), + dp_rank: 0, + arrival_timestamp_ms, + } + } + + #[test] + fn test_next_timestamp_matches_current_choice_logic() { + assert_eq!(next_timestamp(Some(1.0), Some(2.0)), Some(1.0)); + assert_eq!(next_timestamp(Some(2.0), Some(1.0)), Some(1.0)); + assert_eq!(next_timestamp(Some(3.0), None), Some(3.0)); + assert_eq!(next_timestamp(None, Some(4.0)), Some(4.0)); + assert_eq!(next_timestamp(None, None), None); + } + + #[test] + fn test_pop_next_trace_ready_releases_only_arrivals_at_or_before_now() { + let mut pending = VecDeque::from(vec![ + direct_request(1, Some(1.0)), + direct_request(2, Some(1.1)), + direct_request(3, Some(2.0)), + ]); + + let (request_1, arrival_1) = pop_next_trace_ready(&mut pending, 1.0).unwrap(); + assert_eq!(request_1.uuid, Some(Uuid::from_u128(1))); + assert_eq!(arrival_1, 1.0); + + assert!(pop_next_trace_ready(&mut pending, 1.0).is_none()); + + let (request_2, arrival_2) = pop_next_trace_ready(&mut pending, 1.1).unwrap(); + assert_eq!(request_2.uuid, Some(Uuid::from_u128(2))); + assert_eq!(arrival_2, 1.1); + assert_eq!(pending.len(), 1); + } + + #[test] + fn test_pop_next_concurrency_ready_stops_at_max_in_flight() { + let mut pending = VecDeque::from(vec![direct_request(1, None), direct_request(2, None)]); + + assert!(pop_next_concurrency_ready(&mut pending, 5.0, 2, 2).is_none()); + + let (request, arrival_ms) = pop_next_concurrency_ready(&mut pending, 5.0, 1, 2).unwrap(); + assert_eq!(request.uuid, Some(Uuid::from_u128(1))); + assert_eq!(arrival_ms, 5.0); + assert_eq!(pending.len(), 1); + } + + #[test] + fn test_worker_completion_helpers_preserve_same_time_sequence_ordering() { + let mut events = BinaryHeap::new(); + let mut next_event_seq = 0; + + push_worker_completion( + &mut events, + &mut next_event_seq, + 10.0, + WorkerCompletionPayload { + stage: SimulationWorkerStage::Aggregated, + worker_idx: 7, + completed_requests: 1, + output_signals: vec![OutputSignal { + uuid: Uuid::from_u128(7), + completed: true, + handoff_delay_ms: None, + }], + kv_events: Vec::new(), + }, + ); + push_worker_completion( + &mut events, + &mut next_event_seq, + 10.0, + WorkerCompletionPayload { + stage: SimulationWorkerStage::Aggregated, + worker_idx: 8, + completed_requests: 2, + output_signals: vec![OutputSignal { + uuid: Uuid::from_u128(8), + completed: false, + handoff_delay_ms: None, + }], + kv_events: Vec::new(), + }, + ); + + assert!(pop_ready_worker_completion(&mut events, 9.0).is_none()); + + let first = pop_ready_worker_completion(&mut events, 10.0).unwrap(); + let second = pop_ready_worker_completion(&mut events, 10.0).unwrap(); + assert_eq!(first.stage, SimulationWorkerStage::Aggregated); + assert_eq!(first.worker_idx, 7); + assert_eq!(first.completed_requests, 1); + assert_eq!(second.stage, SimulationWorkerStage::Aggregated); + assert_eq!(second.worker_idx, 8); + assert_eq!(second.completed_requests, 2); + assert!(events.is_empty()); + } +} diff --git a/lib/mocker/src/replay/offline/state.rs b/lib/mocker/src/replay/offline/state.rs index 0acbacbace53..47ba06b8b2f9 100644 --- a/lib/mocker/src/replay/offline/state.rs +++ b/lib/mocker/src/replay/offline/state.rs @@ -1,10 +1,162 @@ // SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 +use anyhow::{Result, anyhow, bail}; + use crate::common::protocols::DirectRequest; use crate::common::protocols::MockEngineArgs; use crate::replay::TraceCollector; use crate::scheduler::{EngineCore, EnginePassResult}; +use uuid::Uuid; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AggRequestPhase { + QueuedAtRouter, + Running, +} + +pub(crate) struct AggRequestState { + request: Option, + phase: AggRequestPhase, + prefill_completed: bool, +} + +impl AggRequestState { + pub(crate) fn new_queued(request: DirectRequest) -> Self { + Self { + request: Some(request), + phase: AggRequestPhase::QueuedAtRouter, + prefill_completed: false, + } + } + + pub(crate) fn new_running() -> Self { + Self { + request: None, + phase: AggRequestPhase::Running, + prefill_completed: false, + } + } + + pub(crate) fn is_queued_at_router(&self) -> bool { + self.phase == AggRequestPhase::QueuedAtRouter + } + + pub(crate) fn take_queued_request(&mut self, uuid: Uuid) -> Result { + if !self.is_queued_at_router() { + bail!("offline replay expected queued request state for {uuid}"); + } + let request = self + .request + .take() + .ok_or_else(|| anyhow!("offline replay missing queued request payload for {uuid}"))?; + self.phase = AggRequestPhase::Running; + Ok(request) + } + + pub(crate) fn prefill_completed(&self) -> bool { + self.prefill_completed + } + + pub(crate) fn mark_prefill_completed(&mut self) { + self.prefill_completed = true; + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum DisaggPhase { + QueuedPrefill, + RunningPrefill, + QueuedDecode, + RunningDecode, + Done, +} + +pub(crate) struct DisaggRequestState { + original: Option, + #[cfg(test)] + arrival_ms: f64, + phase: DisaggPhase, + prefill_worker_idx: Option, + decode_worker_idx: Option, +} + +#[cfg(test)] +#[derive(Debug, Clone, PartialEq)] +pub(crate) struct DisaggRequestSnapshot { + pub(crate) arrival_ms: f64, + pub(crate) phase: DisaggPhase, + pub(crate) prefill_worker_idx: Option, + pub(crate) decode_worker_idx: Option, +} + +impl DisaggRequestState { + pub(crate) fn new(request: DirectRequest, arrival_ms: f64) -> Self { + #[cfg(not(test))] + let _ = arrival_ms; + Self { + original: Some(request), + #[cfg(test)] + arrival_ms, + phase: DisaggPhase::QueuedPrefill, + prefill_worker_idx: None, + decode_worker_idx: None, + } + } + + pub(crate) fn is_queued_prefill(&self) -> bool { + self.phase == DisaggPhase::QueuedPrefill + } + + pub(crate) fn is_queued_decode(&self) -> bool { + self.phase == DisaggPhase::QueuedDecode + } + + pub(crate) fn original_request(&self) -> Result<&DirectRequest> { + self.original + .as_ref() + .ok_or_else(|| anyhow!("offline disagg replay request payload was already released")) + } + + pub(crate) fn build_prefill_request(&self) -> Result { + let mut request = self.original_request()?.clone(); + request.max_output_tokens = 1; + Ok(request) + } + + pub(crate) fn build_decode_request(&self) -> Result { + Ok(self.original_request()?.clone()) + } + + pub(crate) fn start_prefill(&mut self, worker_idx: usize) { + self.phase = DisaggPhase::RunningPrefill; + self.prefill_worker_idx = Some(worker_idx); + } + + pub(crate) fn queue_decode(&mut self) { + self.phase = DisaggPhase::QueuedDecode; + } + + pub(crate) fn start_decode(&mut self, worker_idx: usize) { + self.phase = DisaggPhase::RunningDecode; + self.decode_worker_idx = Some(worker_idx); + } + + pub(crate) fn mark_done(&mut self) { + self.phase = DisaggPhase::Done; + self.original = None; + } + + #[cfg(test)] + pub(crate) fn debug_snapshot(&self) -> DisaggRequestSnapshot { + DisaggRequestSnapshot { + arrival_ms: self.arrival_ms, + phase: self.phase, + prefill_worker_idx: self.prefill_worker_idx, + decode_worker_idx: self.decode_worker_idx, + } + } +} pub(crate) struct OfflineWorkerState { core: EngineCore, @@ -91,6 +243,10 @@ impl OfflineWorkerState { self.core.execute_pass(collector, now_ms) } + pub(crate) fn execute_hidden_pass(&mut self, now_ms: f64) -> EnginePassResult { + self.core.execute_hidden_pass(now_ms) + } + #[cfg(test)] pub(crate) fn debug_snapshot(&self) -> OfflineWorkerSnapshot { OfflineWorkerSnapshot { diff --git a/lib/mocker/src/replay/router/offline.rs b/lib/mocker/src/replay/router/offline.rs index b682d885993f..d2ed7ab25bb7 100644 --- a/lib/mocker/src/replay/router/offline.rs +++ b/lib/mocker/src/replay/router/offline.rs @@ -110,6 +110,7 @@ struct PendingRequest { token_seq: Option>, isl_tokens: usize, overlaps: OverlapScores, + track_prefill_tokens: bool, expected_output_tokens: Option, } @@ -130,6 +131,7 @@ impl PendingRequest { overlaps: self.overlaps.clone(), decode_blocks, prefill_tokens, + track_prefill_tokens: self.track_prefill_tokens, router_config_override: None, update_states: true, lora_name: None, @@ -258,7 +260,6 @@ impl OfflineReplayRouter { self.drain_pending() } - #[cfg(test)] pub(crate) fn pending_count(&self) -> usize { self.pending.len() } @@ -371,6 +372,7 @@ impl OfflineReplayRouter { token_seq, isl_tokens: request.tokens.len(), overlaps, + track_prefill_tokens: self.config.router_track_prefill_tokens, expected_output_tokens: Some( u32::try_from(request.max_output_tokens) .context("max_output_tokens does not fit into u32")?, @@ -379,11 +381,14 @@ impl OfflineReplayRouter { } fn admit_request(&mut self, request: PendingRequest) -> Result { - let (decode_blocks, prefill_tokens) = self.slots.potential_blocks_and_tokens( - request.token_seq.as_deref(), - request.isl_tokens, - request.overlaps.clone(), - ); + let (decode_blocks, prefill_tokens) = self + .slots + .potential_blocks_and_tokens_with_prefill_tracking( + request.token_seq.as_deref(), + request.isl_tokens, + request.overlaps.clone(), + request.track_prefill_tokens, + ); let scheduling_request = request.scheduling_request(decode_blocks, prefill_tokens); let selection = self.selector.select_worker( &self.workers_with_configs, @@ -400,6 +405,7 @@ impl OfflineReplayRouter { token_sequence: request.token_seq, isl: request.isl_tokens, overlap: selection.overlap_blocks, + track_prefill_tokens: request.track_prefill_tokens, expected_output_tokens: request.expected_output_tokens, worker: selection.worker, lora_name: None, diff --git a/lib/mocker/src/replay/router/online.rs b/lib/mocker/src/replay/router/online.rs index b462882682c4..87a94e6e097e 100644 --- a/lib/mocker/src/replay/router/online.rs +++ b/lib/mocker/src/replay/router/online.rs @@ -138,6 +138,7 @@ impl KvReplayRouter { args.block_size as u32, selector, policy, + config.router_track_prefill_tokens, CancellationToken::new(), "replay", false, diff --git a/lib/mocker/src/replay/router/shared.rs b/lib/mocker/src/replay/router/shared.rs index 4601f79a3065..56aa85245322 100644 --- a/lib/mocker/src/replay/router/shared.rs +++ b/lib/mocker/src/replay/router/shared.rs @@ -107,7 +107,7 @@ pub(super) fn replay_selector(config: &KvRouterConfig) -> DefaultWorkerSelector DefaultWorkerSelector::new(Some(config.clone()), "replay") } -pub(super) fn replay_router_config( +pub(crate) fn replay_router_config( args: &MockEngineArgs, router_config: Option, ) -> KvRouterConfig { diff --git a/lib/mocker/src/replay/validate.rs b/lib/mocker/src/replay/validate.rs index 824abb5eb9d8..3ebafe1056ff 100644 --- a/lib/mocker/src/replay/validate.rs +++ b/lib/mocker/src/replay/validate.rs @@ -3,9 +3,45 @@ use anyhow::{Result, bail}; -use super::ReplayRouterMode; +use super::{OfflineDisaggReplayConfig, ReplayArgsMode, ReplayRouterMode}; use crate::common::protocols::{MockEngineArgs, WorkerType}; +pub fn validate_replay_args_mode( + aggregated_args: Option<&MockEngineArgs>, + prefill_args: Option<&MockEngineArgs>, + decode_args: Option<&MockEngineArgs>, + num_workers: usize, + num_prefill_workers: usize, + num_decode_workers: usize, +) -> Result { + if aggregated_args.is_some() && (prefill_args.is_some() || decode_args.is_some()) { + bail!("extra_engine_args cannot be combined with prefill_engine_args/decode_engine_args"); + } + + match (aggregated_args, prefill_args, decode_args) { + (Some(_), None, None) | (None, None, None) => { + if num_prefill_workers != 1 || num_decode_workers != 1 { + bail!( + "num_prefill_workers and num_decode_workers are only used for disagg replay; use num_workers for aggregated replay" + ); + } + Ok(ReplayArgsMode::Aggregated) + } + (None, Some(_), Some(_)) => { + if num_workers != 1 { + bail!( + "num_workers is only used for aggregated replay; use num_prefill_workers and num_decode_workers for disagg replay" + ); + } + Ok(ReplayArgsMode::Disagg) + } + (None, Some(_), None) | (None, None, Some(_)) => { + bail!("prefill_engine_args and decode_engine_args must be provided together") + } + (Some(_), Some(_), _) | (Some(_), _, Some(_)) => unreachable!(), + } +} + fn validate_replay_args(args: &MockEngineArgs, num_workers: usize, mode: &str) -> Result<()> { if num_workers == 0 { bail!("{mode} requires num_workers >= 1"); @@ -75,3 +111,63 @@ pub(super) fn validate_online_concurrency_args( validate_replay_args(args, num_workers, "online replay") } + +fn validate_disagg_args(config: &OfflineDisaggReplayConfig, mode: &str) -> Result<()> { + if config.num_prefill_workers == 0 { + bail!("{mode} requires num_prefill_workers >= 1"); + } + if config.num_decode_workers == 0 { + bail!("{mode} requires num_decode_workers >= 1"); + } + if config.prefill_args.worker_type != WorkerType::Prefill { + bail!( + "{mode} requires prefill_engine_args.worker_type=prefill, got {:?}", + config.prefill_args.worker_type, + ); + } + if config.decode_args.worker_type != WorkerType::Decode { + bail!( + "{mode} requires decode_engine_args.worker_type=decode, got {:?}", + config.decode_args.worker_type, + ); + } + if config.prefill_args.dp_size != 1 { + bail!( + "{mode} only supports prefill data_parallel_size=1, got {}", + config.prefill_args.dp_size, + ); + } + if config.decode_args.dp_size != 1 { + bail!( + "{mode} only supports decode data_parallel_size=1, got {}", + config.decode_args.dp_size, + ); + } + if config.prefill_args.block_size != config.decode_args.block_size { + bail!( + "{mode} requires matching prefill/decode block_size, got {} and {}", + config.prefill_args.block_size, + config.decode_args.block_size, + ); + } + + Ok(()) +} + +pub(super) fn validate_offline_disagg_replay_args( + config: &OfflineDisaggReplayConfig, + _router_mode: ReplayRouterMode, +) -> Result<()> { + validate_disagg_args(config, "trace replay") +} + +pub(super) fn validate_offline_disagg_concurrency_args( + config: &OfflineDisaggReplayConfig, + max_in_flight: usize, + _router_mode: ReplayRouterMode, +) -> Result<()> { + if max_in_flight == 0 { + bail!("concurrency replay requires max_in_flight >= 1"); + } + validate_disagg_args(config, "concurrency replay") +} diff --git a/lib/mocker/src/scheduler/mod.rs b/lib/mocker/src/scheduler/mod.rs index 79f7c3ef2a31..b86a6e5c09e7 100644 --- a/lib/mocker/src/scheduler/mod.rs +++ b/lib/mocker/src/scheduler/mod.rs @@ -89,6 +89,13 @@ impl EngineCore { Self::Sglang(core) => core.execute_pass(collector, now_ms), } } + + pub(crate) fn execute_hidden_pass(&mut self, now_ms: f64) -> EnginePassResult { + match self { + Self::Vllm(core) => core.execute_hidden_pass(now_ms), + Self::Sglang(core) => core.execute_hidden_pass(now_ms), + } + } } #[derive(Clone)] diff --git a/lib/mocker/src/scheduler/sglang/config.rs b/lib/mocker/src/scheduler/sglang/config.rs index a95e4dc1332e..547a3c764def 100644 --- a/lib/mocker/src/scheduler/sglang/config.rs +++ b/lib/mocker/src/scheduler/sglang/config.rs @@ -36,6 +36,8 @@ pub(super) struct SglangConfig { pub(super) decode_speedup_ratio: f64, pub(super) worker_type: WorkerType, pub(super) block_size: usize, + pub(super) kv_bytes_per_token: Option, + pub(super) kv_transfer_bandwidth: Option, } impl SglangConfig { @@ -81,6 +83,8 @@ impl SglangConfig { decode_speedup_ratio: args.decode_speedup_ratio, worker_type: args.worker_type, block_size: args.block_size, + kv_bytes_per_token: args.kv_bytes_per_token, + kv_transfer_bandwidth: args.kv_transfer_bandwidth, } } } diff --git a/lib/mocker/src/scheduler/sglang/core.rs b/lib/mocker/src/scheduler/sglang/core.rs index 192d7e3dbc4d..903cb428005b 100644 --- a/lib/mocker/src/scheduler/sglang/core.rs +++ b/lib/mocker/src/scheduler/sglang/core.rs @@ -101,6 +101,10 @@ impl SglangCore { self.execute_pass_internal(Some(collector), now_ms) } + pub(crate) fn execute_hidden_pass(&mut self, now_ms: f64) -> EnginePassResult { + self.execute_pass_internal(None, now_ms) + } + pub(super) fn execute_pass_internal( &mut self, mut collector: Option<&mut TraceCollector>, diff --git a/lib/mocker/src/scheduler/sglang/decode.rs b/lib/mocker/src/scheduler/sglang/decode.rs index ec44c0acf599..39e62ea8fefa 100644 --- a/lib/mocker/src/scheduler/sglang/decode.rs +++ b/lib/mocker/src/scheduler/sglang/decode.rs @@ -4,6 +4,7 @@ use std::time::Duration; use crate::common::protocols::OutputSignal; +use crate::common::utils::compute_prefill_handoff_delay_ms; use crate::kv_manager::SglangKvManager; use super::config::{SglangConfig, floor_to_block}; @@ -179,6 +180,13 @@ pub(super) fn simulate_decode_step( output_signals.push(OutputSignal { uuid: req.uuid, completed: is_complete, + handoff_delay_ms: compute_prefill_handoff_delay_ms( + config.worker_type, + is_complete, + req.prompt_len(), + config.kv_transfer_bandwidth, + config.kv_bytes_per_token, + ), }); if is_complete { diff --git a/lib/mocker/src/scheduler/sglang/tests.rs b/lib/mocker/src/scheduler/sglang/tests.rs index f3368b5f51d7..fe4b4cb1af51 100644 --- a/lib/mocker/src/scheduler/sglang/tests.rs +++ b/lib/mocker/src/scheduler/sglang/tests.rs @@ -860,7 +860,6 @@ mod router_events { } } } - assert_eq!(seen, expected); drop(scheduler); drop(sink); @@ -871,4 +870,41 @@ mod router_events { assert!(harness.ok_count(METRIC_EVENT_REMOVED) > 0); harness.shutdown(); } + + #[test] + fn test_prefill_completion_emits_handoff_delay() { + let args = MockEngineArgs::builder() + .engine_type(EngineType::Sglang) + .num_gpu_blocks(64) + .block_size(4) + .worker_type(crate::common::protocols::WorkerType::Prefill) + .kv_transfer_bandwidth(Some(1.0)) + .kv_bytes_per_token(Some(1_000_000)) + .speedup_ratio(0.0) + .sglang(Some(SglangArgs { + page_size: Some(4), + chunked_prefill_size: Some(16), + ..Default::default() + })) + .build() + .unwrap(); + let mut core = SglangCore::new(args); + core.receive(DirectRequest { + tokens: vec![1; 8], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(91)), + dp_rank: 0, + arrival_timestamp_ms: None, + }); + + let mut collector = crate::replay::TraceCollector::default(); + let pass = core.execute_pass(&mut collector, 0.0); + let signal = pass + .output_signals + .first() + .expect("prefill pass should emit one completed signal"); + + assert!(signal.completed); + assert_eq!(signal.handoff_delay_ms, Some(8.0)); + } } diff --git a/lib/mocker/src/scheduler/vllm/core.rs b/lib/mocker/src/scheduler/vllm/core.rs index 9057bf4861c3..a00d83a65a2d 100644 --- a/lib/mocker/src/scheduler/vllm/core.rs +++ b/lib/mocker/src/scheduler/vllm/core.rs @@ -14,6 +14,7 @@ use crate::common::protocols::{ WorkerType, }; use crate::common::sequence::ActiveSequence; +use crate::common::utils::compute_prefill_handoff_delay_ms; use crate::kv_manager::KvManager; use crate::replay::TraceCollector; use crate::scheduler::{ @@ -278,6 +279,10 @@ impl VllmCore { self.execute_pass_internal(Some(collector), now_ms, None) } + pub(crate) fn execute_hidden_pass(&mut self, now_ms: f64) -> EnginePassResult { + self.execute_pass_internal(None, now_ms, None) + } + pub(super) fn execute_pass_internal( &mut self, mut collector: Option<&mut TraceCollector>, @@ -641,8 +646,25 @@ impl VllmCore { } if let Some(request) = self.state.requests.get(&uuid) { debug_assert_vllm_request_progress(uuid, request); + let handoff_delay_ms = compute_prefill_handoff_delay_ms( + self.args.worker_type, + completed, + request.sequence.num_input_tokens(), + self.args.kv_transfer_bandwidth, + self.args.kv_bytes_per_token, + ); + output_signals.push(OutputSignal { + uuid, + completed, + handoff_delay_ms, + }); + } else { + output_signals.push(OutputSignal { + uuid, + completed, + handoff_delay_ms: None, + }); } - output_signals.push(OutputSignal { uuid, completed }); if completed { self.state.complete(&uuid); } diff --git a/lib/mocker/src/scheduler/vllm/tests.rs b/lib/mocker/src/scheduler/vllm/tests.rs index 87641c53b707..fdfd132aac9d 100644 --- a/lib/mocker/src/scheduler/vllm/tests.rs +++ b/lib/mocker/src/scheduler/vllm/tests.rs @@ -175,6 +175,40 @@ mod core_behavior { ); } + #[test] + fn test_prefill_completion_emits_handoff_delay() { + let args = MockEngineArgs::builder() + .block_size(4) + .num_gpu_blocks(8) + .max_num_batched_tokens(Some(8)) + .max_num_seqs(Some(1)) + .enable_chunked_prefill(true) + .worker_type(crate::common::protocols::WorkerType::Prefill) + .kv_transfer_bandwidth(Some(1.0)) + .kv_bytes_per_token(Some(1_000_000)) + .speedup_ratio(0.0) + .build() + .unwrap(); + let mut core = VllmCore::new(args); + core.receive(DirectRequest { + tokens: vec![1; 8], + max_output_tokens: 1, + uuid: Some(Uuid::from_u128(81)), + dp_rank: 0, + arrival_timestamp_ms: None, + }); + + let mut collector = crate::replay::TraceCollector::default(); + let pass = core.execute_pass(&mut collector, 0.0); + let signal = pass + .output_signals + .first() + .expect("prefill pass should emit one completed signal"); + + assert!(signal.completed); + assert_eq!(signal.handoff_delay_ms, Some(8.0)); + } + #[test] fn test_first_token_can_arrive_on_prompt_completion_pass() { let mut core = VllmCore::new(make_args()); diff --git a/tests/mocker/test_config.py b/tests/mocker/test_config.py index d8cf92948e5e..06b3fe3be092 100644 --- a/tests/mocker/test_config.py +++ b/tests/mocker/test_config.py @@ -1,10 +1,13 @@ # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import argparse import importlib.util +import json from pathlib import Path from types import SimpleNamespace +import numpy as np import pytest from dynamo.llm import EngineType, EntrypointArgs @@ -133,3 +136,99 @@ def test_entrypoint_args_accept_typed_mocker_engine_args(): ) assert entrypoint_args is not None + + +def test_build_mocker_engine_args_preserves_cli_mapped_fields(tmp_path): + planner_profile_data = tmp_path / "planner_profile_data.npz" + np.savez( + planner_profile_data, + prefill_isl=np.array([128.0, 256.0]), + prefill_ttft_ms=np.array([4.0, 8.0]), + decode_active_kv_tokens=np.array([1024.0, 2048.0]), + decode_context_length=np.array([128.0, 256.0]), + decode_itl=np.array([[1.0, 1.5], [2.0, 2.5]]), + ) + + args = argparse.Namespace( + engine_type="sglang", + num_gpu_blocks=2048, + block_size=128, + max_num_seqs=64, + max_num_batched_tokens=4096, + enable_prefix_caching=False, + enable_chunked_prefill=False, + preemption_mode="fifo", + speedup_ratio=2.0, + decode_speedup_ratio=3.0, + dp_size=4, + startup_time=1.5, + planner_profile_data=planner_profile_data, + is_prefill_worker=True, + is_decode_worker=False, + durable_kv_events=False, + kv_transfer_bandwidth=123.0, + reasoning=json.dumps( + { + "start_thinking_token_id": 11, + "end_thinking_token_id": 12, + "thinking_ratio": 0.25, + } + ), + sglang_schedule_policy="lpm", + sglang_page_size=128, + sglang_max_prefill_tokens=8192, + sglang_chunked_prefill_size=2048, + sglang_clip_max_new_tokens=1024, + sglang_schedule_conservativeness=0.8, + aic_perf_model=True, + aic_system="h200_sxm", + aic_backend_version="0.5.6.post2", + aic_tp_size=8, + model_path="/models/mock", + ) + + engine_args = CONFIG.build_mocker_engine_args(args) + payload = json.loads(engine_args.dump_json()) + + assert payload == { + "engine_type": "sglang", + "num_gpu_blocks": 2048, + "block_size": 128, + "max_num_seqs": 64, + "max_num_batched_tokens": 4096, + "enable_prefix_caching": False, + "enable_chunked_prefill": False, + "speedup_ratio": 2.0, + "decode_speedup_ratio": 3.0, + "dp_size": 4, + "startup_time": 1.5, + "worker_type": "prefill", + "planner_profile_data": str(planner_profile_data), + "aic_backend": "sglang", + "aic_system": "h200_sxm", + "aic_backend_version": "0.5.6.post2", + "aic_tp_size": 8, + "aic_model_path": "/models/mock", + "enable_local_indexer": True, + "bootstrap_port": None, + "kv_bytes_per_token": None, + "kv_transfer_bandwidth": 123.0, + "reasoning": { + "start_thinking_token_id": 11, + "end_thinking_token_id": 12, + "thinking_ratio": 0.25, + }, + "zmq_kv_events_port": None, + "zmq_replay_port": None, + "preemption_mode": "fifo", + "router_queue_policy": None, + "sglang": { + "schedule_policy": "lpm", + "page_size": 128, + "max_prefill_tokens": 8192, + "chunked_prefill_size": 2048, + "clip_max_new_tokens": 1024, + "schedule_conservativeness": 0.8, + }, + "has_perf_model": True, + } diff --git a/tests/router/test_router_e2e_with_mockers.py b/tests/router/test_router_e2e_with_mockers.py index e5c3a5da7717..a1619aca63b5 100644 --- a/tests/router/test_router_e2e_with_mockers.py +++ b/tests/router/test_router_e2e_with_mockers.py @@ -8,14 +8,10 @@ # endpoint tables) that races under concurrent xdist workers. Do not add # @pytest.mark.parallel until DRT endpoint registration is confirmed thread-safe. # -# NOTE: TCP request plane is NOT tested here. These tests use --num-workers > 1 which spawns -# multiple workers in a single process sharing one TCP server. The shared TCP server uses -# endpoint_path (e.g., "generate") as the routing key, causing handler collisions when multiple -# workers register the same endpoint. This is a test-only limitation; production deployments -# with separate processes per worker work correctly with TCP. import asyncio import logging import os +from pathlib import Path from typing import Any, Dict, Optional import aiohttp @@ -64,6 +60,9 @@ BASE_PORT_ZMQ = 11100 # Base port for ZMQ KV event publishing NUM_REQUESTS = 100 BLOCK_SIZE = 16 +PLANNER_PROFILE_DATA_DIR = ( + Path(__file__).resolve().parents[1] / "planner/profiling_results/H200_TP1P_TP1D" +) def get_unique_ports( @@ -172,6 +171,20 @@ def _build_mocker_command( command.extend(["--preemption-mode", str(mocker_args["preemption_mode"])]) if "dp_size" in mocker_args: command.extend(["--data-parallel-size", str(mocker_args["dp_size"])]) + if "planner_profile_data" in mocker_args: + command.extend( + ["--planner-profile-data", str(mocker_args["planner_profile_data"])] + ) + if mocker_args.get("aic_perf_model") is True: + command.append("--aic-perf-model") + if "aic_system" in mocker_args: + command.extend(["--aic-system", str(mocker_args["aic_system"])]) + if "aic_backend_version" in mocker_args: + command.extend( + ["--aic-backend-version", str(mocker_args["aic_backend_version"])] + ) + if "aic_tp_size" in mocker_args: + command.extend(["--aic-tp-size", str(mocker_args["aic_tp_size"])]) # Use --durable-kv-events to enable JetStream mode (local indexer disabled) if mocker_args.get("durable_kv_events") is True: command.append("--durable-kv-events") @@ -641,17 +654,29 @@ def __exit__(self, exc_type, exc_val, exc_tb): @pytest.mark.timeout(120) # bumped for xdist contention (was 42s; ~13.80s serial avg) @pytest.mark.parametrize( - "router_mode,durable_kv_events", + "router_mode,durable_kv_events,mocker_args_override", [ - pytest.param("kv", False, id="kv-nondurable"), - pytest.param("kv", True, id="kv-durable"), - pytest.param("round-robin", False, id="roundrobin"), - pytest.param("random", False, id="random"), - pytest.param("power-of-two", False, id="power-of-two"), + pytest.param("kv", False, {}, id="kv-nondurable"), + pytest.param( + "kv", + False, + {"planner_profile_data": PLANNER_PROFILE_DATA_DIR}, + id="kv-planner", + ), + pytest.param( + "kv", + False, + {"aic_perf_model": True, "aic_system": "h200_sxm"}, + id="kv-aic", + ), + pytest.param("kv", True, {}, id="kv-durable"), + pytest.param("round-robin", False, {}, id="roundrobin"), + pytest.param("random", False, {}, id="random"), + pytest.param("power-of-two", False, {}, id="power-of-two"), ], indirect=["durable_kv_events"], ) -@pytest.mark.parametrize("request_plane", ["nats", "tcp"], indirect=True) +@pytest.mark.parametrize("request_plane", ["tcp"], indirect=True) def test_mocker_router( request, runtime_services_dynamic_ports, @@ -659,6 +684,7 @@ def test_mocker_router( router_mode, request_plane, durable_kv_events, + mocker_args_override, ): """Test router with multiple mocker engine instances across all router modes. @@ -675,6 +701,7 @@ def test_mocker_router( "block_size": BLOCK_SIZE, "durable_kv_events": durable_kv_events, } + mocker_args.update(mocker_args_override) with MockerProcess( request,