[Arsene12358]

Why and what changed

Decode-phase (generation) Mamba2 rows were never used. StateSpaceTable::query_mamba2 returned an error for every decode query on purpose. That parity gate copied a defect of a Python loader (load_mamba2_data) that no longer exists, so every decode Mamba2 kernel fell back to its analytic SOL estimate even when measured rows were packaged.

Mamba2Op::query also ignored the database mode. In SOL and EMPIRICAL modes it read measured rows and labelled them silicon. Before this branch, only the context path leaked this way, because decode rows were never read.

• crates/core/src/perfmodel/perf_database/state_space.rs: remove the parity gate. Decode queries now resolve like context queries, with the same key lookup on the 1-axis batch curve; the query sequence length feeds SOL only. The comments that described the gate are removed.
• crates/core/src/perfmodel/operators/mamba.rs: Mamba2Op::query now follows the database-mode semantics of the GEMM, attention, MoE and mHC operators.
  • SOL and SOL_FULL return the analytic SOL without reading the table.
  • EMPIRICAL also returns the analytic SOL, because Mamba2 has no util-space empirical layer.
  • SILICON and HYBRID read the table and keep the analytic SOL fallback on a miss.
• crates/core/parity_tests/perfmodel/goldens/engine_step.json: refresh, with pin_goldens.py --refresh, of the two records that the decode change moves. See the golden diff below.

