FlashInfer fmha_v2 API has some flaws: it introduces some tiny H2Ds or tiny cuda kernels that cause the perf to drop. Although there has been some effort reduce this, in TRTLLM, these preparation steps are all fused into a kernel that runs before the FMHA kernel itself. We might need to do something similar to get the best performance here.
FlashInfer fmha_v2 API has some flaws: it introduces some tiny H2Ds or tiny cuda kernels that cause the perf to drop. Although there has been some effort reduce this, in TRTLLM, these preparation steps are all fused into a kernel that runs before the FMHA kernel itself. We might need to do something similar to get the best performance here.