Repository navigation
[Bug]: Grouped MLA generation cubins are registered as Dense in kernelMetaInfo.h but named and selected as Causal #20376
Description
System Info
TensorRT-LLM main, cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/cubin/kernelMetaInfo.h (artifacts from #19591). On release/1.3 the table is split per architecture since #20022; the same rows are in kernelMetaInfoSm100f.h and kernelMetaInfoSm107.h (see the location list below). Seen on SM107; the SM100f rows have the same pattern.
Who can help?
@pengbowang-nv @farazkh80 (#15009)
Information
- The official example scripts
- My own modified scripts
Reproduction
All 36 grouped MLA generation cubins (fmhaSm100fKernel_... and fmhaSm107aKernel_Qkv{Bfloat16,E4m3}OBfloat16HQk576HV512{,HVPerCta128,HVPerCta256}PagedKvCausalP32...VarSeqQ64Kv128{Static,Persistent}GroupedKeepsAbForGen; per architecture 9 with BF16 and 9 with E4M3 Q/KV, 18 for kSM_100f, 18 for kSM_107) have Causal in the function name, but their mMaskType field is 0 (TrtllmGenAttentionMaskType::Dense). Locations:
- main @ 13e19d6,
kernelMetaInfo.h: SM100f BF16 lines 19057-19110, SM100f E4M3 19347-19400, SM107 BF16 20903-21019, SM107 E4M3 21907-22023. - release/1.3 @ c9d1169 (after [https://nvbugs/6812429][fix] Add SM107 P128 trtllm-gen FMHA generation kernels #20022):
kernelMetaInfoSm100f.hlines 243-296 (BF16) and 533-586 (E4M3);kernelMetaInfoSm107.hlines 887-1003 (BF16) and 1835-1951 (E4M3). The values are unchanged by the split. The grouped P128 E4M3 H128 cubins added there for SM107 are registered withmMaskType = 1(Causal), so only the HQk576 grouped MLA rows are affected.
mMaskType is part of the hash that the cubin path looks up (hashFromFmhaOptions -> hashID). With FmhaOptions::mSelectsGroupedMla = true, the autotuner selects the grouped kernel with a causal mask for MLA generation with multiple query tokens, so the lookup fails (Trtllm-gen kernels not found), and the only remaining route is NVRTC. On SM107, NVRTC currently fails for these configurations (separate issue). We had to map Causal back to Dense for grouped MLA kernels in algoFilterForCubinPath before the cubins could be used.
With that mapping, the cubins give the expected causal result: compared with the non-grouped kernel on SM107, outputs are bit-identical for KV lengths 16 and 64, and at most 2 bf16 ulp apart from 256 up to 128K. So the kernels behave as their names say, and only the metadata is off. #15009 also mentions that its re-exported rows were changed to maskType=Causal "to match the shipped lib's grouped mask rewrite".
Expected behavior
The registered mask type matches what the kernel implements and what the selector asks for, so that grouped MLA generation cubins can be found without a remapping.
actual behavior
mMaskType = Dense for kernels named ...Causal...; the cubin lookup misses them when grouped MLA selection is enabled.
additional notes
This does not affect main today because mSelectsGroupedMla is never set in TensorRT-LLM. It blocks enabling it (#15009), together with the missing cumSeqLensQ/Kv in the MLA generation path (#20375). If the right fix is on the export side, a regenerated kernelMetaInfo.h (kernelMetaInfoSm100f.h / kernelMetaInfoSm107.h after #20022) would be enough; otherwise I can send the remapping as a small PR.
Related: #20375, #20377, #20382 (FP8-KV grouped MLA hang for head counts that do not divide 64).
Activity
lucifer1004 commentedon Oct 15, 2026
Updated the description for #20022, which split the trtllm-gen FMHA metadata table into per-architecture files on release/1.3. The 36 rows are now in
kernelMetaInfoSm100f.handkernelMetaInfoSm107.hthere, still withmMaskType = 0, and the description lists the line ranges on both main and release/1.3. Main still has the singlekernelMetaInfo.h. I also corrected the cubin pattern: each architecture has 9 BF16 and 9 E4M3 rows.saiyambharara commentedon Oct 20, 2026
@lucifer1004 @pengbowang-nv I'd like to take this if you haven't started. I checked all 8078 rows of
kernelMetaInfo.hon main. Every MLA gen row is registered Dense, while the 720 non-MLA grouped gen kernels named Causal are registered Causal. So it looks like the export applied the MLA gen Dense rule to the grouped variants too.A few questions before I write anything:
algoFilterForCubinPath?I can't run SM100f/SM107 locally, so I'd rely on CI for GPU coverage.
lucifer1004 commentedon Oct 22, 2026
@saiyambharara thanks, please go ahead. I haven't started on a fix.
selectGroupedMlaGenerationKernel()) is in the prebuilt library, so I can only describe what we observed. WithmSelectsGroupedMlaenabled for MLA generation with more than one query token, every lookup we saw asked for the grouped kernel with a Causal mask. That is why the cubin path missed the Dense rows. We never saw a Dense request for a grouped MLA generation kernel. This matches the note in [None][feat] Add test-only grouped MLA generation selection probe #15009 that its rows were changed tomaskType=Causal"to match the shipped lib's grouped mask rewrite". @pengbowang-nv can confirm whether the library ever requests Dense for these kernels.algoFilterForCubinPathwas only a local workaround to get going. These files are generated by the export, though, so the owners may prefer to fix the export and carry the change with the next drop. I'd defer to them on that.saiyambharara commentedon Oct 22, 2026
@lucifer1004 the PR is up: #20838. It registers the 36 rows as Causal and adds a test that checks each row's mask against its kernel name. It builds on your SM107 comparison, thanks.
@pengbowang-nv as lucifer1004 suggested, could you confirm the library never requests Dense for these grouped MLA kernels? And that the binaries were built with the causal mask, so the registration and the binaries agree? The table comes from the export, so the next drop would also need to carry the Causal registration or this fix gets overwritten.
[NVIDIA#20376][fix] Register grouped MLA generation cubins as Causal
pengbowang-nv commentedon Oct 23, 2026
Hi @saiyambharara , thanks for investigating.
Thanks!