fix(llama): preserve BF16 checkpoint value range

families/llama/checkpoint_mapper.py

@@ -26,7 +26,7 @@
 
 def _target_np_dtype(precision: str) -> np.dtype:
     """Map precision string to numpy dtype for weight storage."""
-    if precision in ("fp16", "bf16"):
+    if precision == "fp16":
         return np.float16
     return np.float32
 

families/llama/dual_profile_decoder_builder.py

@@ -70,16 +70,7 @@ def _const_in_work_dtype(
     work_np_dtype: np.dtype,
     work_trt_dtype: trt.DataType,
 ) -> trt.ITensor:
-    """Create a constant in work_np_dtype storage and cast it to work_trt_dtype.
-
-    Needed for bf16 builds: the dual-profile builder stores bf16 weights
-    on disk as fp16 (work_np_dtype = np.float16), but the runtime tensor
-    must be bfloat16 to match the rest of the graph. ``add_constant``
-    alone produces an fp16 constant — we need an explicit cast to
-    bfloat16 so layers like IRotaryEmbeddingLayer (which require all
-    inputs to share a dtype) accept it. fp16 / fp32 builds are no-ops
-    because work_np_dtype maps directly to work_trt_dtype.
-    """
+    """Create a constant in storage dtype and cast it to runtime dtype."""
     const = graph_ops.add_constant(network, shape, values, dtype=work_np_dtype)
     if const.dtype != work_trt_dtype:
         const = network.add_cast(const, work_trt_dtype).get_output(0)
@@ -309,7 +300,7 @@ def build_dual_profile_decoder_engine(
     if precision == "fp16":
         work_np_dtype, work_trt_dtype = np.float16, trt.float16
     elif precision == "bf16":
-        work_np_dtype, work_trt_dtype = np.float16, trt.bfloat16
+        work_np_dtype, work_trt_dtype = np.float32, trt.bfloat16
     else:
         work_np_dtype, work_trt_dtype = np.float32, trt.float32
 
@@ -427,14 +418,11 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False):
                 interleaved=interleaved_rope,
                 rope_scaling=config.raw.get("rope_scaling"),
             )
-            # BF16 must round directly from the FP32 indexed table. Routing
-            # through FP16 storage would introduce FP16 -> BF16 double rounding.
-            rope_np_dtype = np.float32 if work_trt_dtype == trt.bfloat16 else work_np_dtype
             cos_half_table = _const_in_work_dtype(
-                network, cos_half_np.shape, cos_half_np, rope_np_dtype, work_trt_dtype
+                network, cos_half_np.shape, cos_half_np, work_np_dtype, work_trt_dtype
             )
             sin_half_table = _const_in_work_dtype(
-                network, sin_half_np.shape, sin_half_np, rope_np_dtype, work_trt_dtype
+                network, sin_half_np.shape, sin_half_np, work_np_dtype, work_trt_dtype
             )
 
     # Learned position embedding (GPT-2 / OPT / GPT-Neo / XGLM).
@@ -490,8 +478,6 @@ def _add_profile(opt_sq: int, max_sq: int, *, fixed: bool = False):
         )
         hidden_state = pos_add.get_output(0)
 
-    # Make sure the main hidden stream is in the requested runtime dtype
-    # before entering the layer stack (BF16 mode stores fp16 constants).
     if hidden_state.dtype != work_trt_dtype:
         hidden_state = network.add_cast(hidden_state, work_trt_dtype).get_output(0)
 

families/llama/graph_ops.py

@@ -125,13 +125,13 @@ def add_rms_norm(
 ) -> trt.ITensor:
     """RMSNorm: gamma * (x / sqrt(mean(x^2) + eps)).
 
-    FP32 precision boundary: when dtype != float32, casts to FP32 before
+    FP32 precision boundary: when input is not float32, casts to FP32 before
     norm computation for numerical stability, then casts back.
 
     TRT's native normalization API implements mean-centered LayerNorm, not
     RMSNorm, so this remains a manual shared implementation.
     """
-    need_cast = (dtype != np.float32)
+    need_cast = (inp.dtype != trt.float32)
     output_dtype = inp.dtype
     if need_cast:
         inp = network.add_cast(inp, trt.float32).get_output(0)