Who consumes this: op-level estimation of NemotronH-family models on data that ships a Mamba2 table. Today that is the TensorRT-LLM 1.3.0rc29 tables of eight systems (the default current slot since the data campaign of #517); the rc20 tables of the same systems also ship a Mamba2 table. Decode Mamba2 conv1d kernels now use their measured rows in SILICON and HYBRID. In SOL and EMPIRICAL, Mamba2 context and decode now return the analytic SOL on every backend.

Readiness rule split: this branch first also made SILICON readiness report a missing Mamba2 table as a data miss. That rule removed default-mode op-level support for NemotronH-family models on backends without a Mamba2 table. By the owner's decision, it moves to AISIM-2194. Commit a0139351 restores crates/core/src/perfmodel/engine/readiness.rs to its base content and removes that rule's test. The branch history keeps the first commit, whose message still describes the rule; this description supersedes it, so trim the readiness sentences of that commit from the squash message when merging.

Review map

• Risk level: medium. It changes predictions through table use and database-mode semantics.
• Start with:
  • the mode guard at the top of Mamba2Op::query in operators/mamba.rs;
  • the removed gate in StateSpaceTable::query_mamba2 in perf_database/state_space.rs;
  • the golden diff.
• Public or serialized contract changed: no. No schema, API, configuration or default changes. readiness.rs is identical to the base revision.
• Compatibility or rollback concern:
  • Decode predictions of NemotronH-family models change wherever a Mamba2 table resolves (SILICON and HYBRID).
  • Mamba2 context values change in SOL and EMPIRICAL modes on such data, because the context op no longer reads measured rows there.
  • Rollback is a revert. The golden refresh reverts with it.

Evidence

Local runs on the branch head a0139351, from the repository root unless noted. The branch was rebased on 2026-10-30 onto main 7980e950, after the data campaign #517, which moved the default TensorRT-LLM slot to 1.3.0rc29, added the Mamba2 alias table to state_space.rs and re-pinned the goldens; the rebase kept main's alias lookup next to the removed gate and regenerated the two golden records with pin_goldens.py --refresh. The native extension was rebuilt from the same tree, and the HTTP proxy variables were unset. The base for comparisons is 7980e950.

• Tests:
  • cargo test --workspace --no-fail-fast: 2680 passed, 0 failed, 1 ignored. The base has 2677 passed; the difference is the three new tests.
  • cargo test --workspace --features embed-python,replay-bench, with a local Python embedding and AIC_REQUIRE_EMBEDDED_ROUND_TRIP=1: 2772 passed, 0 failed, 1 ignored.
  • cargo fmt --all -- --check is clean. cargo clippy --workspace --all-targets reports no warning on a changed line.
  • python -m pytest -q -rx -n 4 -c python/aisimulate/pytest.ini crates/core/parity_tests/perfmodel/test_engine_step_parity.py: 304 passed.
  • The same command with test_compile_engine_parity.py: 85 passed.
  • python scripts/prediction_regression/check_prediction_numerics.py --output <file>: 16 cases, 0 failures.
  • From python/aisimulate: python -m pytest -q -p no:timeout -n 10 tests/unit tests/golden --ignore=tests/unit/sdk/database/test_moe_dispatch.py. Per-test results are identical to the base revision: 13411 passed and 29 skipped. The same 3 failures occur on both sides and come from the local environment: no sklearn, and collector.fpm_forward is not importable from an editable install.
  • Prediction Regression Gate, reproduced locally with the scripts of .github/workflows/prediction-regression-gate.yml (collect_static.py, run_tier2.py, run_silicon.py, then report.py against the base snapshot): exit 0. It compared 174437 rows across 29 combos and found DRIFT 685, which is report-only. There is no REGRESSION, STATUS_CHANGE, GAIN, ROWS_ADDED, ROWS_REMOVED or RELABELED row. Tier-1 status totals are equal on both sides, and the tier-2 and silicon snapshots are byte-identical.
• Fast CI and Full CI: running on a0139351 after the rebase; the earlier runs were on d4db1d79.
• CodeRabbit reviewed commit: d4db1d79, before the rebase; the review of a0139351 is pending.
• Codex reviewed commit: pending (required for medium risk).
• Negative or boundary cases:
  • A decode kernel whose logical name resolves no row (selective_state_update) keeps the analytic estimate. The rc29 tables ship decode SSM rows under the physical label selective_state_update_flashinfer, and no alias maps the logical name to it yet; that mapping is a follow-up.
  • SOL, SOL_FULL and EMPIRICAL never read rows, for both the context and the decode op.
  • evaluate_ops_sol_json still rejects Mamba2 for having no SOL decomposition (evaluate_ops_sol_json_rejects_unexported_families).
  • Each new test failed before its fix:
    • The state-space test unwrapped the gate error.
    • The decode operator test saw Sol instead of Silicon.
    • The mode test saw measured rows labelled Silicon in EMPIRICAL.
• Expected-value derivation (for exact numerical assertions): the new tests use synthetic fixtures whose rows are hit exactly.
  • Measured results equal the stored fixture value times scale_factor.
  • Analytic results are derived by hand in the test comment: 1536 conv channels; 18,874,368 bytes for context and 36,864 bytes for decode; divided by the fixture's memory bandwidth.
• Before/after output, trace, benchmark, or golden diff: see the golden diff below.

Golden diff

engine_step.json changes six leaves in two records of the TensorRT-LLM Nemotron-H-56B smoke case in SMOKE_CASES, whose key is nemotron-h-56b-<system>-trtllm-isl1024-osl2. The diff shows the full key.

Record	Metric	Change
::static	static_gen	higher
::static	static_total	higher, by exactly the static_gen change
::static	static_ctx	unchanged
::disagg	disagg_tpot	higher
::disagg	disagg_request	higher, by exactly the disagg_tpot change
::disagg	disagg_ttft	unchanged

The two post_freeze_pins entries of these records now name the pin commit 19067832.

Unchanged:

• the ::mixed and ::agg records of the same case;
• every other case, including the other NemotronH cases, whose backends resolve no Mamba2 table;
• compile_engine.json and per_op.json.

This is the intended effect:

• In this case only generation_mamba_conv1d (causal_conv1d_update) moves from the analytic SOL to its measured row.
• generation_mamba_ssm (selective_state_update) stays analytic: the rc29 tables ship its rows under the physical label selective_state_update_flashinfer, which no alias maps to the logical name yet.
• The mixed and aggregated records do not move, because the mixed-step composition takes only decode attention from the generation graph.
• Neither the mode fix (9f8551ba) nor the readiness split (a0139351) moved a golden: no parity case runs a model with Mamba2 rows in SOL or EMPIRICAL mode, and readiness does not change predictions.

To reproduce:

• Run the head code against the base engine_step.json. It fails exactly these two records (302 passed, 2 failed).

• The refresh ran from a clean tree at 19067832 with a matching native build:

  PYTHONPATH=python/aisimulate/src:python/aisimulate python/aisimulate/.venv/bin/python crates/core/parity_tests/perfmodel/pin_goldens.py --refresh $(python3 -c "import json; c=json.load(open('crates/core/parity_tests/perfmodel/goldens/engine_step.json'))['cases']; print(' '.join(f'{k}::{s}' for k in c if k.startswith('nemotron-h-56b-') and '-trtllm-' in k for s in ('static', 'disagg')))")

Prediction gate drift

All 685 drift rows are decode (gen) rows on the TensorRT-LLM 1.3.0rc29 combos of the eight systems that ship the Mamba2 table. They cover nine NemotronH-family models: Nemotron-3 Nano BF16 and NVFP4, Nemotron-3 Super FP8 and NVFP4, Nemotron-3 Ultra BF16, FP8 and NVFP4, Nemotron-3.5 Lightning NVFP4, and Nemotron-H-56B.

652 rows go up and 33 go down. All 33 decreases are on the two systems (L40S and RTX PRO 6000) whose measured decode conv1d rows fall below the bandwidth-only SOL at some batch sizes. That is a data-quality question for the collector, not a property of this code.

Beyond the parity records, these move:

• static and disaggregated decode;
• pure-decode iterations, because decode_step_latency runs the full generation graph, so aggregated serving moves in decode-only steps;
• decode energy where the table has a power column, which one of the eight tables does.

Mixed prefill-and-decode steps do not move.

Modeling or data provenance

No performance data changes. The decode rows were already packaged in mamba2_perf.parquet (TensorRT-LLM 1.3.0rc20 and 1.3.0rc29, eight systems); this pull request stops discarding them. Decode rows resolve today only for causal_conv1d_update; the rc29 decode SSM rows carry a physical label that no alias maps yet. The unit tests use synthetic fixtures. The goldens are regression baselines, not accuracy validation. The silicon-accuracy layer of the gate has no NemotronH reference, so this pull request makes no accuracy claim.

Known limitations, unchanged here:

• Decode now inherits the exact-shape miss fallback to the first table entry with the same d_model, which AISIM-2185 tracks.
• The raw data-plane view view_mamba2 (table_view.rs) still drops decode rows, so the engine and the data plane disagree for Mamba2 decode.
• The GDN and KDA operators still read measured rows in SOL and EMPIRICAL modes.
• For EMPIRICAL, the alternative policy would be to refuse Mamba2, as DeepSeek-V4.1 does. This pull request keeps the analytic SOL, which the unchanged readiness arm also assumes.

Tracking

• Closes: none on GitHub; tracked in Linear.
• Related PRs or issues:
  • AISIM-2194: the SILICON readiness rule for Mamba2, split out of this pull request.
  • AISIM-2185 covers the TP-blind lookup and the d_model fallback.

Linear: AISIM-2184

🤖 Generated with Claude Code

[coderabbitai]

📝 Summary

Risk: Medium. Human attention should focus on mode-specific routing in Mamba2Op::query, generation-row selection and fallback in StateSpaceTable::query_mamba2, and behavior when measured modes have no matching row.

Behavior and contract

• SOL, SOL_FULL, and EMPIRICAL return analytic SOL without reading measured rows. SILICON and HYBRID use measured rows and fall back to analytic SOL on a database miss.
• Generation queries can resolve against Mamba2 rows. The generation curve uses batch as its coordinate; sequence length is used for the SOL estimate.
• The supplied change summary reports no public signature changes. Query results now depend on database mode and, in measured modes, whether a matching row resolves.
• Golden data updates the disaggregated and static Nemotron-H-56B values. The mixed record remains unchanged.

Evidence and quality

• The supplied summary reports tests for context and generation row resolution, measured-row scaling, and outputs across all five database modes.
• The author reports Rust workspace tests, embedded-Python workspace tests, and parity tests. These results were not independently verified.
• The author reports three Python unit-test failures on both the branch and base, attributed to environment issues. Fast CI and full CI were pending. Review severity counts are unavailable.
• Shell results provided only log paths, not source or diff output. They do not provide independent verification of the implementation claims.

Merge readiness
The reported tests and behavior do not establish merge readiness. CI and review completion remain unconfirmed; bot checks are not reviewer approval.

Mamba2 generation queries now resolve measured rows in SILICON and HYBRID modes. SOL, SOL_FULL, and EMPIRICAL modes use analytic SOL results. Tests cover generation-row scaling and query behavior across modes. The parity golden values and selected pins are updated.

Changes: Mamba2 performance queries

Layer / File(s)	Summary
Generation-row resolution crates/core/src/perfmodel/perf_database/state_space.rs	Generation queries proceed through grid loading and key selection. Context queries use batch and sequence coordinates; generation queries use the batch-only curve. A parquet-backed test checks the stored latencies for both phases.
Mode selection and parity updates crates/core/src/perfmodel/operators/mamba.rs, crates/core/parity_tests/perfmodel/goldens/engine_step.json	Mamba2Op::query uses analytic SOL results in SOL, SOL_FULL, and EMPIRICAL modes. SILICON and HYBRID modes use measured rows and fall back to SOL on a miss. Tests cover these modes and generation-row scaling. The Nemotron H 56B TRT-LLM golden values and disagg and static pins are updated.

Priority: Normal

Estimated code review effort: 3 (Moderate) | ~20 minutes

Possibly related PRs

• ai-dynamo/aisimulate#619: Changes Mamba2 generation-query semantics in the same operator path.
• ai-dynamo/aisimulate#517: Adds Mamba2 performance tables that this change consumes.
• ai-dynamo/aisimulate#649: Changes collection of Mamba2 context and decode rows consumed by this change.

Merge Risk: Minimal · up to a0139

Measured decode estimates are enabled in supported modes. No merge-blocking issue was identified.

Pre-merge checks: 6 · 1 · 1

Failed checks: 1 warning · 1 inconclusive

Check name	Status	Explanation	Resolution
Cross-Layer Contract		The runtime Rust path is covered, but the Python serialization consumer is stale. StateSpaceTable::query_mamba2 now resolves generation rows, and Mamba2Op sends SILICON/HYBRID queries to that path…	Update view_mamba2 to serialize generation rows as leaves at [kernel_source][phase][model_key][batch], remove the preserved-bug documentation, and add a table-view/Python consumer test with both context and generation rows. Then rerun t…
Review Evidence		The available description identifies local Rust, embedded-Python, formatting, clippy, and parity commands, and it labels CI as pending. The changed tests cover measured-row, missing-row, and database-…	Provide the complete PR description, especially the remainder of the Evidence section. Confirm the final text reports command results, labels local/parity/fake-runner/production evidence correctly, and does not present workflows under pytho…

Passed checks: 6 passed

Check name	Status	Explanation
Title check		The title precisely identifies the behavioral changes: using decode-phase Mamba2 rows and enforcing database-mode semantics for the Mamba2 operator.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.
Modeling And Data Evidence		PASS. The PR supplies provenance for the TensorRT-LLM Nemotron-H Mamba2 data and an explained, machine-readable golden diff. The diff changes only the expected static and disaggregated decode metrics,…
Compatibility Boundaries		No compatibility-boundary failure is introduced. The PR changes only the Mamba2 Rust query path, Mamba2 table lookup, and two parity-golden records. It does not change Python wrappers, Rust PyO3 bindi…
Description check		The pull request description follows the required template and covers the behavior change, review map, risk, compatibility, test evidence, golden updates, provenance, limitations, and tracking. It ide…

Full details: Modeling And Data Evidence

Explanation

PASS. The PR supplies provenance for the TensorRT-LLM Nemotron-H Mamba2 data and an explained, machine-readable golden diff. The diff changes only the expected static and disaggregated decode metrics, while context, mixed, aggregate, and unrelated backend records remain unchanged. Commit evidence identifies the pin_goldens.py --refresh workflow and reports parity, numerical-sentinel, and prediction-regression checks. The added Rust tests provide reproducible measured-row and database-mode checks.

Full details: Compatibility Boundaries

Explanation

No compatibility-boundary failure is introduced. The PR changes only the Mamba2 Rust query path, Mamba2 table lookup, and two parity-golden records. It does not change Python wrappers, Rust PyO3 bindings, DatabaseMode values, schemas, versions, Cargo manifests, wheel configuration, or documented defaults. The Python Mamba2 shell still delegates evaluation through the unchanged native engine path (operations/mamba.py and operations/base.py), while the Rust binding constructor and serialized fields remain unchanged (perfmodel/py_ops.rs). The workspace still has one crate, and the unchanged maturin configuration still builds one aisimulate wheel from that crate. Existing documentation still states that AISimulate has no Dynamo installation dependency.

Full details: Description check

Explanation

The pull request description follows the required template and covers the behavior change, review map, risk, compatibility, test evidence, golden updates, provenance, limitations, and tracking. It identifies the pending Codex review required for medium risk, but the description is otherwise complete and actionable.

Full details: Cross-Layer Contract

Explanation

The runtime Rust path is covered, but the Python serialization consumer is stale. StateSpaceTable::query_mamba2 now resolves generation rows, and Mamba2Op sends SILICON/HYBRID queries to that path (crates/core/src/perfmodel/perf_database/state_space.rs:270-327, crates/core/src/perfmodel/operators/mamba.rs:33-76). However, view_mamba2 still deliberately drops every generation leaf and only vivifies [model_key][batch] (crates/core/src/perfmodel/perf_database/table_view.rs:796-830). Mamba2Kernel.load_data still exposes this view as _mamba2_data (python/aisimulate/src/aisimulate_core/sdk/operations/mamba.py:73-86), so the Python table-view/serialization layer does not provide the decode rows that the changed Rust consumer now uses. The new Rust tests and refreshed engine-step golden cover the direct engine output, but no test covers the Python table-view consumer.

Resolution

Update view_mamba2 to serialize generation rows as leaves at [kernel_source][phase][model_key][batch], remove the preserved-bug documentation, and add a table-view/Python consumer test with both context and generation rows. Then rerun the Mamba2 engine and table-view contract tests.

Full details: Review Evidence

Explanation

The available description identifies local Rust, embedded-Python, formatting, clippy, and parity commands, and it labels CI as pending. The changed tests cover measured-row, missing-row, and database-mode boundaries. However, the PR description is truncated at the Python test command. The complete results and any later hosted-CI claims, including possible references to inactive nested workflows, cannot be verified.

Resolution

Provide the complete PR description, especially the remainder of the Evidence section. Confirm the final text reports command results, labels local/parity/fake-runner/production evidence correctly, and does not present workflows under python/aisimulate/.github/workflows/ as hosted CI evidence.

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

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

[Arsene12358]

@simone-chen this pull request is ready for your review. It is a consumer-side fix in the Rust perf model, and nothing changes on the collector side.

• Decode predictions now read the decode-phase Mamba2 rows. The parity gate that forced the context rows for decode is removed.
• The Mamba2 operator is gated by database mode. SOL, SOL_FULL and EMPIRICAL stay analytic, SILICON and HYBRID read the measured rows.
• The engine-step goldens are refreshed for the decode rows.

The SILICON readiness rule for Mamba2 is split out to AISIM-2194 and stays out of this pull request. Checks: 61 passed, 8 still running, none failed. Please squash on merge, and trim the readiness sentences of the first commit from the squash message.

[Arsene12358]

@simone-chen the branch is rebased onto the current main (head a013935) and the CI is green. It is ready for your review. The description has the local evidence: the cargo and embedded suites, the parity run, and the consumer gate on main plus this branch.

[Arsene12358]

@simone-chen one more input for this review, now that a complete op-level dataset for a NemotronH model exists on our side. A coverage gate runs aisimulate predict --detail source --format json on that dataset for the model at TP 1, 2 and 4 and classifies every operation by its source tag. On the current main the only entries it flags are the two decode Mamba2 operations, generation_mamba_conv1d and generation_mamba_ssm, which come back tagged sol in every configuration although the mamba2 table holds the decode rows. On a local build of main plus this branch (a013935) the same gate reports every operation served from the tables, with nothing flagged.

The branch is rebased on main and the CI is green, so the review is the only step left. Thanks for taking a look when you have a moment.