Skip to content

feat: consume measured SD rounds in profile-backed replay - #477

Merged
tedzhouhk merged 22 commits into
mainfrom
hzhou-codex/fpm-sd-round-consumer
Oct 21, 2026
Merged

tedzhouhk merged 22 commits into
mainfrom
hzhou-codex/fpm-sd-round-consumer

Conversation

@tedzhouhk

@tedzhouhk tedzhouhk commented Oct 14, 2026 •

Copy link
Copy Markdown
Contributor

Summary

Depends on #476. Enable measured whole-SD-round timing through canonical ForwardPassPerfModel.best_available, profile compilation and ordinary replay.

  • Retain nextn for scheduling/acceptance, but price one complete measured round at the actual request count and past-KV total. Do not multiply lookup batch by K + 1 or charge draft operations again.
  • Require matching P/D execution identity, matching depth, direct SILICON interpolation, denied fallback and observed linear-cache runtime memory.
  • Preserve the recorded usable KV pool without subtracting draft/graph allocations again; retain identity across serialization and capacity materialization.
  • Fail explicitly for unsupported low-level operations or malformed bindings, including nested SD-selected operations without an execution binding.

Updated integration

Current head: 69c3ccd1b384b4c5f001e1abb9cd9766e9c4fe93, with main 00f0eff747801c00af98388dd0028df9eb95c53f merged into the existing branch using a signed, history-preserving commit.

The runtime conflict retains main's new GLM KV readiness check in the shared validated constructor, as well as this PR's explicitly bound whole-SD-round validation. The original 14-file consumer delta is unchanged. Independent full-tree comparison verified preservation of main's #512 HF evidence plumbing, #514 GLM changes, #523 catalog and merged AgentX warmup/speedup behavior. No force-push, version bump or data/golden rewrite was made.

Validation

  • Fresh exact-head Rust 1.96.1 embed-python native; immutable module identity verified.
  • 1,829 focused Python tests passed, including measured SD/memory/ad-hoc guards, GLM readiness/precision, catalog and current-main AgentX warmup/CLI paths.
  • 1,009 Rust performance-model, 55 native Python-API, and 200 loadgen tests passed; one existing R17 artifact-dependent test ignored.
  • Both frozen-golden parity suites: 390 passed, with goldens unchanged.
  • Independent native review: 126 probes passed, covering SD guards, memory topology, qualified replay, and GLM Hopper/Blackwell boundaries. Independent whole-tree proof found no lost main changes or unintended PR changes.
  • Source/legal, documentation, strict/generated CODEOWNERS, Rust formatting, scoped Ruff, whitespace, GPG and DCO checks passed.
  • Fresh hosted Full CI and forward-prediction advisory belong to this exact head. Historical checks are not substituted for current qualification.

Boundaries

vLLM EAGLE3, pure TP, DP1/PP1/CP1, linear cache and HBM-only. No ngram/generic speculation, recorded DCP, grouped-cache SD or offload SD support is added. Pure P/D raw FPM queries are supported; mixed replay timing remains separate P+D composition, not measured mixed-round fidelity. Partial-width rounds still use configured-depth timing, and the existing prefill callback still rounds ragged batches to integer mean lengths.

Guided collection and independent AgentX accuracy qualification are separate work. No new GPU runs, performance-table changes or golden refreshes were performed.

…ract

Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
@tedzhouhk
tedzhouhk requested review from a team as code owners October 14, 2026 23:56
@coderabbitai

coderabbitai Bot commented Oct 14, 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: f52543bc-5487-43b8-b9d7-f36862ce7583
📥 Commits

Reviewing files that changed from the base of the PR and between 1c7aaf0 and 69c3ccd.

📒 Files selected for processing: 4
  • crates/core/src/perfmodel/engine/runtime.rs
  • crates/core/src/python.rs
  • python/aisimulate/src/aisimulate/capacity.py
  • python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py
