[Bug]: FP8-KV MLA generation hangs on SM107 with the grouped trtllm-gen E4M3 kernels when heads per rank do not divide 64 (e.g. 12) #20382

System Info

• GPU: SM107, 8 GPUs over 2 nodes (TP8 attention)
• TensorRT-LLM main @ 69a455a (code and cubin list checked there), PyTorch backend, TRTLLM attention backend
• Model: MLA model with 96 query heads -> 12 heads per rank at TP8 (kv_lora_rank 512, qk_rope_head_dim 64)
• kv_cache_config.dtype: fp8, tokens_per_block 32, MTP speculative decoding with 8 query tokens per step
• Grouped (tokens x heads in one Q tile) MLA generation kernel selection enabled locally (FmhaOptions::mSelectsGroupedMla; main does not set it, see [None][feat] Add test-only grouped MLA generation selection probe #15009)

Who can help?

@pengbowang-nv @farazkh80

Information

• The official example scripts
• My own modified scripts

Reproduction

1. Serve the model with kv_cache_config.dtype: fp8 and grouped MLA generation selection enabled.
2. Startup never finishes: every rank stops in the first generation-shape attention call of the warmup (num_tokens=8, num_gen_requests=1, KV length 8). GPU utilization stays at 100 %, nothing is logged after Attention workspace size is not enough, increase the size from 0 bytes to .... TRTLLM_ENABLE_PDL=0 does not change it.
3. A host-side backtrace of a rank shows the main thread inside the next cuBLAS call (cublasGemmStridedBatchedEx -> cuKernelSetAttribute), i.e. queued behind a kernel that never completes.
4. With synchronizing checkpoints in MLA.forward_absorption_generation: the q_nope bmm and mla_rope_generation (FP8 Q quantization, FP8 KV append, bmm scales) complete; the attention op does not. A sync right after the FMHA launch in TllmGenFmhaRunner::run never returns. The selected cubin is fmhaSm107aKernel_QkvE4m3OBfloat16HQk576HV512HVPerCta256PagedKvCausalP32VarSeqQ64Kv128StaticGroupedKeepsAbForGen (grid 2x2x1, numHeadsQ 12, maxSeqLenQ 8, multiCtasKv off for this KV length).

Standalone reproduction with the FMHA runner only (no model): TllmGenFmhaRunner(E4M3, E4M3, E4M3, BF16), paged KV page 32, MLA head dims 576/512, 1 KV head, generation kernel, grouped selection enabled, seqLenQ 8 (12 heads also hangs at seqLenQ 1, 4, 5, 6), seqLenKv 2048:

The FP8 kernel completes only when numHeadsQ divides the 64-row Q tile. BF16 grouped kernels with 12 heads are fine.

Expected behavior

The FP8 grouped MLA generation kernel works for any head count, or is not selected for head counts it cannot handle.

actual behavior

The FMHA kernel never completes; the process hangs with the GPU at 100 %.

additional notes

Analysis. In makeTmaShapeStrideQ (cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/trtllmGen_fmha_export/KernelParams.h), for mGroupsTokensHeadsQ:

numTokensPerTileQ = numTokensPerTileQ / numGroupedHeads;   // 64 / 12 = 5
tileShapes = {clampedHeadDimQ, numGroupedHeads, 1, numTokensPerTileQ};  // 5 tokens x 12 heads = 60 rows

so the Q TMA box covers 60 of the 64 tile rows. The comment there says padding at the end of the tile is allowed. The BF16 cubins handle it; the E4M3 cubins never finish, which looks like the kernel waits for a full 64-row tile (a transaction byte count or row bookkeeping derived from tileSizeQ instead of from the box). Unchanged on current main.

Reachability on main. Main never sets mSelectsGroupedMla, so this hang is latent there today. But the precompiled SM107 E4M3 MLA generation cubins are KeepsAb only (Dense / Sliding / Sparse, plus the grouped ones above); there is no SwapsAb E4M3 MLA generation cubin. Non-grouped FP8 MLA generation on SM107 therefore goes to NVRTC and fails with Array.cpp:376 Out-of-bound access (#20377, case 2, including the 2048-query-token shape). I think the two are likely the same coverage gap seen from two sides: shapes without an E4M3 cubin fail in NVRTC, and the existing grouped E4M3 cubins hang when 64 % numHeadsQ != 0. The only FP8 MLA generation path that works on SM107 right now seems to be grouped selection with 4/8/16/32 heads per rank.

Workaround (removes the hang in our runs). In AttentionOp::mlaGeneration (TRTLLM-Gen path), when the selected kernel is grouped E4M3, Q is decode-size (<= 16 tokens) and 64 % numHeadsQ != 0: copy Q into a zero-padded [numTokens, nextPow2(numHeadsQ), 576] FP8 buffer, run the FMHA with mNumHeadsQ = mNumHeadsQPerKv = nextPow2(numHeadsQ) and a padded BF16 output, then copy the real heads of O back (two small row-copy kernels; MLA heads are independent, so padded heads only produce output that is dropped). Not used when softmax stats, attention sinks or an output scale-factor buffer are requested. I have not opened a PR for it: it only matters once grouped selection is enabled (#15009, #20375, #20376), so it would be dead code on main as is; it could go in together with that change.

Possible proper fixes: (a) make the E4M3 grouped kernels handle a partial last tile like the BF16 ones; (b) build the Q TMA box from whole tokens rounded up, with OOB fill and row masking; (c) exclude such head counts from grouped selection for FP8 (on SM107 this currently lands on the NVRTC failure in #20377, so it needs that fixed too).

End-to-end with the workaround (one serving run, TP8 over 2 nodes, 8 query tokens per step, batch 1; the build also includes the split-KV reduction change from #20380; bf16 KV vs fp8 KV):

• Without the workaround: all ranks hang in the first generation warmup, as above. With it: startup completes, including CUDA graph capture for batch 16..1; no new errors in the logs.
• Decode step period (median, ms), bf16 KV -> fp8 KV: 2K 12.57 -> 12.64, 32K 12.87 -> 12.85, 128K 13.62 -> 13.39.
• Per MLA call (8 query tokens, 1 sequence): BF16 grouped with 12 heads vs FP8 grouped with 16 padded heads: 2K 12.16 vs 11.32 us, 32K 21.50 vs 17.01 us, 128K 48.99 vs 34.97 us.
• Accuracy checks: 113 -> 112 correct (prompt tokens identical, one answer differs).
• This was a single end-to-end validation. Batch 16, long-context concurrency and >128K contexts were not exercised with FP8 KV.

Related: #20375, #20376, #20377, #15009.