perf(llama): reduce build host-memory retention

families/llama/checkpoint_mapper.py

@@ -116,6 +116,8 @@ def load_standard_weights(
         f"Embedding shape {embedding.shape} != ({vocab}, {hidden})"
     )
     weights["embedding"] = embedding.astype(target_dtype)
+    # Jiaxin Deng: keep only the stored embedding while loading decoder layers.
+    del embedding
 
     def _load_layer(layer_idx: int) -> tuple[int, WeightDict, int, int]:
         prefix = f"layer.{layer_idx}"
@@ -226,7 +228,11 @@ def _projection(stem: str, hf_name: str) -> np.ndarray:
         )
     else:
         # Tied embeddings
-        weights["w_out"] = _transpose_2d(embedding.copy(), "embedding_tied", precision=precision)
+        weights["w_out"] = _transpose_2d(
+            weights["embedding"], "embedding_tied", precision=precision
+        )
+        if np.shares_memory(weights["w_out"], weights["embedding"]):
+            weights["w_out"] = weights["w_out"].copy()
 
     weights["_attention_size"] = attention_size  # type: ignore[assignment]
     weights["_kv_attention_size"] = kv_attention_size  # type: ignore[assignment]

families/llama/model.py

@@ -283,6 +283,9 @@ def _build_native(request: "BuildRequest", writer: "BundleWriter") -> None:
                 verbose=bool(request.verbose),
                 quant_ctx=quant_ctx,
             )
+            # Note (Jiaxin Deng): avoid retaining both plans during decode compilation.
+            writer.add_bytes("prefill.plan", prefill)
+            del prefill
             config.raw["_decoder_engine_role"] = "decode"
             decode = _build_engine(
                 config,
@@ -293,7 +296,6 @@ def _build_native(request: "BuildRequest", writer: "BundleWriter") -> None:
                 quant_ctx=quant_ctx,
             )
             writer.add_bytes("engine.plan", decode)
-            writer.add_bytes("prefill.plan", prefill)
             layout = "split"
         config.raw.pop("_decoder_engine_role", None)
 

families/llama/tests/test_checkpoint_cache.py

@@ -0,0 +1,44 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Exercise staged checkpoint lookup without Hub metadata or network access."""
+
+import pytest
+from huggingface_hub import constants
+from huggingface_hub.errors import LocalEntryNotFoundError
+
+from .test_e2e import _checkpoint
+
+
+@pytest.fixture
+def offline_cache(tmp_path, monkeypatch):
+    monkeypatch.setattr(constants, "HF_HUB_CACHE", str(tmp_path))
+    monkeypatch.setattr(constants, "HF_HUB_OFFLINE", True)
+    return tmp_path
+
+
+def test_pinned_checkpoint_uses_snapshot_without_cached_tree(offline_cache):
+    revision = "a" * 40
+    snapshot = offline_cache / "models--trtmc-test--llama" / "snapshots" / revision
+    snapshot.mkdir(parents=True)
+    (snapshot / "config.json").write_text("{}", encoding="utf-8")
+
+    assert _checkpoint({"hf_id": "trtmc-test/llama", "hf_revision": revision}) == snapshot
+
+
+def test_missing_checkpoint_is_not_replaced_by_another_revision(offline_cache):
+    cached = offline_cache / "models--trtmc-test--llama" / "snapshots" / ("a" * 40)
+    cached.mkdir(parents=True)
+    (cached / "config.json").write_text("{}", encoding="utf-8")
+
+    with pytest.raises(LocalEntryNotFoundError):
+        _checkpoint({"hf_id": "trtmc-test/llama", "hf_revision": "b" * 40})
+
+
+def test_staged_snapshot_still_requires_config(offline_cache):
+    revision = "a" * 40
+    snapshot = offline_cache / "models--trtmc-test--llama" / "snapshots" / revision
+    snapshot.mkdir(parents=True)
+
+    with pytest.raises(AssertionError):
+        _checkpoint({"hf_id": "trtmc-test/llama", "hf_revision": revision})