[milesial]

Why and what changed

NemotronH inherits attention KV accounting that multiplies by every layer in the model. Its hybrid stack includes Mamba, MoE, and MLP layers that do not allocate attention KV. This overstates the token-dependent memory cost and understates KV token capacity.

Count only attention layers from the hybrid pattern while preserving KV-head sharding and replication across tensor-parallel ranks. For the bundled Nemotron 3 Super configuration at TP1, this changes FP8 attention KV from 45,056 to 4,096 bytes per token: 8 attention layers × 2 KV heads × 128 elements × K/V.

Review map

• Risk level: medium.
• Start with NemotronHModel.get_kvcache_elements_per_token.
• Public or serialized contract changed: no schema change. NemotronH attention KV estimates and derived token capacities change.
• Compatibility or rollback concern: Mamba recurrent-state reservation remains separate. This does not add a total hybrid-memory estimator.

Evidence

• Regression coverage exercises TP1, TP2, and TP4 with FP8 and BF16 KV, including KV-head replication at TP4.
• Assertions cover element counts, sequence bytes, and the inverse byte-budget-to-token calculation.
• Expected values come from the bundled model geometry: eight attention layers, two KV heads, and a head dimension of 128.
• CI results are pending.

Modeling or data provenance

Attention tensor geometry from the bundled nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8 configuration. This is a memory-accounting correction, not a measured latency calibration.

[coderabbitai]

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: 5ec25265-7f6a-4eb2-a5da-460616180f89

📥 Commits

Reviewing files that changed from the base of the PR and between a0a35e7 and 84c34f4.

📒 Files selected for processing: 2

• python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py
• python/aisimulate/tests/unit/sdk/models/test_model_config.py

🔗 Linked repositories identified

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

• ai-dynamo/dynamo (manual)
• ai-dynamo/aiconfigurator (manual)

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

📜 Recent review details

🧰 Additional context used

📚 Code guidelines: 5

python/aisimulate/.claude/rules/rust-core/parity.md — configured
python/aisimulate/.claude/rules/perfmodel-api.md — auto-discovered
AGENTS.md — configured
REVIEW.md — configured
DEVELOPMENT.md — auto-discovered

📓 Path-based instructions: 7

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/models/nemotron_h.py

Read REVIEW.md before commenting.

⚙️ CodeRabbit configuration file

Files:

• python/aisimulate/tests/unit/sdk/models/test_model_config.py
• python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py

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/models/nemotron_h.py

Applicability reference: AGENTS.md: # AGENTS Before changing a performance model, its configuration, or a caller in Rust, Python, CLI, Sweeper, Replay, or Planner, MUST read and follow [`perfmodel-api.md`](python/aisimulat...

📄 CodeRabbit inference engine (python/aisimulate/.claude/rules/perfmodel-api.md)

Files:

• python/aisimulate/tests/unit/sdk/models/test_model_config.py

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/tests/unit/sdk/models/test_model_config.py
• python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py

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/tests/unit/sdk/models/test_model_config.py
• python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py

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/models/nemotron_h.py

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

Linked repositories findings

ai-dynamo/aiconfigurator

• The frozen Nemotron-H implementation explicitly classifies M (Mamba), E (MoE), * (attention), and - (MLP) layers via _count_layer_types(). [::ai-dynamo/aiconfigurator::aic-core/src/aiconfigurator_core/sdk/models/nemotron_h.py:15-23] [::ai-dynamo/aiconfigurator::aic-core/src/aiconfigurator_core/sdk/models/nemotron_h.py:87-95]
• Only * layers instantiate ContextAttention and GenerationAttention; Mamba and MoE layers use separate operation paths. This corroborates counting attention layers only for attention KV sizing. [::ai-dynamo/aiconfigurator::aic-core/src/aiconfigurator_core/sdk/models/nemotron_h.py:183-205] [::ai-dynamo/aiconfigurator::aic-core/src/aiconfigurator_core/sdk/models/nemotron_h.py:476-496]
• Nemotron-H configuration supports both hybrid pattern strings and explicit layer-block types, so sizing logic should remain compatible with both parsed forms. [::ai-dynamo/aiconfigurator::tests/unit/sdk/test_utils.py:118-159] [::ai-dynamo/aiconfigurator::tests/unit/sdk/test_utils.py:193-225]

ai-dynamo/dynamo

