feat(agentic): D→P snapshot orchestration in reseed path + CLI flag

Phase 3 — wires the SGLang-side snapshot RPCs (committed in 86412bb)
into the agentic reseed slow-path. On _invoke_kvcache_seeded_router:

  1. POST {prefill_url}/_snapshot/prepare_receive   alloc P-side slots
  2. POST {old_decode_url}/_snapshot/dump           RDMA push session KV
  3. POST {prefill_url}/_snapshot/finalize_ingest   insert into P radix

After step 3 P's radix tree has the session prefix cached; the subsequent
SGLang router-driven prefill on P hits cache instead of re-computing.

Any RPC failure short-circuits to the existing seeded_router fallback
(re-prefill from scratch). All steps are best-effort and structurally
logged for post-hoc analysis.

Flag plumbing:
  cli.py             --enable-d-to-p-sync          (replay + benchmark)
  topology.py        SingleNodeTopology.enable_d_to_p_sync
  stack.py           SGLANG_SNAPSHOT_LINK_ENABLE=1 injection per worker
  replay.py          ReplayConfig.enable_d_to_p_sync +
                     _attempt_d_to_p_sync helper

Snapshot port per worker derives from disaggregation_bootstrap_port +
1000 (set in third_party/.../snapshot/controller.py), so different
workers get distinct mooncake snapshot engines on the same node.

Smoke (next): scripts/smoke_snapshot_sglang_integration.py spawns one
D + one P, exercises the 3 RPCs end-to-end, checks cache_tokens on a
follow-up generate request.

See docs/D_TO_P_SYNC_DESIGN_ZH.md for the full design.

src/agentic_pd_hybrid/cli.py

@@ -283,6 +283,17 @@ def main() -> None:
             "See docs/E1_E2_FIX_DESIGN_ZH.md §Q2."
         ),
     )
+    replay.add_argument(
+        "--enable-d-to-p-sync",
+        action="store_true",
+        help=(
+            "Enable D→P RDMA KV snapshot push for reseed fast-path. "
+            "When set, on _invoke_kvcache_seeded_router agentic will probe D's "
+            "session_aware_cache, RDMA-dump session KV to P's snapshot link, "
+            "and insert into P's radix tree so the upcoming P prefill hits "
+            "cache. See docs/D_TO_P_SYNC_DESIGN_ZH.md."
+        ),
+    )
 
     sample = subparsers.add_parser(
         "sample-sessions",
@@ -547,6 +558,14 @@ def main() -> None:
             "See docs/E1_E2_FIX_DESIGN_ZH.md §Q2."
         ),
     )
