Skip to content

feat: enable vLLM AgentX local and shared G2 replay - #381

Merged
Harrilee merged 5 commits into
mainfrom
harrli/agentx-vllm-g2-api
Oct 5, 2026
Merged

Harrilee merged 5 commits into
mainfrom
harrli/agentx-vllm-g2-api

Conversation

@Harrilee

@Harrilee Harrilee commented Oct 5, 2026 •

Copy link
Copy Markdown
Contributor

AgentX can use existing vLLM local/shared G2 through CLI, Python ReplaySpec and native JSON on one aggregated worker or 1P1D, with attention DP=1 per role. Python and native enforce the same support matrix, preserve explicit runner capabilities, and reject G3/speculation. Existing HBM-only behavior stays compatible.

This PR focuses on opening and qualifying the AgentX G2 guards, with original restoration/lifecycle fixtures and single-worker/1P1D examples. It uses the existing execution entrypoints. The proposed generic JSON/composition bridge, public core-version export and paired source-install recipe have been removed following review. There are no Cargo, package-version or consumer-build changes.

Dynamo integration belongs to #15626, built on #15647's existing MockEngineArgs and replay wiring. Its installation and end-to-end qualification are separate from this Engine acceptance.

Validation on current 9fb66903, including main through b0ae2eda: 1,965 Rust library tests passed (1 existing ignore), 23 AgentX tests and 9 host-scope tests passed. A freshly built and installed locked wheel passed 528 Python runner/G2/CLI tests. Both documented CLI examples completed and restored eight input tokens from host; the shared example reports one eight-block pool. Rust formatting, generated CODEOWNERS and packaged legal-file checks pass. Hosted CI for this update is pending. Qualification remains functional_only.

@coderabbitai

coderabbitai Bot commented Oct 5, 2026 •

Copy link
Copy Markdown

Review in Change Stack →

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info
⚙️ Run configuration
  • Configuration used: Repository: ai-dynamo/aisimulate/.coderabbit.yaml
  • Review profile: ASSERTIVE
  • Plan: Enterprise
  • Run ID: 2fb2fb34-8687-4677-8371-c18379236c79
📥 Commits

Reviewing files that changed from the base of the PR and between 9fb6690 and 64d51e7.

📒 Files selected for processing (1)
  • docs/g2-cache-scope.md
🔗 Linked repositories identified

CodeRabbit considers these linked repositories for cross-repo context during reviews:

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 10 remain after this review.

📜 Recent review details
⏰ Context from checks skipped due to timeout. (25)
  • GitHub Check: Prediction Regression / Collect prediction snapshot (old)
  • GitHub Check: Prediction Regression / Collect prediction snapshot (new)
  • GitHub Check: Platform Wheels / Build wheels (macosx_arm64)
  • GitHub Check: Collector Data / Check collector data
  • GitHub Check: Collector Data / Perf data sanity (informational)
  • GitHub Check: Platform Wheels / Build wheels (manylinux_2_28_x86_64)
  • GitHub Check: Platform Wheels / Build wheels (manylinux_2_28_aarch64)
  • GitHub Check: Public API Rust (arm64)
  • GitHub Check: readme-commands / README / dynamo
  • GitHub Check: readme-commands / README / source
  • GitHub Check: Release Artifact Contract (amd64)
  • GitHub Check: Python 3.13 compatibility
  • GitHub Check: Release Artifact Contract (arm64)
  • GitHub Check: Python 3.11 compatibility
  • GitHub Check: Application Test Wheel (arm64)
  • GitHub Check: Application Test Wheel (amd64)
  • GitHub Check: Public API Rust (amd64)
  • GitHub Check: Engine Golden Regression
  • GitHub Check: Rust (amd64)
  • GitHub Check: Rust feature modes
  • GitHub Check: Rust (arm64)
  • GitHub Check: Python Dependency Licenses
  • GitHub Check: Build simulation wheel (base)
  • GitHub Check: Build simulation wheel (head)
  • GitHub Check: Forward Prediction Performance (advisory)
🧰 Additional context used
📚 Code guidelines (2)
AGENTS.md — configured
REVIEW.md — configured
📓 Path-based instructions (4)
Check commands, defaults, supported runtimes, public names, and claims against executable behavior.

