From 2ed6621836e98457c82ffb25222255e63f23615f Mon Sep 17 00:00:00 2001 From: JiaxinD Date: Sun, 27 Sep 2026 18:12:36 -0700 Subject: [PATCH 1/5] perf(llama): release prefill plan before decode build Signed-off-by: JiaxinD --- families/llama/model.py | 4 +- families/llama/tests/test_support.py | 99 ++++++++++++++++++++++++++++ 2 files changed, 102 insertions(+), 1 deletion(-) create mode 100644 families/llama/tests/test_support.py diff --git a/families/llama/model.py b/families/llama/model.py index f6a895ce40..036da5d8a1 100644 --- a/families/llama/model.py +++ b/families/llama/model.py @@ -247,6 +247,9 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: precision=precision, verbose=bool(request.verbose), ) + # Jiaxin Deng: stage the plan before decode compilation to avoid retaining both engines. + writer.add_bytes("prefill.plan", prefill) + del prefill config.raw["_decoder_engine_role"] = "decode" decode = _build_engine( config, @@ -256,7 +259,6 @@ def build(request: "BuildRequest", writer: "BundleWriter") -> None: verbose=bool(request.verbose), ) writer.add_bytes("engine.plan", decode) - writer.add_bytes("prefill.plan", prefill) layout = "split" config.raw.pop("_decoder_engine_role", None) diff --git a/families/llama/tests/test_support.py b/families/llama/tests/test_support.py new file mode 100644 index 0000000000..fdfbb9884b --- /dev/null +++ b/families/llama/tests/test_support.py @@ -0,0 +1,99 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU coverage of serialized-plan lifetime and atomic bundle publication.""" + +import gc +import importlib.util +import json +from pathlib import Path +import struct +import sys +import types + +import pytest + +from tensorrt_model_connect import BuildRequest +from tensorrt_model_connect.bundle_writer import BundleWriter + +@pytest.fixture +def request_and_writer(tmp_path, monkeypatch): + # Jiaxin Deng: stub only GPU compilation; exercise the real build and writer on CPU. + with monkeypatch.context() as imports: + for module_name, function_name in ( + ("dual_profile_decoder_builder", "build_dual_profile_decoder_engine"), + ("standard_decoder_builder", "build_standard_decoder_engine"), + ): + module = types.ModuleType(f"families.llama.{module_name}") + setattr(module, function_name, None) + imports.setitem(sys.modules, module.__name__, module) + spec = importlib.util.spec_from_file_location( + "families.llama._plan_lifetime_model", Path(__file__).parents[1] / "model.py" + ) + model = importlib.util.module_from_spec(spec) + spec.loader.exec_module(model) + checkpoint = tmp_path / "checkpoint" + checkpoint.mkdir() + (checkpoint / "config.json").write_text(json.dumps({ + "model_type": "llama", "hidden_size": 8, "intermediate_size": 16, + "num_hidden_layers": 1, "num_attention_heads": 2, + "num_key_value_heads": 1, "vocab_size": 32, + "max_position_embeddings": 32, "eos_token_id": 2, + })) + monkeypatch.setattr(model, "load_standard_weights", lambda *args, **kwargs: {}) + destination = tmp_path / "model.bundle" + request = BuildRequest( + model_dir=checkpoint, output_path=destination, family="llama", + task="text_generation", precision="fp16", max_sequence_length=16, + ) + writer = BundleWriter(destination) + yield model, request, writer + writer.abort() + + +def test_prefill_plan_is_released_before_decode_build(request_and_writer, monkeypatch): + model, request, writer = request_and_writer + released = [] + + class Plan(bytes): + def __del__(self): + released.append(True) + + def build_engine(config, *args, **kwargs): + if config.raw["_decoder_engine_role"] == "prefill": + return Plan(b"prefill bytes") + gc.collect() + assert released, "prefill plan is still resident during decode build" + return b"decode bytes" + + monkeypatch.setattr(model, "_build_engine", build_engine) + model.build(request, writer) + writer.finish() + data = request.output_path.read_bytes() + header_size = struct.unpack(" Date: Mon, 28 Sep 2026 09:06:09 -0700 Subject: [PATCH 2/5] perf(llama): release projection scratch between loads Signed-off-by: JiaxinD --- families/llama/checkpoint_mapper.py | 57 ++++++++----------------- families/llama/tests/test_support.py | 63 +++++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 41 deletions(-) diff --git a/families/llama/checkpoint_mapper.py b/families/llama/checkpoint_mapper.py index 7fb1fb8b58..eb1a944801 100644 --- a/families/llama/checkpoint_mapper.py +++ b/families/llama/checkpoint_mapper.py @@ -128,36 +128,23 @@ def _load_layer(layer_idx: int) -> tuple[int, WeightDict, int, int]: layer[f"{prefix}.input_norm"] = input_norm.astype(np.float32) layer[f"{prefix}.post_attn_norm"] = post_norm.astype(np.float32) - # Q/K/V/O projections - q_raw = _load_tensor( - readers, _layer_key(layer_idx, "self_attn.q_proj.weight", model_prefix) - ) - k_raw = _load_tensor( - readers, _layer_key(layer_idx, "self_attn.k_proj.weight", model_prefix) - ) - v_raw = _load_tensor( - readers, _layer_key(layer_idx, "self_attn.v_proj.weight", model_prefix) - ) - o_raw = _load_tensor( - readers, _layer_key(layer_idx, "self_attn.o_proj.weight", model_prefix) - ) - - q_hidden = q_raw.shape[0] - gate_raw = _load_tensor( - readers, _layer_key(layer_idx, "mlp.gate_proj.weight", model_prefix) - ) - layer_mlp_size = gate_raw.shape[0] - - # Transpose all projections [out, in] -> [in, out] - q_t = _transpose_2d(q_raw, "q_proj", precision=layer_precision) - k_t = _transpose_2d(k_raw, "k_proj", precision=layer_precision) - v_t = _transpose_2d(v_raw, "v_proj", precision=layer_precision) - o_t = _transpose_2d(o_raw, "o_proj", precision=layer_precision) - - layer[f"{prefix}.w_q"] = q_t - layer[f"{prefix}.w_k"] = k_t - layer[f"{prefix}.w_v"] = v_t - layer[f"{prefix}.w_o"] = o_t + # Jiaxin Deng: release each FP32 source before loading the next projection. + for source, target in ( + ("self_attn.q_proj", "w_q"), + ("self_attn.k_proj", "w_k"), + ("self_attn.v_proj", "w_v"), + ("self_attn.o_proj", "w_o"), + ("mlp.gate_proj", "w_gate"), + ("mlp.up_proj", "w_up"), + ("mlp.down_proj", "w_down"), + ): + raw = _load_tensor(readers, _layer_key(layer_idx, f"{source}.weight", model_prefix)) + layer[f"{prefix}.{target}"] = _transpose_2d(raw, source, precision=layer_precision) + if target == "w_q": + q_hidden = raw.shape[0] + elif target == "w_gate": + layer_mlp_size = raw.shape[0] + del raw # Optional QKV biases (Qwen2 style) q_bias_key = _layer_key(layer_idx, "self_attn.q_proj.bias", model_prefix) @@ -182,16 +169,6 @@ def _load_layer(layer_idx: int) -> tuple[int, WeightDict, int, int]: _load_tensor(readers, k_norm_key).astype(np.float32), num_kv_heads ) - # MLP projections - up_raw = _load_tensor(readers, _layer_key(layer_idx, "mlp.up_proj.weight", model_prefix)) - down_raw = _load_tensor( - readers, _layer_key(layer_idx, "mlp.down_proj.weight", model_prefix) - ) - - layer[f"{prefix}.w_gate"] = _transpose_2d(gate_raw, "gate_proj", precision=layer_precision) - layer[f"{prefix}.w_up"] = _transpose_2d(up_raw, "up_proj", precision=layer_precision) - layer[f"{prefix}.w_down"] = _transpose_2d(down_raw, "down_proj", precision=layer_precision) - return layer_idx, layer, q_hidden, layer_mlp_size layer_results: list[tuple[int, WeightDict, int, int] | None] = [None] * num_layers diff --git a/families/llama/tests/test_support.py b/families/llama/tests/test_support.py index fdfbb9884b..97151862e7 100644 --- a/families/llama/tests/test_support.py +++ b/families/llama/tests/test_support.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""CPU coverage of serialized-plan lifetime and atomic bundle publication.""" +"""CPU coverage of build memory lifetime and bundle publication.""" import gc import importlib.util @@ -10,12 +10,18 @@ import struct import sys import types +import weakref +import numpy as np import pytest +from safetensors.numpy import save_file from tensorrt_model_connect import BuildRequest from tensorrt_model_connect.bundle_writer import BundleWriter +from .. import checkpoint_mapper +from ..config import ModelConfig + @pytest.fixture def request_and_writer(tmp_path, monkeypatch): # Jiaxin Deng: stub only GPU compilation; exercise the real build and writer on CPU. @@ -97,3 +103,58 @@ def build_engine(config, *args, **kwargs): writer.abort() assert request.output_path.read_bytes() == b"previous bundle" assert not list(request.output_path.parent.glob(".model.bundle.sections.*")) + + +@pytest.mark.parametrize("precision,fp32_layers", [("fp16", ()), ("fp32", ()), ("fp16", (1,))]) +def test_projection_scratch_is_released_without_changing_weights( + tmp_path, monkeypatch, precision, fp32_layers +): + config = ModelConfig( + model_type="llama", hidden_size=4, vocab_size=8, num_hidden_layers=2, + num_attention_heads=2, num_key_value_heads=1, intermediate_size=6, + ) + tensors = {"model.embed_tokens.weight": np.arange(32, dtype=np.float16).reshape(8, 4)} + shapes = { + "self_attn.q_proj": (4, 4), "self_attn.k_proj": (2, 4), + "self_attn.v_proj": (2, 4), "self_attn.o_proj": (4, 4), + "mlp.gate_proj": (6, 4), "mlp.up_proj": (6, 4), "mlp.down_proj": (4, 6), + } + output_names = ("w_q", "w_k", "w_v", "w_o", "w_gate", "w_up", "w_down") + expected = {} + for index in range(2): + for projection_index, (source, output) in enumerate(zip(shapes, output_names)): + shape = shapes[source] + values = ( + np.arange(np.prod(shape), dtype=np.float32).reshape(shape) / 7 + + index * 3 + projection_index / 9 - 2 + ).astype(np.float16) + tensors[f"model.layers.{index}.{source}.weight"] = values + dtype = np.float32 if precision == "fp32" or index in fp32_layers else np.float16 + expected[f"layer.{index}.{output}"] = values.astype(np.float32).T.astype(dtype) + for name in ("input_layernorm", "post_attention_layernorm"): + tensors[f"model.layers.{index}.{name}.weight"] = np.ones(4, dtype=np.float16) + save_file(tensors, tmp_path / "model.safetensors") + load_tensor = checkpoint_mapper._load_tensor + previous = {} + + def tracked_load(readers, name): + if ".layers." in name and name.endswith("_proj.weight"): + layer = name.split(".")[2] + if layer in previous: + assert previous[layer]() is None, "previous projection scratch is still live" + value = load_tensor(readers, name) + previous[layer] = weakref.ref(value) + return value + return load_tensor(readers, name) + + monkeypatch.setattr(checkpoint_mapper, "_load_tensor", tracked_load) + weights = checkpoint_mapper.load_standard_weights( + tmp_path, config, precision=precision, fp32_layers=fp32_layers + ) + for name, values in expected.items(): + np.testing.assert_array_equal(weights[name], values) + assert weights[name].dtype == values.dtype + assert weights[name].flags.c_contiguous + assert weights["_attention_size"] == 4 + assert weights["_kv_attention_size"] == 2 + assert weights["_mlp_size"] == 6 From b2f1dc39a5fa90da613dfbaaf89245c9a6b1b256 Mon Sep 17 00:00:00 2001 From: JiaxinD Date: Mon, 28 Sep 2026 18:06:34 -0700 Subject: [PATCH 3/5] perf(llama): release embedding scratch before layers Signed-off-by: JiaxinD --- families/llama/checkpoint_mapper.py | 6 +++++- families/llama/tests/test_support.py | 21 +++++++++++++++++++-- 2 files changed, 24 insertions(+), 3 deletions(-) diff --git a/families/llama/checkpoint_mapper.py b/families/llama/checkpoint_mapper.py index eb1a944801..17a74ff53b 100644 --- a/families/llama/checkpoint_mapper.py +++ b/families/llama/checkpoint_mapper.py @@ -111,6 +111,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}" @@ -209,7 +211,9 @@ def _load_layer(layer_idx: int) -> tuple[int, WeightDict, int, int]: ) else: # Tied embeddings - weights["w_out"] = _transpose_2d(embedding.copy(), "embedding_tied", precision=precision) + weights["w_out"] = _transpose_2d( + weights["embedding"].copy(), "embedding_tied", precision=precision + ) weights["_attention_size"] = attention_size # type: ignore[assignment] weights["_kv_attention_size"] = kv_attention_size # type: ignore[assignment] diff --git a/families/llama/tests/test_support.py b/families/llama/tests/test_support.py index 97151862e7..762eb964b7 100644 --- a/families/llama/tests/test_support.py +++ b/families/llama/tests/test_support.py @@ -106,14 +106,18 @@ def build_engine(config, *args, **kwargs): @pytest.mark.parametrize("precision,fp32_layers", [("fp16", ()), ("fp32", ()), ("fp16", (1,))]) +@pytest.mark.parametrize("tied", [False, True]) def test_projection_scratch_is_released_without_changing_weights( - tmp_path, monkeypatch, precision, fp32_layers + tmp_path, monkeypatch, precision, fp32_layers, tied ): config = ModelConfig( model_type="llama", hidden_size=4, vocab_size=8, num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1, intermediate_size=6, ) - tensors = {"model.embed_tokens.weight": np.arange(32, dtype=np.float16).reshape(8, 4)} + embedding = np.arange(32, dtype=np.float32).reshape(8, 4) / 7 - 2 + tensors = {"model.embed_tokens.weight": embedding} + if not tied: + tensors["lm_head.weight"] = embedding + 3 shapes = { "self_attn.q_proj": (4, 4), "self_attn.k_proj": (2, 4), "self_attn.v_proj": (2, 4), "self_attn.o_proj": (4, 4), @@ -136,8 +140,15 @@ def test_projection_scratch_is_released_without_changing_weights( save_file(tensors, tmp_path / "model.safetensors") load_tensor = checkpoint_mapper._load_tensor previous = {} + embedding_source = None def tracked_load(readers, name): + nonlocal embedding_source + if name == "model.embed_tokens.weight": + value = load_tensor(readers, name) + embedding_source = weakref.ref(value) + return value + assert embedding_source() is None, "embedding scratch is still live" if ".layers." in name and name.endswith("_proj.weight"): layer = name.split(".")[2] if layer in previous: @@ -151,6 +162,12 @@ def tracked_load(readers, name): weights = checkpoint_mapper.load_standard_weights( tmp_path, config, precision=precision, fp32_layers=fp32_layers ) + dtype = np.float32 if precision == "fp32" else np.float16 + np.testing.assert_array_equal(weights["embedding"], embedding.astype(dtype)) + output = embedding if tied else tensors["lm_head.weight"] + np.testing.assert_array_equal(weights["w_out"], output.T.astype(dtype)) + assert weights["w_out"].dtype == dtype + assert weights["w_out"].flags.c_contiguous for name, values in expected.items(): np.testing.assert_array_equal(weights[name], values) assert weights[name].dtype == values.dtype From 1154d0a7b7a44a02da38e54c9262b9a792c5893b Mon Sep 17 00:00:00 2001 From: JiaxinD Date: Tue, 29 Sep 2026 09:06:13 -0700 Subject: [PATCH 4/5] fix(llama): resolve staged E2E checkpoints offline Reuse the checkpoint-cache prerequisite from #1424 so Dev CI can validate the host-memory changes without Hub metadata access. Preserve exact revisions and config checks. Signed-off-by: JiaxinD --- families/llama/tests/test_checkpoint_cache.py | 44 +++++++++++++++++++ families/llama/tests/test_e2e.py | 1 + 2 files changed, 45 insertions(+) create mode 100644 families/llama/tests/test_checkpoint_cache.py diff --git a/families/llama/tests/test_checkpoint_cache.py b/families/llama/tests/test_checkpoint_cache.py new file mode 100644 index 0000000000..9f765a0ca5 --- /dev/null +++ b/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}) diff --git a/families/llama/tests/test_e2e.py b/families/llama/tests/test_e2e.py index fb32605d8d..9610758321 100644 --- a/families/llama/tests/test_e2e.py +++ b/families/llama/tests/test_e2e.py @@ -113,6 +113,7 @@ def _checkpoint(manifest: dict) -> Path: snapshot_download( repo_id=manifest["hf_id"], revision=manifest.get("hf_revision"), + local_files_only=True, ) ) assert (path / "config.json").is_file(), path From d9ef21e2b2222b4f4262587b603b88660813ed88 Mon Sep 17 00:00:00 2001 From: JiaxinD Date: Tue, 20 Oct 2026 09:04:25 -0700 Subject: [PATCH 5/5] perf(llama): avoid tied-head input copy Reuse the stored embedding for transposition and copy only an aliased singleton output. Retain independent contiguous storage, target values and canonical validation while avoiding the full input allocation. Signed-off-by: JiaxinD --- families/llama/checkpoint_mapper.py | 4 +- families/llama/tests/test_support.py | 67 ++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+), 1 deletion(-) diff --git a/families/llama/checkpoint_mapper.py b/families/llama/checkpoint_mapper.py index 302c8e86d5..00f5a10bc2 100644 --- a/families/llama/checkpoint_mapper.py +++ b/families/llama/checkpoint_mapper.py @@ -229,8 +229,10 @@ def _projection(stem: str, hf_name: str) -> np.ndarray: else: # Tied embeddings weights["w_out"] = _transpose_2d( - weights["embedding"].copy(), "embedding_tied", precision=precision + 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] diff --git a/families/llama/tests/test_support.py b/families/llama/tests/test_support.py index 8eab3c06d7..0bacb43246 100644 --- a/families/llama/tests/test_support.py +++ b/families/llama/tests/test_support.py @@ -9,6 +9,7 @@ from pathlib import Path import struct import sys +import tracemalloc import types from types import SimpleNamespace import weakref @@ -281,3 +282,69 @@ def test_quantized_placeholders_skip_projection_reads(tmp_path, precision, fp32_ 2, 6, ) + + +def _write_tied_checkpoint(path, vocab, hidden): + config = ModelConfig( + model_type="llama", + hidden_size=hidden, + vocab_size=vocab, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + intermediate_size=hidden, + ) + embedding = (np.arange(vocab * hidden, dtype=np.float32).reshape(vocab, hidden) % 29 - 14) / 7 + tensors = {"model.embed_tokens.weight": embedding} + for norm in ("input_layernorm", "post_attention_layernorm"): + tensors[f"model.layers.0.{norm}.weight"] = np.ones(hidden, dtype=np.float32) + for projection in ( + "self_attn.q_proj", + "self_attn.k_proj", + "self_attn.v_proj", + "self_attn.o_proj", + "mlp.gate_proj", + "mlp.up_proj", + "mlp.down_proj", + ): + tensors[f"model.layers.0.{projection}.weight"] = np.zeros( + (hidden, hidden), dtype=np.float32 + ) + save_file(tensors, path / "model.safetensors") + return config, embedding + + +@pytest.mark.parametrize("precision", ["fp16", "bf16", "fp32"]) +@pytest.mark.parametrize("vocab,hidden", [(8, 4), (8, 1), (1, 4), (1, 1)]) +def test_tied_output_owns_contiguous_storage(tmp_path, precision, vocab, hidden): + config, source = _write_tied_checkpoint(tmp_path, vocab, hidden) + weights = checkpoint_mapper.load_standard_weights(tmp_path, config, precision=precision) + dtype = np.float32 if precision == "fp32" else np.float16 + np.testing.assert_array_equal(weights["w_out"], source.astype(dtype).T) + assert weights["w_out"].dtype == dtype + assert weights["w_out"].flags.c_contiguous + assert not np.shares_memory(weights["w_out"], weights["embedding"]) + original = weights["embedding"].copy() + weights["w_out"][0, 0] = 99 + np.testing.assert_array_equal(weights["embedding"], original) + + +def test_tied_output_allocation_does_not_copy_input(tmp_path, monkeypatch): + config, _ = _write_tied_checkpoint(tmp_path, 4096, 256) + has_tensor = checkpoint_mapper._has_tensor + + def begin_tracking(readers, name): + found = has_tensor(readers, name) + if name == "lm_head.weight": + assert not found + tracemalloc.start() + return found + + monkeypatch.setattr(checkpoint_mapper, "_has_tensor", begin_tracking) + try: + weights = checkpoint_mapper.load_standard_weights(tmp_path, config, precision="fp16") + _, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + assert weights["w_out"].nbytes == 2 * 4096 * 256 + assert peak < weights["w_out"].nbytes * 1.5