perf(llama): reduce build host-memory retention

[JiaxinD]

Background

Llama compilation retains avoidable host allocations: FP32 embedding/projection sources during layer loading and serialized prefill bytes during decode compilation. These lifetimes create memory pressure without contributing to engine computation.

Exit Criteria

Release each temporary source before the next load and stage/release prefill bytes before decode compilation. Preserve mapped weights, named engine payloads and atomic publication. No GPU performance or whole-process RSS claim.

Implementation

• Release the embedding source after storage conversion and each projection source after transposition within existing layer workers.
• Stage prefill.plan before decode compilation and release its in-memory bytes.
• Merge upstream 41c5552c, applying the lifetime change inside the new native-builder callback while retaining complete Edge/paired dispatch. Keep the upstream mode-aware offline checkpoint resolver and our cache regressions.
• Changes remain in families/llama. Physical bundle section order changes, but names, payloads, reading by offset, ABI and bundle format remain unchanged.

Change categories

• Model or runtime behavior

Validation

Commands and Results

Command	Result
python -m pytest families/llama/tests/test_support.py families/llama/tests/test_checkpoint_cache.py core/builder/tests/test_bundle_writer.py core/builder/tests/test_bundle_reader.py tools/tests/test_architecture.py tools/tests/test_family_impact.py -q -p no:cacheprovider	105 passed
python -m tools.model_ci validate, python -m tools.test_impact --validate, focused Ruff and git diff --check	Passed

Existing regressions exercise the real BundleWriter/safetensors paths with compilation stubs: release before decode, unchanged distinct weights/overrides/tied embeddings, and retention of the published bundle after decode failure. The lifetime and cache regressions previously failed before their fixes; this run checks integration with the new upstream builder.

Earlier synthetic CPU tracemalloc experiments reduced the two-plan peak from 33567185 to 16789012 bytes and the projection fixture from 16543495 to 10769185 bytes, with identical mapped-output SHA256. Those historical allocation experiments are not real-model RSS/GPU benchmarks for this merged head.

Hardware, Environment, and Revisions

Head 2df0010a87f39b56496349256a430553bcf73ef2, tested tree 4726bb5216dac054b29513d20d5f0116cecbfb90, based on upstream 41c5552ce998f6d46a083e15da105258a78778c7. Linux/WSL, Python 3.12, CPU only. Commands use PYTHONPATH=core/builder:apps/benchmark:.. No model weights or datasets downloaded.

Not Run / Remaining Gaps

No local real-checkpoint TensorRT compilation, GPU parity, throughput/latency or whole-process RSS measurement. Native Edge/speculative execution was retained from upstream, not newly hardware-qualified here. Previous-head Stable/Dev passed; new-head CI remains separate and must be checked. Radix authentication currently blocks GPU work.

Contributor Self-Review

• I have completed a self-review of this change.

Notes For Future Readers

Converted outputs remain owned by the weight dictionary. BundleWriter stages synchronously and its existing caller aborts temporary sections on failure. The synthetic tests verify numerical conversion and release timing; they do not establish GPU resource/performance equivalence. Scoped implementation and applicable CPU regression checks are complete; Draft remains due to the repository's unchanged Ready permission/capacity restriction.

Risk level

• Low

Allocation lifetimes and physical section order change. Numerical conversions, engine graphs and names remain unchanged, with focused output/publication coverage.

[coderabbitai]

Navigate logical layers of code changes, visualize relationships, and explore their blast radius.

🧰 Additional context used

📚 Code guidelines (1)

REVIEW.md — configured

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/TensorRT-Model-Connect/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: 23ee9aa2-6f0b-43b9-b350-d3bd623f3e14

📥 Commits

Reviewing files that changed from the base of the PR and between 41c5552 and 2df0010.

📒 Files selected for processing (4)

• families/llama/checkpoint_mapper.py
• families/llama/model.py
• families/llama/tests/test_checkpoint_cache.py
• families/llama/tests/test_support.py

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

---

📝 Walkthrough

⚠️ A high-level summary could not be generated for this review. CodeRabbit will regenerate it on the next update, or you can request a refresh with @coderabbitai summary.

Walkthrough

The Llama changes add offline checkpoint-cache tests, release temporary tensors during checkpoint loading, and publish the prefill plan before compiling the decode plan.

Changes

Llama checkpoint and build resource handling

Layer / File(s)	Summary
Offline checkpoint-cache resolution tests families/llama/tests/test_checkpoint_cache.py	Adds offline tests for resolving a pinned revision to a staged snapshot, rejecting a missing requested revision, and rejecting a snapshot without config.json.
Checkpoint tensor loading and scratch release families/llama/checkpoint_mapper.py, families/llama/tests/test_support.py	The loader deletes embedding and projection source arrays after storing converted weights. Tests check scratch-array release and verify loaded weight values, dtypes, contiguity, and attention dimensions across the tested configurations.
Split-engine plan publication families/llama/model.py, families/llama/tests/test_support.py	The builder writes the prefill plan and deletes its in-memory value before compiling the decode plan. Tests check plan publication and verify that a decode-build failure preserves the existing bundle and removes temporary section files.

Priority: ⬇️ Low

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

Change: Refactor

Suggested reviewers: yifeif-nv

Merge Risk: ⚪ Minimal · up to 2df00

The reviewed changes have no established merge-blocking issue. The reported tests and checks passed, though real-checkpoint TensorRT compilation was not run for this head.

🚥 Pre-merge checks | ✅ 8 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 18.75% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 16 functions across 4 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (8 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: reducing host-memory retention during Llama builds.
Description check	✅ Passed	The description covers the motivation, exit criteria, implementation, change category, validation, environment, remaining gaps, self-review, and risk rationale. It reports the relevant limitations and…
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.
Family Ownership Boundary	✅ Passed	No family-boundary failure was introduced. The reviewed diff changes only four files under families/llama; it changes no other family or central registry. The new checkpoint test imports `_checkpoin…
Shared Semantic Neutrality	✅ Passed	The check does not identify a shared-code change. The reviewed diff changes only families/llama/checkpoint_mapper.py, families/llama/model.py, and Llama tests. These are model-owned Python and fam…
Benchmark Validation Integrity	✅ Passed	The PR changes Llama checkpoint and plan memory lifetimes. The reviewed diff changes no benchmark, reference, metric, gate, workload, or report implementation. The added tests check plan release and b…
Shared Change Blast Radius	✅ Passed	The check is not triggered. The authoritative diff changes only families/llama/checkpoint_mapper.py, families/llama/model.py, and Llama tests. The production changes adjust Llama checkpoint loadin…

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

---

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