⚙️ CodeRabbit configuration file

Files:

  • docs/g2-cache-scope.md
Read REVIEW.md before commenting.

⚙️ CodeRabbit configuration file

Files:

  • docs/g2-cache-scope.md
Before making any change under: `python/aisimulate/src/aisimulate/generator/**` MUST read: `python/aisimulate/.claude/rules/generator-development.md` Before making any change under `python/aisimulate/collector/**` MUST read: `python/aisimul...

📄 CodeRabbit inference engine (AGENTS.md)

Files:

  • docs/g2-cache-scope.md
Source excerpt: Only workflows under the repository-root `.github/workflows/` run for this repository.

📄 CodeRabbit inference engine (REVIEW.md)

Files:

  • docs/g2-cache-scope.md
🪛 LanguageTool
docs/g2-cache-scope.md

[style] ~219-~219: The double modal “requires aggregated” is nonstandard (only accepted in certain dialects). Consider “to be aggregated”.
Context: ...shared` are rejected. G3 still requires aggregated replay with attention DP 1. Agentic sn...

(NEEDS_FIXED)

🔀 Multi-repo context ai-dynamo/dynamo, ai-dynamo/aiconfigurator

Linked repositories findings

ai-dynamo/dynamo

Inspected the branch associated with open PR #15626 at commit 59bf8dc50e0bf4362424e75499c18f292105c894—not the default branch.

  • AgentX G2 integration contract: Offline replay validates that agentic host offload uses one worker per role, static pools, vLLM on every role, attention DP=1 on every role, and no speculative decoding. The disaggregated entrypoint validates both prefill and decode roles. This matches the PR’s documented single-aggregated-worker or 1P1D scope. [lib/mocker/src/replay/offline/entrypoints.rs:47-95, 297-305][::ai-dynamo/dynamo::]
  • Shared-pool boundary: cluster_shared native host offload is explicitly limited to offline replay; the grouped scheduler rejects it outside that path. [lib/mocker/src/grouped_scheduler.rs:293-300][::ai-dynamo/dynamo::]
  • Consumer tests: Dynamo’s AgentX G2 tests exercise dp_rank_local and cluster_shared, test H2D delay at different bandwidths, and verify that a 1P1D shared pool has eight blocks—not sixteen. These tests are present on the PR #15626 branch. [components/src/dynamo/replay/tests/test_agentx_g2.py:200-235, 239-250][::ai-dynamo/dynamo::]
  • Qualification caveat: The Dynamo runner advertises agentic host-offload support but labels its qualification functional_only; its source says full AgentX runtime conformance remains a separate checkpoint. The integration test checks this metadata. [components/src/dynamo/replay/simulation.py:87-101; components/src/dynamo/replay/tests/test_simulation_integration.py:383-387][::ai-dynamo/dynamo::]

ai-dynamo/aiconfigurator

The broad search found no matching AgentX, ReplaySpec, or host-offload contract references in this frozen compatibility reference. [::ai-dynamo/aiconfigurator::]

🔇 Additional comments (1)
docs/g2-cache-scope.md (1)

221-221: 🎯 Functional Correctness

The inspected source supports a generic Agentic scope, not an AgentX-only runtime scope. The documentation links to the AgentX guide for qualification limits, so the wording does not claim that every supported Agentic path has separate qualification. No actionable scope error remains.


📝 Summary

Risk: Medium. The three areas needing the most human attention are:

  1. Python and native validation parity for G2 eligibility and host-offload scope.
  2. Restore, cache-lifetime, and cleanup behavior for local G2 and shared 1P1D replay.
  3. Runner capabilities and CLI configuration, especially rejection of G3, speculation, multi-worker layouts, and unsupported DP settings.

Behavior and contract changes

  • AgentX replay now accepts vLLM local or cluster-shared G2 host offload for a static, single aggregated worker or one prefill/decode pair. Attention DP must be 1 on every active role.
  • Python and native validation reject unsupported backends, speculative decoding, G3, invalid scopes, scaling hooks, and incompatible worker layouts. The Engine runner now declares agentic host-offload support.
  • CLI scope preservation, documentation, and examples cover local and shared G2 configurations. HBM-only replay retains its existing multi-worker and DP behavior.
  • Tests cover host-KV reuse, transfer delay, warmup and profile lifecycle, cancellation during restore, and Python/native/CLI configuration behavior.

