📌 Description

Summary

• Port TRT-LLM's SM120 CUTLASS FP4 GEMM template parameters for better kernel performance
  • TRT-LLM reference file: nvfp4_nvfp4_gemm_template_sm120.h
• Enable PDL (programmatic dependent launch) for SM120/SM121 FP4 GEMM

Details

TRT-LLM kernel parameter port

Updated the SM120 FP4 GEMM kernel template (fp4_gemm_template_sm120.h) to match TRT-LLM's optimized configuration:

1. KernelScheduleAuto → KernelTmaWarpSpecializedCooperative — Removes auto-resolution ambiguity. Explicitly selects the cooperative warp-specialized mainloop where dedicated warps handle TMA loads while others run MMA.
2. StageCount<2> → StageCountAutoCarveout<sizeof(EpilogueSharedStorage)> — Instead of a hardcoded 2-stage pipeline, dynamically computes how many stages fit in shared memory after reserving space for the epilogue. More stages = better latency hiding of TMA loads behind MMA compute.
3. void (scheduler) → StaticPersistentScheduler — The old void scheduler launches one CTA per output tile (data-parallel). The persistent scheduler launches fewer CTAs that loop over multiple tiles, reducing kernel launch overhead — most impactful at small M where the kernel is short.
4. EpilogueScheduleAuto → TmaWarpSpecialized — Explicitly selects TMA-based epilogue with warp specialization for output writes, rather than relying on auto-resolution.
5. OpClassBlockScaledTensorOp → OpClassTensorOp (in epilogue builder only) — The epilogue doesn't need the block-scaled op class (that's only for the mainloop MMA). Using OpClassTensorOp matches what TRT-LLM uses and avoids potential misrouting in the epilogue collective builder.

The persistent scheduler reduces launch overhead (most impactful for small M), dynamic stage carveout adapts pipeline depth to available smem, and explicit cooperative warp specialization avoids auto-resolution ambiguity.

PDL enablement

Changed enablePDL=false → enablePDL=true in runFp4GemmImpl. The CUTLASS_ENABLE_GDC_FOR_SM100=1 compile flag is already set (since PR #2780), and both SM100 FP4 GEMM and SM120 MXFP8 GEMM already run with PDL enabled. The false was a stale leftover.

Performance Numbers on RTX 5090 (SM120) and DGX Spark (SM121)

Performance changes most relevant to SM120. Very minor on Spark

🔍 Related Issues

#3013

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• Tests have been added or updated as needed.
• All tests are passing (unittest, etc.).

Reviewer Notes

Summary by CodeRabbit

• Optimizations
  • Optimized FP4 GEMM kernel execution by refining scheduler configuration, epilogue fusion operations, and memory staging parameters to improve performance and resource utilization.