feat(agentx): allow multi-replica and shared G2 for AgentX replay [AISIM-1893, AISIM-1894]

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"),
                     (

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.
 

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.
 

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

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:
 

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 |

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

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.

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; "

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:

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:

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)