fix(phi4): stage engine plans and export AIPerf tokenizer

families/phi4_multimodal/model.py

@@ -343,6 +343,9 @@ def native(request: "BuildRequest", writer: "BundleWriter") -> None:
         config.raw["_model_dir"] = str(model_dir)
         model = _Phi4MultimodalModel()
         weights = model.load_weights(str(model_dir), config)
+        writer.set_header(family="phi4_multimodal", task=request.task, backend=request.backend)
+        # Stage each serialized plan before building the next network. The
+        # bundle writer keeps sections on disk and publishes only at finish.
         config.raw["_decoder_engine_role"] = "prefill"
         prefill = model.build_engine(
             config,
@@ -352,6 +355,8 @@ def native(request: "BuildRequest", writer: "BundleWriter") -> None:
             quant_ctx=None,
             verbose=request.verbose,
         )
+        writer.add_bytes("prefill.plan", prefill)
+        del prefill
         config.raw["_decoder_engine_role"] = "decode"
         decode = model.build_engine(
             config,
@@ -361,12 +366,16 @@ def native(request: "BuildRequest", writer: "BundleWriter") -> None:
             quant_ctx=None,
             verbose=request.verbose,
         )
+        writer.add_bytes("engine.plan", decode)
+        del decode
         config.raw.pop("_decoder_engine_role", None)
         vision = model.build_vision_engine(
             str(model_dir), config, weights, precision=precision, verbose=request.verbose
         )
         if vision is None:
             raise RuntimeError("Phi-4 Multimodal vision build returned no engine")
+        writer.add_bytes("vision.plan", vision)
+        del vision
         vl = model.get_vl_config(config) or {}
         runtime = {
             "tensor_parallel_size": 1,
@@ -386,10 +395,6 @@ def native(request: "BuildRequest", writer: "BundleWriter") -> None:
             },
         }
         runtime.update(vl)
