[Bug]: Grouped MLA generation kernels read out of bounds at batch > 1: the MLA generation path never sets cumSeqLensQPtr / cumSeqLensKvPtr #20375

System Info

SM107, TensorRT-LLM main (C++ attention op, trtllm-gen FMHA, BF16 MLA generation with HQk 576 / HV 512, paged KV with 32 tokens per page).

Who can help?

@farazkh80 (#15009) @pengbowang-nv

Information

• The official example scripts
• My own modified scripts

Reproduction

AttentionOp fills TllmGenFmhaRunnerParams for MLA generation (cpp/tensorrt_llm/common/attentionOp.cpp, the trtllm-gen MLA generation branch) with seqLensKvPtr = params.cache_seq_lens, but it never sets cumSeqLensQPtr or cumSeqLensKvPtr. The SwapsAbForGen kernels selected today do not read them, so nothing fails on main.

The grouped tokens x heads MLA generation kernels (...VarSeqQ64Kv128{Static,Persistent}GroupedKeepsAbForGen, selected when FmhaOptions::mSelectsGroupedMla is set, as #15009 proposed) do read them. With mSelectsGroupedMla = true and the runner parameters exactly as AttentionOp builds them:

• In a standalone TllmGenFmhaRunner microbenchmark (Kimi K3 MLA decode shape with MTP, KV length 64): batch 1 output is correct, but batch 4 and batch 16 are wrong (output sums -209.5 and -218.6, expected -243.7). Passing correct cumSeqLensQ = [0, T, 2T, ...] and cumSeqLensKv = prefix sum of seqLensKv makes all batch sizes match the non-grouped kernel.
• In serving (Kimi K3 with MTP, CUDA graphs), capture of the batch-4 graph fails with CUDA_ERROR_ILLEGAL_ADDRESS right after the grouped cubin is loaded. Captures of batches 5..16 happened to succeed.

#15009 only measured batch 1, where the offsets are never needed, which may be why this did not show up there.

Expected behavior

Every kernel that the MLA generation path can select gets the metadata it reads. Either AttentionOp provides cumSeqLensQPtr / cumSeqLensKvPtr whenever a grouped MLA generation kernel can be selected, or the selection refuses grouped kernels when they are absent.

actual behavior

With grouped selection enabled, batch > 1 MLA generation reads past the end of its offsets: the output is silently wrong or there is an illegal address.

additional notes

Workaround we validated on SM107 (Kimi K3 serving, all graph batch sizes 1..16, unchanged accuracy): when grouped selection is enabled and batch > 1, AttentionOp launches a one-CTA kernel per MLA generation call. That kernel writes cumSeqLensQ[i] = i * maxSeqLenQ and cumSeqLensKv = the prefix sum of cache_seq_lens into a device buffer that is allocated before graph capture, and passes both pointers to the runner. Batch 1 skips the kernel. I can turn this into a PR together with whatever change enables grouped MLA selection, or separately behind the same switch, whichever maintainers prefer.

Related: the grouped MLA generation cubins are also registered with a mask type that does not match their names (#20376), the NVRTC fallback fails for these kernels on SM107 (#20377), and with FP8 KV the grouped E4M3 cubins hang when heads per rank do not divide 64 (#20382).