Skip to content

perf: Port TRT-LLM SM120/SM121 FP4 CUTLASS GEMM optimizations. Add PDL #3026

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bkryu merged 6 commits into from
Apr 10, 2026
+22 −22
Merged

perf: Port TRT-LLM SM120/SM121 FP4 CUTLASS GEMM optimizations. Add PDL #3026

Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e81e707
Port commits from TRTLLM
bkryu Apr 8, 2026
ca74b80
Pingpong and PDL
bkryu Apr 8, 2026
7a9523b
Merge branch 'main' into sm120_mm_fp4_opt
bkryu Apr 9, 2026
a4802a6
Remove ping pong kernel
bkryu Apr 9, 2026
35a7808
Remove ping pong
bkryu Apr 9, 2026
589990a
Remove unnecessary changes to reduce diff
bkryu Apr 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 22 additions & 22 deletions include/flashinfer/gemm/fp4_gemm_template_sm120.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -185,8 +185,8 @@ inline size_t runFp4GemmImpl(void* D, void const* A, void const* B, void const*
std::string(cutlass::cutlassGetStatusString(initStatus)));
}

// Disable PDL since GDC flag is not set
auto runStatus = gemm.run(args, workspace, stream, nullptr, /*enablePDL=*/false);
// Enable PDL — GDC flag (CUTLASS_ENABLE_GDC_FOR_SM100) is set at compile time
auto runStatus = gemm.run(args, workspace, stream, nullptr, /*enablePDL=*/true);
if (runStatus != cutlass::Status::kSuccess) {
throw std::runtime_error(std::string("[FP4 gemm Runner") + scheduler_name + "] " +
"Failed to run cutlass FP4 gemm on sm120/sm121. Error: " +
Expand Down Expand Up @@ -240,39 +240,39 @@ inline size_t runFp4GemmImpl(void* D, void const* A, void const* B, void const*
using ElementCompute = float; \
using ElementAccumulator = float; \
using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; \
using EpilogueTileType = \
cutlass::epilogue::collective::EpilogueTileAuto; /* Always use auto for SM100 */ \
using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; \
using MainloopSchedule = cutlass::gemm::collective::KernelScheduleAuto; \
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto; \
using FusionOperation = \
cutlass::epilogue::fusion::LinearCombination<OutElementType, float, void, float>; \
using ThreadBlockShape = cute::Shape<cute::Int<CTA_M_>, cute::Int<CTA_N_>, cute::Int<CTA_K_>>; \
/* Epilogue: explicit TmaWarpSpecialized schedule (matches TRT-LLM SM120 pattern) */ \
using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< \
Arch, OperatorClass, ThreadBlockShape, ClusterShape, EpilogueTileType, ElementAccumulator, \
ElementCompute, ElementC, LayoutC, AlignmentC, OutElementType, LayoutC, AlignmentC, \
EpilogueSchedule, \
cutlass::epilogue::fusion::LinearCombination<OutElementType, float, void, \
float>>::CollectiveOp; \
Arch, cutlass::arch::OpClassTensorOp, ThreadBlockShape, ClusterShape, \
cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, ElementCompute, \
ElementC, LayoutC, AlignmentC, OutElementType, LayoutC, AlignmentC, \
cutlass::epilogue::TmaWarpSpecialized, FusionOperation>::CollectiveOp; \
\
/* SM120/SM121 BlockScaled - Use nv_float4_t without tuples like example 79 */ \
/* Use fixed 2 stages for SM120/SM121 to meet minimum requirement */ \
/* Dynamic stage carveout adapts pipeline depth to available smem after epilogue */ \
using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< \
Arch, OperatorClass, ElementA, LayoutA, AlignmentA, ElementB, LayoutB, AlignmentB, \
ElementAccumulator, ThreadBlockShape, ClusterShape, \
cutlass::gemm::collective::StageCount<2>, /* Fixed 2 stages for SM120 */ \
cutlass::gemm::collective::KernelScheduleAuto>::CollectiveOp; \
cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>( \
sizeof(typename CollectiveEpilogue::SharedStorage))>, \
cutlass::gemm::KernelTmaWarpSpecializedCooperative>::CollectiveOp; \
Comment thread
bkryu marked this conversation as resolved.
\
/* Two scheduler options for different workloads */ \
/* See: https://github.com/NVIDIA/cutlass/blob/main/examples/79_blackwell_geforce_gemm */ \
using TileSchedulerTag = cutlass::gemm::StaticPersistentScheduler; \
\
/* Option 1: Default scheduler (void) - Data Parallel, good for regular shapes */ \
using GemmKernelDefault = cutlass::gemm::kernel::GemmUniversal< \
cute::Shape<int, int, int, int>, /* Indicates ProblemShape */ \
CollectiveMainloop, CollectiveEpilogue, void>; /* Default DP scheduler */ \
/* Option 1: Persistent scheduler - reduced launch overhead, good default */ \
using GemmKernelDefault = \
cutlass::gemm::kernel::GemmUniversal<cute::Shape<int, int, int, int>, CollectiveMainloop, \
CollectiveEpilogue, TileSchedulerTag>; \
Comment thread
coderabbitai[bot] marked this conversation as resolved.
\
/* Option 2: StreamK scheduler - better load balancing for small M/N, large K */ \
using GemmKernelStreamK = cutlass::gemm::kernel::GemmUniversal< \
cute::Shape<int, int, int, int>, /* Indicates ProblemShape */ \
CollectiveMainloop, CollectiveEpilogue, \
cutlass::gemm::StreamKScheduler>; /* StreamK scheduler */ \
using GemmKernelStreamK = \
cutlass::gemm::kernel::GemmUniversal<cute::Shape<int, int, int, int>, CollectiveMainloop, \
CollectiveEpilogue, cutlass::gemm::StreamKScheduler>; \
\
using GemmDefault = typename cutlass::gemm::device::GemmUniversalAdapter<GemmKernelDefault>; \
using GemmStreamK = typename cutlass::gemm::device::GemmUniversalAdapter<GemmKernelStreamK>; \
Expand Down
Loading