Evidence supplied

The author reports 1,965 Rust library tests passed, with one existing ignore; 23 AgentX tests and 9 host-scope tests passed; and a freshly built, installed locked wheel passed 528 Python runner/G2/CLI tests. The author also reports both documented CLI examples completed and restored eight host input tokens. Formatting and packaging checks reportedly passed.

These results are author-reported and not independently verified here. Hosted CI was pending, and qualification remains functional_only.

Readiness

Functional qualification evidence is supplied, but hosted CI status and performance equivalence are not established. Dynamo integration and its end-to-end qualification are outside this change. Review severity counts are unavailable because no current review findings were supplied. The test reports do not constitute bot review approval.

Walkthrough

AgentX replay now supports qualified vLLM G2 host offload for specified static worker layouts. The change adds role validation, cache-reuse and lifecycle tests, functional fixtures, and documentation for configuration and qualification limits.

Changes

AgentX G2 replay

Layer / File(s) Summary
Validate supported replay configurations
python/aisimulate/src/aisimulate/sweeper/replay.py, python/aisimulate/src/aisimulate/runner.py, crates/core/src/python.rs, tests/test_runner.py, tests/test_host_offload_cli.py
Validation now allows qualified G2 host offload for supported vLLM layouts and continues to reject unsupported settings, including G3 offload and speculative decoding. Tests cover role constraints, capability checks, and native payload scope preservation.
Exercise G2 cache reuse and replay lifecycle
crates/core/src/engine/scheduler/rank.rs, crates/core/tests/agentx_qualification.rs
Native replay tests cover decode-owned host KV reuse by prefill, bandwidth and capacity effects, warmup state, recycled plays, and cancellation during a pending restore.
Add functional fixtures and cross-interface checks
examples/cli/agentx-g2-*.yaml, tests/e2e/configs/unified_cli/fixtures/traces/weka-g2-restore.jsonl, tests/test_agentx_g2.py, CODEOWNERS, .github/codeowners/areas.yaml
Local and shared-PD fixtures support acceptance checks across Python, native, and CLI replay. Tests check reuse evidence, timing controls, shared-pool compatibility, and invalid configurations. Ownership rules include the new acceptance test.
Document supported G2 configurations
docs/agentx-g2.md, docs/agentic-profile.md, docs/agentic-warmup.md, docs/cli/user-guide.md, docs/sweeper/traffic.md, docs/g2-cache-scope.md, python/aisimulate/src/aisimulate/cli_args.py
Documentation describes supported worker layouts, attention DP requirements, unsupported configurations, fixture use, report interpretation, and qualification scope.

Priority: ➖ Normal

Estimated code review effort: 4 (Complex) | ~50 minutes

Merge Risk: ⚪ Minimal · up to 64d51

The documentation change does not establish a merge-blocking issue; it is ready for normal checks.

🚥 Pre-merge checks | ✅ 6 | ❌ 2

❌ Failed checks (2 warnings)

Check name Status Explanation Resolution
Description check ⚠️ Warning The description explains the behavior, scope, and reported test results, but it omits several required template sections and evidence details. Add a Review map with risk level, starting points, contract impact, and compatibility concerns. Add the missing Evidence fields, including exact test commands, CI status, reviewed commits, numerical-assertion derivation, and before/after ou…
Cross-Layer Contract ⚠️ Warning The nested-rank Python ReplaySpec path is accepted by the new AgentX G2 capability gate but cannot execute cluster_shared unless callers add a native-only field themselves. tests/test_runner.py … Derive kv_layout_id for nested as well as flat Python rank input when cluster_shared is configured and the ID is absent, while preserving explicit IDs. Add an execution-level Python ReplaySpec test for nested-rank shared G2 that verif…
✅ Passed checks (6 passed)
Check name Status Explanation
Title check ✅ Passed The title precisely states the change: enabling vLLM AgentX local and shared G2 replay.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
Modeling And Data Evidence ✅ Passed The changed qualification evidence is traceable and reproducible. docs/agentx-g2.md identifies the Weka fixture as repository-authored and states that fixed timing makes it a functional test, not me…
Compatibility Boundaries ✅ Passed Python and Rust keep the replay schema at version 1 (REPLAY_SPEC_API_VERSION and CURRENT_REPLAY_SPEC_VERSION). The changed Python and Rust guards align on G2 deployment limits, G3/speculation reje…
Review Evidence ✅ Passed The PR description reports test and CLI-example results, states that hosted CI is pending, and separates Dynamo integration qualification from Engine acceptance. It does not cite an inactive nested wo…
Full details: Description check