🔗 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 (new)
  • GitHub Check: Prediction Regression / Collect prediction snapshot (old)
  • GitHub Check: Platform Wheels / Build wheels (manylinux_2_28_x86_64)
  • GitHub Check: Collector Data / Perf data sanity (informational)
  • GitHub Check: Platform Wheels / Build wheels (macosx_arm64)
  • GitHub Check: Collector Data / Check collector data
  • GitHub Check: Platform Wheels / Build wheels (manylinux_2_28_aarch64)
  • GitHub Check: Release Artifact Contract (arm64)
  • GitHub Check: readme-commands / README / dynamo
  • GitHub Check: readme-commands / README / source
  • GitHub Check: Public API Rust (amd64)
  • GitHub Check: Public API Rust (arm64)
  • GitHub Check: Release Artifact Contract (amd64)
  • GitHub Check: Rust feature modes
  • GitHub Check: Python 3.13 compatibility
  • GitHub Check: Application Test Wheel (amd64)
  • GitHub Check: Application Test Wheel (arm64)
  • GitHub Check: Engine Golden Regression
  • GitHub Check: Python 3.11 compatibility
  • GitHub Check: Rust (arm64)
  • GitHub Check: Rust (amd64)
  • 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: 4
python/aisimulate/.claude/rules/rust-core/parity.md — configured
AGENTS.md — configured
REVIEW.md — configured
DEVELOPMENT.md — auto-discovered
📓 Path-based instructions: 9
Preserve the Rust single oracle: Python may describe operations, load raw data, orchestrate, and present results, but must not compute per-op performance values.

⚙️ CodeRabbit configuration file

Files:

  • python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py
Check unified CLI, Replay, Sweeper, and orchestration behavior together.

⚙️ CodeRabbit configuration file

Files:

  • python/aisimulate/src/aisimulate/capacity.py
Enforce the single-oracle and golden-diff rules in python/aisimulate/.claude/rules/rust-core/parity.md.

⚙️ CodeRabbit configuration file

Files:

  • crates/core/src/perfmodel/engine/runtime.rs
Treat top-level exports and bindings as public and release boundaries.

⚙️ CodeRabbit configuration file

Files:

  • crates/core/src/python.rs
Read REVIEW.md before commenting.

⚙️ CodeRabbit configuration file

Files:

  • python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py
  • crates/core/src/python.rs
  • python/aisimulate/src/aisimulate/capacity.py
  • crates/core/src/perfmodel/engine/runtime.rs
Source excerpt: Any PR that changes what the compiled engine computes (operators, loaders, interpolation, selection rules, engine composition) MUST in the same PR: Keep `test_engine_step_parity.py` / `test_compile_engine_parity.py` green —...

📄 CodeRabbit inference engine (python/aisimulate/.claude/rules/rust-core/parity.md)

Files:

  • python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py
  • crates/core/src/perfmodel/engine/runtime.rs
Source excerpt: Keep run logs, investigation notes, and benchmark results out of `docs/`; measured evidence belongs under `benchmarks/evidence/`.

📄 CodeRabbit inference engine (AGENTS.md)

Files:

  • python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py
  • crates/core/src/python.rs
  • python/aisimulate/src/aisimulate/capacity.py
  • crates/core/src/perfmodel/engine/runtime.rs
Source excerpt: Binary parquet changes require a machine-readable summary of keys, shapes, units, row counts, coverage, and anomalies plus proof that a real consumer reaches the new data.

📄 CodeRabbit inference engine (REVIEW.md)

Files:

  • python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py
  • crates/core/src/python.rs
  • python/aisimulate/src/aisimulate/capacity.py
  • crates/core/src/perfmodel/engine/runtime.rs
Applicability reference: AGENTS.md: # AGENTS **Lint:** `python/aisimulate/.venv/bin/ruff check --config python/aisimulate/pyproject.toml python/aisimulate tests` and `python/aisimulate/.venv/bin/ruff format --check --config python/aisimulat...

📄 CodeRabbit inference engine (DEVELOPMENT.md)

Files:

  • python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py
  • python/aisimulate/src/aisimulate/capacity.py
🔀 Multi-repo context ai-dynamo/dynamo, ai-dynamo/aiconfigurator

Linked repositories findings

