Open-sourcing custom attention kernels in TensorRT-LLM

[juney-nvidia]

Dear Community,

We're excited to share significant updates regarding our custom attention kernels:

• Decoding Phase Kernels (XQA) Now Available
  The high-performance kernels for decoding-phase attention have been fully open-sourced in PR #3762. This fulfills numerous community requests to enable deeper customization of attention mechanisms.

• Prefill Phase Kernels Coming Soon
  Our team is actively working to open-source the optimized prefill-phase attention kernels, with completion targeted by mid-May 2025.
  Update on May 14: the high-performance kernels for prefill-phase attention have been fully open-sourced in PR #4185!

These releases empower developers to:

• Implement novel attention optimizations
• Experiment with model architecture modifications
• Contribute performance improvements back to the core project

We encourage all community members to explore these components and share your enhancements. Your contributions will help advance LLM inference efficiency across NVIDIA GPUs.

Let's build the future of accelerated LLMs together!

Best regards,
The TensorRT-LLM Team

[sfzhu93]

Thanks for sharing the XQA kernel! This is more straightforward than CUTLASS to understand some of the low-level details. Thus, I got one probably noob question to ask: why there are only wrappers for mma.* and wgmma.* PTX instructions, but wmma.* instructions are not used?

[juney-nvidia]

@lowsfer for vis of this community question :)

June

[sfzhu93]

Thanks! I think I got the answer somewhere else: mma.* and wgmma.* can generally cover all usages of gemm. wmma.* is mostly a legacy instruction.

[lowsfer]

Both mma (all arch since Volta) and wgmma (Hopper-only) are direct mapping of hardware instructions. wmma is a wrapper of mma (and other helper instructions) which gives us less control than the lower-level mma. It was exposed before mma went public.

[Aya-ZIbra]

Thanks for sharing these kernels! Are the kernels compatible with Blackwell? Any plans for Blackwell-specific optimizations?