@@ -166,12 +166,12 @@ def add_rms_norm_per_head(
 ) -> trt.ITensor:
     """Per-head RMSNorm for [Sq, num_heads * head_dim] tensors.
 
-    FP32 precision boundary: when dtype != float32, casts to FP32 before
+    FP32 precision boundary: when input is not float32, casts to FP32 before
     norm computation for numerical stability, then casts back.
     ``sequence_length=None`` means runtime-dynamic Sq.
     ``gamma`` may be [num_heads * head_dim] or [head_dim] broadcast to heads.
     """
-    need_cast = (dtype != np.float32)
+    need_cast = (inp.dtype != trt.float32)
     output_dtype = inp.dtype
     seq_dim = -1 if sequence_length is None else sequence_length
     reshape_in = network.add_shuffle(inp)
@@ -223,10 +223,10 @@ def add_layer_norm(
 ) -> trt.ITensor:
     """LayerNorm: gamma * ((x - mean) / sqrt(var + eps)) + beta.
 
-    FP32 precision boundary: when dtype != float32, casts to FP32 before
+    FP32 precision boundary: when input is not float32, casts to FP32 before
     norm computation for numerical stability, then casts back.
     """
-    need_cast = (dtype != np.float32)
+    need_cast = (inp.dtype != trt.float32)
     output_dtype = inp.dtype
     if need_cast:
         inp = network.add_cast(inp, trt.float32).get_output(0)