ai-dynamo/dynamo

  • Dynamo pins aisimulate and aisimulate-core to exact versions, so the measured-round changes require coordinated package updates. [::ai-dynamo/dynamo::pyproject.toml:17] [::ai-dynamo/dynamo::Cargo.toml:62]
  • Planner speculative-depth normalization only checks existing speculation/nextn fields and then injects config["nextn"]; it does not add execution-profile identity. [::ai-dynamo/dynamo::components/src/dynamo/planner/core/perf_model/ais_adapter.py:265-276] Profiles using positive nextn therefore need identity already present, or the new SDK validation will reject them.

ai-dynamo/aiconfigurator

  • The frozen Rust engine explicitly rejects FPM specs with nextn > 0, confirming it cannot consume the new measured speculative-round format. [::ai-dynamo/aiconfigurator::aic-core/rust/aiconfigurator-core/src/engine/runtime.rs:336-372]
  • The frozen public compile_engine contract exposes nextn and forward_model but no execution-profile identity parameter. [::ai-dynamo/aiconfigurator::tests/cross_package/test_core_public_api.py:51-78]
🔇 Additional comments: 4
python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py: 1

18-18: LGTM!

crates/core/src/perfmodel/engine/runtime.rs: 1

416-416: LGTM!

Also applies to: 447-465, 494-494, 528-547, 570-643, 680-680, 726-726, 806-811, 855-858, 925-925, 1200-1200, 1524-1524, 1701-1701, 1893-1893, 1989-1989, 2091-2091, 2125-2125, 2160-2160, 2199-2199, 2248-2248, 2303-2303, 2382-2386, 2480-2488, 4055-4310

crates/core/src/python.rs: 1

1041-1046: 🎯 Functional Correctness

Do not gate nextn on execution_profile here.

The Python profile-backed capacity path rejects nextn > 0 when the selected deployment has no execution_profile, so an unbound hybrid profile does not reach the native estimator with the forwarded value. The claimed failure or block-count divergence is therefore not established, and the proposed change would replace an explicit validation error with a different capacity request.

python/aisimulate/src/aisimulate/capacity.py: 1

221-228: 🎯 Functional Correctness

Do not change this depth lookup.

Measured execution profiles reject generic speculation and require nextn to match execution_profile.num_draft_tokens. The canonical Rust configuration also serializes its nextn field, including its default value. The proposed speculation-only and missing-nextn paths do not reach this lookup.

📝 Summary

Risk: High

Human attention

  1. Verify complete-round pricing uses actual request count and past-KV total without duplicate lookup or draft-operation charges.
  2. Verify strict execution identity, nextn, topology, interpolation, fallback, and cache-memory validation.
  3. Verify serialization and capacity materialization preserve the recorded usable KV pool and execution identity.

Changed behavior and contracts

  • Measured speculative-round profiles now use FpmExecutionBinding validation.
  • ForwardPassPerfModel.best_available can consume canonical measured full-round timings.
  • nextn remains the scheduling and acceptance depth.
  • Measured profiles require direct SILICON interpolation, denied fallback, matching prefill/decode identities, and vLLM pure-TP topology.
  • Unsupported generic speculation, mixed telemetry, grouped-cache SD, DCP, offload SD, and malformed bindings are rejected.
  • SDK and capacity APIs propagate speculative execution identity and nextn.
  • FpmSpeculativeExecutionProfile is exported.
  • Documentation defines profile schema, replay restrictions, and complete-round charging.

Technical quality

  • Supplied evidence reports 1,454 focused Python tests, 1,001 Rust performance-model tests, two frozen-golden suites with 390 passes each, and 112 exact-native probes.
  • Supplied evidence reports offline catalog discovery for 55 profiles across 9 manifests.
  • The evidence does not verify downloading or re-benchmarking external datasets.
  • One pre-existing real-artifact test was ignored.
  • Current review findings and severity counts are unavailable. Bot summaries are not approval.
  • Repository inspection completed without reported errors, but returned no diff or source excerpts. It therefore adds no verifiable implementation detail.

