Skip to content

fix(trtllm): restrict routed-MoE backends to supported architectures (release port for #4107) #4230

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
kahyunnam merged 1 commit into from
Jul 29, 2026
Merged

fix(trtllm): restrict routed-MoE backends to supported architectures (release port for #4107) #4230

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
87 changes: 70 additions & 17 deletions csrc/trtllm_batched_gemm_runner.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
* limitations under the License.
*/

#include <algorithm>
#include <cstring>
#include <vector>

Expand All @@ -36,6 +37,49 @@ using namespace batchedGemm::trtllm::gen;

static BatchedGemmInterface::ModuleCache globalTrtllmGenBatchedGemmModuleCache;

namespace {

// The trtllm-gen cubin manifest is a downloaded artifact, so which architectures
// it actually covers is not knowable at compile time. Encode only the
// cubin-arch -> SM-version compatibility rules here and let `config.mSm` decide
// what is available. Unknown cubin families are rejected so that a newly shipped
// one fails loudly instead of being mis-dispatched onto hardware that cannot run
// it (see #4107).
bool isArchCompatible(int smVersion, tg::CudaArch cubinArch) {
switch (cubinArch) {
case tg::CudaArch::Sm100a:
return smVersion == 100;
case tg::CudaArch::Sm100f:
return smVersion == 100 || smVersion == 103;
case tg::CudaArch::Sm103a:
return smVersion == 103;
#ifdef TLLM_RUBIN_FEATURES
// CudaArch::Sm107a only exists in the Rubin cubin pin's generated headers,
// which is also the only build that defines TLLM_RUBIN_FEATURES. sm107 is
// unreachable in the non-Rubin module: get_trtllm_moe_sm100_module() picks
// the Rubin variant by device compute capability before this runs.
case tg::CudaArch::Sm107a:
return smVersion == 107;
#endif
default:
return false;
}
}

void checkPassingConfigIndex(std::vector<int64_t> const& passingConfigIndices,
int32_t configIndex) {
auto const it = std::find(passingConfigIndices.begin(), passingConfigIndices.end(), configIndex);
if (it == passingConfigIndices.end()) {
std::ostringstream msg;
msg << "Config index " << configIndex
<< " is not in this runner's compatible config set (device architecture or GEMM options "
"mismatch)";
FLASHINFER_ERROR(msg.str());
}
}

} // namespace

std::vector<int64_t> prioritizePredefinedConfigs(
int m, int n, int k, std::vector<int64_t> const& sortedIndices,
batchedGemm::batchedGemm::BatchedGemmConfig const* configs) {
Expand Down Expand Up @@ -141,23 +185,9 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(
if ((int64_t)options.mEltwiseActType != (int64_t)mOptions.eltwiseActType) {
continue;
}
#ifdef TLLM_RUBIN_FEATURES
// CudaArch::Sm107a only exists in the Rubin cubin pin's generated
// headers (TRTLLM_GEN_BMM_RUBIN). The default pin has no such enumerator,
// so this comparison must not be compiled into the non-Rubin module.
// sm_version == 107 is unreachable there anyway: the Rubin module is
// selected by device compute capability before this runs.
if (sm_version == 107) {
if (config.mSm != tg::CudaArch::Sm107a) continue;
}
#endif
if (sm_version == 103) {
if (config.mSm != tg::CudaArch::Sm103a && config.mSm != tg::CudaArch::Sm100f) continue;
if (options.mPatchF2fp && config.mSm != tg::CudaArch::Sm103a) continue;
}
if (sm_version == 100) {
if (config.mSm != tg::CudaArch::Sm100a && config.mSm != tg::CudaArch::Sm100f) continue;
}
if (!isArchCompatible(sm_version, config.mSm)) continue;
// Sm100f cubins miss the f2fp patch, so sm103 must fall back to Sm103a for it.
if (sm_version == 103 && options.mPatchF2fp && config.mSm != tg::CudaArch::Sm103a) continue;
if (mOptions.transposeMmaOutput && options.mEpilogueTileM == mOptions.epilogueTileM) {
// Skip cubins with clusterZ > 1 due to correctness issues described in
// https://github.com/flashinfer-ai/flashinfer/issues/3197
Expand All @@ -167,6 +197,23 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(
}
}

if (mPassingConfigIndices.empty()) {
// Distinguish "this GPU has no cubins at all" from "no cubin matches these GEMM
// options". The former is the common failure on unsupported hardware, and the
// option dump below would send users looking in entirely the wrong place.
bool anyArchCompatible = false;
for (size_t i = 0; i < bmm.getNumBatchedGemmConfigs(); ++i) {
if (isArchCompatible(sm_version, configs[i].mSm)) {
anyArchCompatible = true;
break;
}
}
std::ostringstream arch_msg;
arch_msg << "The trtllm-gen batched GEMM cubin manifest contains no kernels runnable on sm"
<< sm_version << "; this backend currently ships cubins for sm100, sm103 and sm107.";
FLASHINFER_CHECK(anyArchCompatible, arch_msg.str());
}

std::ostringstream error_msg;
error_msg << "No kernel found for the given options: "
<< "mDtypeA: " << tg::dtypeToString(mOptions.dtypeA)
Expand All @@ -189,6 +236,7 @@ TrtllmGenBatchedGemmRunner::TrtllmGenBatchedGemmRunner(
size_t TrtllmGenBatchedGemmRunner::getWorkspaceSizeInBytes(
int32_t m, int32_t n, int32_t k, std::vector<int32_t> const& batchedTokens, int32_t numTokens,
int32_t numBatches, int32_t maxNumCtasInBatchDim, int32_t configIndex) const {
checkPassingConfigIndex(mPassingConfigIndices, configIndex);
BatchedGemmData gemmData{};
gemmData.mProblemDimensions.mNumBatches = numBatches;
gemmData.mProblemDimensions.mNumTokens = numTokens;
Expand Down Expand Up @@ -232,7 +280,12 @@ void TrtllmGenBatchedGemmRunner::run(
BatchedGemmData gemmData{};

auto const configs = bmm.getBatchedGemmConfigs();
auto const numConfigs = bmm.getNumBatchedGemmConfigs();

FLASHINFER_CHECK(
configIndex >= 0 && static_cast<size_t>(configIndex) < static_cast<size_t>(numConfigs),
"Config index", configIndex, "is out of range; the manifest has", numConfigs, "configs");
checkPassingConfigIndex(mPassingConfigIndices, configIndex);
auto const& config = configs[configIndex];
// printf("running config %d: %s\n", configIndex, config.mFunctionName);

Expand Down
65 changes: 65 additions & 0 deletions csrc/trtllm_gemm_runner.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@

#include <cuda.h>

#include <algorithm>
#include <sstream>
#include <string>

#include "flashinfer/exception.h"
Expand All @@ -33,6 +35,38 @@ static thread_local gemm::gemm::GemmInterface::ModuleCache globalTrtllmGenGemmMo

namespace flashinfer {

namespace {

// The trtllm-gen cubin manifest is a downloaded artifact, so which architectures
// it actually covers is not knowable at compile time. Encode only the
// cubin-arch -> SM-version compatibility rules here and let `config.mSm` decide
// what is available. Unknown cubin families are rejected so that a newly shipped
// one fails loudly instead of being mis-dispatched onto hardware that cannot run
// it (see #4107).
bool isArchCompatible(int smVersion, gemm::trtllm::gen::CudaArch cubinArch) {
using CudaArch = gemm::trtllm::gen::CudaArch;
switch (cubinArch) {
case CudaArch::Sm100a:
return smVersion == 100;
case CudaArch::Sm100f:
return smVersion == 100 || smVersion == 103;
case CudaArch::Sm103a:
return smVersion == 103;
#ifdef TLLM_RUBIN_FEATURES
// CudaArch::Sm107a only exists in the Rubin cubin pin's generated headers,
// which is also the only build that defines TLLM_RUBIN_FEATURES. sm107 is
// unreachable in the non-Rubin module: get_trtllm_gemm_module() picks the
// Rubin variant by device compute capability before this runs.
case CudaArch::Sm107a:
return smVersion == 107;
#endif
default:
return false;
}
}

} // namespace

struct TrtllmGenGemmRunnerOptions {
gemm::trtllm::gen::Dtype eltType;
gemm::trtllm::gen::Dtype outputType;
Expand Down Expand Up @@ -110,6 +144,7 @@ class TrtllmGenGemmRunner {
auto const configs = gemm.getGemmConfigs();

mPassingConfigIndices.clear();
int const sv = getSMVersion();

for (size_t i = 0; i < gemm.getNumGemmConfigs(); ++i) {
auto const options = configs[i].mOptions;
Expand All @@ -118,10 +153,30 @@ class TrtllmGenGemmRunner {
options.mTransposeMmaOutput == mOptions.transposeMmaOutput &&
options.mSfLayoutB == mOptions.sfLayoutB &&
options.mLayoutA == mOptions.layoutA) { // FIXME(siyuanf): expose matrix layout to user
if (!isArchCompatible(sv, configs[i].mSm)) continue;
mPassingConfigIndices.push_back(i);
}
}

if (mPassingConfigIndices.empty()) {
// Distinguish "this GPU has no cubins at all" from "no cubin matches these GEMM
// options". The former is the common failure on unsupported hardware, and the
// option dump below would send users looking in entirely the wrong place.
bool anyArchCompatible = false;
for (size_t i = 0; i < gemm.getNumGemmConfigs(); ++i) {
if (isArchCompatible(sv, configs[i].mSm)) {
anyArchCompatible = true;
break;
}
}
if (!anyArchCompatible) {
std::ostringstream arch_msg;
arch_msg << "The trtllm-gen GEMM cubin manifest contains no kernels runnable on sm" << sv
<< "; this backend currently ships cubins for sm100, sm103 and sm107.";
FLASHINFER_ERROR(arch_msg.str());
}
}

FLASHINFER_CHECK(mPassingConfigIndices.size() > 0,
"No valid tactic found for the given options",
"mDtypeA: ", gemm::trtllm::gen::dtypeToString(mOptions.eltType),
Expand All @@ -130,11 +185,20 @@ class TrtllmGenGemmRunner {
"mSfLayoutB: ", gemm::trtllm::gen::sfLayoutToString(mOptions.sfLayoutB));
}

void checkPassingConfigIndex(int64_t tactic) const {
auto it = std::find(mPassingConfigIndices.begin(), mPassingConfigIndices.end(), tactic);
TVM_FFI_ICHECK(it != mPassingConfigIndices.end())
<< "Tactic " << tactic
<< " is not in this runner's compatible config set (device architecture or GEMM options "
"mismatch)";
}

int64_t getWorkspaceSizeInBytes(int64_t m, int64_t n, int64_t k, int64_t tactic) {
auto gemm = gemm::gemm::GemmInterface();
auto const configs = gemm.getGemmConfigs();
FLASHINFER_CHECK(tactic >= 0 && tactic < gemm.getNumGemmConfigs(),
"Invalid tactic in getWorkspaceSizeInBytes");
checkPassingConfigIndex(tactic);
auto const config = configs[tactic];

gemm::gemm::GemmData gemmData;
Expand All @@ -156,6 +220,7 @@ class TrtllmGenGemmRunner {
auto gemm = gemm::gemm::GemmInterface();
auto const configs = gemm.getGemmConfigs();
TVM_FFI_ICHECK(tactic >= 0 && tactic < gemm.getNumGemmConfigs()) << "Invalid tactic id in run";
checkPassingConfigIndex(tactic);
auto const& config = configs[tactic];
TVM_FFI_ICHECK(config.mOptions.mSfLayoutB == mOptions.sfLayoutB) << "Invalid sf layout in run";

Expand Down
32 changes: 20 additions & 12 deletions flashinfer/fused_moe/api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,18 +223,25 @@ def __repr__(self) -> str:
# Backend configs — each declares hardware preconditions
# ---------------------------------------------------------------------------

# Architectures the TRT-LLM routed-MoE cubin manifest ships kernels for. The
# manifest is a downloaded artifact, so this cannot be derived at import time and
# has to be kept in sync with the cubin-arch compatibility rules in
# csrc/trtllm_batched_gemm_runner.cu. SM110/120/121 need upstream cubins first;
# claiming them here makes the batched-GEMM runner abort at dispatch (#4107).
_TRTLLM_ROUTED_ARCHS = (100, 103, 107)

# The FP8 kernels are validated on the SM100 family only — the outer JIT module
# compiles for major 12 as well, but those cubins fail at runtime on SM120/121.
_TRTLLM_ROUTED_FP8_ARCHS = (100, 103)


@dataclass(frozen=True)
class TrtllmFp4Config:
"""TensorRT-LLM FP4 block-scale backend."""

@classmethod
def supported(cls, arch: int) -> bool:
# SM100+ only: the routed runner delegates to the trtllm-gen sm100
# module, which core.is_trtllm_moe_supported() gates on major >= 10.
# Returning True on SM90 would mark the backend available on H100 and
# then fail at dispatch.
return arch >= 100
return arch in _TRTLLM_ROUTED_ARCHS

@staticmethod
def prepare_weights(
Expand Down Expand Up @@ -274,10 +281,7 @@ class TrtllmFp8BlockConfig:

@classmethod
def supported(cls, arch: int) -> bool:
# The available TRTLLM block-FP8 BMM cubins are validated only on the
# SM100 family. The outer JIT can compile for major 12, but its FP8
# kernels currently fail at runtime on SM120/121.
return arch in (100, 103)
return arch in _TRTLLM_ROUTED_FP8_ARCHS

@staticmethod
def prepare_weights(
Expand Down Expand Up @@ -327,7 +331,7 @@ class TrtllmFp8PerTensorConfig:

@classmethod
def supported(cls, arch: int) -> bool:
return arch in (100, 103)
return arch in _TRTLLM_ROUTED_FP8_ARCHS

@staticmethod
def prepare_weights(
Expand Down Expand Up @@ -375,7 +379,7 @@ class TrtllmBf16Config:

@classmethod
def supported(cls, arch: int) -> bool:
return arch >= 100
return arch in _TRTLLM_ROUTED_ARCHS

@staticmethod
def prepare_weights(
Expand Down Expand Up @@ -415,7 +419,11 @@ class TrtllmMxInt4Config:

@classmethod
def supported(cls, arch: int) -> bool:
return arch >= 100
# Same trtllm-gen routed batched-GEMM path as the FP4/BF16 backends, so it
# inherits the same manifest coverage. Whether sm107a MxE2m1 cubins exist
# is not verifiable from this repo; this only narrows the previous
# ``arch >= 100``, leaving SM100/103/107 behaviour unchanged.
return arch in _TRTLLM_ROUTED_ARCHS

def __repr__(self) -> str:
return "TrtllmMxInt4Config()"
Expand Down
Loading
Loading