diff --git a/crates/core/src/python.rs b/crates/core/src/python.rs index a085833593..2273bc1bec 100644 --- a/crates/core/src/python.rs +++ b/crates/core/src/python.rs @@ -250,16 +250,8 @@ fn validate_public_agentic_engine( .iter() .any(|role| role.rank.native_host_offload.is_some()) { - let single_worker_per_role = match &spec.topology { - ReplayTopology::Aggregated { workers } => workers.initial_workers == 1, - ReplayTopology::Disaggregated { - prefill, decode, .. - } => prefill.initial_workers == 1 && decode.initial_workers == 1, - }; - ensure!( - single_worker_per_role, - "agentic host offload requires one aggregated worker or one prefill and one decode worker" - ); + // Worker count is unrestricted: placement consumes each worker's native + // G2 KV events, and per-play token namespaces keep G2 identities play-local. ensure!( spec.adapters.scaling.provider == "none", "agentic host offload requires static worker pools without a scaling policy" @@ -3948,17 +3940,16 @@ mod tests { if scope.is_none() && other.is_none() { continue; } + let mut multi_replica = pd.clone(); + for pointer in [ + "/spec/topology/prefill/initial_workers", + "/spec/topology/decode/initial_workers", + ] { + *multi_replica.pointer_mut(pointer).unwrap() = serde_json::json!(2); + } + let report = run(&multi_replica); + assert_eq!(report["completed_requests"], 1, "{scope:?}/{other:?}"); for (pointer, value, message) in [ - ( - "/spec/topology/prefill/initial_workers", - serde_json::json!(2), - "one aggregated worker or one prefill and one decode worker", - ), - ( - "/spec/topology/decode/initial_workers", - serde_json::json!(2), - "one aggregated worker or one prefill and one decode worker", - ), ("/spec/engine/prefill/dp_size", serde_json::json!(2), "DP=1"), ("/spec/engine/decode/dp_size", serde_json::json!(2), "DP=1"), ( diff --git a/docs/adapters/runner-abi.md b/docs/adapters/runner-abi.md index 46da94aa37..a9dbeb239e 100644 --- a/docs/adapters/runner-abi.md +++ b/docs/adapters/runner-abi.md @@ -80,8 +80,8 @@ schemes' draft costs. A stack must explicitly declare generic speculation support and retain those costs through its timing provider. `require_compatible(spec)` checks versions and coarse combinations before -execution. AgentX G2 additionally requires offline vLLM, a static single -aggregated worker or 1P1D, and DP1 on every role. G3 and speculation cannot be +execution. AgentX G2 additionally requires offline vLLM, static worker +pools, and DP1 on every role. G3 and speculation cannot be combined with AgentX G2. The [feature matrix](../replay/features.md) records the built-in stack's actual boundary, including limitations stricter than a capability flag. diff --git a/docs/replay/agentic/continuous-profiles.md b/docs/replay/agentic/continuous-profiles.md index 11758fe01f..7095deb333 100644 --- a/docs/replay/agentic/continuous-profiles.md +++ b/docs/replay/agentic/continuous-profiles.md @@ -38,8 +38,8 @@ The built-in engine supports offline vLLM and SGLang with aggregated or separate prefill/decode workers and HBM-only KV cache, including speculative decoding with explicit acceptance assumptions and a compatible timing model. vLLM also supports local or shared [G2 host offload](../engine/kv-cache.md#agentic-g2) on a static -single aggregated worker or 1P1D, with attention DP=1 on every role and -speculative decoding disabled. +aggregated or P/D deployment with any number of workers per role, attention DP=1 on +every role and speculative decoding disabled. Weka, Agentic Mooncake, and agentic Dynamo traces share this path. Online execution, other backends, and Dynamo-owned routing need their own qualified integration. diff --git a/docs/replay/agentic/warmup.md b/docs/replay/agentic/warmup.md index 9417bbd3ef..913bb98457 100644 --- a/docs/replay/agentic/warmup.md +++ b/docs/replay/agentic/warmup.md @@ -66,7 +66,7 @@ vLLM and SGLang replay with either aggregated workers or separate prefill/decode pools and HBM-only KV cache, including [speculative decoding](../engine/speculation.md) with explicit acceptance assumptions; ngram requires vLLM. vLLM also supports local or shared -[G2 host offload](../engine/kv-cache.md#agentic-g2) on a static single aggregated worker or 1P1D, +[G2 host offload](../engine/kv-cache.md#agentic-g2) on static aggregated or P/D pools of any size, with attention DP1 on every role and speculative decoding disabled. For example: ```yaml diff --git a/docs/replay/engine/kv-cache.md b/docs/replay/engine/kv-cache.md index 56236a8d27..4e229e998a 100644 --- a/docs/replay/engine/kv-cache.md +++ b/docs/replay/engine/kv-cache.md @@ -471,10 +471,37 @@ Rejected: `recommend`, G2/G3 offload, and speculative decoding with sub-block ## Agentic traffic with G2 AgentX snapshot, warmup and continuous-profile traffic supports G2 on vLLM with -one aggregated worker, or one prefill and one decode worker, all with -`attention_data: 1`, static workers and no speculative decoding. Either G2 scope -works; G3 is not supported. `kvcr` is accepted, but with one rank per role there -is no peer, so it stays idle. Results are qualified `functional_only`. +aggregated or prefill/decode pools of any size, all with `attention_data: 1`, +static workers and no speculative decoding. Either G2 scope works; G3 is not +supported. Each worker publishes its G2 residency as native `HostPinned` KV +events, so a KV-aware placement policy can follow a G2-resident prefix to its +owner, or to any replica of a `cluster_shared` pool. Snapshot and profile plays +use play-local token identities, so G1 and G2 hits never cross a recycled play. +`kvcr` is accepted but is not part of the AgentX qualification. Results are +qualified `functional_only`. + +A `cluster_shared` pool under AgentX follows this contract: + +- **Ownership and discovery.** The replay deployment owns one pool. Every + participating rank, across replicas and P/D roles, joins it; the first + participant fixes layout, capacity, and shared bandwidth, and capacity is + counted once in `g2_domains`. A block becomes visible when its D2H store + completes, and the pool publishes a `HostPinned` store to every subscribed + rank. Pool-wide LRU eviction publishes the matching removal to every rank. +- **Identity.** Conversations of one play share G2 prefixes. Snapshot and + profile plays use play-local token identities and a distinct `cache_id`, so + equal source prompts in recycled plays occupy separate entries and never hit + each other. +- **Contention.** Concurrent restores of the same entries count residency once + and split the pool's directional bandwidth max-min fairly, deterministically + from the replay clock. +- **Stale entries.** If a placement scored residency that is evicted before the + request is admitted, the request recomputes the missing prefix. Its routing + overlap then exceeds the realized `host_reused_input_tokens`. +- **Membership.** AgentX G2 requires static pools, so worker failure and + ownership handoff are outside its support boundary. +- **Disabling sharing.** `scope: dp_rank_local` reproduces the private-G2 + result exactly, even when shared-only bandwidth fields remain set. Two small offline fixtures use fixed timing and need no GPU: diff --git a/docs/replay/features.md b/docs/replay/features.md index 5642869337..55774de354 100644 --- a/docs/replay/features.md +++ b/docs/replay/features.md @@ -22,7 +22,7 @@ is a separate [performance-model question](../perf-model/support-matrix.md). | Prefix reuse | Native backend cache | G1 with prefix caching enabled | Actual reuse differs from Router overlap; complete block and backend admission rules apply | | Grouped FPM cache | vLLM | Cold aggregated, PP1/CP1 | HBM only, no speculation, prefix caching off, no scalar block capacity, no P/D | | G2 host offload | vLLM | Aggregated and token-only P/D, any attention-DP size | Prefix caching, no native MTP/Belady/state-cache offload; local or layout-compatible shared pool | -| Agentic G2 | vLLM | One aggregated worker or 1P1D, attention DP1 on every role | Static pools, no speculation/G3; each role may omit G2 or use local/shared scope | +| Agentic G2 | vLLM | Aggregated or P/D with any worker count, attention DP1 on every role | Static pools, no speculation/G3; each role may omit G2 or use local/shared scope | | G3 offload | vLLM | Aggregated, attention DP1; fixed or scaled pools | Requires G2; no speculation, AgentX, P/D, recommendation or hardware integration | | KVCR peer fetch | vLLM | Aggregated with any attention-DP size, or P/D prefill ranks; fixed or scaled pools | Requires private G2 and hints from a placement policy or ReplaySpec requests (CLI traffic has none); no G3, decode-role KVCR or recommendation; idle under AgentX | | Best-effort Belady | All three native backends | Fixed aggregated, DP1, complete open-loop input trace | Native descriptor; prefix caching required; rejects closed-loop/generated/agentic/delta inputs, scaling and offload | diff --git a/docs/replay/workloads.md b/docs/replay/workloads.md index 47526b2409..3697427db8 100644 --- a/docs/replay/workloads.md +++ b/docs/replay/workloads.md @@ -290,7 +290,7 @@ workers and HBM-only KV cache, including [speculative decoding](engine/speculation.md) with explicit acceptance assumptions; ngram requires vLLM. vLLM additionally supports [G2 host offload](engine/kv-cache.md#agentic-g2) with static workers, -either one aggregated worker or 1P1D, `attention_data: 1` on every role and +aggregated or P/D with any worker count, `attention_data: 1` on every role and speculative decoding disabled. TensorRT-LLM, G3 offload and online P/D are rejected. Results are qualified `functional_only`, not hardware accuracy. [Start an AgentX simulation](agentic/quickstart.md) walks through a diff --git a/docs/sweeper/sdk.md b/docs/sweeper/sdk.md index e2790f88df..661247aace 100644 --- a/docs/sweeper/sdk.md +++ b/docs/sweeper/sdk.md @@ -242,8 +242,8 @@ shared synthetic knobs carried by `ReplaySpec.workload`. Profile controls are concrete workload values, not search dimensions. The built-in Engine runner supports profiles on offline aggregated or P/D vLLM/SGLang replay with HBM-only KV cache. -vLLM also supports local or shared [G2 host offload](../replay/engine/kv-cache.md) on a static single -aggregated worker or 1P1D, with attention DP1 on every role. G3 remains unsupported; +vLLM also supports local or shared [G2 host offload](../replay/engine/kv-cache.md) on static +aggregated or P/D pools of any size, with attention DP1 on every role. G3 remains unsupported; speculative decoding follows the narrower [engine SD constraints](../replay/engine/speculation.md). Warmup is optional: `repeated_prefixes` retains the saved snapshot frontier, while `live_trajectory_v1` advances it. diff --git a/python/aisimulate/src/aisimulate/cli_args.py b/python/aisimulate/src/aisimulate/cli_args.py index ca79547b88..4b3d2ae1ff 100644 --- a/python/aisimulate/src/aisimulate/cli_args.py +++ b/python/aisimulate/src/aisimulate/cli_args.py @@ -76,8 +76,8 @@ def build_parser() -> argparse.ArgumentParser: subparsers.choices["predict"].epilog = ( "AgentX replay: use weka, agentic_mooncake, or agentic Dynamo traces with " "trace_timestamps and positive agentic_lanes. The offline engine stack supports " - "aggregated and P/D vLLM/SGLang with HBM-only KV cache, plus vLLM G2 on a static " - "single aggregated worker or 1P1D, with DP1 on every role and speculative decoding disabled for G2. " + "aggregated and P/D vLLM/SGLang with HBM-only KV cache, plus vLLM G2 on static " + "aggregated or P/D pools of any size, with DP1 on every role and speculative decoding disabled for G2. " "HBM-only AgentX supports speculative decoding with explicit acceptance assumptions. " "agentic_snapshot selects seeded starts; agentic_warmup primes saved prefixes; " "agentic_profile enables duration controls on the offline engine stack; " diff --git a/python/aisimulate/src/aisimulate/sweeper/replay.py b/python/aisimulate/src/aisimulate/sweeper/replay.py index b9a3db7375..f3e3ace417 100644 --- a/python/aisimulate/src/aisimulate/sweeper/replay.py +++ b/python/aisimulate/src/aisimulate/sweeper/replay.py @@ -312,8 +312,8 @@ def supports(self, hook: RuntimeHookSpec) -> bool: def _require_agentic_g2_scope(spec: ReplaySpec, roles: list[tuple[str, Mapping[str, Any]]]) -> None: """Constrain agentic offload without narrowing ordinary HBM-only replay. - Inspect every role, including roles without G2: the supported P/D fleet is - one DP1 worker per role. Canonical timing and flat topology aliases must + Inspect every role, including roles without G2: every role of the supported + fleet is DP1, with any number of workers. Canonical timing and flat topology aliases must receive the same early gate before a runner materializes either form. """ @@ -331,8 +331,8 @@ def _require_agentic_g2_scope(spec: ReplaySpec, roles: list[tuple[str, Mapping[s counts = (deployment.num_prefill_workers, deployment.num_decode_workers) else: raise ValueError("agentic host offload supports only aggregated or disaggregated workers") - if any(type(count) is not int or count != 1 for count in counts): - raise ValueError("agentic host offload requires one aggregated worker or one prefill and one decode worker") + if any(type(count) is not int or count < 1 for count in counts): + raise ValueError("agentic host offload requires at least one worker on every role") if any(hook.kind == "scaling_policy" for hook in spec.runtime_hooks): raise ValueError("agentic host offload requires static worker pools without a scaling policy") for role, args, rank in ranks: diff --git a/tests/test_agentx_g2.py b/tests/test_agentx_g2.py index b149780b4e..79d970b11e 100644 --- a/tests/test_agentx_g2.py +++ b/tests/test_agentx_g2.py @@ -135,6 +135,43 @@ def test_weka_g2_restore_has_causal_bandwidth_and_capacity_controls(name) -> Non assert fast["committed_prefill_tokens"] < hbm["committed_prefill_tokens"] +@pytest.mark.parametrize("name", ["agentx-g2-local.yaml", "agentx-g2-shared-pd.yaml"]) +def test_weka_g2_runs_on_multiple_replicas(name) -> None: + config = _config(name) + single = _run(config) + for worker in config["engine"]["workers"].values(): + worker["parallelism"] = {**worker["parallelism"], "replicas": 2} + report = _run(config) + assert report["completed_requests"] == single["completed_requests"] + assert len({route["logical_worker_id"] for row in report["per_request"] for route in row["routing_history"]}) == 2 + # One cluster-shared pool, counted once whatever the replica count. + assert report.get("g2_domains", []) == single.get("g2_domains", []) + + +def test_weka_shared_g2_contract_across_replicas() -> None: + config = _config("agentx-g2-shared-pd.yaml") + for worker in config["engine"]["workers"].values(): + worker["parallelism"] = {**worker["parallelism"], "replicas": 2} + shared = _run(config) + # Four participants join one deployment-owned pool, counted once. + assert [domain["capacity_blocks"] for domain in shared["g2_domains"]] == [8] + assert shared["completed_requests"] == _run(_config("agentx-g2-shared-pd.yaml"))["completed_requests"] + + disabled, private = copy.deepcopy(config), copy.deepcopy(config) + for worker in disabled["engine"]["workers"].values(): + worker["kv_cache"]["host_offload"] = {**worker["kv_cache"]["host_offload"], "scope": "dp_rank_local"} + for worker in private["engine"]["workers"].values(): + offload = dict(worker["kv_cache"]["host_offload"], scope="dp_rank_local") + for field in ("shared_d2h_bandwidth_gbps", "shared_h2d_bandwidth_gbps"): + offload.pop(field, None) + worker["kv_cache"]["host_offload"] = offload + disabled_report, private_report = _run(disabled), _run(private) + # Disabling sharing reproduces private G2 even with shared-only fields left set. + assert "g2_domains" not in disabled_report or not disabled_report["g2_domains"] + assert disabled_report["per_request"] == private_report["per_request"] + assert disabled_report["committed_prefill_tokens"] == private_report["committed_prefill_tokens"] + + @pytest.mark.parametrize("field,value", [("num_host_blocks", 16), ("shared_h2d_bandwidth_gbps", 2)]) def test_weka_g2_rejects_inconsistent_shared_pool(field, value) -> None: config = _config("agentx-g2-shared-pd.yaml") @@ -156,7 +193,6 @@ def test_weka_g2_rejects_incompatible_shared_layout() -> None: @pytest.mark.parametrize( "fault,message", [ - ("workers", "agentic host offload requires one aggregated worker or one prefill and one decode worker"), ("dp", "agentic host offload requires attention DP=1 on every role"), ("backend", "agentic host offload requires backend=vllm on every role"), ("aic_nextn", "agentic replay requires speculative decoding disabled"), @@ -183,11 +219,7 @@ def test_agentic_invalid_configuration_matrix_rejects_python_and_native(role, fa args = getattr(spec.backend_deployment, args_field) engine = payload["spec"]["engine"] if aggregated else payload["spec"]["engine"][role] rank = engine["rank"] - if fault == "workers": - workers_field = "num_workers" if aggregated else f"num_{role}_workers" - spec = replace(spec, backend_deployment=replace(spec.backend_deployment, **{workers_field: 2})) - payload["spec"]["topology"]["workers" if aggregated else role]["initial_workers"] = 2 - elif fault == "dp": + if fault == "dp": args["aic_attention_dp_size"] = 2 engine["dp_size"] = 2 if not aggregated: diff --git a/tests/test_runner.py b/tests/test_runner.py index 351a82630e..05de473c0e 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -492,17 +492,29 @@ def test_agentic_g2_runner_capability_remains_explicit(nested_rank, scopes, work "scopes,field,value", [ (("dp_rank_local",), "num_workers", 2), + (("dp_rank_local", None), "num_prefill_workers", 2), + ((None, "cluster_shared"), "num_decode_workers", 3), + ], +) +def test_agentic_g2_accepts_multiple_replicas(nested_rank, scopes, field, value): + spec = _agentic_g2_spec(scopes=scopes, nested_rank=nested_rank) + spec = replace(spec, backend_deployment=replace(spec.backend_deployment, **{field: value})) + EngineReplayRunnerFactory().capabilities().require_compatible(spec) + + +@pytest.mark.parametrize("nested_rank", [False, True]) +@pytest.mark.parametrize( + "scopes,field,value", + [ (("cluster_shared",), "num_workers", 0), (("dp_rank_local",), "num_workers", True), - (("dp_rank_local", None), "num_prefill_workers", 2), - ((None, "cluster_shared"), "num_decode_workers", 2), (("dp_rank_local", "cluster_shared"), "num_prefill_workers", 0), ], ) -def test_agentic_g2_rejects_unqualified_worker_counts(nested_rank, scopes, field, value): +def test_agentic_g2_rejects_invalid_worker_counts(nested_rank, scopes, field, value): spec = _agentic_g2_spec(scopes=scopes, nested_rank=nested_rank) spec = replace(spec, backend_deployment=replace(spec.backend_deployment, **{field: value})) - with pytest.raises(ValueError, match="one aggregated worker or one prefill and one decode"): + with pytest.raises(ValueError, match="at least one worker on every role"): EngineReplayRunnerFactory().capabilities().require_compatible(spec)