-        writer.set_header(family="phi4_multimodal", task=request.task, backend=request.backend)
-        writer.add_bytes("engine.plan", decode)
-        writer.add_bytes("prefill.plan", prefill)
-        writer.add_bytes("vision.plan", vision)
         runtime.update(_tokenizer_runtime_contract(model_dir))
         writer.add_json("runtime.json", runtime)
         for filename in (

families/phi4_multimodal/tests/aiperf/prepare_tokenizer.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Export Phi4's standard tokenizer without its model-only remote configuration."""
+
+import argparse
+import importlib.metadata
+import json
+from pathlib import Path
+
+
+def export_tokenizer(model_dir: Path, output_dir: Path) -> None:
+    from transformers import AutoTokenizer, GPT2TokenizerFast, PreTrainedTokenizerFast
+
+    config = json.loads((model_dir / "config.json").read_text())
+    if config.get("model_type") not in {"phi4mm", "phi4_multimodal", "phi4multimodal"}:
+        raise ValueError("model directory must contain the Phi4 multimodal checkpoint")
+    tokenizer = GPT2TokenizerFast.from_pretrained(model_dir, local_files_only=True)
+    probes = ("hello", "What is the main object in this image? Answer in one word.",
+              "<|user|>\n<|image_1|>\nQuestion?<|end|>\n<|assistant|>\n", "找到图片中的狗。")
+    identities = [
+        {"text": text, "add_special_tokens": special,
+         "token_ids": tokenizer.encode(text, add_special_tokens=special)}
+        for text in probes for special in (False, True)
+    ]
+    vocabulary = tokenizer.get_vocab()
+    # Keep the checkpoint's GPT-4o regex/backend unchanged. A model-specific
+    # GPT2 wrapper in a newer client can rebuild it with a different regex.
+    exported = PreTrainedTokenizerFast(
+        tokenizer_object=tokenizer.backend_tokenizer,
+        model_max_length=tokenizer.model_max_length,
+        padding_side=tokenizer.padding_side,
+        truncation_side=tokenizer.truncation_side,
+        clean_up_tokenization_spaces=tokenizer.clean_up_tokenization_spaces,
+        chat_template=tokenizer.chat_template,
+        **tokenizer.special_tokens_map,
+    )
+    output_dir.mkdir(parents=True, exist_ok=False)
+    exported.save_pretrained(output_dir)
+    restored = AutoTokenizer.from_pretrained(output_dir, use_fast=True, local_files_only=True)
+    if restored.get_vocab() != vocabulary:
+        raise RuntimeError("exported tokenizer changed the vocabulary")
+    if restored.chat_template != tokenizer.chat_template:
+        raise RuntimeError("exported tokenizer changed the chat template")
+    for probe in identities:
+        if restored.encode(probe["text"], add_special_tokens=probe["add_special_tokens"]) != probe["token_ids"]:
+            raise RuntimeError("exported tokenizer changed token IDs or framing")
+    audit = {"source": str(model_dir.resolve()),
+             "transformers": importlib.metadata.version("transformers"),
+             "tokenizer_class": type(tokenizer).__name__, "probes": identities}
+    (output_dir / "trtmc-tokenizer-export.json").write_text(json.dumps(audit, indent=2) + "\n")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--model-dir", type=Path, required=True, help="local Phi4 multimodal checkpoint")
+    parser.add_argument("--output-dir", type=Path, required=True, help="new directory for client tokenizer files")
+    args = parser.parse_args()
+    export_tokenizer(args.model_dir, args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

families/phi4_multimodal/tests/conftest.py

@@ -0,0 +1,13 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Resource configuration for the original FP32 eager family reference."""
+
+import os
+
+
+def pytest_configure(config):
+    # The eager reference's varying allocations fragment a nearly full GPU.
+    # Keep its precision/operators unchanged and preserve caller configuration.
+    if "PYTORCH_ALLOC_CONF" not in os.environ:
+        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

families/phi4_multimodal/tests/test_plan_lifetime.py

@@ -0,0 +1,78 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Keep serialized engines out of RAM while building the next network."""
+
+import gc
+import json
+from types import SimpleNamespace
+
+import pytest
+
+
+def test_each_plan_is_staged_and_released_before_the_next_build(tmp_path, monkeypatch):
+    pytest.importorskip("tensorrt")
+    from families.phi4_multimodal import model
+    from families.phi4_multimodal.edge_llm import dispatch
+    from tensorrt_model_connect.bundle_writer import BundleWriter
+
+    live = set()
+    built = []
+
+    class Plan(bytes):
+        def __new__(cls, role):
+            value = super().__new__(cls, role.encode())
+            value.role = role
+            live.add(role)
+            return value
+
+        def __del__(self):
+            live.remove(self.role)
+
+    def next_plan(role):
+        gc.collect()
+        assert not live, f"previous serialized engine retained during {role}: {live}"
+        built.append(role)
+        return Plan(role)
+
+    class Model:
+        @staticmethod
+        def load_weights(*_args):
+            return {}
+
+        @staticmethod
+        def build_engine(config, *_args, **_kwargs):
+            return next_plan(config.raw["_decoder_engine_role"])
+
+        @staticmethod
+        def build_vision_engine(*_args, **_kwargs):
+            return next_plan("vision")
+
+        @staticmethod
+        def get_vl_config(_config):
+            return {"image_token_id": 3, "vision_output_dim": 4}
+
+    monkeypatch.setattr(model, "_Phi4MultimodalModel", Model)
+    monkeypatch.setattr(model, "_tokenizer_runtime_contract", lambda _directory: {})
+    monkeypatch.setattr(dispatch, "build", lambda request, writer, native: native(request, writer))
+    (tmp_path / "config.json").write_text(json.dumps({
+        "model_type": "phi4mm", "hidden_size": 4, "num_hidden_layers": 1,
+        "vocab_size": 8, "num_attention_heads": 1, "num_key_value_heads": 1,
+    }))
+    request = SimpleNamespace(
+        mtp_seq_len=None, dynamic_kv_cache=False, image_height=None, image_width=None,
+        video_num_frames=None, max_batch_size=1, context_parallel_size=1,
+        task="vision_language_generation", tensor_parallel_size=1, quantization=None,
+        fp32_layers=(), model_dir=tmp_path, precision="fp16", max_sequence_length=768,
+        verbose=False, backend="tensorrt",
+    )
+    writer = BundleWriter(tmp_path / "model.bundle")
+    try:
+        model.build(request, writer)
+        writer.finish()
+    finally:
+        writer.abort()
+    gc.collect()
+    assert built == ["prefill", "decode", "vision"]
+    assert not live
+    assert (tmp_path / "model.bundle").is_file()