[JiaxinD]

Background

Llama BF16 checkpoints currently pass through FP16 host storage before the runtime BF16 cast. Valid BF16 values such as ±2^20 and ±2^-30 become infinities or zero. Both decoder builders repeat this narrowing. This is an independently reproduced correctness bug in the existing family, not new model support.

Exit Criteria

• Preserve BF16 checkpoint values through host staging and both decoder builders.
• Keep FP32 normalization computation and the requested runtime/cache dtype.
• Preserve FP16 and FP32 controls; do not claim full-model qualification or performance equivalence.

Implementation

• Use FP32 host storage for BF16 checkpoint loading and standard/dual-profile constants.
• Select RMSNorm, per-head RMSNorm and LayerNorm compute casts from the input tensor dtype rather than storage dtype.
• Add real BF16 safetensors, builder-prefix and strict normalization regressions within Llama ownership.

Change categories

• Model or runtime behavior
• Public API
• ABI
• Bundle or artifact format
• Dependencies
• Documentation only
• CI or developer tooling

No public API, ABI, format or dependency changes. Rebuild affected BF16 plans to obtain corrected constants.

Validation

Commands and Results

Command	Result
``$env:PYTHONPATH='core/builder;.'; python -m pytest families/llama/tests/test_bf16_storage.py families/llama/tests/test_runtime_receipt.py families/llama/tests/test_speculative_contract.py -q`	37 passed. Before the fix: 7 new regression failures, 14 controls passed.
PYTHONPATH=core/builder:. python -m pytest -c /dev/null --rootdir=. --confcutdir=. families/llama/tests/test_builder_variants.py families/llama/tests/test_bf16_storage.py -q	H200: 63 passed, two unregistered-marker warnings from using an empty pytest config.
python native-bf16-probe.py --source <tree> --output <result.json>	Archived standalone diagnostic: actual build, deserialize and execution of a tiny one-layer zero-projection decoder. Baseline: four FP16/FP32 controls passed, both BF16 builders failed. Fixed: all six passed with unchanged rtol=1e-6, atol=1e-7.
python -m pytest tools/tests/test_architecture.py tools/tests/test_family_impact.py tools/tests/test_public_source_hygiene.py tools/tests/test_pr_metadata.py -q	99 passed.

Ruff, legal headers, catalog, impact and git diff --check passed.

Hardware, Environment, and Revisions

Base 59fe695bdd5cd328bfc276a534cf762beeede787, tested head f179f32c6db3cdda6bc32d5cf90ee83d3ad8ffcf. Windows/Python 3.13 CPU checks; Linux/Python 3.12.3, NVIDIA H200, TensorRT 11.4.0.106 CUDA 13 package for native checks. Five normalized source-file hashes matched the tested archive plus patch. Synthetic BF16 safetensors, no external checkpoint or dataset.

Not Run / Remaining Gaps

Dev CI39636402874 on L4/TRT11.1 passed FP8 E2E but NVFP4 failed engine deserialization with CUDA OOM; three later cases were not run. Attribution needs a clean base/head memory comparison. TP, native-KV dual-profile execution and performance qualification remain unverified. The tiny native probe exercises the dense dual-profile path. Windows collection of two additional TensorRT-dependent modules failed because TensorRT is absent there; the native builder suite above ran on H200.

Contributor Self-Review

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

Manual call-chain review plus independent review found no blocking defect.

Notes For Future Readers

Review mapper, builders, then normalization boundaries. BF16 host weight ndarray elements increase from two to four bytes; peak RSS, GPU memory and latency were not measured. No resource-equivalence claim. If #1427 lands first, its BF16 host-storage expectations in test_support.py must follow this FP32 storage contract. No runtime dependency or required merge order.

Risk level

• Low
• Medium
• High

Corrects shared Llama-family BF16 staging and normalization boundaries, with increased host storage and incomplete full-model qualification. Draft retained for the unresolved NVFP4 qualification failure and unchanged three-slot review capacity.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

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

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

[JiaxinD]

@yifeif-nv Community CPU passed for f179f32; Dev39636402874 passed FP8 but NVFP4 hit CUDA OOM during engine deserialization on L4/TRT11.1.
Please run protected CI for this head; a clean base/head L4 memory comparison would help attribute the failure.