Bug: default Llama 3.1 builds fail to build (fp32) or cannot run (fp16) at the 131072-row default length

[LiRunGuo]

Affected area

Model resolution or TensorRT engine build

TensorRT-Model-Connect version

main at 393ab02f; the cited code is unchanged at ccc3579c.

Installation method

Development container

Model and revision

meta-llama/Llama-3.1-8B-Instruct (local snapshot, max_position_embeddings=131072).

Environment

NVIDIA H200 (141 GB), driver 580.159.03, container built from Dockerfile.dev.x86 (TensorRT 11.1.0.106, CUDA 13.3).

Minimal reproducer

# 1. Defaults (no --precision -> fp32, no --max-sequence-length -> 131072)
python -m tensorrt_model_connect build meta-llama/Llama-3.1-8B-Instruct --output default.bundle

# 2. fp16 with the default length
python -m tensorrt_model_connect build meta-llama/Llama-3.1-8B-Instruct --precision fp16 --output fp16.bundle
trtmc run fp16.bundle --runtime-root "$TRTMC_BUILD_DIR" --prompt "Explain what TensorRT does in one sentence." --max-new-tokens 8

Observed behavior

1. The default build runs for several minutes and then fails without writing a bundle:
   IBuilder::buildSerializedNetwork: Error Code 10: Internal Error (Could not find any implementation for node {ForeignNode[(Unnamed Layer* 9) [Shuffle]...(Unnamed Layer* 2842) [ElementWise]]} ...) followed by RuntimeError: dual-profile decoder engine build failed.
2. The fp16 build succeeds and writes a 32.2 GB bundle, but every run fails at load:
   [trtmc] KV cache rows=131072 (bundle max=131072) then Error: Unable to allocate TensorRT input buffer for 'attention_mask'.
3. --precision bf16 works: 26.94 ms/token at the default length, 11.26 ms/token with --max-sequence-length 4096.

Cause, as far as I can trace it:

• The native-KV route requires BF16 (families/llama/build_routing.py:283). Llama's support.py declares no default_precision, so the CLI uses the shared default fp32 (build_cli.py:118, model_support.py:58), and fp32/fp16 builds take the dense dual-profile path.
• families/llama/model.py:187 picks the default length from the architecture's native-KV eligibility, independent of precision, so fp32/fp16 builds still get max_position_embeddings (131072).
• families/llama/dual_profile_decoder_builder.py:248-257 sets max_prefill_length = max_cache_length and clamps it to the 64-token explicit-attention chunk only for the native-KV and runtime-sized paths ("so they cannot bypass this safety bound"). The dense path keeps 131072.
• The dense path declares attention_mask as float32 (-1, -1) (line 317) with profile maximum (max_sq, max_cache_length + max_sq) (line 358), that is [131072, 262144] float32 ≈ 137 GB, which the runtime cannot allocate (core/runtime/tensorrt/trt_module_impl.cpp:359).

families/llama/tests/manifests/minitron-4b-width.json also builds fp16 at 131072, so this path is only exercised where that allocation fits.

Expected behavior

A default build of a supported checkpoint should either produce a bundle that runs on the target GPU or fail early with a clear message. Possible fixes for the family owners to choose from:

• default Llama to BF16 when the native route is eligible, or
• choose the default length from the route the requested precision actually takes (the ineligible branch already falls back to 256), and/or
• bound the dense path's prefill profile like the native paths, or reject builds whose mask profile cannot fit.

Relevant logs

$ python -m tensorrt_model_connect build .../Llama-3.1-8B-Instruct --output runs/llama31-8b-default.bundle
[TRT] [E] IBuilder::buildSerializedNetwork: Error Code 10: Internal Error (Could not find any implementation for node {ForeignNode[...]} ...)
RuntimeError: dual-profile decoder engine build failed

$ trtmc run runs/llama31-8b.bundle ...   # --precision fp16 build, 32243009680 bytes
[trtmc] KV cache rows=131072 (bundle max=131072)
Error: Unable to allocate TensorRT input buffer for 'attention_mask'

Submission checks

• I searched open and closed issues and found no duplicate.
• This is not a security vulnerability.
• I removed secrets, credentials, private URLs, and restricted artifacts.