@@ -273,8 +273,8 @@ def add_gelu_new(
     """GELU (tanh approximation): 0.5*x*(1+tanh(sqrt(2/pi)*(x+0.044715*x^3))).
 
     Constants are cast to ``inp.dtype`` so the elementwise ops are valid in
-    a STRONGLY_TYPED network when ``inp`` is bf16 (storage np_dtype is
-    fp16, runtime trt_dtype is bfloat16) or any other non-matching combo.
+    a STRONGLY_TYPED network when ``inp`` is bf16 or any other
+    non-matching combo.
     """
     target_dtype = inp.dtype
     const_shape = (1,) * max(1, len(tuple(inp.shape)))

families/llama/standard_decoder_builder.py

@@ -181,7 +181,7 @@ def build_standard_decoder_engine(
         work_np_dtype = np.float16
         work_trt_dtype = trt.float16
     elif precision == "bf16":
-        work_np_dtype = np.float16  # stored as float16, TRT uses bfloat16
+        work_np_dtype = np.float32
         work_trt_dtype = trt.bfloat16
     else:
         work_np_dtype = np.float32
@@ -341,9 +341,6 @@ def _cast_work_dtype(tensor: trt.ITensor) -> trt.ITensor:
         )
         hidden_state = pos_add.get_output(0)
 
-    # In BF16 mode many embedding/position constants are still materialized from
-    # float16 storage. Normalize the decoder's main hidden stream back to the
-    # requested runtime dtype before entering the layer stack.
     if hidden_state.dtype != work_trt_dtype:
         hidden_state = network.add_cast(hidden_state, work_trt_dtype).get_output(0)
 

families/llama/tests/test_bf16_storage.py

@@ -0,0 +1,242 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""CPU checkpoint and typed graph-prefix checks, without TensorRT execution."""
+
+import importlib
+from pathlib import Path
+import sys
+from types import ModuleType, SimpleNamespace
+
+import ml_dtypes
+import numpy as np
+import pytest
+from safetensors.numpy import save_file
+
+from ..checkpoint_mapper import load_standard_weights
+from ..config import ModelConfig
+
+
+class _Tensor:
+    def __init__(self, values, dtype="fp32", shape=None):
+        storage = {"fp32": np.float32, "fp16": np.float16, "bf16": ml_dtypes.bfloat16}
+        with np.errstate(over="ignore", under="ignore", invalid="ignore"):
+            self.values = np.asarray(values).astype(storage.get(dtype, np.int32)).astype(np.float32)
+        self.dtype = dtype
+        self.shape = self.values.shape if shape is None else shape
+
+
+def _layer(tensor):
+    return SimpleNamespace(get_output=lambda index: tensor)
+
+
+class _PrefixComplete(Exception):
+    pass
+
+
+class _Network:
+    def __init__(self, stop_at_epsilon=False):
+        self.constants = []
+        self.inputs = {}
+        self.stop_at_epsilon = stop_at_epsilon
+
+    def add_constant(self, shape, values):
+        values = values.reshape(shape)
+        self.constants.append(values.copy())
+        if self.stop_at_epsilon and shape == (1, 1):
+            raise _PrefixComplete
+        dtype = "fp16" if values.dtype == np.float16 else "fp32"
+        return _layer(_Tensor(values, dtype))
+
+    def add_input(self, name, dtype, shape):
+        tensor = _Tensor(np.zeros(tuple(max(1, n) for n in shape)), dtype, shape)
+        self.inputs[name] = tensor
+        return tensor
+
+    def add_cast(self, tensor, dtype):
+        return _layer(_Tensor(tensor.values, dtype, tensor.shape))
+
+    def add_elementwise(self, lhs, rhs, operation):
+        assert lhs.dtype == rhs.dtype, "mixed elementwise dtypes"
+        function = {"prod": np.multiply, "sum": np.add, "sub": np.subtract}[operation]
+        return _layer(_Tensor(function(lhs.values, rhs.values), lhs.dtype))
+
+    def add_reduce(self, tensor, operation, axes, keep_dims):
+        assert operation == "avg"
+        dimensions = tuple(i for i in range(tensor.values.ndim) if axes & (1 << i))
+        return _layer(
+            _Tensor(tensor.values.mean(axis=dimensions, keepdims=keep_dims), tensor.dtype)
+        )
+
+    def add_unary(self, tensor, operation):
+        function = {"sqrt": np.sqrt, "recip": np.reciprocal}[operation]
+        return _layer(_Tensor(function(tensor.values), tensor.dtype))
+
+    def add_shuffle(self, tensor):
+        shuffle = SimpleNamespace(reshape_dims=tensor.shape)
+        shuffle.get_output = lambda index: _Tensor(
+            tensor.values.reshape(shuffle.reshape_dims), tensor.dtype
+        )
+        return shuffle
+
+
+@pytest.fixture
+def family(monkeypatch):
+    network = _Network()
+    config = SimpleNamespace(add_optimization_profile=lambda profile: None)
+    builder = SimpleNamespace(
+        create_network=lambda flags: network,
+        create_builder_config=lambda: config,
+        create_optimization_profile=lambda: SimpleNamespace(set_shape=lambda *args: None),
+    )
+    logger = type("Logger", (), {"VERBOSE": 1, "WARNING": 2, "__init__": lambda self, level: None})
+    trt = SimpleNamespace(
+        float16="fp16",
+        float32="fp32",
+        bfloat16="bf16",
+        int32="int32",
+        Weights=lambda values: values,
+        Logger=logger,
+        Builder=lambda logger: builder,
+        NetworkDefinitionCreationFlag=SimpleNamespace(STRONGLY_TYPED=1),
+        ElementWiseOperation=SimpleNamespace(PROD="prod", SUM="sum", SUB="sub"),
+        ReduceOperation=SimpleNamespace(AVG="avg"),
+        UnaryOperation=SimpleNamespace(SQRT="sqrt", RECIP="recip"),
+    )
+    name = "_llama_bf16_storage_test"
+    package = ModuleType(name)
+    package.__path__ = [str(Path(__file__).resolve().parents[1])]
+    monkeypatch.setitem(sys.modules, "tensorrt", trt)
+    monkeypatch.setitem(sys.modules, name, package)
+    try:
+        yield SimpleNamespace(
+            graph=importlib.import_module(name + ".graph_ops"),
+            standard=importlib.import_module(name + ".standard_decoder_builder"),
+            dual=importlib.import_module(name + ".dual_profile_decoder_builder"),
+            network=network,
+        )
+    finally:
+        for key in tuple(sys.modules):
+            if key.startswith(name + "."):
+                del sys.modules[key]
+
+
+def _checkpoint(tmp_path, tied, *, wide=True):
+    config = ModelConfig(
+        model_type="llama",
+        hidden_size=4,
+        vocab_size=4,
+        num_hidden_layers=1,
+        num_attention_heads=1,
+        num_key_value_heads=1,
+        intermediate_size=4,
+        rms_norm_eps=2**-30,
+        raw={"_decoder_engine_role": "decode"},
+    )
+    values = np.array(
+        [2**20, -(2**20), 2**-30, -(2**-30)] if wide else [1024, -1024, 0.125, -0.125],
+        dtype=ml_dtypes.bfloat16,
+    )
+    matrix = np.tile(values, (4, 1))
+    tensors = {"model.embed_tokens.weight": matrix}
+    if not tied:
+        tensors["lm_head.weight"] = matrix
+    for norm in ["input_layernorm", "post_attention_layernorm"]:
+        tensors[f"model.layers.0.{norm}.weight"] = np.ones(4, dtype=ml_dtypes.bfloat16)
+    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"] = matrix
+    save_file(tensors, tmp_path / "model.safetensors")
+    return config, matrix.astype(np.float32)
+
+
+@pytest.mark.parametrize("precision", ["bf16", "fp16", "fp32"])
+@pytest.mark.parametrize("tied", [False, True])
+def test_checkpoint_preserves_requested_storage_range(tmp_path, precision, tied):
+    config, source = _checkpoint(tmp_path, tied, wide=precision != "fp16")
+    dtype = np.float16 if precision == "fp16" else np.float32
+    with np.errstate(over="ignore", under="ignore"):
+        weights = load_standard_weights(tmp_path, config, precision=precision)
+        expected = source.astype(dtype)
+    for key in [
+        "embedding",
+        "w_out",
+        "layer.0.w_q",
+        "layer.0.w_k",
+        "layer.0.w_v",
+        "layer.0.w_o",
+        "layer.0.w_gate",
+        "layer.0.w_up",
+        "layer.0.w_down",
+    ]:
+        assert weights[key].dtype == dtype
+        np.testing.assert_array_equal(weights[key], expected if key == "embedding" else expected.T)
+
+
+@pytest.mark.parametrize("precision", ["bf16", "fp16", "fp32"])
+@pytest.mark.parametrize("builder", ["standard", "dual"])
+def test_decoder_prefix_preserves_constant_storage(family, tmp_path, precision, builder):
+    config, values = _checkpoint(tmp_path, True, wide=precision != "fp16")
+    weights = {
+        "embedding": values,
+        "position_embedding": values,
+        "final_norm": np.ones(4),
+        "w_out": values.T,
+    }
+    family.network.stop_at_epsilon = True
+    module = getattr(family, builder)
+    function = getattr(
+        module, f"build_{'standard' if builder == 'standard' else 'dual_profile'}_decoder_engine"
+    )
+    kwargs = {} if builder == "standard" else {"native_kv_cache": False, "profile_mode": "decode"}
+    with np.errstate(over="ignore", under="ignore"):
+        with pytest.raises(_PrefixComplete):
+            function(config, weights, 4, precision=precision, position_type="learned", **kwargs)
+        dtype = np.float16 if precision == "fp16" else np.float32
+        expected = values.astype(dtype)
+    for constant in family.network.constants[:2]:
+        assert constant.dtype == dtype
+        np.testing.assert_array_equal(constant, expected)
+    assert (
+        family.network.inputs["cache_k_0"].dtype
+        == {"bf16": "bf16", "fp16": "fp16", "fp32": "fp32"}[precision]
+    )
+    if precision == "bf16":
+        np.testing.assert_array_equal(family.network.constants[-1], [[2**-30]])
+
+
+@pytest.mark.parametrize(
+    "dtype,storage", [("bf16", np.float32), ("fp16", np.float16), ("fp32", np.float32)]
+)
+@pytest.mark.parametrize("norm", ["rms", "head", "layer"])
+def test_norm_uses_fp32_compute_with_independent_storage(family, dtype, storage, norm):
+    inputs = _Tensor([[1.0, 2.0, 3.0, 4.0]], dtype)
+    epsilon = _Tensor([[1e-5]], "fp32")
+    gamma = np.ones(4, dtype=np.float32)
+    kwargs = {"dtype": storage}
+    if norm == "rms":
+        result = family.graph.add_rms_norm(family.network, inputs, 4, gamma, epsilon, **kwargs)
+        expected = inputs.values / np.sqrt(np.mean(inputs.values**2) + 1e-5)
+    elif norm == "head":
+        result = family.graph.add_rms_norm_per_head(
+            family.network, inputs, 2, 2, gamma, epsilon, **kwargs
+        )
+        reshaped = inputs.values.reshape(1, 2, 2)
+        expected = (reshaped / np.sqrt(np.mean(reshaped**2, axis=2, keepdims=True) + 1e-5)).reshape(
+            1, 4
+        )
+    else:
+        result = family.graph.add_layer_norm(
+            family.network, inputs, 4, gamma, np.zeros(4), epsilon, **kwargs
+        )
+        centered = inputs.values - inputs.values.mean(axis=1, keepdims=True)
+        expected = centered / np.sqrt(np.mean(centered**2) + 1e-5)
+    assert result.dtype == dtype
+    np.testing.assert_allclose(result.values, _Tensor(expected, dtype).values, rtol=1e-6, atol=1e-7)