Qwen3-TTS CodePredictor (MTP) crashes with "No available kernel for the GQA" on Jetson AGX Orin

[WonderW-hub]

Title: [Bug] Qwen3-TTS CodePredictor (MTP) crashes with "No available kernel for the GQA" on Jetson AGX Orin

Environment:

• Hardware: Jetson AGX Orin (Ampere / SM87 architecture)
• Model: Qwen3-TTS-0.6B (from HuggingFace/ModelScope)
• Framework: TensorRT-Edge-LLM (Runtime version: 0.10.0)
• Context: Testing via NVIDIA Developer Program SDKs and TensorRT-Edge-LLM toolkit.
• Reporter: @WonderW-hub

Bug Description:
When running the end-to-end qwen3_tts_inference pipeline, the Talker engine successfully executes and outputs the first codec token. However, the pipeline immediately crashes when passing the tensor to the CodePredictor (MTP) engine for step 2 decoding. The error explicitly points to a missing GQA kernel for this specific architecture.

Steps to Reproduce:

1. Export Qwen3-TTS-0.6B ONNX models and compile the engines (llm, code_predictor, and code2wav) using llm_build and audio_build.
2. Configure input.json with standard ChatML formatting and roles mapped in processed_chat_template.json.
3. Execute the C++ omni inference binary with the compiled engines:

EDGELLM_PLUGIN_PATH="./libNvInfer_edgellm_plugin.so" \
LD_PRELOAD="./libNvInfer_edgellm_plugin.so" \
./examples/omni/qwen3_tts_inference \
    --inputFile "input.json" \
    --talkerEngineDir "..." \
    --code2wavEngineDir "..." \
    --tokenizerDir "..."

4. Attempted workarounds that failed:

• Rebuilding the MTP engine with --maxBatchSize 32 (to match RVQ layers) and lowered --maxKVCacheCapacity 2048.
• Applying --specDraft during llm_build (fails due to engine_role mismatch).
• Forcing attention fallback via environment variables: TRTLLM_ENABLE_XQA=0, TRTLLM_ENABLE_MMHA=0, NVTE_FUSED_ATTN=0.

Error Logs:

[05:28:48.205] [INFO] [qwen3OmniTTSRuntime.cpp:1963:handleAudioGeneration] Batch 0: first codec token: 1350
[05:28:48.206] [INFO] [TensorRT] Switching optimization profile from: 1 to 0...
[05:28:48.212] [ERROR] [attentionPlugin.cpp:1093:enqueue] AttentionPlugin: enqueue failed: No available kernel available for the GQA
[05:28:48.212] [ERROR] [TensorRT] [pluginV3Runner.cpp::execute::252] Error Code 2: Internal Error (Assertion pluginUtils::isSuccess(status) failed...)
[05:28:48.212] [ERROR] [qwen3OmniTTSRuntime.cpp:2749:runCodePredictorGenerationForFrame] CP decode execute failed (step=2)
[05:28:48.212] [ERROR] [qwen3OmniTTSRuntime.cpp:2353:runTalkerGenerationLoop] CodePredictor failed at frame 0
[05:28:48.212] [WARNING] [qwen3_tts_inference.cpp:601:main] TTS generation failed for batch group [0, 1)

Expected Behavior:
The CodePredictor should either successfully map to a supported Ampere attention kernel (or cuBLAS fallback) to process the remaining RVQ layers, or llm_build should expose a specific flag to compile the MTP engine safely for Jetson Orin devices.

---