• Dynamo’s AISimulate estimator derives maximum KV tokens from the model’s maximum batch size and (isl + osl). Any Nemotron-H per-token memory correction therefore changes the reported planner capacity, even though Dynamo does not directly call the new method. [::ai-dynamo/dynamo::components/src/dynamo/planner/monitoring/aic_estimator.py:220-240]
• DGD materialization independently writes the selected KV-cache settings and max_num_batched_tokens into runtime configuration, so changed capacity recommendations affect generated deployment arguments rather than a public Dynamo schema. [::ai-dynamo/dynamo::components/src/dynamo/aisimulate/output/dgd/renderers/direct/materializer.py:157-169]
• Existing DGD tests cover backend KV-cache defaults, fixed-block overrides, and mutually exclusive capacity settings; these are configuration-contract tests, not hard-coded Nemotron-H byte expectations. [::ai-dynamo/dynamo::components/src/dynamo/aisimulate/tests/dgd/test_renderer_direct.py:189-264]

📝 Summary

Risk: Medium. The change alters NemotronH KV-memory estimates and derived token capacity.

Human attention should focus on:

1. Correctly counting only attention layers in the hybrid NemotronH pattern.
2. Preserving KV-head sharding and replication across TP1, TP2, and TP4.
3. Confirming that Mamba state remains outside the attention KV estimate.

Changed behavior and contracts

• Adds NemotronHModel.get_kvcache_elements_per_token().
• The method estimates attention KV elements per token from attention-layer count, TP-adjusted KV heads, head size, and K/V storage.
• Mamba, MoE, and MLP layers are excluded from the attention KV calculation.
• NemotronH attention KV estimates and derived token capacities change.
• The public model API gains the new method.
• No schema change is reported.
• Mamba recurrent-state reservation remains separate.

For the bundled Nemotron 3 Super configuration at TP1, the reported FP8 attention KV estimate changes from 45,056 to 4,096 bytes per token. The estimate uses eight attention layers, two KV heads, 128 elements per head, and K/V storage.

Evidence supplied

• Regression coverage is reported for TP1, TP2, and TP4.
• Coverage includes FP8 and BF16 KV-cache dtypes.
• Tests reportedly verify per-token element counts, dtype-based byte scaling, and inverse byte-budget-to-token calculations.
• TP4 coverage includes KV-head replication.

Evidence still missing

• CI results are pending.
• The supplied evidence does not establish merge readiness.
• No current review severity counts were supplied.
• No total hybrid-memory estimator was added.

Nemotron-H now estimates KV-cache elements per token from attention-layer count, tensor-parallel KV-head allocation, and head size. Regression tests cover multiple tensor-parallel sizes, cache dtypes, byte scaling, and token capacity.

Changes: Nemotron-H KV-cache sizing

Layer / File(s)	Summary
KV-cache calculation and validation python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py, python/aisimulate/tests/unit/sdk/models/test_model_config.py	NemotronHModel.get_kvcache_elements_per_token calculates attention KV-cache elements per token. Tests cover tensor-parallel sizes 1, 2, and 4, FP8 and bfloat16 cache dtypes, byte scaling, and capacity for a 64-token budget.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 84c34

The Nemotron-H KV estimate now counts only attention KV storage and preserves the intended tensor-parallel behavior; no actionable merge risk remains.

