CUDA_ERROR_ILLEGAL_ADDRESS in XQA on Jetson Orin Nx for Gemma4

Describe the bug

Prefill fails when running Gemma4 E2B and E4B for i4 awq and i8 sq using 0.9.0.

#11 0x0000ffffd45992f4 in trt_edgellm::check::_checkCudaDriver (result=CUDA_ERROR_ILLEGAL_ADDRESS,
    func=0xffffd465d7a0 "cuLaunchKernel(kernelInfo.mDeviceFunction, dimGrid.x, dimGrid.y, dimGrid.z, dimCta.x, dimCta.y, dimCta.z, kernelInfo.mSharedMemBytes, stream, kernelParams, nullptr)",
    file=0xffffd465d358 "TensorRT-Edge-LLM/cpp/kernels/decodeAttentionKernels/decoderXQARunner.cpp", line=563)
    at TensorRT-Edge-LLM/cpp/common/checkMacros.h:85
#12 0x0000ffffd45a72ec in trt_edgellm::DecoderXQARunner::dispatchXQAKernel (this=0xffff63ffa5a8, params=..., stream=@0xffff63ffa4c8: 0xffff5c000e30)
    at TensorRT-Edge-LLM/cpp/kernels/decodeAttentionKernels/decoderXQARunner.cpp:563
#13 0x0000ffffd4503f64 in trt_edgellm::plugins::AttentionPlugin::enqueue (this=0xaaaaac541930, inputDesc=0xaaaaae57c680, outputDesc=0xaaaaae57c8e8, inputs=0xaaaaae57c998, outputs=0xaaaaae57c9d0,
    workspace=0x27a1d4000, stream=0xffff5c000e30) at TensorRT-Edge-LLM/cpp/plugins/attentionPlugin/attentionPlugin.cpp:1157
--Type <RET> for more, q to quit, c to continue without paging--#14 0x0000ffffdbc37474 in ?? () from /lib/aarch64-linux-gnu/libnvinfer.so.10
#15 0x0000ffffdbbe4c04 in ?? () from /lib/aarch64-linux-gnu/libnvinfer.so.10
#16 0x0000ffffdbbe672c in ?? () from /lib/aarch64-linux-gnu/libnvinfer.so.10
#17 0x0000aaaaaaad1664 in std::_Function_handler<bool (sample::TrtCudaStream&), sample::(anonymous namespace)::EnqueueExplicit>::_M_invoke(std::_Any_data const&, sample::TrtCudaStream&) ()
#18 0x0000aaaaaaad9b70 in void sample::(anonymous namespace)::inferenceExecution<nvinfer1::IExecutionContext>(sample::InferenceOptions const&, sample::InferenceEnvironment&, sample::(anonymous namespace)::SyncS
truct&, int, int, int, std::vector<sample::InferenceTrace, std::allocator<sample::InferenceTrace> >&) ()

Steps/Code to reproduce bug

1. Build TensorRT-Edge-LLM
2. Quantize Gemma4
3. Convert to onnx
4. Create engine
5. Try to run inference

Build configuration:

cmake .. \
    -DCMAKE_BUILD_TYPE=Debug \
    -DTRT_PACKAGE_DIR=/usr \
    -DCMAKE_TOOLCHAIN_FILE=../cmake/aarch64_linux_toolchain.cmake \
    -DEMBEDDED_TARGET=jetson-orin \
    -DBUILD_PYTHON_BINDINGS=ON \
    -DCUDA_CTK_VERSION=12.6 \
    -DENABLE_CUTE_DSL='gemm'
cmake --build .>

Runtime command used:

tensorrt-edgellm-quantize llm --model_dir google/gemma-4-E2B --quantization int8_sq --output_dir gemma4_e2b_i8_sq
tensorrt-edgellm-export gemma4_e2b_i8_sq trt_gemma4_e2b_i8_sq

from experimental.server import LLM
LLM(onnx_dir="trt_gemma4_e2b_i8_sq/llm")

/usr/src/tensorrt/bin/trtexec --loadEngine=llm.engine --dumpLayerInfo --staticPlugins=TensorRT-Edge-LLM/build/libNvInfer_edgellm_plugin.so --verbose --profilingVerbosity=detailed

Expected behavior

No error.

System information (Edge Device)

• Platform : Jetson Orin NX 16GB
• Software release : Jetpack 6.2.1
• CPU architecture: aarch64
• GPU compute capability: SM87
• Total device memory: 16GB
• Build type: Release and Debug
• Library versions:
  • TensorRT Edge-LLM version or commit hash: 0.9.0
  • CUDA: 12.6.68
  • TensorRT: 10.3.0.30
  • C++ compiler (e.g., GCC 11.4): GCC 11.4.0
• CMake options used:
  • CMAKE_TOOLCHAIN_FILE: ../cmake/aarch64_linux_toolchain.cmake
  • EMBEDDED_TARGET: jetson-orin
  • TRT_PACKAGE_DIR: /usr

Start with cpp/kernels/decodeAttentionKernels/decoderXQARunner.cpp around dispatchXQAKernel line 563, then inspect the AttentionPlugin enqueue call at cpp/plugins/attentionPlugin/attentionPlugin.cpp line 1157 and the CUDA check in cpp/common/checkMacros.h. Reproduce the Gemma4 i4 AWQ or i8 SQ run on Jetson Orin NX using the listed build and runtime commands; done means prefill completes without CUDA_ERROR_ILLEGAL_ADDRESS.