Merge readiness

  • The supplied evidence supports targeted coverage for profile validation, identity preservation, capacity materialization, and replay.
  • Full-workspace validation and current review findings are not supplied.
  • Merge requires confirmation of the three high-risk areas and any required broader validation.

This change adds validated speculative-round FPM profiles across Rust and Python. The engine prices measured prefill and decode rounds, while capacity estimation forwards the configured draft depth. The change also adds replay documentation and validation tests.

Changes: Speculative-round FPM

Layer / File(s) Summary
Execution identity and profile contract
crates/core/src/perfmodel/fpm/*, python/aisimulate/src/aisimulate_core/sdk/engine.py, python/aisimulate/src/aisimulate_core/sdk/fpm_profile.py, python/aisimulate/tests/unit/sdk/*
FPM profiles accept nonzero nextn only with a validated execution profile. Direct FPM specifications serialize execution identity. Tests cover validation, serialization, normalization, and profile selection.
Bound engine execution and query constraints
crates/core/src/perfmodel/engine/runtime.rs
Engine construction validates execution bindings and phase selectors. Full-round FPM uses measured prefill and decode totals, requires one attention-DP rank, and rejects unsupported query surfaces.
Python capacity and cache estimation
python/aisimulate/src/aisimulate/capacity.py, python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py, crates/core/src/python.rs, python/aisimulate/tests/unit/sdk/test_fpm_profile.py
Capacity and KV-cache estimation forward speculative depth for execution-profile deployments. Tests cover observed runtime memory, supported depth, and GPU-block estimates.
Replay behavior and end-to-end validation
docs/perf-model/methods/whole-forward.md, docs/replay/engine/speculation.md, tests/test_fpm_profile_workflow.py
Documentation defines measurement, lookup, topology, and replay constraints. The workflow test checks identity propagation, observed KV capacity, speculative-depth handling, and replay metrics.

Priority: ➖ Normal

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

Merge Risk: ⚪ Minimal · up to 69c3c

This change enables measured speculative-decoding rounds in profile-backed replay and capacity estimation. The one concern was that a hybrid speculative profile could pass the wrong draft depth to capacity estimation. That configuration is rejected earlier in both the old and new code. No outstanding issue blocks merging.

Pre-merge checks: Passed 5 · Failed 3

Failed Failed checks: 3 warnings

Check name Status Explanation Resolution
Cross-Layer Contract Warning The compatibility layer remains stale. This PR adds FpmSpeculativeExecutionProfile to aisimulate_core.sdk.fpm_profile at line 18, but aisimulate.fpm_profile still exports only the older profile … Add FpmSpeculativeExecutionProfile to python/aisimulate/src/aisimulate/fpm_profile.py from aisimulate_core.fpm_profile. Extend test_fpm_profile_alias_preserves_module_and_type_identity to include the new type and verify its pickle i…
Compatibility Boundaries Warning Python and Rust public surfaces remain unsynchronized. The PR adds the measured-profile path in python/aisimulate/src/aisimulate_core/sdk/engine.py:750-753, which calls `speculative_execution_identi… Add def _normalize_fpm_execution_profile(model_revision: str, execution_profile_json: str) -> str: ... to python/aisimulate/src/aisimulate_core/_native.pyi. Add a cross-package test that compares this stub declaration with the native me…
Review Evidence Warning The description reports counts and broad validation categories, plus negative boundaries, parity results, native probes, and hosted CI links. It does not name the actual local commands, identify fake-… Update the Validation section with the exact commands and results. Label each result as local, fake-runner/synthetic, parity, or hosted/production evidence. State explicitly when production or GPU evidence was not run, and keep the active t…
Passed Passed checks: 5 passed
Check name Status Explanation
Title check Passed The title precisely describes the main behavioral change: consuming measured speculative-decoding rounds during profile-backed replay. It avoids vague wording.
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 PR provides reproducible numeric oracle evidence and provenance. crates/core/src/perfmodel/engine/runtime.rs:4070-4178 builds a hand-authored schema-7/native-CUDA fixture with explicit selectors…
Description check Passed The description clearly explains the behavior change, scope, compatibility boundaries, integration details, and extensive validation evidence. It does not use the required section headings for Review …
Full details: Modeling And Data Evidence

Explanation

The PR provides reproducible numeric oracle evidence and provenance. crates/core/src/perfmodel/engine/runtime.rs:4070-4178 builds a hand-authored schema-7/native-CUDA fixture with explicit selectors, real_kv, and known latencies, then checks exact prefill/decode totals, batch-width behavior, and no draft recharge. tests/test_fpm_profile_workflow.py:230-320 writes hashed Parquet metadata, marks the timing rows as synthetic, and verifies capacity, replay identity, TTFT, round counts, and ITL. The changed-file inventory contains no golden or performance-data refresh. The PR description also reports frozen-golden parity with no golden updates and independent native probes.

Full details: Description check

Explanation

The description clearly explains the behavior change, scope, compatibility boundaries, integration details, and extensive validation evidence. It does not use the required section headings for Review map, Modeling or data provenance, or Tracking, and it omits explicit CodeRabbit and Codex reviewed-commit fields. These omissions are non-critical because the required technical context is otherwise substantially present.

Full details: Cross-Layer Contract

Explanation

The compatibility layer remains stale. This PR adds FpmSpeculativeExecutionProfile to aisimulate_core.sdk.fpm_profile at line 18, but aisimulate.fpm_profile still exports only the older profile types at lines 6-11. The cross-package identity test at lines 158-164 checks only three older types, so it does not cover the missing compatibility import or pickle identity.

Resolution

Add FpmSpeculativeExecutionProfile to python/aisimulate/src/aisimulate/fpm_profile.py from aisimulate_core.fpm_profile. Extend test_fpm_profile_alias_preserves_module_and_type_identity to include the new type and verify its pickle identity through aisimulate.fpm_profile.

Full details: Compatibility Boundaries

Explanation

Python and Rust public surfaces remain unsynchronized. The PR adds the measured-profile path in python/aisimulate/src/aisimulate_core/sdk/engine.py:750-753, which calls speculative_execution_identity() and then the native _normalize_fpm_execution_profile method at python/aisimulate/src/aisimulate_core/sdk/fpm_identity.py:16-23. Rust exposes that static method with two string parameters and a string result at crates/core/src/perfmodel/py.rs:2016-2032, but the shipped python/aisimulate/src/aisimulate_core/_native.pyi:166-205 omits it. The pull request does not change _native.pyi, and the existing cross-package tests do not check this method. The artifact manifests are unchanged, so no deliberate artifact-contract update addresses the mismatch.

Resolution

Add def _normalize_fpm_execution_profile(model_revision: str, execution_profile_json: str) -&gt; str: ... to python/aisimulate/src/aisimulate_core/_native.pyi. Add a cross-package test that compares this stub declaration with the native method signature. Re-run the one-wheel import and typing-contract tests.

Full details: Review Evidence

Explanation

The description reports counts and broad validation categories, plus negative boundaries, parity results, native probes, and hosted CI links. It does not name the actual local commands, identify fake-runner coverage, or clearly separate local, fake-runner, parity, and production evidence. The cited Full CI and Forward Prediction Performance workflows are active top-level workflows (.github/workflows/ci.yml:6-22 and .github/workflows/performance.yml:4-17), so the inactive-nested-workflow condition is not the issue.

Resolution

Update the Validation section with the exact commands and results. Label each result as local, fake-runner/synthetic, parity, or hosted/production evidence. State explicitly when production or GPU evidence was not run, and keep the active top-level workflow run links with their status and exact commit.

  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>

@diffractiongoo diffractiongoo 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.

One P2 finding in the standalone measured-profile memory API: unsupported DCP is dropped before canonical validation. Details inline.

Comment thread python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Base automatically changed from hzhou-codex/fpm-sd-round-identity to main October 20, 2026 20:33
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
Signed-off-by: hongkuanz <hongkuanz@nvidia.com>
@tedzhouhk
tedzhouhk merged commit a0a35e7 into main Oct 21, 2026
73 checks passed
@tedzhouhk
tedzhouhk deleted the hzhou-codex/fpm-sd-round-consumer branch October 21, 2026 00:16
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