Describe the bug
The llm_build C++ tool fails to serialize an engine for the Gemma-4-31B model. The failure stems from Gemma 4's heterogeneous, asymmetric architecture (fluctuating attention head dimensions).
During compilation, the AttentionPlugin initially attempts to map the model using FMHA but fails when it encounters headSize=512, falling back to naive attention. Immediately after, the engine builder crashes with an API Usage Error due to a dimension mismatch because TensorRT expects a static number of KV heads across all layers, but Gemma 4 dynamically alters its KV heads (e.g., dropping to 4 heads at Layer 5).
Impact: Blocker - Gemma 4 cannot be compiled or run using the C++ TensorRT-Edge-LLM runtime on Jetson Thor.
Steps/Code to reproduce bug
[WARNING] [attentionPlugin.cpp:359:AttentionPlugin] AttentionPlugin: FMHA/FFPA unsupported for headSize=512, using naive attention for prefill + XQA for decode.
[INFO] [attentionPlugin.cpp:327:AttentionPlugin] AttentionPlugin: FMHA supported for headSize=256, using FMHA for prefill + XQA for decode.
[ERROR] [contextFMHARunner.cpp:435:canImplement] ContextFMHARunner::canImplement() unsupported headSize=512. Supported head sizes are 64, 72, 80, 128, 256.
...
[ERROR] [TensorRT] IBuilder::buildSerializedNetwork: Error Code 4: API Usage Error (Dimension mismatch for tensor past_key_values_5 and profile 0. At dimension axis 2, profile has min=16, opt=16, max=16 but tensor has 4. In validate at optimizer/api/network.cpp:3788)
[ERROR] [builderUtils.cpp:307:buildAndSerializeEngine] Failed to build serialized engine
[ERROR] [llm_build.cpp:243:main] Failed to build LLM engine.
Build configuration:
cmake ..
-DTRT_PACKAGE_DIR=/usr
-DCMAKE_TOOLCHAIN_FILE=cmake/aarch64_linux_toolchain.cmake
-DEMBEDDED_TARGET=jetson-thor
-DCMAKE_CUDA_ARCHITECTURES=110
make -j$(nproc)
Runtime command used:
EDGELLM_PLUGIN_PATH=./build/libNvInfer_edgellm_plugin.so \
./examples/llm/llm_build \
--onnxDir ~/onnx_models/gemma4/llm \
--engineDir ~/engines/gemma4 \
--maxBatchSize 1 \
--maxInputLen 262144 \
--maxKVCacheCapacity 262144
Expected behavior
The TensorRT builder should dynamically support Gemma 4's shifting attention layer dimensions and asymmetric KV cache heads without crashing. If dynamic dimension mapping is physically impossible within the static constraints of an .engine file, the builder should fail fast during the initial ONNX parsing phase with a clear "Unsupported Model Architecture" error, rather than crashing midway through graph validation.
System information (Edge Device)
System information (Edge Device)
Platform: NVIDIA Jetson AGX Thor Developer Kit
Software release: JetPack 7.1 (L4T 38.4.0)
CPU architecture: aarch64
GPU compute capability: SM110
Total device memory: 128GB LPDDR5X
Build type: Release
Library versions:
TensorRT Edge-LLM version or commit hash: (Latest main branch)
CUDA: 13.0
TensorRT: 10.13.3.9
C++ compiler: GCC 11.4
CMake options used:
CMAKE_TOOLCHAIN_FILE: cmake/aarch64_linux_toolchain.cmake
EMBEDDED_TARGET: jetson-thor
TRT_PACKAGE_DIR: /usr
CMAKE_CUDA_ARCHITECTURES: 110
Describe the bug
The llm_build C++ tool fails to serialize an engine for the Gemma-4-31B model. The failure stems from Gemma 4's heterogeneous, asymmetric architecture (fluctuating attention head dimensions).
During compilation, the AttentionPlugin initially attempts to map the model using FMHA but fails when it encounters headSize=512, falling back to naive attention. Immediately after, the engine builder crashes with an API Usage Error due to a dimension mismatch because TensorRT expects a static number of KV heads across all layers, but Gemma 4 dynamically alters its KV heads (e.g., dropping to 4 heads at Layer 5).
Impact: Blocker - Gemma 4 cannot be compiled or run using the C++ TensorRT-Edge-LLM runtime on Jetson Thor.
Steps/Code to reproduce bug
[WARNING] [attentionPlugin.cpp:359:AttentionPlugin] AttentionPlugin: FMHA/FFPA unsupported for headSize=512, using naive attention for prefill + XQA for decode.
[INFO] [attentionPlugin.cpp:327:AttentionPlugin] AttentionPlugin: FMHA supported for headSize=256, using FMHA for prefill + XQA for decode.
[ERROR] [contextFMHARunner.cpp:435:canImplement] ContextFMHARunner::canImplement() unsupported headSize=512. Supported head sizes are 64, 72, 80, 128, 256.
...
[ERROR] [TensorRT] IBuilder::buildSerializedNetwork: Error Code 4: API Usage Error (Dimension mismatch for tensor past_key_values_5 and profile 0. At dimension axis 2, profile has min=16, opt=16, max=16 but tensor has 4. In validate at optimizer/api/network.cpp:3788)
[ERROR] [builderUtils.cpp:307:buildAndSerializeEngine] Failed to build serialized engine
[ERROR] [llm_build.cpp:243:main] Failed to build LLM engine.
Build configuration:
cmake ..
-DTRT_PACKAGE_DIR=/usr
-DCMAKE_TOOLCHAIN_FILE=cmake/aarch64_linux_toolchain.cmake
-DEMBEDDED_TARGET=jetson-thor
-DCMAKE_CUDA_ARCHITECTURES=110
make -j$(nproc)
Runtime command used:
Expected behavior
The TensorRT builder should dynamically support Gemma 4's shifting attention layer dimensions and asymmetric KV cache heads without crashing. If dynamic dimension mapping is physically impossible within the static constraints of an .engine file, the builder should fail fast during the initial ONNX parsing phase with a clear "Unsupported Model Architecture" error, rather than crashing midway through graph validation.
System information (Edge Device)
System information (Edge Device)
Platform: NVIDIA Jetson AGX Thor Developer Kit
Software release: JetPack 7.1 (L4T 38.4.0)
CPU architecture: aarch64
GPU compute capability: SM110
Total device memory: 128GB LPDDR5X
Build type: Release
Library versions:
TensorRT Edge-LLM version or commit hash: (Latest main branch)
CUDA: 13.0
TensorRT: 10.13.3.9
C++ compiler: GCC 11.4
CMake options used:
CMAKE_TOOLCHAIN_FILE: cmake/aarch64_linux_toolchain.cmake
EMBEDDED_TARGET: jetson-thor
TRT_PACKAGE_DIR: /usr
CMAKE_CUDA_ARCHITECTURES: 110