+    benchmark.add_argument(
+        "--enable-d-to-p-sync",
+        action="store_true",
+        help=(
+            "Enable D→P RDMA KV snapshot push for reseed fast-path. "
+            "See docs/D_TO_P_SYNC_DESIGN_ZH.md."
+        ),
+    )
     benchmark.add_argument(
         "--sample-profile",
         choices=["default", "small-append"],
@@ -634,6 +653,7 @@ def main() -> None:
             backpressure_max_pause_s=args.backpressure_max_pause_s,
             kvcache_migration_reject_threshold=args.kvcache_migration_reject_threshold,
             kvcache_load_floor_bonus=args.kvcache_load_floor_bonus,
+            enable_d_to_p_sync=args.enable_d_to_p_sync,
         )
         results = asyncio.run(replay_trace(config))
         print(
@@ -876,6 +896,7 @@ def _topology_from_args(args: argparse.Namespace):
         force_rdma=args.force_rdma,
         trust_remote_code=not args.no_trust_remote_code,
         ib_device=args.ib_device,
+        enable_d_to_p_sync=getattr(args, "enable_d_to_p_sync", False),
         prefill_extra_server_args=("--disable-overlap-schedule",),
         decode_extra_server_args=("--disable-overlap-schedule",),
         direct_extra_server_args=("--enable-streaming-session",),

src/agentic_pd_hybrid/replay.py

@@ -116,6 +116,11 @@ class ReplayConfig:
     # with shared cross-session prefix. 0 disables. See
     # docs/E1_E2_FIX_DESIGN_ZH.md §Q2.
     kvcache_load_floor_bonus: int = 0
+    # D→P snapshot push: when True and reseed fires, agentic will RDMA-dump
+    # the session's KV from the D-side worker that last held it onto the P
+    # worker and insert into P's radix tree, so the subsequent P prefill
+    # hits cache. See docs/D_TO_P_SYNC_DESIGN_ZH.md.
+    enable_d_to_p_sync: bool = False
     structural_log_dir: Path | None = None
 
 
@@ -2104,6 +2109,119 @@ async def _invoke_plain_router(
     )
 
 
+async def _attempt_d_to_p_sync(
+    *,
+    client: httpx.AsyncClient,
+    request: TraceRequest,
+    config: ReplayConfig,
+    prefill_url: str,
+    decode_session: DirectSessionState,
+) -> dict | None:
+    """Try to RDMA-dump session KV from the D that last held it to ``prefill_url``.
+
+    Returns a dict with status info on success/skip, or ``None`` on a
+    non-recoverable error. The caller falls back to normal re-prefill on
+    any failure.
+    """
+    if not config.enable_d_to_p_sync:
+        return None
+    source_d_url = decode_session.server_url
+    if not source_d_url:
+        return {"status": "skipped-no-source-d"}
+    if not decode_session.opened:
+        return {"status": "skipped-d-closed"}
+    # Compose token list for radix insert: we don't have the actual token_ids
+    # on the agentic side in a stable form; use the request's prompt_token_ids
+    # via the residency bookkeeping. For now we use a length proxy.
+    target_tokens = max(0, int(_estimate_session_resident_tokens(request)))
+    if target_tokens <= 0:
+        return {"status": "skipped-zero-tokens"}
+
+    try:
+        prep_resp = await client.post(
+            f"{prefill_url}/_snapshot/prepare_receive",
+            json={
+                "session_id": request.session_id,
+                "num_tokens": target_tokens,
+            },
+            timeout=30.0,
+        )
+        prep_resp.raise_for_status()
+        prep = prep_resp.json()
+    except Exception as exc:
+        return {"status": "prepare-failed", "error": repr(exc)}
+    if not prep.get("ok"):
+        return {"status": "prepare-not-ok", "reason": prep.get("reason")}
+
+    try:
+        dump_resp = await client.post(
+            f"{source_d_url}/_snapshot/dump",
+            json={
+                "session_id": request.session_id,
+                "target_snapshot_session_id": prep["snapshot_session_id"],
+                "target_k_base_ptrs": prep["k_base_ptrs"],
+                "target_v_base_ptrs": prep["v_base_ptrs"],
+                "target_slot_indices": prep["slot_indices"],
+                "target_stride_k_bytes": prep["stride_k_bytes"],
+                "target_stride_v_bytes": prep["stride_v_bytes"],
+            },
+            timeout=60.0,
+        )
+        dump_resp.raise_for_status()
+        dump = dump_resp.json()
+    except Exception as exc:
+        return {"status": "dump-failed", "error": repr(exc)}
+    if not dump.get("ok"):
+        return {"status": "dump-not-ok", "reason": dump.get("reason"),
+                "bytes_pushed": dump.get("bytes_pushed", 0)}
+
+    # We need token_ids for radix insert. The caller has request.input_token_ids
+    # for the first N — use that as best-available approximation.
+    tokens = list(getattr(request, "input_token_ids", []) or [])
+    if not tokens:
+        # No token_ids available — can't insert into radix. P will fall back
+        # to normal prefill but will have wasted slots. Discard.
+        try:
+            await client.post(
+                f"{prefill_url}/_snapshot/finalize_ingest",
+                json={
+                    "session_id": request.session_id,
+                    "token_ids": [],
+                    "slot_indices": prep["slot_indices"],
+                },
+                timeout=15.0,
+            )
+        except Exception:
+            pass
+        return {"status": "no-tokens-discard", "bytes_pushed": dump.get("bytes_pushed", 0)}
+
+    n = min(len(tokens), len(prep["slot_indices"]))
+    try:
+        fin_resp = await client.post(
+            f"{prefill_url}/_snapshot/finalize_ingest",
+            json={
+                "session_id": request.session_id,
+                "token_ids": tokens[:n],
+                "slot_indices": prep["slot_indices"][:n],
+            },
+            timeout=30.0,
+        )
+        fin_resp.raise_for_status()
+        fin = fin_resp.json()
+    except Exception as exc:
+        return {"status": "finalize-failed", "error": repr(exc),
+                "bytes_pushed": dump.get("bytes_pushed", 0)}
+    if not fin.get("ok"):
+        return {"status": "finalize-not-ok", "reason": fin.get("reason"),
+                "bytes_pushed": dump.get("bytes_pushed", 0)}
+    return {
+        "status": "ok",
+        "bytes_pushed": int(dump.get("bytes_pushed", 0)),
+        "inserted_prefix_len": int(fin.get("inserted_prefix_len", 0)),
+        "snapshot_session_id": prep.get("snapshot_session_id"),
+    }
+
+
 async def _invoke_kvcache_seeded_router(
     *,
     client: httpx.AsyncClient,
@@ -2155,6 +2273,31 @@ async def _invoke_kvcache_seeded_router(
             decode_session.prefill_server_url = prefill_url
             prefill_session_newly_opened = True
 
+    # D→P snapshot push (Phase 3) — best-effort; on any failure we silently
+    # fall back to the existing re-prefill path. The result is logged for
+    # post-hoc analysis but does not affect correctness.
+    if config.enable_d_to_p_sync:
+        sync_result = await _attempt_d_to_p_sync(
+            client=client,
+            request=request,
+            config=config,
+            prefill_url=prefill_url,
+            decode_session=decode_session,
+        )
+        if sync_result is not None and sync_result.get("status") != "ok":
+            logger.info(
+                "d_to_p_sync sid=%s rid=%s skipped: %s",
+                request.session_id, request.request_id, sync_result,
+            )
+        elif sync_result and sync_result.get("status") == "ok":
+            logger.info(
+                "d_to_p_sync sid=%s rid=%s pushed=%d ingested_prefix=%d",
+                request.session_id,
+                request.request_id,
+                sync_result.get("bytes_pushed", 0),
+                sync_result.get("inserted_prefix_len", 0),
+            )
+
     decode_session_newly_opened = False
     try:
         prefill_priority = _prefill_priority_for_router_request(

src/agentic_pd_hybrid/stack.py

@@ -209,6 +209,15 @@ def _build_process_env(topology: SingleNodeTopology) -> dict[str, str]:
     if topology.transfer_backend == "mooncake":
         env.setdefault("MC_TRANSFER_TIMEOUT", "1800")
 
+    # D→P snapshot link (Phase 2). Each worker reads its own
+    # `disaggregation_bootstrap_port` and binds at `bootstrap_port + 1000`
+    # for the snapshot mooncake engine (see
+    # third_party/sglang/.../disaggregation/snapshot/controller.py).
+    if topology.enable_d_to_p_sync:
+        env["SGLANG_SNAPSHOT_LINK_ENABLE"] = "1"
+        if topology.ib_device:
+            env.setdefault("SGLANG_SNAPSHOT_LINK_IB_DEVICE", topology.ib_device)
+
     repo_root = Path(__file__).resolve().parents[2]
     python_paths = [
         str(repo_root / "src"),

src/agentic_pd_hybrid/topology.py

@@ -46,6 +46,7 @@ class SingleNodeTopology:
     trust_remote_code: bool
     force_rdma: bool = False
     ib_device: str | None = None
+    enable_d_to_p_sync: bool = False
     extra_server_args: tuple[str, ...] = ()
     prefill_extra_server_args: tuple[str, ...] = ()
     decode_extra_server_args: tuple[str, ...] = ()
@@ -95,6 +96,7 @@ def build_single_node_topology(
     force_rdma: bool = False,
     trust_remote_code: bool = True,
     ib_device: str | None = None,
+    enable_d_to_p_sync: bool = False,
     extra_server_args: tuple[str, ...] = (),
     prefill_extra_server_args: tuple[str, ...] = (),
     decode_extra_server_args: tuple[str, ...] = (),
@@ -238,6 +240,7 @@ def build_single_node_topology(
         trust_remote_code=trust_remote_code,
         force_rdma=force_rdma,
         ib_device=ib_device,
+        enable_d_to_p_sync=enable_d_to_p_sync,
         extra_server_args=extra_server_args,
         prefill_extra_server_args=prefill_extra_server_args,
         decode_extra_server_args=decode_extra_server_args,