[None][perf] Optimize MiniMax-M3 MXFP8 GEMMs

cpp/tensorrt_llm/thop/mxfp8Gemm.cpp

@@ -25,6 +25,11 @@
 
 #include <cstdint>
 #include <cuda_fp16.h>
+#include <map>
+#include <mutex>
+#include <optional>
+#include <shared_mutex>
+#include <tuple>
 #include <vector>
 
 namespace tkc = tensorrt_llm::cutlass_extensions;
@@ -39,11 +44,103 @@ namespace torch_ext
 namespace
 {
 
+constexpr int64_t kMxfp8LargeMMin = 6553;
+constexpr int64_t kMxfp8M8kBucket = 8192;
+constexpr int64_t kMxfp8M16kMin = 13106;
+constexpr int64_t kMxfp8M16kBucket = 16384;
+constexpr int64_t kMxfp8M32kMin = 19659;
+constexpr int64_t kMxfp8M32kBucket = 32768;
+constexpr int64_t kMxfp8TacticCacheMiss = -2;
+
+int getMxfp8SmVersion()
+{
+    // PyExecutor binds one GPU architecture per rank.
+    static int const smVersion = tensorrt_llm::common::getSMVersion();
+    return smVersion;
+}
+
+int64_t getMxfp8TuningBucket(int64_t const m)
+{
+    if (m < kMxfp8LargeMMin)
+    {
+        return m;
+    }
+    if (m <= kMxfp8M8kBucket)
+    {
+        return kMxfp8M8kBucket;
+    }
+    if (m >= kMxfp8M16kMin && m <= kMxfp8M16kBucket)
+    {
+        return kMxfp8M16kBucket;
+    }
+    if (m >= kMxfp8M32kMin && m <= kMxfp8M32kBucket)
+    {
+        return kMxfp8M32kBucket;
+    }
+    return m;
+}
+
+using Mxfp8TacticCacheKey = std::tuple<int, at::ScalarType, int64_t, int64_t, int64_t>;
+
+struct Mxfp8TacticCacheEntry
+{
+    tkc::CutlassGemmConfig config;
+    int64_t tactic;
+};
+
+using Mxfp8TacticCache = std::map<Mxfp8TacticCacheKey, Mxfp8TacticCacheEntry>;
+
+Mxfp8TacticCache& getMxfp8TacticCache()
+{
+    static Mxfp8TacticCache cache;
+    return cache;
+}
+
+std::shared_mutex& getMxfp8TacticCacheMutex()
+{
+    static std::shared_mutex mutex;
+    return mutex;
+}
+
+Mxfp8TacticCacheKey makeMxfp8TacticCacheKey(
+    int64_t const m, int64_t const n, int64_t const k, at::ScalarType const outputDtype)
+{
+    return {getMxfp8SmVersion(), outputDtype, getMxfp8TuningBucket(m), n, k};
+}
+
+std::optional<Mxfp8TacticCacheEntry> findMxfp8TacticCacheEntry(
+    int64_t const m, int64_t const n, int64_t const k, at::ScalarType const outputDtype)
+{
+    std::shared_lock<std::shared_mutex> lock(getMxfp8TacticCacheMutex());
+    auto const& cache = getMxfp8TacticCache();
+    auto const iterator = cache.find(makeMxfp8TacticCacheKey(m, n, k, outputDtype));
+    if (iterator == cache.end())
+    {
+        return std::nullopt;
+    }
+    return iterator->second;
+}
+
+void cacheMxfp8Tactic(int64_t const m, int64_t const n, int64_t const k, at::ScalarType const outputDtype,
+    tkc::CutlassGemmConfig const& config, int64_t const tactic)
+{
+    std::unique_lock<std::shared_mutex> lock(getMxfp8TacticCacheMutex());
+    getMxfp8TacticCache().insert_or_assign(
+        makeMxfp8TacticCacheKey(m, n, k, outputDtype), Mxfp8TacticCacheEntry{config, tactic});
+}
+
+void clearMxfp8CachedTactics()
+{
+    std::unique_lock<std::shared_mutex> lock(getMxfp8TacticCacheMutex());
+    getMxfp8TacticCache().clear();
+}
+
 tkc::CutlassGemmConfig getDefaultMxfp8GemmConfig()
 {
     // Reuse the same default tile/cluster as MXFP8xMXFP4 -- the B operand is
     // 2x wider in MXFP8xMXFP8, but the same 4x4 cluster/256x256 tile shape is
-    // a reasonable starting point on B200.
+    // a reasonable Blackwell fallback before startup tuning populates the
+    // native tactic cache.
     return tkc::CutlassGemmConfig(tkc::CutlassTileConfigSM100::CtaShape128x256x256B, tkc::MainloopScheduleType::AUTO,
         tkc::EpilogueScheduleType::AUTO, tkc::ClusterShape::ClusterShape_4x4x1);
 }
@@ -63,25 +160,9 @@ void runMxfp8Gemm(at::Tensor& out, at::Tensor const& act, at::Tensor const& weig
         reinterpret_cast<char*>(workspace.data_ptr()), wsBytes, at::cuda::getCurrentCUDAStream(act.get_device()));
 }
 
