Skip to content

Update moe gemm #3239

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
May 8, 2026
Merged

Update moe gemm #3239

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
acf982f
update artifacts
IwakuraRein May 6, 2026
1912390
update batched gemm runner
IwakuraRein May 1, 2026
c72926c
update low-latency-gemm-runner
IwakuraRein May 6, 2026
52d54c9
clean up
IwakuraRein May 6, 2026
28ff015
update bmm cubins
IwakuraRein May 6, 2026
faaf038
fix select_kernel_fp8
IwakuraRein May 7, 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
12 changes: 10 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 @@ -91,9 +91,11 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(
auto const configs = bmm.getBatchedGemmConfigs();

mPassingConfigIndices.clear();
auto sm_version = getSMVersion();

Comment on lines +94 to 95

@coderabbitai coderabbitai Bot May 6, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Do not make SM-specific config admission depend on the current device at construction time.

Line 94 snapshots getSMVersion() from whichever CUDA device is current during construction, but the runner later executes against an explicit device argument in run(). That means a runner created on one GPU and reused on another can permanently admit the wrong patchF2fp configs into mPassingConfigIndices, which later drives getValidConfigIndices(). Please either bind the runner to a single device and assert that on use, or move this SM filtering to a path that already knows the target device.

Also applies to: 124-130

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@csrc/trtllm_batched_gemm_runner.cu` around lines 94 - 95, The code snapshots
SM via getSMVersion() at construction which can differ from the device used in
run(), causing incorrect entries in mPassingConfigIndices for patchF2fp; either
(A) bind the runner to a single device by storing the intended device ID in the
constructor and assert that the device argument passed to run() matches that
stored ID before using sm_version (so the original SM-based filter remains
safe), or (B) remove the SM-based filtering from construction and instead
perform the SM filtering of patchF2fp when computing valid configs in
getValidConfigIndices() or inside run() where the explicit device argument is
available (call getSMVersion() for that device there) and then populate
mPassingConfigIndices based on that device-specific SM version. Ensure you
update references to getSMVersion(), mPassingConfigIndices, run(),
getValidConfigIndices(), and patchF2fp accordingly.

for (size_t i = 0; i < bmm.getNumBatchedGemmConfigs(); ++i) {
auto const options = configs[i].mOptions;
auto const config = configs[i];
auto const options = config.mOptions;
auto const tileSize = mOptions.transposeMmaOutput ? options.mTileN : options.mTileM;
// When we include low-latency kernels we can set transposeMmaOutput via constructor
if (options.mDtypeA == mOptions.dtypeA && options.mDtypeB == mOptions.dtypeB &&
Expand All @@ -119,7 +121,13 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(
if ((int64_t)options.mEltwiseActType != (int64_t)mOptions.eltwiseActType) {
continue;
}

// if patchF2fp is enabled, sm100f cubins cannot be used for sm103
if (options.mPatchF2fp && sm_version == 103) {
if (config.mSm != tg::CudaArch::Sm103a) continue;
}
if (options.mPatchF2fp && sm_version == 100) {
if (config.mSm != tg::CudaArch::Sm100a && config.mSm != tg::CudaArch::Sm100f) continue;
}
if (mOptions.transposeMmaOutput && options.mEpilogueTileM == mOptions.epilogueTileM) {
mPassingConfigIndices.push_back(i);
}
Expand Down
16 changes: 8 additions & 8 deletions csrc/trtllm_gemm_runner.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,21 +44,21 @@ struct TrtllmGenGemmRunnerOptions {
int64_t select_kernel_fp8(int32_t M, int32_t N, int32_t K,
const gemm::gemm::GemmInterface& interface) {
static constexpr const char* KERNEL_NAME_HIGH_N_K_RATIO =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s6_et64x8_m64x8x32_c1x1x1_16dp256b_rM_TN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x8x128u2_s6_et64x8_m64x8x32_c1x1x1_rM_TN_"
"transOut_"
"noShflA_dsFp8_schPd2x2x1x3_sm100f";
"noShfl_dsFp8_schPd2x2x1x3_sm100f";

static constexpr const char* KERNEL_NAME_LOW_N_K_RATIO =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s6_et64x32_m64x32x32_c1x1x1_16dp256b_rM_TN_"
"transOut_noShflA_dsFp8_schedS_sm100f";
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s6_et64x32_m64x32x32_c1x1x1_rM_TN_"
"transOut_noShfl_dsFp8_schedS_sm100f";

static constexpr const char* KERNEL_NAME_LARGE_N =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s6_et64x32_m64x32x32_c1x1x1_16dp256b_rM_TN_"
"transOut_noShflA_dsFp8_schPd2x2x1x3_sm100f";
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128u2_s6_et64x32_m64x32x32_c1x1x1_rM_TN_"
"transOut_noShfl_dsFp8_schPd2x2x1x3_sm100f";

static constexpr const char* KERNEL_NAME_DEFAULT =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_c1x1x1_16dp256b_rM_TN_"
"transOut_noShflA_dsFp8_schedS_sm100f";
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x16x128u2_s6_et64x16_m64x16x32_c1x1x1_rM_TN_"
"transOut_noShfl_dsFp8_schedS_sm100f";

double const n_k_ratio = static_cast<double>(N) / static_cast<double>(K);

Expand Down
16 changes: 8 additions & 8 deletions csrc/trtllm_low_latency_gemm_runner.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,28 +63,28 @@ gemm::gemm::GemmData createGemmData(int64_t m, int64_t n, int64_t k) {
*/
int64_t select_kernel(int32_t m, int32_t n, int32_t k, const gemm::gemm::GemmInterface& interface) {
static constexpr const char* KERNEL_MMAN_8_TILEK_128 =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s7_et128x8_m128x8x32_c1x1x1_16dp256b_rM_BN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x8x128_s7_et128x8_m128x8x32_c1x1x1_rM_BN_"
"transOut_schedS_sm100f";
static constexpr const char* KERNEL_MMAN_8_TILEK_256 =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_c1x1x1_16dp256b_rM_BN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x8x256_s4_et128x8_m128x8x32_c1x1x1_rM_BN_"
"transOut_schedS_sm100f";
static constexpr const char* KERNEL_MMAN_16_TILEK_128 =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s7_et128x32_m128x64x32_c1x1x1_16dp256b_rM_BN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x64x128_s7_et128x32_m128x64x32_c1x1x1_rM_BN_"
"transOut_schedS_sm100f";
static constexpr const char* KERNEL_MMAN_16_TILEK_256 =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x64x256_s3_et128x32_m128x64x32_c1x1x1_16dp256b_rM_BN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x64x256_s3_et128x32_m128x64x32_c1x1x1_rM_BN_"
"transOut_schedS_sm100f";
static constexpr const char* KERNEL_MMAN_32_TILEK_128 =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s9_et128x32_m128x32x32_c1x1x1_16dp256b_rM_BN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x128_s9_et128x32_m128x32x32_c1x1x1_rM_BN_"
"transOut_schedS_sm100f";
static constexpr const char* KERNEL_MMAN_32_TILEK_256 =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_c1x1x1_16dp256b_rM_BN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x32x256_s5_et128x32_m128x32x32_c1x1x1_rM_BN_"
"transOut_schedS_sm100f";
static constexpr const char* KERNEL_MMAN_64_TILEK_128 =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s7_et128x16_m128x16x32_c1x1x1_16dp256b_rM_BN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x16x128_s7_et128x16_m128x16x32_c1x1x1_rM_BN_"
"transOut_schedS_sm100f";
static constexpr const char* KERNEL_MMAN_64_TILEK_256 =
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_c1x1x1_16dp256b_rM_BN_"
"gemm_Bfloat16_E4m3E4m3_Fp32_t128x16x256_s5_et128x16_m128x16x32_c1x1x1_rM_BN_"
"transOut_schedS_sm100f";

std::string kernel_name;
Expand Down
8 changes: 4 additions & 4 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 = "1d876ee612888821b168c25ffa75a9dcbb963aaa/fmha/trtllm-gen/"
TRTLLM_GEN_BMM: str = (
"3d9dd08b1691e63e298a7b862d74fd7af3daf594/batched_gemm-4fc8a68-6743435/"
"c21ddd11585c1eea5764927465d0be15dd957e45/batched_gemm-91e0ba0-da44fdf/"
)
TRTLLM_GEN_GEMM: str = (
"31e75d429ff3f710de1251afdd148185f53da44d/gemm-4daf11e-1fddea2/"
"10f64528a1172dae8e29601a3b99ab9dc78d37be/gemm-91e0ba0-2710384/"
)
CUDNN_SDPA: str = "a72d85b019dc125b9f711300cb989430f762f5a6/fmha/cudnn/"
# For DEEPGEMM, we also need to update KernelMap.KERNEL_MAP_HASH in flashinfer/deep_gemm.py
Expand All @@ -160,11 +160,11 @@ class CheckSumHash:
"1abeea012a8779c6df5b84332fad43c6cfc3b257fe5ab883c8ea501464010d16"
)
TRTLLM_GEN_BMM: str = (
"44174e2a08bb427088f5b5443bf0108bb6fb6cb0812ff6018f6418b3d2273824"
"4a3ed9c3dc6547ea3eed01ebda75b0e4322f6c01fc40cd2a4978e4deaba2732a"
)
DEEPGEMM: str = "1a2a166839042dbd2a57f48051c82cd1ad032815927c753db269a4ed10d0ffbf"
TRTLLM_GEN_GEMM: str = (
"64b7114a429ea153528dd4d4b0299363d7320964789eb5efaefec66f301523c7"
"f97f90f9ce1dab73eb3d7c90fca4bbd52687642dd87a79dd10b77d7802b25c33"
)
# SHA256 of the checksums.txt manifest file per cpu-arch/sm-arch,
# NOT hashes of individual kernel .so files.
Expand Down
Loading