Skip to content

perf: Update trtllm-gen batched GEMM kernels - faster, more NVFP4 tile dims, MXFP8 with relu2 act #2667

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
aleozlx merged 6 commits into from
Mar 3, 2026
Merged

perf: Update trtllm-gen batched GEMM kernels - faster, more NVFP4 tile dims, MXFP8 with relu2 act #2667

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
8b0cf4b
Update trtllm-gen batched GEMM kernels with newer & faster ones, more…
amitz-nv Mar 2, 2026
b3568ef
upstream header file deletion
jimmyzho Mar 2, 2026
0a10131
Merge branch 'main' into update-batched-gemm-faster-gemms
jimmyzho Mar 2, 2026
109b6b3
add /
jimmyzho Mar 2, 2026
351a374
Merge branch 'update-batched-gemm-faster-gemms' of github.com:amitz-n…
jimmyzho Mar 2, 2026
ca601c4
Update csrc/trtllm_batched_gemm_runner.cu
aleozlx Mar 3, 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
5 changes: 3 additions & 2 deletions csrc/trtllm_batched_gemm_runner.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -258,8 +258,9 @@ void TrtllmGenBatchedGemmRunner::run(
// FIXME once we start using all-reduce in the epilogue of the bmm this can be moved elsewhere
bmm.runInitBeforeWorldSync(config, gemmData, static_cast<void*>(stream));

auto const err = bmm.run(config, workspace, gemmData, static_cast<void*>(stream),
multiProcessorCount, enable_pdl, globalTrtllmGenBatchedGemmModuleCache);
auto const err =
bmm.run(config, workspace, gemmData, static_cast<void*>(stream), multiProcessorCount,
enable_pdl, /*pinnedHostBuffer=*/nullptr, globalTrtllmGenBatchedGemmModuleCache);

FLASHINFER_CHECK(err == 0,
"Error occurred when running GEMM!"
Expand Down
6 changes: 3 additions & 3 deletions flashinfer/artifacts.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,10 +137,10 @@ class ArtifactPath:

TRTLLM_GEN_FMHA: str = "e86f0e45764555d070c3d143b4caaea61a45b777/fmha/trtllm-gen/"
TRTLLM_GEN_BMM: str = (
"de418500652abd6325552396fc931d5165f12804/batched_gemm-2a674db-3a84a12/"
"b55211623be7f5697c5262ffd8361fc06c147bc9/batched_gemm-b3c1646-c111d7c/"
)
TRTLLM_GEN_GEMM: str = (
"1fddc48b7b48af33914d040051b3e2ee9ba4701e/gemm-145d1b1-9b113e3"
"1fddc48b7b48af33914d040051b3e2ee9ba4701e/gemm-145d1b1-9b113e3/"
)
CUDNN_SDPA: str = "a72d85b019dc125b9f711300cb989430f762f5a6/fmha/cudnn/"
# For DEEPGEMM, we also need to update KernelMap.KERNEL_MAP_HASH in flashinfer/deep_gemm.py
Expand All @@ -158,7 +158,7 @@ class CheckSumHash:
"c4c93904a4c72b8a3d0d5c525c6decb71c835b477d7d75651ecaaa7007c5a3ef"
)
TRTLLM_GEN_BMM: str = (
"69ade639fa6640ea90261599581c8cda50f8f6b9ac9be07a3ca38f3bf5458452"
"0af823880730c4f0b3832d2208fab035946694b83444410b9309db5613d60195"
)
DEEPGEMM: str = "1a2a166839042dbd2a57f48051c82cd1ad032815927c753db269a4ed10d0ffbf"
TRTLLM_GEN_GEMM: str = (
Expand Down
Loading