Pre-merge checks: 7 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Cross-Layer Contract		The native Python path is updated and covered, but the alternative public KV-estimation path remains stale. NemotronHModel.get_kvcache_elements_per_token now counts pattern.count("*") at `python/a…	Update NaiveKVCacheEstimator to carry NemotronH attention-layer geometry from hybrid_override_pattern or layers_block_type and use that count for standard KV bytes and its inverse. Cover the forced `estimate_kv_cache(..., allow_naive_…

Passed checks: 7 passed

Check name	Status	Explanation
Title check		The title precisely states the behavioral change: NemotronH KV-memory accounting now counts only attention layers.
Description check		The description covers the problem, behavior change, review map, contract impact, compatibility concern, regression coverage, expected-value derivation, and modeling provenance. It omits the Tracking …
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		The PR supplies an explained and reproducible golden diff. The prior base formula used all layers (base.py:566-589), while the changed NemotronHModel formula counts "*" attention layers and pres…
Compatibility Boundaries		The change stays within the Python model implementation and its unit test. The reviewed diff changes only python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py and `python/aisimulate/tests/…
Review Evidence		PASS. The PR description states that CI results are pending and does not cite any inactive nested workflow as hosted CI evidence. It describes TP1/TP2/TP4 regression coverage and boundary-style KV-hea…

Full details: Description check

Explanation

The description covers the problem, behavior change, review map, contract impact, compatibility concern, regression coverage, expected-value derivation, and modeling provenance. It omits the Tracking section and does not provide exact test commands, CI results, reviewed commits, or explicit before/after output.

Full details: Modeling And Data Evidence

Explanation

The PR supplies an explained and reproducible golden diff. The prior base formula used all layers (base.py:566-589), while the changed NemotronHModel formula counts "*" attention layers and preserves TP KV-head rounding (nemotron_h.py:97-101). The PR description gives the before/after FP8 values and geometry. The added parameterized test records expected element counts for TP1/2/4 and both FP8 and BF16, then checks byte scaling and the inverse token budget (test_model_config.py:912-927). The bundled config provenance is also machine-readable: it specifies head dimension 128, eight * layers, and two KV heads (aiconfigurator commit f254959eb89e2f206b8f9a77051644d7c1cbdb89, aic-core/src/aiconfigurator_core/model_configs/nvidia--NVIDIA-Nemotron-3-Super-120B-A12B-FP8_config.json:16-20,38-40). This is not binary-only validation or an unexplained refresh.

Full details: Compatibility Boundaries

Explanation

The change stays within the Python model implementation and its unit test. The reviewed diff changes only python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py and python/aisimulate/tests/unit/sdk/models/test_model_config.py; it changes no Rust source, PyO3 binding, schema, version, or manifest. NemotronHModel already has a public package re-export at models/__init__.py:313-325, and the new override follows the existing BaseModel KV-cache method contract at models/base.py:566-594. The changed module imports only aisimulate_core and has no Dynamo dependency. Packaging remains unified: pyproject.toml:113-123 points the wheel build to crates/core and includes both Python packages; the release verifier still enforces one wheel and one Rust crate at verify_release_wheels.py:87-91 and 130-152. Versions remain 0.13.0 in the Python and Rust manifests.

Full details: Review Evidence

Explanation

PASS. The PR description states that CI results are pending and does not cite any inactive nested workflow as hosted CI evidence. It describes TP1/TP2/TP4 regression coverage and boundary-style KV-head replication checks, with no false hosted-CI claim.

Full details: Cross-Layer Contract

Explanation

The native Python path is updated and covered, but the alternative public KV-estimation path remains stale. NemotronHModel.get_kvcache_elements_per_token now counts pattern.count("*") at python/aisimulate/src/aisimulate_core/sdk/models/nemotron_h.py:97-101, and the added test covers get_model, sequence bytes, and inverse capacity at python/aisimulate/tests/unit/sdk/models/test_model_config.py:912-927. However, NaiveKVCacheEstimator.kv_bytes_per_token still computes standard KV as 2 * ceil(num_kv_heads / tp) * head_dim * layers * dtype_bytes at python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py:832-863. Its parsed geometry stores total layers and does not retain the Nemotron hybrid attention-layer count. estimate_kv_cache can select this estimator when the native build fails and fallback is enabled at python/aisimulate/src/aisimulate_core/sdk/memory/kv_cache.py:1314-1349. The fallback tests cover generic fixtures but no Nemotron case. Therefore a CLI/API caller using the documented naive fallback still receives the old all-layer estimate.

Resolution

Update NaiveKVCacheEstimator to carry NemotronH attention-layer geometry from hybrid_override_pattern or layers_block_type and use that count for standard KV bytes and its inverse. Cover the forced estimate_kv_cache(..., allow_naive_fallback=True) path for Nemotron at TP1, TP2, and TP4 with FP8 and BF16 scaling. Add coverage for any CLI or capacity consumer that uses the naive estimator, or document and enforce that such consumers cannot estimate NemotronH.

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

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

[tedzhouhk]

Reviewed 84c34f45; no new correctness issue found. Counting attention layers from the normalized hybrid pattern fixes attention-KV sizing while preserving KV-head sharding/replication (including TP greater than the KV-head count).

Validation on the exact head: 547 model/config/memory tests passed, plus 72 independent model/topology/dtype cases and 12 public estimate_kv_cache() cases. The cited Super 120B FP8 model gives 4,096 bytes/token at TP1 and 2,048 at TP2/TP4, with the inverse capacity calculation consistent.

The 11x correction is attention KV bytes/token for that model, not total GPU memory. Mamba recurrent-state sizing remains separate; this PR does not add automatic Nemotron state sizing. No end-to-end GPU accuracy validation was performed.