Resolution

Add a Review map with risk level, starting points, contract impact, and compatibility concerns. Add the missing Evidence fields, including exact test commands, CI status, reviewed commits, numerical-assertion derivation, and before/after output or trace. Add Modeling or data provenance with N/A if not applicable, and complete Tracking with issue or PR references, or N/A.

Full details: Cross-Layer Contract

Explanation

The nested-rank Python ReplaySpec path is accepted by the new AgentX G2 capability gate but cannot execute cluster_shared unless callers add a native-only field themselves. tests/test_runner.py builds nested cluster_shared inputs without kv_layout_id and asserts require_compatible accepts them (lines 216–269), but does not execute them. In runner.py, _materialize_engine_role adds a derived kv_layout_id only when nested_rank is None (lines 1815–1822), so nested rank input reaches Rust without one. Rust requires a nonempty ID for cluster_shared (crates/core/src/engine/config.rs lines 410–416) and calls rank validation during role construction (crates/core/src/replay/engine.rs lines 384–389). The PR enables AgentX host offload in the Engine capability (runner.py line 412), exposing this previously rejected AgentX path. The new end-to-end test covers flat CLI-compiled configurations, not nested-rank execution. This also conflicts with docs/agentx-g2.md lines 64–72, which says Python supplies the ID and limits the caller-supplied requirement to direct native JSON callers.

Resolution

Derive kv_layout_id for nested as well as flat Python rank input when cluster_shared is configured and the ID is absent, while preserving explicit IDs. Add an execution-level Python ReplaySpec test for nested-rank shared G2 that verifies materialization and successful native replay. If nested rank input is intended to require an explicit ID instead, reject the missing ID at the Python boundary and document that requirement separately from direct native JSON.

  • Fix all pre-merge checks with AI
  • Autopilot · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

Signed-off-by: Harry Lee <harrli@nvidia.com>
@Harrilee
Harrilee changed the base branch from harrli/aic-1817-routing to main October 5, 2026 04:50
@Harrilee
Harrilee force-pushed the harrli/agentx-vllm-g2-api branch from 22d8ad0 to 7f47667 Compare October 5, 2026 04:50
Signed-off-by: Harry Lee <harrli@nvidia.com>
Comment thread crates/core/src/python.rs Outdated
Signed-off-by: Harry Lee <harrli@nvidia.com>
…-01a109e3

Signed-off-by: Harry Lee <harrli@nvidia.com>
@Harrilee
Harrilee marked this pull request as ready for review October 5, 2026 21:40
@Harrilee
Harrilee requested review from a team as code owners October 5, 2026 21:40
@Harrilee
Harrilee requested a review from dreamtalen October 5, 2026 21:40

@dreamtalen dreamtalen left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] Update the G2 unsupported-mode documentation

docs/g2-cache-scope.md:218 still says agentic snapshot/warmup host offload
is rejected. This PR enables qualified configurations. Replace the blanket
rejection with the supported limits or link to docs/agentx-g2.md.

Signed-off-by: Harry Lee <harrli@nvidia.com>
@Harrilee
Harrilee enabled auto-merge (squash) October 5, 2026 23:11
@Harrilee
Harrilee disabled auto-merge October 5, 2026 23:52
@Harrilee
Harrilee merged commit 9093044 into main Oct 5, 2026
73 checks passed
@Harrilee
Harrilee deleted the harrli/agentx-vllm-g2-api branch October 5, 2026 23:53
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants