[Bug][FMHA] Incorrect Q Sequence Start Calculation in FMHA Kernels for Cross-Attention

System Info

Ubuntu 22.04, x86_64
NVIDIA A10-24Q, 24G VRAM
Tensorrt-LLM v1.1.0rc5

Who can help?

The fused multi-head attention (FMHA) kernels incorrectly calculate q_sequence_start and kv_mask_loop_start by using actual_kv_seqlen to offset the query position. This assumption is valid for self-attention where query and key/value sequences are aligned, but breaks in cross-attention scenarios where the encoder output (KV) and decoder input (Q) have independent, arbitrary sequence lengths.

https://github.com/NVIDIA/TensorRT-LLM/blob/c1cfb61b1b0940e9212b68e7ee72d42c6126e242/cpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_noloop.h#L124

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

Run the following test scripts:

Test configuration:

• input_len=64, kv_input_len=128
• num_heads=16, head_size=32
• Cross-attention with padding mask
  • Q: all-1s
  • cross_k: all-1s
  • cross_v: sequential values (0, 1, 2, ...) for each position
  • padding_mask: first 1/4 valid, others masked

https://github.com/toothache/TensorRT-LLM/blob/c3fe137363e04e0cb6a6783ea60a42baf9e07e98/tests/unittest/trt/attention/test_cross_attention.py

Expected behavior

TRT-LLM results should match with Torch/ORT version result.

actual behavior

FAIL: test_plugin_cross_attention_float16_64_128_ContextFMHAType_enabled (main.TestPluginCrossAttention)

Launching FMHA kernel: fmha_v2_flash_attention_fp16_64_64_S_q_paged_kv_32_sm86_kernel_nl

TRT-LLM vs ONNX - Mean diff: 16.0, Max diff: 16.0
TRT-LLM vs PyTorch - Mean diff: 16.0, Max diff: 16.0
ONNX vs PyTorch - Mean diff: 0.0, Max diff: 0.0

additional notes

This is a correctness issue that produces incorrect results for cross-attention, a core component in many production models.

I've verified that fixing the incorrect q_sequence_start and kv_mask_loop_start offset calculations resolves the correctness issue.

However, I'm not very familiar with the FMHA kernel codebase and its design considerations. I'm happy to contribute a proper fix if guidance is provided on:

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.

Start with cpp/kernels/fmha_v2/src/fused_multihead_flash_attention_kernel_noloop.h around line 124, then run tests/unittest/trt/attention/test_cross_attention.py using the reported 64/128 cross-attention configuration. Compare the TensorRT-LLM output with the Torch and ONNX Runtime results; done means the test passes with matching results.