feat: add SM120 fmha_v2 kernels to AOT pip wheel builds (#2885)
## Summary `gen_trtllm_fmha_v2_sm120_module()` exists in `jit/attention/modules.py` and the JIT runtime path (`generate_kernels.py`) already dispatches to it correctly. However, `aot.py`'s `gen_all_modules()` — which drives the pip wheel AOT build — was missing it from the `has_sm120 or has_sm121` section. This means SM120/SM121 devices using a pip wheel would never get the fmha_v2 SM120 kernels compiled into the wheel, and would have to fall back to slower paths. **Fix:** Add `gen_trtllm_fmha_v2_sm120_module()` to the `has_sm120 or has_sm121` block in `aot.py`, alongside the other SM120 modules (fused MOE, GEMM, FP4 quantization). No behavior change for JIT users; only affects AOT pip wheel builds. Addresses the AOT gap noted in #2555. Contributed by Second Nature Computing (https://joinsecondnature.com) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Chores** * Expanded optimized inference module support for SM120 and SM121 GPUs to include attention kernels in addition to existing fused MoE and GEMM optimizations. * Increased runtime coverage and readiness for attention-heavy workloads on those architectures, improving performance consistency for models using attention. [](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/2885) <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
B
Blake Ledden committed
4dba29fbd626fc0ac28eea22101db8f8588d78b7
Parent: 6885e76
Committed by GitHub <noreply@github.com>
on 5/11/2026, 10:43:10 PM