-} // namespace
-
-// MXFP8 (e4m3 + UE8M0 1x32 block scales) x MXFP8 (e4m3 + UE8M0 1x32 block
-// scales) GEMM on Blackwell sm_100/103.
-//
-// Operands (matching the CUTLASS block-scaled tensor-op convention):
-//   act:          [M, K] Float8_e4m3fn, row-major.
-//   actScale:     1D uint8 (UE8M0), swizzled layout produced by
-//                 torch.ops.trtllm.mxfp8_quantize(input, swizzedLayout=True).
-//   weight:       [N, K] Float8_e4m3fn, expected to be column-major in memory.
-//                 The caller is responsible for ensuring the weight tensor is
-//                 contiguous in the column-major sense that CUTLASS expects.
-//   weightScale:  1D uint8 (UE8M0), swizzled layout produced by
-//                 torch.ops.trtllm.block_scale_interleave(scale).
-//   globalScale:  [1] float -- alpha multiplier baked into the epilogue.
-//                 For pure MXFP8xMXFP8 this is usually [1.0].
-//   out_dtype:    fp16 / bf16 / fp32 output element type.
-at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight,
-    at::Tensor const& weightScale, at::Tensor const& globalScale, std::optional<c10::ScalarType> out_dtype)
+at::Tensor mxfp8Mxfp8GemmImpl(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight,
+    at::Tensor const& weightScale, at::Tensor const& globalScale, std::optional<c10::ScalarType> outDtype,
+    tkc::CutlassGemmConfig const* gemmConfig, bool const useTacticCache)
 {
     CHECK_INPUT(act, torch::kFloat8_e4m3fn);
     CHECK_INPUT(weight, torch::kFloat8_e4m3fn);
@@ -105,14 +186,17 @@ at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, a
     constexpr int kAlignmentN = 32;
     TORCH_CHECK(n % kAlignmentN == 0, "N (", n, ") must be divisible by ", kAlignmentN);
 
-    auto chosen_dtype = out_dtype.value_or(torch::kBFloat16);
-    TORCH_CHECK(chosen_dtype == torch::kFloat || chosen_dtype == torch::kHalf || chosen_dtype == torch::kBFloat16,
+    auto const chosenDtype = outDtype.value_or(torch::kBFloat16);
+    TORCH_CHECK(chosenDtype == torch::kFloat || chosenDtype == torch::kHalf || chosenDtype == torch::kBFloat16,
         "out_dtype must be one of fp16/bf16/fp32 (default bf16).");
 
-    at::Tensor out = at::detail::empty_cuda({m, n}, chosen_dtype, act.device(), std::nullopt);
+    at::Tensor out = at::detail::empty_cuda({m, n}, chosenDtype, act.device(), std::nullopt);
 
-    auto const config = getDefaultMxfp8GemmConfig();
-    switch (chosen_dtype)
+    auto const cachedEntry = useTacticCache ? findMxfp8TacticCacheEntry(m, n, k, chosenDtype) : std::nullopt;
+    auto const config = gemmConfig != nullptr
+        ? *gemmConfig
+        : (cachedEntry.has_value() ? cachedEntry->config : getDefaultMxfp8GemmConfig());
+    switch (chosenDtype)
     {
     case at::ScalarType::Half:
         runMxfp8Gemm<half>(out, act, weight, actScale, weightScale, globalScale, m, n, k, config);
@@ -132,12 +216,97 @@ at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, a
     return out;
 }
 
+} // namespace
+
+// MXFP8 (e4m3 + UE8M0 1x32 block scales) x MXFP8 (e4m3 + UE8M0 1x32 block
+// scales) GEMM on Blackwell sm_100/103.
+//
+// Operands (matching the CUTLASS block-scaled tensor-op convention):
+//   act:          [M, K] Float8_e4m3fn, row-major.
+//   actScale:     1D uint8 (UE8M0), swizzled layout produced by
+//                 torch.ops.trtllm.mxfp8_quantize(input, swizzedLayout=True).
+//   weight:       [N, K] Float8_e4m3fn, expected to be column-major in memory.
+//                 The caller is responsible for ensuring the weight tensor is
+//                 contiguous in the column-major sense that CUTLASS expects.
+//   weightScale:  1D uint8 (UE8M0), swizzled layout produced by
+//                 torch.ops.trtllm.block_scale_interleave(scale).
+//   globalScale:  [1] float -- alpha multiplier baked into the epilogue.
+//                 For pure MXFP8xMXFP8 this is usually [1.0].
+//   out_dtype:    fp16 / bf16 / fp32 output element type.
+at::Tensor mxfp8_mxfp8_gemm(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight,
+    at::Tensor const& weightScale, at::Tensor const& globalScale, std::optional<c10::ScalarType> outDtype)
+{
+    return mxfp8Mxfp8GemmImpl(act, actScale, weight, weightScale, globalScale, outDtype, /*gemmConfig=*/nullptr,
+        /*useTacticCache=*/true);
+}
+
+class MXFP8GemmRunner : public torch::CustomClassHolder
+{
+public:
+    explicit MXFP8GemmRunner(at::ScalarType outputDtype)
+        : mOutputDtype(outputDtype)
+    {
+        TORCH_CHECK(outputDtype == torch::kFloat || outputDtype == torch::kHalf || outputDtype == torch::kBFloat16,
+            "output_dtype must be one of fp16/bf16/fp32.");
+        mConfigs = CutlassFp4GemmRunner<half, FP4GemmType::W8A8_MXFP8_MXFP8>{}.getConfigs();
+    }
+
+    at::Tensor runGemm(at::Tensor const& act, at::Tensor const& actScale, at::Tensor const& weight,
+        at::Tensor const& weightScale, at::Tensor const& globalScale, int64_t configIdx) const
+    {
+        auto const config = configIdx == -1 ? getDefaultMxfp8GemmConfig() : getConfig(configIdx);
+        return mxfp8Mxfp8GemmImpl(
+            act, actScale, weight, weightScale, globalScale, mOutputDtype, &config, /*useTacticCache=*/false);
+    }
+
+    void registerTactic(int64_t const m, int64_t const n, int64_t const k, int64_t const configIdx) const
+    {
+        tkc::CutlassGemmConfig const config = configIdx == -1 ? getDefaultMxfp8GemmConfig() : getConfig(configIdx);
+        cacheMxfp8Tactic(m, n, k, mOutputDtype, config, configIdx);
+    }
+
+    int64_t getCachedTactic(int64_t const m, int64_t const n, int64_t const k) const
+    {
+        auto const entry = findMxfp8TacticCacheEntry(m, n, k, mOutputDtype);
+        return entry.has_value() ? entry->tactic : kMxfp8TacticCacheMiss;
+    }
+
+    void clearTacticCache() const
+    {
+        clearMxfp8CachedTactics();
+    }
+
+    int64_t getNumConfigs() const
+    {
+        return static_cast<int64_t>(mConfigs.size());
+    }
+
+private:
+    tkc::CutlassGemmConfig const& getConfig(int64_t const configIdx) const
+    {
+        TORCH_CHECK(configIdx >= 0 && configIdx < getNumConfigs(), "MXFP8 config index ", configIdx,
+            " is out of range [0, ", getNumConfigs(), ").");
+        return mConfigs.at(configIdx);
+    }
+
+    at::ScalarType mOutputDtype;
+    std::vector<tkc::CutlassGemmConfig> mConfigs;
+};
+
 } // namespace torch_ext
 
 TRTLLM_NAMESPACE_END
 
 TORCH_LIBRARY_FRAGMENT(trtllm, m)
 {
+    m.class_<tensorrt_llm::torch_ext::MXFP8GemmRunner>("MXFP8GemmRunner")
+        .def(torch::init<at::ScalarType>())
+        .def("run_gemm", &tensorrt_llm::torch_ext::MXFP8GemmRunner::runGemm)
+        .def("get_num_configs", &tensorrt_llm::torch_ext::MXFP8GemmRunner::getNumConfigs)
+        .def("register_tactic", &tensorrt_llm::torch_ext::MXFP8GemmRunner::registerTactic)
+        .def("get_cached_tactic", &tensorrt_llm::torch_ext::MXFP8GemmRunner::getCachedTactic)
+        .def("clear_tactic_cache", &tensorrt_llm::torch_ext::MXFP8GemmRunner::clearTacticCache);
+
     m.def(
         "mxfp8_mxfp8_gemm(Tensor act, Tensor actScale, Tensor weight, Tensor weightScale, "
         "Tensor globalScale, ScalarType? out_dtype=None) -> Tensor");

docs/source/deployment-guide/deployment-guide-for-minimax-m3-on-trtllm.md

@@ -101,6 +101,19 @@ If you don't have access to the source code locally, you can manually create the
 
 The configuration uses Data-Expert Parallelism (DEP): `enable_attention_dp: true` runs the attention layers data-parallel across ranks while the MoE experts run expert-parallel, which favors high-throughput / large-batch serving on MiniMax-M3.
 
+For MXFP8 checkpoints, TensorRT LLM selects the GEMM backend automatically.
+`TRTLLM_MXFP8_GEMM_BACKEND` is an advanced override for debugging and
+performance experiments:
+
+* `trtllm` uses the native TensorRT LLM GEMM for eager execution and CUDA graphs.
+* `flashinfer` forces FlashInfer for both eager execution and CUDA graphs; it
+  requires the pinned `flashinfer-python` package and Blackwell MXFP8 support.
+* `auto` keeps eager execution on the native GEMM and uses FlashInfer in
+  captured decode CUDA graphs after startup tuning.
+
+Leave the variable unset for normal deployments. An explicit value disables
+MiniMax-M3's automatic backend selection and uses the requested policy.
+
 ### Launch the TensorRT LLM Server
 
 MiniMax-M3 is launched through the `trtllm-llmapi-launch` wrapper, which sets up the multi-rank (MPI/Slurm) environment that the parallel server requires. The wrapper is run once per rank by Slurm (`srun`), with one task (rank) per GPU. The example below launches the server across 2 nodes (`-N 2`), 4 GPUs per node (`--ntasks-per-node 4`, 8 ranks total), using the curated YAML to drive parallelism, batching, and the MiniMax-M3 sparse-attention backend: