/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION &
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#pragma once

#include <optional>
#include <set>
#include <sstream>

#include "Enums.h"
#include "KernelParams.h"
#include "KernelTraits.h"
#include "trtllm/gen/CudaArchDecl.h"
#include "trtllm/gen/DtypeDecl.h"
#include "trtllm/gen/SfLayoutDecl.h"
#include "trtllm/gen/SparsityDecl.h"
#include "trtllm/gen/MmaDecl.h"
#ifndef TLLM_GEN_EXPORT_INTERFACE
#include "trtllm/gen/GenCtx.h"
#include "trtllm/gen/CudaRunner.h"
#else
#ifdef TLLM_GEN_EXPORT_FLASHINFER
#include <string>
namespace flashinfer::trtllm_cubin_loader {
std::string getCubin(const std::string& kernelName, const std::string& sha256);
}
#endif // TLLM_GEN_EXPORT_FLASHINFER
#include <iostream>
namespace batchedGemm {


template <typename T> void printArgs(T arg) {
#ifdef TLLM_GEN_DEBUG
  std::cout << arg;
#endif
}

template <typename T, typename... Args> void printArgs(T first, Args... args) {
  printArgs(first);
  if constexpr (sizeof...(args) > 0) {
    printArgs(", ");
    printArgs(args...);
  }
}

#define TLLM_CHECK_ERROR(cond, ...)                                                                \
  if (!(cond)) {                                                                                   \
    printArgs(__VA_ARGS__);                                                                        \
    printArgs("\n");                                                                               \
    return false;                                                                                  \
  }

#define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__)

#define TLLM_CHECK_ERROR_FMT(cond, ...) TLLM_CHECK_ERROR(cond, __VA_ARGS__)

#define TLLM_CHECK_WARNING(cond, ...)                                                              \
  if (!(cond)) {                                                                                   \
    printArgs(__VA_ARGS__);                                                                        \
    printArgs("\n");                                                                               \
    return false;                                                                                  \
  }

#define TLLM_LOG_WARNING(...) TLLM_CHECK_WARNING(false, __VA_ARGS__)

#define TLLM_LOG_INFO(...) TLLM_CHECK_WARNING(false, __VA_ARGS__)

#endif // TLLM_GEN_EXPORT_INTERFACE

#define GEMM_UPDATE_OR_ERROR(OPTION, VALUE)                                                        \
  if (updateOptions) {                                                                             \
    OPTION = VALUE;                                                                                \
  } else                                                                                           \
    return false

namespace trtllm {
namespace gen {
class CudaRunner;
class GenCfg;
} // namespace gen
} // namespace trtllm

namespace gemm {

////////////////////////////////////////////////////////////////////////////////////////////////////

namespace tg = trtllm::gen;

////////////////////////////////////////////////////////////////////////////////////////////////////

// NOTE: when adding new parameters, please update the dumpOptions function and
// gemm_export_config.json for cubin export.
struct GemmOptions {
#ifndef TLLM_GEN_EXPORT_INTERFACE
  // allow safely down-casting
  virtual ~GemmOptions() = default;
#endif

  GemmOptions() = default;
  GemmOptions(AllReduceAlgo allReduceAlgo,
              tg::Dtype biasDtype,
              BiasType biasType,
              int blockK,
              bool clcFastDrain,
              int clusterDimX,
              int clusterDimY,
              int clusterDimZ,
              CtaSwizzleType ctaSwizzleType,
              tg::Dtype dtypeAcc,
              tg::Dtype dtypeA,
              tg::Dtype dtypeB,
              tg::Dtype dtypeC,
              tg::Dtype dtypeMmaA,
              tg::Dtype dtypeMmaB,
              tg::Dtype dtypeSfC,
              EltwiseActType eltwiseActType,
              bool enablesEarlyExit,
              bool enablesDelayedEarlyExit,
              int epilogueLdtmDps,
              int epilogueLdtmBits,
              int epilogueTileM,
              int epilogueTileN,
              int fallbackClusterDimX,
              int fallbackClusterDimY,
              int fallbackClusterDimZ,
              FusedBiasShuffleMode fusedBiasShuffleMode,
              bool fuseLoadSfTask,
              bool fuseUtccpWithUtcmma,
              bool gridTriggerSecondaryA,
              bool gridTriggerSecondaryB,
              bool gridWaitForPrimaryEarlyExit,
              bool gridWaitForPrimaryA,
              bool gridWaitForPrimaryB,
              bool hoistLoadTaskInit,
              bool hoistMmaTaskTryWaits,
              int k,
              KernelTraits kernelTraits,
              bool fineGrainedConsumerA,
              bool fineGrainedConsumerB,
              bool fineGrainedForceValid,
              bool fineGrainedProducer,
              MatrixLayout layoutA,
              MatrixLayout layoutB,
              int m,
              int mmaK,
              tg::MmaKind mmaKind,
              int mmaM,
              int mmaN,
              int mmaTileK,
              bool mockAllReduce,
              int n,
              int numEpilogueWarps,
              int numRegsCastAWarps,
              int numRegsCopySfLdsSttm,
              int numRegsCopySparsityInfo,
              int numRegsPerThreadEpilogueWarp,
              int numRegsPerThreadNonEpilogueWarp,
              int numSlicesForSplitK,
              int numStagesA,
              int numStagesB,
              int numStagesMma,
              int numStagesMmaWithinWorkTile,
              int numStagesMmaAcrossWorkTile,
              int numStagesSmemSfA,
              int numStagesSmemSfB,
              int numStagesTmemSfA,
              int numStagesTmemSfB,
              int numStagesWorkId,
              bool outputDebugTensors,
              bool patchF2fp,
              tg::Dtype perTokenSfDtype,
              SchedHostTask schedHostTask,
              int32_t sfBlockSizeA,
              int32_t sfBlockSizeB,
              int32_t sfBlockSizeC,
              tg::SfLayout sfLayoutA,
              tg::SfLayout sfLayoutB,
              tg::SfLayout sfLayoutC,
              int sfReshapeFactor,
              tg::Sparsity sparsityA,
              SplitK splitK,
              int tileK,
              int tileM,
              int tileN,
              TileScheduler tileScheduler,
              bool transposeMmaOutput,
              bool useCustomizedMma3xNvFp4,
              bool useCustomMmaSchedule,
              bool useDeepSeekFp8,
              bool useFlexibleClusterDims,
              bool useHoistTryWaitForCustomMmaSchedule,
              bool useMaxTmemOverlap,
              bool usePerTokenSfA,
              bool usePerTokenSfB,
              bool useShuffledMatrix,
              bool useTmaStore,
              bool useTwoTmaLoadWarps,
              bool useTwoMmaWarps,
              bool useUnrollLoop2xForMma,
              int validM,
              int validN,
              int validK,
              int worldSize)
    : mAllReduceAlgo{allReduceAlgo}
    , mBiasDtype{biasDtype}
    , mBiasType{biasType}
    , mBlockK(blockK)
    , mClcFastDrain{clcFastDrain}
    , mClusterDimX{clusterDimX}
    , mClusterDimY{clusterDimY}
    , mClusterDimZ{clusterDimZ}
    , mCtaSwizzleType{ctaSwizzleType}
    , mDtypeAcc{dtypeAcc}
    , mDtypeA{dtypeA}
    , mDtypeB{dtypeB}
    , mDtypeC{dtypeC}
    , mDtypeMmaA{dtypeMmaA}
    , mDtypeMmaB{dtypeMmaB}
    , mDtypeSfC{dtypeSfC}
    , mEltwiseActType{eltwiseActType}
    , mEnablesEarlyExit{enablesEarlyExit}
    , mEnablesDelayedEarlyExit{enablesDelayedEarlyExit}
    , mEpilogueLdtmDps{epilogueLdtmDps}
    , mEpilogueLdtmBits{epilogueLdtmBits}
    , mEpilogueTileM{epilogueTileM}
    , mEpilogueTileN{epilogueTileN}
    , mFallbackClusterDimX{fallbackClusterDimX}
    , mFallbackClusterDimY{fallbackClusterDimY}
    , mFallbackClusterDimZ{fallbackClusterDimZ}
    , mFusedBiasShuffleMode{fusedBiasShuffleMode}
    , mFuseLoadSfTask{fuseLoadSfTask}
    , mFuseUtccpWithUtcmma{fuseUtccpWithUtcmma}
    , mGridTriggerSecondaryA{gridTriggerSecondaryA}
    , mGridTriggerSecondaryB{gridTriggerSecondaryB}
    , mGridWaitForPrimaryEarlyExit{gridWaitForPrimaryEarlyExit}
    , mGridWaitForPrimaryA{gridWaitForPrimaryA}
    , mGridWaitForPrimaryB{gridWaitForPrimaryB}
    , mHoistLoadTaskInit{hoistLoadTaskInit}
    , mHoistMmaTaskTryWaits{hoistMmaTaskTryWaits}
    , mK{k}
    , mKernelTraits{kernelTraits}
    , mFineGrainedConsumerA(fineGrainedConsumerA)
    , mFineGrainedConsumerB(fineGrainedConsumerB)
    , mFineGrainedForceValid(fineGrainedForceValid)
    , mFineGrainedProducer(fineGrainedProducer)
    , mLayoutA{layoutA}
    , mLayoutB{layoutB}
    , mM{m}
    , mMmaK{mmaK}
    , mMmaKind{mmaKind}
    , mMmaM{mmaM}
    , mMmaN{mmaN}
    , mMmaTileK{mmaTileK}
    , mMockAllReduce{mockAllReduce}
    , mN{n}
    , mNumEpilogueWarps{numEpilogueWarps}
    , mNumRegsCastAWarps(numRegsCastAWarps)
    , mNumRegsCopySfLdsSttm(numRegsCopySfLdsSttm)
    , mNumRegsCopySparsityInfo(numRegsCopySparsityInfo)
    , mNumRegsPerThreadEpilogueWarp(numRegsPerThreadEpilogueWarp)
    , mNumRegsPerThreadNonEpilogueWarp(numRegsPerThreadNonEpilogueWarp)
    , mNumSlicesForSplitK{numSlicesForSplitK}
    , mNumStagesA{numStagesA}
    , mNumStagesB{numStagesB}
    , mNumStagesMma{numStagesMma}
    , mNumStagesMmaWithinWorkTile{numStagesMmaWithinWorkTile}
    , mNumStagesMmaAcrossWorkTile{numStagesMmaAcrossWorkTile}
    , mNumStagesSmemSfA{numStagesSmemSfA}
    , mNumStagesSmemSfB{numStagesSmemSfB}
    , mNumStagesTmemSfA{numStagesTmemSfA}
    , mNumStagesTmemSfB{numStagesTmemSfB}
    , mNumStagesWorkId{numStagesWorkId}
    , mOutputDebugTensors{outputDebugTensors}
    , mPatchF2fp{patchF2fp}
    , mPerTokenSfDtype{perTokenSfDtype}
    , mSchedHostTask{schedHostTask}
    , mSfBlockSizeA{sfBlockSizeA}
    , mSfBlockSizeB{sfBlockSizeB}
    , mSfBlockSizeC{sfBlockSizeC}
    , mSfLayoutA{sfLayoutA}
    , mSfLayoutB{sfLayoutB}
    , mSfLayoutC{sfLayoutC}
    , mSfReshapeFactor{sfReshapeFactor}
    , mSparsityA{sparsityA}
    , mSplitK{splitK}
    , mTileK{tileK}
    , mTileM{tileM}
    , mTileN{tileN}
    , mTileScheduler{tileScheduler}
    , mTransposeMmaOutput{transposeMmaOutput}
    , mUseCustomizedMma3xNvFp4{useCustomizedMma3xNvFp4}
    , mUseCustomMmaSchedule{useCustomMmaSchedule}
    , mUseDeepSeekFp8{useDeepSeekFp8}
    , mUseFlexibleClusterDims{useFlexibleClusterDims}
    , mUseHoistTryWaitForCustomMmaSchedule{useHoistTryWaitForCustomMmaSchedule}
    , mUseMaxTmemOverlap{useMaxTmemOverlap}
    , mUsePerTokenSfA{usePerTokenSfA}
    , mUsePerTokenSfB{usePerTokenSfB}
    , mUseShuffledMatrix{useShuffledMatrix}
    , mUseTmaStore{useTmaStore}
    , mUseTwoTmaLoadWarps{useTwoTmaLoadWarps}
    , mUseTwoMmaWarps{useTwoMmaWarps}
    , mUseUnrollLoop2xForMma{useUnrollLoop2xForMma}
    , mValidM{validM}
    , mValidN{validN}
    , mValidK{validK}
    , mWorldSize{worldSize} {
  }
  // The all-reduce algorithm.
  AllReduceAlgo mAllReduceAlgo{AllReduceAlgo::None};
  // The data type of bias in global memory.
  tg::Dtype mBiasDtype{tg::Dtype::Fp32};
  // The type of bias.
  BiasType mBiasType{BiasType::None};
  // Block size in the K dimension
  int mBlockK{-1};
  // Whether to enable CLC fast drain for early exit in SM100 CLC-based scheduler.
  bool mClcFastDrain{true};
  // Cluster size in X dim.
  int mClusterDimX{1};
  // Cluster size in Y dim.
  int mClusterDimY{1};
  // Cluster size in Z dim.
  int mClusterDimZ{1};
  // The type of CTA swizzle.
  CtaSwizzleType mCtaSwizzleType{CtaSwizzleType::RasterizeAlongM};
  // Data type of the accumulators.
  tg::Dtype mDtypeAcc{tg::Dtype::Fp32};
  // Data type of the A matrix.
  tg::Dtype mDtypeA{tg::Dtype::Fp16};
  // Data type of the B matrix.
  tg::Dtype mDtypeB{tg::Dtype::Void};
  // Data type of the outputs.
  tg::Dtype mDtypeC{tg::Dtype::Void};
  // Data type of the A matrix for the MMA, if different from the input type.
  tg::Dtype mDtypeMmaA{tg::Dtype::Void};
  // Data type of the B matrix for the MMA, if different from the input type.
  tg::Dtype mDtypeMmaB{tg::Dtype::Void};
  // Data type of the block scaling factors of the output C matrix. When Void, it is auto-deduced
  // from mDtypeC by checkAndUpdateGemmOptions (e.g., E4m3 for E2m1, UE8m0 for Mx formats).
  // Currently supported override: Fp32 (to keep the high-precision scaling factors).
  tg::Dtype mDtypeSfC{tg::Dtype::Void};
  // The type of activation.
  EltwiseActType mEltwiseActType{EltwiseActType::None};
  // Whether to enable early exit.
  bool mEnablesEarlyExit{false};
  // Whether to enable delayed early exit to overlap
  // numNonExitingCtas loading with the other instructions.
  bool mEnablesDelayedEarlyExit{false};
  // The epilogue supports multiple LDTM shapes, although not every shape is applicable in every
  // case. In particular:
  // - On Hopper: must be 16dp256bit.
  // - Transposed output: must be 16dp256bit.
  // - Non-transposed output:
  //     - NvFp4 with fused activation: must be 32dp32bit.
  //     - Else it can be either 16dp256bit or 32dp32bit.
  // The number of DP lanes in the epilogue LDTM.
  int mEpilogueLdtmDps{16};
  // The number of bits in the epilogue LDTM.
  int mEpilogueLdtmBits{256};
  // Tile size for the epilogue in M dimension.
  int mEpilogueTileM{128};
  // Tile size for the epilogue in N dimension.
  int mEpilogueTileN{32};
  // Fallback Cluster size in X dim.
  int mFallbackClusterDimX{1};
  // Fallback Cluster size in Y dim.
  int mFallbackClusterDimY{1};
  // Fallback Cluster size in Z dim.
  int mFallbackClusterDimZ{1};
  // Controls which BiasType::Mn preprocessing steps are fused into the kernel instead of done on
  // the host.
  FusedBiasShuffleMode mFusedBiasShuffleMode{FusedBiasShuffleMode::None};
  // Fuse LoadSfA into LoadA and LoadSfB into LoadB when pipeline stages match.
  // Requires useTwoTmaLoadWarps. Saves LoadSfA/LoadSfB warp groups.
  bool mFuseLoadSfTask{false};
  // Whether fuse UTCCP with UTC*MMA.
  bool mFuseUtccpWithUtcmma{false};
  // Whether load task A triggers the next grid.
  bool mGridTriggerSecondaryA{false};
  // Whether load task B triggers the next grid.
  bool mGridTriggerSecondaryB{false};
  // Whether the loads that check for an early exit should wait on a grid dependency.
  bool mGridWaitForPrimaryEarlyExit{true};
  // Whether the load of A should wait on a grid dependency.
  bool mGridWaitForPrimaryA{true};
  // Whether the load of B should wait on a grid dependency.
  bool mGridWaitForPrimaryB{true};
  // Whether to hoist the initialization of the loading tasks.
  bool mHoistLoadTaskInit{true};
  // Whether to hoist the mbarrier try_waits (e.g., mma.prodAcq, smemAb.consWait) in the MMA task.
  bool mHoistMmaTaskTryWaits{false};
  // The K dimension of GEMM.
  int mK{16 * 16};
  // Traits of the kernel.
  KernelTraits mKernelTraits{};
  // Whether the load task A uses a FineGrained pipeline.
  bool mFineGrainedConsumerA{false};
  // Whether the load task B uses a FineGrained pipeline.
  bool mFineGrainedConsumerB{false};
  // Whether the kernel forces the output to be valid values (not -0.0 for FP and NAN for UE8m0).
  bool mFineGrainedForceValid{false};
  // Whether the kernel is a FineGrained producer (invalidates output buffer prior to writing
  // results).
  bool mFineGrainedProducer{false};
  // Layout of A matrix
  MatrixLayout mLayoutA{MatrixLayout::MajorK};
  // Layout of B matrix
  MatrixLayout mLayoutB{MatrixLayout::MajorK};
  // The M dimension of GEMM.
  int mM{128 * 2};
  // Size of the MMA instruction in the K dimension.
  int mMmaK{16};
  // The kind of MMA instruction to use.
  tg::MmaKind mMmaKind{tg::MmaKind::Auto};
  // Size of the MMA instruction in the M dimension.
  int mMmaM{64};
  // Size of the MMA instruction in the N dimension.
  int mMmaN{16};
  // Size of the MMA tileK. When mUseCustomizedMma3xNvFp4 is true, the mTileK is the tileK for A/B
  // and mMmaTileK is the tileK for SfA/B and MMA.
  int mMmaTileK{0};
  // Whether to mock all-reduce code for single-GPU debugging.
  bool mMockAllReduce{false};
  // The N dimension of GEMM.
  int mN{64 * 4};
  // Number of Epilogue Warps
  int mNumEpilogueWarps{4};
  // Number of registers for the cast A warps.
  int mNumRegsCastAWarps{0};
  // Number of registers for the LDS+STTM warps.
  int mNumRegsCopySfLdsSttm{0};
  // Number of registers per thread to copy sparsity info with LDS+STTM.
  int mNumRegsCopySparsityInfo{0};
  // Number of registers per thread for epilogue warps
  int mNumRegsPerThreadEpilogueWarp{0};
  // Number of registers per thread for non-epilogue warps
  int mNumRegsPerThreadNonEpilogueWarp{0};
  // Number of partitions along the K dimension. When mNumSlicesForSplitK > 1,
  // the problem is distributed across several SMs, where each CTA works on its local K slice.
  // Partial results are accumulated afterwards using either GMEM or DSMEM (in CGA)
  // to exchange the data between CTAs.
  int mNumSlicesForSplitK{1};
  // The depth of the pipeline for A.
  int mNumStagesA{0};
  // The depth of the pipeline for B.
  int mNumStagesB{0};
  // The depth of the mma pipeline. Equals numStagesMmaWithinWorkTile * numStagesMmaAcrossWorkTile.
  int mNumStagesMma{1};
  // The depth of the mma pipeline within work tile. Only GmemC classes with "WithAccInReg" suffix
  // are allowed to be greater than 1.
  int mNumStagesMmaWithinWorkTile{-1};
  // The depth of the mma pipeline across work tiles in the persistent loop.
  int mNumStagesMmaAcrossWorkTile{-1};
  // The depth of the pipeline for the SF of A in SMEM.
  int mNumStagesSmemSfA{0};
  // The depth of the pipeline for the SF of B in SMEM.
  int mNumStagesSmemSfB{0};
  // The depth of the pipeline for the SF of A in TMEM.
  int mNumStagesTmemSfA{0};
  // The depth of the pipeline for the SF of B in TMEM.
  int mNumStagesTmemSfB{0};
  // The depth of the work id pipeline and the work throttle pipeline.
  int mNumStagesWorkId{3};
  // Whether to output debug tensors.
  bool mOutputDebugTensors{false};
  // Patch float conversions.
  bool mPatchF2fp{false};
  // The dtype of scale factors when using per-token scaling.
  // When not provided, it will be Fp32 if usePerTokenSfA and usePerTokenSfB else Bfloat16.
  tg::Dtype mPerTokenSfDtype{tg::Dtype::Void};
  // Which host task the persistent scheduler is fused with (Self = dedicated scheduler task).
  SchedHostTask mSchedHostTask{SchedHostTask::Self};
  // Block size of A, for block-scaled types.
  int mSfBlockSizeA{-1};
  // Block size of B, for block-scaled types.
  int mSfBlockSizeB{-1};
  // Block size of C, for block-scaled types.
  int mSfBlockSizeC{-1};
  // Scale factors layout for A.
  tg::SfLayout mSfLayoutA{tg::SfLayout::R128c4};
  // Scale factors layout for B.
  tg::SfLayout mSfLayoutB{tg::SfLayout::R128c4};
  // Scale factors layout for C.
  tg::SfLayout mSfLayoutC{tg::SfLayout::R128c4};
  // Number of "repeats", i.e. reshaping factor, to fold hidden dimension into SfBlock dimension.
  // As result, the hidden dimension of the SF tensor must be a multiple of NumRepeats *
  // numEltsPerSf * 4. This reduces the problem shape space that the kernel is able to run.
  // But it reduces the number of L2 requests under the hood and potentially improves perf.
  // Applies to layout 8x4 only.
  int mSfReshapeFactor{1};
  // Sparsity of A.
  tg::Sparsity mSparsityA{tg::Sparsity::Dense};
  // The location of the exchange for split-K (it's None when split-K is disabled).
  SplitK mSplitK{SplitK::None};
  // K tile dimension of GEMM.
  int mTileK{16};
  // M tile dimension of GEMM.
  int mTileM{128};
  // N tile dimension of GEMM.
  int mTileN{32};
  // Tile scheduler type.
  TileScheduler mTileScheduler{TileScheduler::Static};
  // Save output of MMA in M-major format.
  bool mTransposeMmaOutput{false};
  // Use custom MMA schedule for 3xNvFp4
  bool mUseCustomizedMma3xNvFp4{false};
  // Use custom MMA schedule optimized for low-latency.
  bool mUseCustomMmaSchedule{false};
  // Use DeepSeek Fp8.
  bool mUseDeepSeekFp8{false};
  // Use flexible cluster dims.
  bool mUseFlexibleClusterDims{false};
  // The purpose of hoisting trywaits is to opportunistically peek at the availability of the next
  // k-block. It benefits when the next k-block is already available and thus sustaining the
  // momentum, but it adds latency to the first k-block for smaller k-loop.
  bool mUseHoistTryWaitForCustomMmaSchedule{false};
  // Whether use the max Tmem overlap trick.
  bool mUseMaxTmemOverlap{false};
  // Apply per-token scales from A
  bool mUsePerTokenSfA{false};
  // Apply per-token scales from B
  bool mUsePerTokenSfB{false};
  // Reorder rows/cols in the A matrix (when TransposeMmaOutput is true, otherwise B matrix) for the
  // better memory accesses in the M-major epilogue.
  bool mUseShuffledMatrix{false};
  // Use TMA to store the result.
  bool mUseTmaStore{true};
  // Use two different warps for A and B matrix load.
  bool mUseTwoTmaLoadWarps{false};
  // Use two different warps for MMA tasks. Applicable only to DeepSeek FP8.
  bool mUseTwoMmaWarps{false};
  // Whether to unroll the loop by 2x.
  bool mUseUnrollLoop2xForMma{true};
  // The valid range of M/N/K dimension of GEMM without padding values.
  // Used to opportunistically remove memory traffic from the padding due to rigid SF shape
  // constraint or TMA constraint. Such as:
  // 1. outputDim % (4 * sfBlockSize) == 0; as 4x SFs are packed into 4 bytes
  // 2. MxFp4 x Fp8 mmaType requires bespoke TMA load which requires hiddenDim % 128 == 0
  // 3. TMA requires 16B alignment for each row
  int mValidM{-1};
  int mValidN{-1};
  int mValidK{-1};
  // World size for all-reduce.
  int mWorldSize{1};
};

////////////////////////////////////////////////////////////////////////////////////////////////////

using SmVersion = tg::CudaArch;

////////////////////////////////////////////////////////////////////////////////////////////////////
//
// GemmConfig
//
////////////////////////////////////////////////////////////////////////////////////////////////////

struct GemmConfig {
  uint8_t const* mData{nullptr};
  uint32_t mSize{0};
  uint32_t mSharedMemSize{0};
  char const* mFunctionName{nullptr};
  uint32_t mNumThreadsPerCTA{0};
  char const* mHash{nullptr};
  std::string mGenCfgJsonStr{""};
  char const* mExecPath{nullptr};
  trtllm::gen::CudaRunner* mCudaRunner{nullptr};
  trtllm::gen::GenCfg* mGenCfg{nullptr};
  int32_t mInstanceIdx{0};

  GemmOptions mOptions{};
  tg::CudaArch mSm{tg::CudaArch::Sm100a};
};

////////////////////////////////////////////////////////////////////////////////////////////////////

// Serialization helpers.
template <typename T> inline std::string toString(T e) {
  return std::to_string(e);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <> inline std::string toString(trtllm::gen::Dtype e) {
  return trtllm::gen::dtypeToString(e);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <> inline std::string toString(trtllm::gen::MmaKind e) {
  return trtllm::gen::mmaKindToString(e);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <> inline std::string toString(CtaSwizzleType e) {
  switch (e) {
  case CtaSwizzleType::RasterizeAlongM:
    return "RasterizeAlongM";
  case CtaSwizzleType::RasterizeAlongN:
    return "RasterizeAlongN";
  case CtaSwizzleType::ZigZagAlongM2:
    return "ZigZagAlongM2";
  case CtaSwizzleType::ZigZagAlongN2:
    return "ZigZagAlongN2";
  case CtaSwizzleType::ZigZagAlongM4:
    return "ZigZagAlongM4";
  case CtaSwizzleType::ZigZagAlongN4:
    return "ZigZagAlongN4";
  default:
    return std::to_string(static_cast<int32_t>(e));
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <> inline std::string toString(EltwiseActType e) {
  switch (e) {
  case EltwiseActType::None:
    return "None";
  case EltwiseActType::Gelu:
    return "Gelu";
  case EltwiseActType::Relu2:
    return "Relu2";
  case EltwiseActType::Silu:
    return "Silu";
  default:
    return std::to_string(static_cast<int32_t>(e));
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <> inline std::string toString(FusedBiasShuffleMode e) {
  switch (e) {
  case FusedBiasShuffleMode::None:
    return "None";
  case FusedBiasShuffleMode::Shuffle:
    return "Shuffle";
  case FusedBiasShuffleMode::ReorderAndShuffle:
    return "ReorderAndShuffle";
  default:
    return std::to_string(static_cast<int32_t>(e));
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <> inline std::string toString(SchedHostTask e) {
  switch (e) {
  case SchedHostTask::Self:
    return "Self";
  case SchedHostTask::LoadA:
    return "LoadA";
  case SchedHostTask::LoadB:
    return "LoadB";
  case SchedHostTask::LoadSfA:
    return "LoadSfA";
  case SchedHostTask::LoadSfB:
    return "LoadSfB";
  default:
    return std::to_string(static_cast<int32_t>(e));
  }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

inline std::string dumpOptions(GemmOptions const& options, bool dumpRuntimeParams = true) {
  std::stringstream ss;
  ss << "mAllReduceAlgo=" << "gemm::AllReduceAlgo(" << static_cast<int32_t>(options.mAllReduceAlgo)
     << ")" << "," << std::endl;
  ss << "mBiasDtype=" << "trtllm::gen::Dtype(" << static_cast<int32_t>(options.mBiasDtype) << ")"
     << "," << std::endl;
  ss << "mBiasType=" << "gemm::BiasType(" << static_cast<int32_t>(options.mBiasType) << ")" << ","
     << std::endl;
  ss << "mBlockK=" << options.mBlockK << "," << std::endl;
  ss << "mClcFastDrain=" << options.mClcFastDrain << "," << std::endl;
  ss << "mClusterDimX=" << options.mClusterDimX << "," << std::endl;
  ss << "mClusterDimY=" << options.mClusterDimY << "," << std::endl;
  ss << "mClusterDimZ=" << options.mClusterDimZ << "," << std::endl;
  ss << "mCtaSwizzleType=" << "gemm::CtaSwizzleType("
     << static_cast<int32_t>(options.mCtaSwizzleType) << ")" << "," << std::endl;
  ss << "mDtypeAcc=" << "trtllm::gen::Dtype(" << static_cast<int32_t>(options.mDtypeAcc) << ")"
     << "," << std::endl;
  ss << "mDtypeA=" << "trtllm::gen::Dtype(" << static_cast<int32_t>(options.mDtypeA) << ")" << ","
     << std::endl;
  ss << "mDtypeB=" << "trtllm::gen::Dtype(" << static_cast<int32_t>(options.mDtypeB) << ")" << ","
     << std::endl;
  ss << "mDtypeC=" << "trtllm::gen::Dtype(" << static_cast<int32_t>(options.mDtypeC) << ")" << ","
     << std::endl;
  ss << "mDtypeMmaA=" << "trtllm::gen::Dtype(" << static_cast<int32_t>(options.mDtypeMmaA) << ")"
     << "," << std::endl;
  ss << "mDtypeMmaB=" << "trtllm::gen::Dtype(" << static_cast<int32_t>(options.mDtypeMmaB) << ")"
     << "," << std::endl;
  ss << "mDtypeSfC=" << "trtllm::gen::Dtype(" << static_cast<int32_t>(options.mDtypeSfC) << ")"
     << "," << std::endl;
  ss << "mEltwiseActType=" << "gemm::EltwiseActType("
     << static_cast<int32_t>(options.mEltwiseActType) << ")" << "," << std::endl;
  ss << "mEnablesEarlyExit=" << options.mEnablesEarlyExit << "," << std::endl;
  ss << "mEnablesDelayedEarlyExit=" << options.mEnablesDelayedEarlyExit << "," << std::endl;
  ss << "mEpilogueLdtmDps=" << options.mEpilogueLdtmDps << "," << std::endl;
  ss << "mEpilogueLdtmBits=" << options.mEpilogueLdtmBits << "," << std::endl;
  ss << "mEpilogueTileM=" << options.mEpilogueTileM << "," << std::endl;
  ss << "mEpilogueTileN=" << options.mEpilogueTileN << "," << std::endl;
  ss << "mFallbackClusterDimX=" << options.mFallbackClusterDimX << "," << std::endl;
  ss << "mFallbackClusterDimY=" << options.mFallbackClusterDimY << "," << std::endl;
  ss << "mFallbackClusterDimZ=" << options.mFallbackClusterDimZ << "," << std::endl;
  ss << "mFusedBiasShuffleMode=" << "gemm::FusedBiasShuffleMode("
     << static_cast<int32_t>(options.mFusedBiasShuffleMode) << ")" << "," << std::endl;
  ss << "mFuseLoadSfTask=" << options.mFuseLoadSfTask << "," << std::endl;
  ss << "mFuseUtccpWithUtcmma=" << options.mFuseUtccpWithUtcmma << "," << std::endl;
  ss << "mGridTriggerSecondaryA=" << options.mGridTriggerSecondaryA << "," << std::endl;
  ss << "mGridTriggerSecondaryB=" << options.mGridTriggerSecondaryB << "," << std::endl;
  ss << "mGridWaitForPrimaryEarlyExit=" << options.mGridWaitForPrimaryEarlyExit << "," << std::endl;
  ss << "mGridWaitForPrimaryA=" << options.mGridWaitForPrimaryA << "," << std::endl;
  ss << "mGridWaitForPrimaryB=" << options.mGridWaitForPrimaryB << "," << std::endl;
  ss << "mHoistLoadTaskInit=" << options.mHoistLoadTaskInit << "," << std::endl;
  ss << "mHoistMmaTaskTryWaits=" << options.mHoistMmaTaskTryWaits << "," << std::endl;
  if (dumpRuntimeParams) {
    ss << "mK=" << options.mK << "," << std::endl;
  }
  ss << "mKernelTraits={}" << "," << std::endl;
  ss << "mFineGrainedConsumerA=" << options.mFineGrainedConsumerA << "," << std::endl;
  ss << "mFineGrainedConsumerB=" << options.mFineGrainedConsumerB << "," << std::endl;
  ss << "mFineGrainedForceValid=" << options.mFineGrainedForceValid << "," << std::endl;
  ss << "mFineGrainedProducer=" << options.mFineGrainedProducer << "," << std::endl;
  ss << "mLayoutA=gemm::MatrixLayout(" << static_cast<int32_t>(options.mLayoutA) << ")" << ","
     << std::endl;
  ss << "mLayoutB=gemm::MatrixLayout(" << static_cast<int32_t>(options.mLayoutB) << ")" << ","
     << std::endl;
  if (dumpRuntimeParams) {
    ss << "mM=" << options.mM << "," << std::endl;
  }
  ss << "mMmaK=" << options.mMmaK << "," << std::endl;
  ss << "mMmaKind=" << "trtllm::gen::MmaKind(" << static_cast<int32_t>(options.mMmaKind) << ")"
     << "," << std::endl;
  ss << "mMmaM=" << options.mMmaM << "," << std::endl;
  ss << "mMmaN=" << options.mMmaN << "," << std::endl;
  ss << "mMmaTileK=" << options.mMmaTileK << "," << std::endl;
  ss << "mMockAllReduce=" << options.mMockAllReduce << "," << std::endl;
  if (dumpRuntimeParams) {
    ss << "mN=" << options.mN << "," << std::endl;
  }
  ss << "mNumEpilogueWarps=" << options.mNumEpilogueWarps << "," << std::endl;
  ss << "mNumRegsCastAWarps=" << options.mNumRegsCastAWarps << "," << std::endl;
  ss << "mNumRegsCopySfLdsSttm=" << options.mNumRegsCopySfLdsSttm << "," << std::endl;
  ss << "mNumRegsCopySparsityInfo=" << options.mNumRegsCopySparsityInfo << "," << std::endl;
  ss << "mNumRegsPerThreadEpilogueWarp=" << options.mNumRegsPerThreadEpilogueWarp << ","
     << std::endl;
  ss << "mNumRegsPerThreadNonEpilogueWarp=" << options.mNumRegsPerThreadNonEpilogueWarp << ","
     << std::endl;
  ss << "mNumSlicesForSplitK=" << options.mNumSlicesForSplitK << "," << std::endl;
  ss << "mNumStagesA=" << options.mNumStagesA << "," << std::endl;
  ss << "mNumStagesB=" << options.mNumStagesB << "," << std::endl;
  ss << "mNumStagesMma=" << options.mNumStagesMma << "," << std::endl;
  ss << "mNumStagesMmaWithinWorkTile=" << options.mNumStagesMmaWithinWorkTile << "," << std::endl;
  ss << "mNumStagesMmaAcrossWorkTile=" << options.mNumStagesMmaAcrossWorkTile << "," << std::endl;
  ss << "mNumStagesSmemSfA=" << options.mNumStagesSmemSfA << "," << std::endl;
  ss << "mNumStagesSmemSfB=" << options.mNumStagesSmemSfB << "," << std::endl;
  ss << "mNumStagesTmemSfA=" << options.mNumStagesTmemSfA << "," << std::endl;
  ss << "mNumStagesTmemSfB=" << options.mNumStagesTmemSfB << "," << std::endl;
  ss << "mNumStagesWorkId=" << options.mNumStagesWorkId << "," << std::endl;
  ss << "mOutputDebugTensors=" << options.mOutputDebugTensors << "," << std::endl;
  ss << "mPatchF2fp=" << options.mPatchF2fp << "," << std::endl;
  ss << "mPerTokenSfDtype=" << "trtllm::gen::Dtype("
     << static_cast<int32_t>(options.mPerTokenSfDtype) << ")" << "," << std::endl;
  ss << "mSchedHostTask=" << "gemm::SchedHostTask(" << static_cast<int32_t>(options.mSchedHostTask)
     << ")" << "," << std::endl;
  ss << "mSfBlockSizeA=" << options.mSfBlockSizeA << "," << std::endl;
  ss << "mSfBlockSizeB=" << options.mSfBlockSizeB << "," << std::endl;
  ss << "mSfBlockSizeC=" << options.mSfBlockSizeC << "," << std::endl;
  ss << "mSfLayoutA=" << "trtllm::gen::SfLayout(" << static_cast<int32_t>(options.mSfLayoutA) << ")"
     << "," << std::endl;
  ss << "mSfLayoutB=" << "trtllm::gen::SfLayout(" << static_cast<int32_t>(options.mSfLayoutB) << ")"
     << "," << std::endl;
  ss << "mSfLayoutC=" << "trtllm::gen::SfLayout(" << static_cast<int32_t>(options.mSfLayoutC) << ")"
     << "," << std::endl;
  ss << "mSfReshapeFactor=" << options.mSfReshapeFactor << "," << std::endl;
  ss << "mSparsityA=" << "trtllm::gen::Sparsity(" << static_cast<int32_t>(options.mSparsityA) << ")"
     << "," << std::endl;
  ss << "mSplitK=" << "gemm::SplitK(" << static_cast<int32_t>(options.mSplitK) << ")" << ","
     << std::endl;
  ss << "mTileK=" << options.mTileK << "," << std::endl;
  ss << "mTileM=" << options.mTileM << "," << std::endl;
  ss << "mTileN=" << options.mTileN << "," << std::endl;
  ss << "mTileScheduler=" << "gemm::TileScheduler(" << static_cast<int32_t>(options.mTileScheduler)
     << ")" << "," << std::endl;
  ss << "mTransposeMmaOutput=" << options.mTransposeMmaOutput << "," << std::endl;
  ss << "mUseCustomizedMma3xNvFp4=" << options.mUseCustomizedMma3xNvFp4 << "," << std::endl;
  ss << "mUseCustomMmaSchedule=" << options.mUseCustomMmaSchedule << "," << std::endl;
  ss << "mUseDeepSeekFp8=" << options.mUseDeepSeekFp8 << "," << std::endl;
  ss << "mUseFlexibleClusterDims=" << options.mUseFlexibleClusterDims << "," << std::endl;
  ss << "mUseHoistTryWaitForCustomMmaSchedule=" << options.mUseHoistTryWaitForCustomMmaSchedule
     << "," << std::endl;
  ss << "mUseMaxTmemOverlap=" << options.mUseMaxTmemOverlap << "," << std::endl;
  ss << "mUsePerTokenSfA=" << options.mUsePerTokenSfA << "," << std::endl;
  ss << "mUsePerTokenSfB=" << options.mUsePerTokenSfB << "," << std::endl;
  ss << "mUseShuffledMatrix=" << options.mUseShuffledMatrix << "," << std::endl;
  ss << "mUseTmaStore=" << options.mUseTmaStore << "," << std::endl;
  ss << "mUseTwoTmaLoadWarps=" << options.mUseTwoTmaLoadWarps << "," << std::endl;
  ss << "mUseTwoMmaWarps=" << options.mUseTwoMmaWarps << "," << std::endl;
  ss << "mUseUnrollLoop2xForMma=" << options.mUseUnrollLoop2xForMma << "," << std::endl;
  if (dumpRuntimeParams) {
    ss << "mValidM=" << options.mValidM << "," << std::endl;
    ss << "mValidN=" << options.mValidN << "," << std::endl;
    ss << "mValidK=" << options.mValidK << "," << std::endl;
    ss << "mWorldSize=" << options.mWorldSize << std::endl;
  }
  return ss.str();
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename T> inline T divUp(T a, T b) {
  return (a + b - 1) / b;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename T> inline T divUpMul(T a, T b) {
  return gemm::divUp(a, b) * b;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// clang-format off
inline std::vector<int> srcToDstBlk16RowMap = 
      {
        0,  8,
        1,  9,
        2, 10,
        3, 11,
        4, 12,
        5, 13,
        6, 14,
        7, 15
      };
inline std::vector<int> srcToDstBlk32RowMap = 
      {
        0,  8, 16, 24,
        1,  9, 17, 25,
        2, 10, 18, 26,
        3, 11, 19, 27,
        4, 12, 20, 28,
        5, 13, 21, 29,
        6, 14, 22, 30,
        7, 15, 23, 31
      };
// clang-format on

////////////////////////////////////////////////////////////////////////////////////////////////////

inline int32_t getShuffleBlockSize(int epilogueTileM) {
  int shuffleBlockSize = 16;
  if (epilogueTileM % 128 == 0) {
    shuffleBlockSize = 32;
  }
  return shuffleBlockSize;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

inline std::vector<int> const& getShuffleIndices(int epilogueTileM) {
  auto const shuffleBlockSize = getShuffleBlockSize(epilogueTileM);
  return shuffleBlockSize == 16 ? srcToDstBlk16RowMap : srcToDstBlk32RowMap;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// Check if the options are valid or not.
inline bool checkAndUpdateGemmOptions(GemmOptions& options,
                                      tg::CudaArch cudaArch,
                                      int tpGrpSize,
                                      bool updateOptions = true) {
  options.mWorldSize = tpGrpSize;

  bool isBlackwell = tg::isArchBlackwell(cudaArch);
  bool const isRubin = cudaArch == tg::CudaArch::Sm107a;

  // If dtypeB is unspecified (Dtype::Void), assign to dtypeA.
  if (options.mDtypeB == tg::Dtype::Void) {
    if (updateOptions) {
      options.mDtypeB = options.mDtypeA;
    } else {
      return false;
    }
  }
  // If dtypeC is unspecified (Dtype::Void), assign to dtypeA.
  if (options.mDtypeC == tg::Dtype::Void) {
    TLLM_LOG_INFO("Setting dtypeC to ", tg::dtypeToString(options.mDtypeA));
    if (updateOptions) {
      options.mDtypeC = options.mDtypeA;
    } else {
      return false;
    }
  }

  // If not specified, used the input dtypes as MMA dtypes (no cast required).
  if (options.mDtypeMmaA == tg::Dtype::Void) {
    if (updateOptions) {
      options.mDtypeMmaA = options.mDtypeA;
    } else {
      return false;
    }
  }
  if (options.mDtypeMmaB == tg::Dtype::Void) {
    if (updateOptions) {
      options.mDtypeMmaB = options.mDtypeB;
    } else {
      return false;
    }
  }

  // If validM/N/K is not specified, then assume the full range of the dimension is valid.
  if (options.mValidM < 0 || options.mValidN < 0 || options.mValidK < 0) {
    if (updateOptions) {
      options.mValidM = options.mValidM < 0 ? options.mM : options.mValidM;
      options.mValidN = options.mValidN < 0 ? options.mN : options.mValidN;
      options.mValidK = options.mValidK < 0 ? options.mK : options.mValidK;
    } else {
      return false;
    }
  }

  // It must not exceed the padded dimensions.
  if (options.mValidM > options.mM || options.mValidN > options.mN ||
      options.mValidK > options.mK) {
    TLLM_LOG_WARNING(
      "ValidM, ValidN, and ValidK must be less than or equal to M, N, and K respectively.");
    if (updateOptions) {
      options.mValidM = std::min(options.mValidM, options.mM);
      options.mValidN = std::min(options.mValidN, options.mN);
      options.mValidK = std::min(options.mValidK, options.mK);
    } else {
      return false;
    }
  }

  // BlockMajorK layout does not support validM, validN, validK parameters
  if (options.mLayoutA == gemm::MatrixLayout::BlockMajorK ||
      options.mLayoutB == gemm::MatrixLayout::BlockMajorK) {
    bool hasValidParams = (options.mValidM != -1 && options.mValidM != options.mM) ||
                          (options.mValidN != -1 && options.mValidN != options.mN) ||
                          (options.mValidK != -1 && options.mValidK != options.mK);
    TLLM_CHECK_ERROR(
      !hasValidParams,
      "BlockMajorK layout does not support validM/validN/validK parameters due to swizzled layout. "
      "Found validM=",
      options.mValidM,
      " validN=",
      options.mValidN,
      " validK=",
      options.mValidK);
  }

#ifndef TLLM_GEN_EXPORT_INTERFACE
  // The exported interface also defines TLLM_PUBLIC_RELEASE but dispatches prebuilt cubins, so the
  // JIT-only rejection must not be baked into it.
  if (options.mDtypeA == tg::Dtype::E2m1 && options.mDtypeMmaA == tg::Dtype::E4m3) {
    TLLM_CHECK_ERROR(false, "E2m1 x E4m3 is not supported for JIT compile. Use cubins instead.");
  }
#endif // TLLM_GEN_EXPORT_INTERFACE

  // Check that the A cast is supported.
  // Currently, we only support {MxFp4, NvFp4, MxInt4} -> Bf16.
  TLLM_CHECK_ERROR((options.mDtypeA == options.mDtypeMmaA) ||
                     ((options.mDtypeA == tg::Dtype::MxE2m1 || options.mDtypeA == tg::Dtype::E2m1 ||
                       options.mDtypeA == tg::Dtype::MxInt4) &&
                      options.mDtypeMmaA == tg::Dtype::Bfloat16) ||
                     (options.mDtypeA == tg::Dtype::E2m1 && options.mDtypeMmaA == tg::Dtype::E4m3),
                   "Unsupported cast for A: ",
                   tg::dtypeToString(options.mDtypeA),
                   " -> ",
                   tg::dtypeToString(options.mDtypeMmaA));

  // Check that the B cast is supported.
  // Currently, we only support Fp8 -> MxFp8.
  // TODO: add same support for A (no transpose)
  TLLM_CHECK_ERROR(
    (options.mDtypeB == options.mDtypeMmaB) ||
      (options.mDtypeB == tg::Dtype::E4m3 && options.mDtypeMmaB == tg::Dtype::MxE4m3),
    "Unsupported cast for B: ",
    tg::dtypeToString(options.mDtypeB),
    " -> ",
    tg::dtypeToString(options.mDtypeMmaB));

  if (options.mDtypeA != options.mDtypeMmaA) {
    TLLM_CHECK_ERROR(options.mTileM == 128,
                     "TileM must be 128 when casting the input matrix A before the MMA.");
  }

  if (options.mPatchF2fp) {
    TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::MxE2m1 &&
                       options.mDtypeMmaA == tg::Dtype::Bfloat16,
                     "PatchF2fp is only supported for MxFp4 to Bf16 casts.");
  }

  // FIXME: We do not support different dtypes for A and B when not on Blackwell.
  if (!isBlackwell) {
    TLLM_CHECK_ERROR(options.mDtypeMmaA == options.mDtypeMmaB,
                     "For non-Blackwell, A and B must have the same dtype.");
  }

  // Check that the different dtypes for A and B are supported by the tensor core
  // kind::f8f6f4
  if (options.mDtypeMmaA == tg::Dtype::E4m3 || options.mDtypeMmaA == tg::Dtype::E2m1) {
    TLLM_CHECK_ERROR(options.mDtypeMmaB == tg::Dtype::E4m3 || options.mDtypeMmaB == tg::Dtype::E2m1,
                     "For dtypeMmaA = E4m3/E2m1 A, dtypeMmaB must also be E4m3/E2m1.");
  }

  // kind::mxf8f6f4
  if (options.mDtypeMmaA == tg::Dtype::MxE4m3 || options.mDtypeMmaA == tg::Dtype::MxE2m1) {
    TLLM_CHECK_ERROR(options.mDtypeMmaB == tg::Dtype::MxE4m3 ||
                       options.mDtypeMmaB == tg::Dtype::MxE2m1,
                     "For dtypeMmaA = MxE4m3 or MxE2m1, dtypeMmaB must also be MxE4m3 or MxE2m1.");
  }
  if (options.mDtypeMmaB == tg::Dtype::MxE4m3 || options.mDtypeMmaB == tg::Dtype::MxE2m1) {
    TLLM_CHECK_ERROR(options.mDtypeMmaA == tg::Dtype::MxE4m3 ||
                       options.mDtypeMmaA == tg::Dtype::MxE2m1,
                     "For dtypeMmaB = MxE4m3 or MxE2m1, dtypeMmaA must also be MxE4m3 or MxE2m1.");
  }

  // kind::f16
  if (options.mDtypeMmaA == tg::Dtype::Fp16 || options.mDtypeMmaA == tg::Dtype::Bfloat16) {
    TLLM_CHECK_ERROR(options.mDtypeMmaB == options.mDtypeMmaA,
                     "For dtypeMmaA = Fp16/Bfloat16, dtypeMmaB must be the same as dtypeMmaA.");
  }

  // When one of the inputs needs to be cast, we must use two load warps.
  if ((options.mDtypeMmaA != options.mDtypeA || options.mDtypeMmaB != options.mDtypeB) &&
      !options.mUseTwoTmaLoadWarps) {
    TLLM_LOG_WARNING("Two TMA load warps must be enabled if any of the inputs needs to be cast.");
  }

  // When different dtypes are used for A and B, we must use different tiles to do the loading.
  // It is not strictly required, but current implementation of SmemAb requires that.
  if (options.mDtypeA != options.mDtypeB) {
    TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps,
                     "Two TMA load warps must be enabled for different input types of A and B.");
  }

  // Number of stages.
  {
    // If updateOptions is false, all the numStages parameters must be set.
    if (!updateOptions && (options.mNumStagesA == 0 || options.mNumStagesB == 0 ||
                           options.mNumStagesSmemSfA == 0 || options.mNumStagesSmemSfB == 0 ||
                           options.mNumStagesTmemSfA == 0 || options.mNumStagesTmemSfB == 0)) {
      TLLM_CHECK_WARNING(options.mNumStagesA != 0, "numStagesA is not set.");
      TLLM_CHECK_WARNING(options.mNumStagesB != 0, "numStagesB is not set.");
      TLLM_CHECK_WARNING(options.mNumStagesSmemSfA != 0, "numStagesSmemSfA is not set.");
      TLLM_CHECK_WARNING(options.mNumStagesSmemSfB != 0, "numStagesSmemSfB is not set.");
      TLLM_CHECK_WARNING(options.mNumStagesTmemSfA != 0, "numStagesTmemSfA is not set.");
      TLLM_CHECK_WARNING(options.mNumStagesTmemSfB != 0, "numStagesTmemSfB is not set.");
      return false;
    }

    // If neither numStagesA nor numStagesB are set, default to 2.
    // Else, if one is set and the other is not, set them to the same value.
    if (options.mNumStagesA == 0 && options.mNumStagesB == 0) {
      TLLM_LOG_WARNING("Setting numStagesA and numStagesB to default value 2.");
      options.mNumStagesA = 2;
      options.mNumStagesB = 2;
    } else if (options.mNumStagesB == 0) {
      TLLM_LOG_WARNING("Setting numStagesB to numStagesA=", options.mNumStagesA);
      options.mNumStagesB = options.mNumStagesA;
    } else if (options.mNumStagesA == 0) {
      TLLM_LOG_WARNING("Setting numStagesA to numStagesB=", options.mNumStagesB);
      options.mNumStagesA = options.mNumStagesB;
    }

    // If neither numStagesSmemSfA nor numStagesSmemSfB is set, default to numStagesA.
    // Else, if one is set and the other is not, set them to the same value.
    if (options.mNumStagesSmemSfA == 0 && options.mNumStagesSmemSfB == 0) {
      TLLM_LOG_WARNING("Setting numStagesSmemSfA and numStagesSmemSfB to numStagesA=",
                       options.mNumStagesA);
      options.mNumStagesSmemSfA = options.mNumStagesA;
      options.mNumStagesSmemSfB = options.mNumStagesA;
    } else if (options.mNumStagesSmemSfB == 0) {
      TLLM_LOG_WARNING("Setting numStagesSmemSfB to numStagesSmemSfA=", options.mNumStagesSmemSfA);
      options.mNumStagesSmemSfB = options.mNumStagesSmemSfA;
    } else if (options.mNumStagesSmemSfA == 0) {
      TLLM_LOG_WARNING("Setting numStagesSmemSfA to numStagesSmemSfB=", options.mNumStagesSmemSfB);
      options.mNumStagesSmemSfA = options.mNumStagesSmemSfB;
    }

    // If fusedUtccpWithUtcmma is used, numStagesTmemSf for UTCCP-copyable layouts must be 1.
    if (options.mFuseUtccpWithUtcmma) {
      if (tg::sfLayoutCanUseUtccp(options.mSfLayoutA) && options.mNumStagesTmemSfA != 1) {
        TLLM_LOG_WARNING("Using fusedUtccpWithUtcmma and sfLayoutA=",
                         tg::sfLayoutToString(options.mSfLayoutA),
                         " requires numStagesTmemSfA=1 (got ",
                         options.mNumStagesTmemSfA,
                         ").");
        GEMM_UPDATE_OR_ERROR(options.mNumStagesTmemSfA, 1);
      }
      if (tg::sfLayoutCanUseUtccp(options.mSfLayoutB) && options.mNumStagesTmemSfB != 1) {
        TLLM_LOG_WARNING("Using fusedUtccpWithUtcmma and sfLayoutB=",
                         tg::sfLayoutToString(options.mSfLayoutB),
                         " requires numStagesTmemSfB=1 (got ",
                         options.mNumStagesTmemSfB,
                         ").");
        GEMM_UPDATE_OR_ERROR(options.mNumStagesTmemSfB, 1);
      }
    }

    // If numStagesTmemSfA or numStagesTmemSfB is not set, set to the respective numStagesSmemSf.
    if (options.mNumStagesTmemSfA == 0) {
      TLLM_LOG_WARNING("Setting numStagesTmemSfA to numStagesSmemSfA=", options.mNumStagesSmemSfA);
      options.mNumStagesTmemSfA = options.mNumStagesSmemSfA;
    }
    if (options.mNumStagesTmemSfB == 0) {
      TLLM_LOG_WARNING("Setting numStagesTmemSfB to numStagesSmemSfB=", options.mNumStagesSmemSfB);
      options.mNumStagesTmemSfB = options.mNumStagesSmemSfB;
    }

    // Check that the numbers of stages are valid.
    if (!options.mUseTwoTmaLoadWarps) {
      TLLM_CHECK_ERROR(options.mNumStagesA == options.mNumStagesB,
                       "numStagesA (",
                       options.mNumStagesA,
                       ") must equal numStagesB (",
                       options.mNumStagesB,
                       ") when not using separate load warps for A and B.");
    }
    if (tg::sfLayoutCanUseUtccp(options.mSfLayoutA) &&
        tg::sfLayoutCanUseUtccp(options.mSfLayoutB)) {
      TLLM_CHECK_ERROR(options.mNumStagesSmemSfA == options.mNumStagesSmemSfB,
                       "numStagesSmemSfA (",
                       options.mNumStagesSmemSfA,
                       ") must equal numStagesSmemSfB (",
                       options.mNumStagesSmemSfB,
                       ") when both SF layouts are UTCCP-copyable.");
      TLLM_CHECK_ERROR(options.mNumStagesTmemSfA == options.mNumStagesTmemSfB,
                       "numStagesTmemSfA (",
                       options.mNumStagesTmemSfA,
                       ") must equal numStagesTmemSfB (",
                       options.mNumStagesTmemSfB,
                       ") when both SF layouts are UTCCP-copyable.");
    }
  }

  // Get the mma kind for the input types.
  if (options.mMmaKind == tg::MmaKind::Auto) {
    if (updateOptions) {
      options.mMmaKind = dtypeGetMmaKind(options.mDtypeMmaA, options.mDtypeMmaB);
    } else {
      return false;
    }
  }

  // Check that the sparsity mode of A is supported, and compatible with the MMA kind.
  // Note: trtllm-gen currently does not support sparsity with tf32, fp16, bf16.
  switch (options.mSparsityA) {
  case tg::Sparsity::Dense:
    // Always supported.
    break;
  case tg::Sparsity::Any_1_2:
    TLLM_LOG_ERROR("1:2 sparsity is not supported.");
    break;
  case tg::Sparsity::Any_2_4: {
    bool isSupported_2_4 =
      (options.mMmaKind == tg::MmaKind::Fp8Fp6Fp4 || options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4);
    TLLM_CHECK_ERROR(isSupported_2_4,
                     "2:4 sparsity is not supported for MMA kind ",
                     tg::mmaKindToString(options.mMmaKind),
                     " on target ",
                     tg::cudaArchToString(cudaArch));
    break;
  }
  case tg::Sparsity::Pairwise_4_8:
    TLLM_CHECK_ERROR(options.mMmaKind == tg::MmaKind::MxFp4NvFp4,
                     "Pairwise 4:8 sparsity is only supported for MMA kind MxFp4NvFp4.");
    break;
  default:
    TLLM_CHECK_ERROR(false, "Unsupported sparsityA: ", tg::sparsityToString(options.mSparsityA));
    break;
  }

  // Is A sparse?
  bool const isSparseA = tg::isSparse(options.mSparsityA);

  // Requirements specific to sparsity, and compatibility with other features.
  if (isSparseA) {
    TLLM_CHECK_ERROR(isBlackwell, "Sparsity is only supported on Blackwell");
    // The following requirement is for TMA load: the box width must be a multiple of 16B.
    TLLM_CHECK_ERROR(
      tg::getNumBytesSparsityInfo(options.mSparsityA, options.mTileK) % 16 == 0,
      "The sparsity information for one tile row must be a multiple of 16B. Use larger tileK.");
    TLLM_CHECK_ERROR(options.mDtypeA == options.mDtypeMmaA,
                     "Sparsity is not supported with on-the-fly upcasting.");
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, "Sparsity is not supported with DeepSeek Fp8.");
    // Note: we could probably support FineGrained + sparsity, but it is not necessary if we use
    // sparsity only for weights, and FineGrained for activations.
    TLLM_CHECK_ERROR(!options.mFineGrainedConsumerA,
                     "Sparsity is not supported with FineGrained consumer.");
  }

  if (options.mMmaKind == tg::MmaKind::Fp8Fp6Fp4) {
    int mmaK = isSparseA ? 64 : 32;
    // SM107 supports 2x-mmaK for FP8.
    if (options.mMmaK == (isSparseA ? 128 : 64)) {
      TLLM_CHECK_ERROR(isRubin,
                       "MmaK == 64 (/128 sparse) is only supported on SM 107a for ",
                       tg::mmaKindToString(options.mMmaKind));
      TLLM_CHECK_ERROR(options.mMmaM % 128 == 0,
                       "MmaM must be a multiple of 128 for MmaK == 64 (/128 sparse) and ",
                       tg::mmaKindToString(options.mMmaKind),
                       ". Got mmaM == ",
                       options.mMmaM);
      mmaK = isSparseA ? 128 : 64;
    }

    if (options.mMmaK != mmaK) {
      TLLM_LOG_WARNING("Unsupported MmaK (",
                       options.mMmaK,
                       ") for MmaKind=",
                       gemm::toString(options.mMmaKind),
                       " and sparsity=",
                       tg::sparsityToString(options.mSparsityA),
                       ". Setting MmaK to ",
                       mmaK);
      if (updateOptions) {
        options.mMmaK = mmaK;
        options.mTileK = std::max(options.mMmaK, options.mTileK);
      } else {
        return false;
      }
    }
  }

  // Check LDTM shape.
  if (isBlackwell) {
    TLLM_CHECK_ERROR((options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256) ||
                       (options.mEpilogueLdtmDps == 32 && options.mEpilogueLdtmBits == 32),
                     "Unsupported LDTM shape: ",
                     options.mEpilogueLdtmDps,
                     "dp",
                     options.mEpilogueLdtmBits,
                     "bit.");
    if (options.mEpilogueTileM == 64) {
      TLLM_CHECK_ERROR(options.mEpilogueLdtmDps == 16,
                       "Unsupported LDTM shape for epilogueTileM=64: ",
                       options.mEpilogueLdtmDps,
                       "dp",
                       options.mEpilogueLdtmBits,
                       "bit.");
    }
    if (options.mTransposeMmaOutput) {
      // We can't use 32dp32bit LDTM for transposed outputs because we need each thread to own
      // multiple consecutive output elements.
      TLLM_CHECK_ERROR((options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256),
                       "Only 16dp256bit LDTM is supported for transposed outputs.");
    }
  } else {
    TLLM_CHECK_ERROR(
      options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256,
      "Hopper does not use TMEM. The register layout corresponds to 16dp256bit. Got ",
      options.mEpilogueLdtmDps,
      "dp",
      options.mEpilogueLdtmBits,
      "bit.");
  }

  // Constraints for NvFp4, MxFp8, and MxFp4.
  if (
    (options.mMmaKind == tg::MmaKind::MxFp4NvFp4 || options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4 ||
     options.mDtypeC == tg::Dtype::MxE4m3) &&
    options.mMmaM != 128) {

    if (options.mClusterDimX == 1) {
      // MMA M must be 128 when the input uses block scaling, or when the output is an Mx format.
      int newTileM = 128 * divUp(options.mTileM, 128);
      TLLM_LOG_WARNING("Unsupported MmaM (",
                       options.mMmaM,
                       ") for MmaKind=",
                       gemm::toString(options.mMmaKind),
                       ". Setting MmaM to 128 and TileM to ",
                       newTileM);
      if (updateOptions) {
        options.mMmaM = 128;
        options.mTileM = newTileM;
      } else {
        return false;
      }
    } else {
      TLLM_CHECK_ERROR(options.mMmaM == 256 && options.mTileM == 128,
                       "2CTA UTCxMMA only supports mmaM = 256 and tileM = 128.");
    }
  }
  if (options.mMmaTileK == 0) {
    options.mMmaTileK = options.mTileK;
  } else {
    if (options.mMmaTileK != options.mTileK) {
      TLLM_CHECK_ERROR(options.mUseCustomizedMma3xNvFp4,
                       "mMmaTileK is only supported for mUseCustomizedMma3xNvFp4");
    }
  }
  if (options.mMmaKind == tg::MmaKind::MxFp4NvFp4 || options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) {
    TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell");

    int mmaK = isSparseA ? 64 : 32;
    if (options.mMmaKind == tg::MmaKind::MxFp4NvFp4) {
      mmaK = isSparseA ? 128 : 64;
      // SM107 supports 2x-mmaK for FP4.
      if (options.mMmaK == 128 && !isSparseA) {
        TLLM_CHECK_ERROR(isRubin,
                         "MmaK == 128 dense is only supported on SM 107a for ",
                         tg::mmaKindToString(options.mMmaKind));
        // The SF layout for k=128 depends on whether mmaN is <= or > 128. To keep it simple in smem
        // -> tmem data movement, we don't want to have multiple possible layouts for the same tileN
        // (i.e. we keep that code agnostic of mmaN for now).
        TLLM_CHECK_ERROR(options.mTileN == options.mMmaN,
                         "When mmaK == 128 dense, TileN must be equal to MmaN.");
        mmaK = 128;
      }
      if (options.mMmaK == 96 && !isSparseA) {
        mmaK = 96;
      }
    }
    // SM107 supports 2x-mmaK for MxFP8.
    if (options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4 && options.mMmaK == (isSparseA ? 128 : 64)) {
      TLLM_CHECK_ERROR(isRubin,
                       "MmaK == 64 (/128 sparse) is only supported on SM 107a for ",
                       tg::mmaKindToString(options.mMmaKind));
      mmaK = isSparseA ? 128 : 64;
    }
    if (options.mMmaK != mmaK) {
      int newTileK = mmaK * divUp(options.mTileK, mmaK);
      TLLM_LOG_WARNING("Unsupported MmaK (",
                       options.mMmaK,
                       ") for MmaKind=",
                       gemm::toString(options.mMmaKind),
                       ". Setting MmaK to ",
                       mmaK,
                       " and TileK to ",
                       newTileK);
      if (updateOptions) {
        options.mMmaK = mmaK;
        options.mTileK = newTileK;
      } else {
        return false;
      }
    }

    // The MMA N may only be smaller than 64 if it is equal to the tile N.
    TLLM_CHECK_ERROR(options.mMmaN >= 64 || options.mMmaN == options.mTileN,
                     "MmaN (",
                     options.mMmaN,
                     ") must be >= 64 or equal to TileN (",
                     options.mTileN,
                     ")");
  }

  if (options.mUseCustomizedMma3xNvFp4) {
    TLLM_CHECK_ERROR(options.mMmaKind == tg::MmaKind::MxFp4NvFp4,
                     "mMmaKind needs to be tg::MmaKind::MxFp4NvFp4 for mUseCustomizedMma3xNvFp4.");
    TLLM_CHECK_ERROR(options.mMmaTileK == 192,
                     "mMmaTileK needs to be 192 for mUseCustomizedMma3xNvFp4.");
    TLLM_CHECK_ERROR(options.mMmaK == 96, "mMmaK needs to be 96 for mUseCustomizedMma3xNvFp4.");
    TLLM_CHECK_ERROR(options.mTileK == 256, "mTileK needs to be 256 for mUseCustomizedMma3xNvFp4");
    TLLM_CHECK_ERROR(options.mUseMaxTmemOverlap,
                     "mUseCustomizedMma3xNvFp4 requires mUseMaxTmemOverlap");
  }

  // Note: the logic for selecting/checking the correct block size based on dtypes and sparsity is
  // centralized here, to avoid error-prone code duplication and make it a more explicit "contract"
  // with the user who is providing inputs in this format.
  // Additionally, in some cases, multiple values are possible:
  // - When we use type casting before the MMA (e.g. e2m1 x e4m3).
  // - For output C, based on whether the consumer will use sparsity.

  // SF block size for A.
  if (options.mDtypeA == tg::Dtype::E2m1 && options.mDtypeB == tg::Dtype::E4m3) {
    // Note that the type conversion needs to happen before TCs.
    // For example, convert e2m1 to e4m3 inside TmemCastA.
    if (!(options.mSfBlockSizeA == 16 || options.mSfBlockSizeA == 32)) {
      TLLM_LOG_WARNING("sfBlockSizeA must be 16 or 32 for e2m1 x e4m3, got ",
                       options.mSfBlockSizeA);
      GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeA, 16);
    }
  } else if (options.mDtypeA == tg::Dtype::E2m1) {
    if (!((options.mSfBlockSizeA == 16 && !isSparseA) ||
          (options.mSfBlockSizeA == 32 && isSparseA))) {
      TLLM_LOG_WARNING("sfBlockSizeA must be 16 (dense) or 32 (sparse) for dtypeA=e2m1, got ",
                       options.mSfBlockSizeA);
      GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeA, isSparseA ? 32 : 16);
    }
  } else if (options.mDtypeA == tg::Dtype::MxE2m1 || options.mDtypeA == tg::Dtype::MxE4m3 ||
             options.mDtypeA == tg::Dtype::MxInt4) {
    if (!((options.mSfBlockSizeA == 32 && !isSparseA) ||
          (options.mSfBlockSizeA == 64 && isSparseA))) {
      TLLM_LOG_WARNING(
        "sfBlockSizeA must be 32 (dense) or 64 (sparse) for dtypeA=mx{e2m1,e4m3,int4}, got ",
        options.mSfBlockSizeA);
      GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeA, isSparseA ? 64 : 32);
    }
  } else if (options.mSfBlockSizeA > 0) {
    TLLM_LOG_WARNING("Got sfBlockSizeA=",
                     options.mSfBlockSizeA,
                     " but dtypeA=",
                     tg::dtypeToString(options.mDtypeA),
                     " does not use block scales");
    GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeA, -1);
  }
  // SF block size for B.
  if (options.mDtypeB == tg::Dtype::E2m1) {
    if (!((options.mSfBlockSizeB == 16 && !isSparseA) ||
          (options.mSfBlockSizeB == 32 && isSparseA))) {
      TLLM_LOG_WARNING("sfBlockSizeB must be 16 (dense) or 32 (sparse) for dtypeB=e2m1, got ",
                       options.mSfBlockSizeB);
      GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeB, isSparseA ? 32 : 16);
    }
  } else if (options.mDtypeB == tg::Dtype::MxE2m1 || options.mDtypeB == tg::Dtype::MxE4m3 ||
             (options.mDtypeB == tg::Dtype::E4m3 && options.mDtypeMmaB == tg::Dtype::MxE4m3)) {
    if (!((options.mSfBlockSizeB == 32 && !isSparseA) ||
          (options.mSfBlockSizeB == 64 && isSparseA))) {
      TLLM_LOG_WARNING(
        "sfBlockSizeB must be 32 (dense) or 64 (sparse) for dtypeB=mx{e2m1,e4m3}, got ",
        options.mSfBlockSizeB);
      GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeB, isSparseA ? 64 : 32);
    }
  } else if (options.mSfBlockSizeB > 0) {
    TLLM_LOG_WARNING("Got sfBlockSizeB=",
                     options.mSfBlockSizeB,
                     " but dtypeB=",
                     tg::dtypeToString(options.mDtypeB),
                     " does not use block scales");
    GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeB, -1);
  }
  // SF block size for C.
  if (options.mDtypeC == tg::Dtype::E2m1) {
    if (!(options.mSfBlockSizeC == 16 || options.mSfBlockSizeC == 32)) {
      TLLM_LOG_WARNING("sfBlockSizeC must be 16 or 32 for dtypeC=e2m1, got ",
                       options.mSfBlockSizeC);
      GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeC, 16);
    }
  } else if (options.mDtypeC == tg::Dtype::MxE2m1 || options.mDtypeC == tg::Dtype::MxE4m3) {
    if (!(options.mSfBlockSizeC == 32 || options.mSfBlockSizeC == 64)) {
      TLLM_LOG_WARNING("sfBlockSizeC must be 32 or 64 for dtypeC=mx{e2m1,e4m3}, got ",
                       options.mSfBlockSizeC);
      GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeC, 32);
    }
  } else if (options.mSfBlockSizeC > 0) {
    TLLM_LOG_WARNING("Got sfBlockSizeC=",
                     options.mSfBlockSizeC,
                     " but dtypeC=",
                     tg::dtypeToString(options.mDtypeC),
                     " does not use block scales");
    GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeC, -1);
  }


  if (tg::dtypeIsBlockFmt(options.mDtypeA)) {
    int sfATileK = 4;
    int numEltsPerSfA = options.mSfBlockSizeA;
    TLLM_CHECK_ERROR(options.mTileK % (sfATileK * numEltsPerSfA) == 0,
                     "TileK (",
                     options.mTileK,
                     ") must be a multiple of ",
                     (sfATileK * numEltsPerSfA),
                     " for numEltsPerSfA=",
                     numEltsPerSfA,
                     " and SF layout ",
                     tg::sfLayoutToString(options.mSfLayoutA));
    auto const numEltsPerSfAInK = divUp(options.mK, numEltsPerSfA);
    TLLM_CHECK_ERROR(numEltsPerSfAInK % sfATileK == 0,
                     "K dimension of scaling factors for A (",
                     numEltsPerSfAInK,
                     ") must be a multiple of ",
                     sfATileK,
                     " for SF layout ",
                     tg::sfLayoutToString(options.mSfLayoutA));
  }
  if (tg::dtypeIsBlockFmt(options.mDtypeB)) {
    TLLM_CHECK_ERROR(options.mSfLayoutB == tg::SfLayout::R128c4 ||
                       options.mSfLayoutB == tg::SfLayout::R8c4 ||
                       options.mSfLayoutB == tg::SfLayout::Linear,
                     "Only the 128x4, 8x4 and linear SF layouts are supported for B, got ",
                     tg::sfLayoutToString(options.mSfLayoutB));

    // TileN must be a multiple of the number of rows per SF tile.
    int const numSfTileRowsB = options.mSfLayoutB == tg::SfLayout::R128c4 ? 128 : 8;
    TLLM_CHECK_ERROR(options.mTileN % numSfTileRowsB == 0,
                     "TileN (",
                     options.mTileN,
                     ") must be a multiple of ",
                     numSfTileRowsB,
                     " for B SF layout ",
                     tg::sfLayoutToString(options.mSfLayoutB));

    int numEltsPerSfB = options.mSfBlockSizeB;
    TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSfB) == 0,
                     "TileK (",
                     options.mTileK,
                     ") must be a multiple of ",
                     (4 * numEltsPerSfB),
                     " for typeB ",
                     gemm::toString(options.mDtypeB));
    auto const numEltsPerSfBInK = options.mK / numEltsPerSfB;
    TLLM_CHECK_ERROR(numEltsPerSfBInK % 4 == 0,
                     "K dimension of scaling factors for B (",
                     numEltsPerSfBInK,
                     ") must be a multiple of 4");
  }

  int32_t padMultiplierA = 1;
  int32_t padMultiplierB = 1;
  if (options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) {
    if (options.mDtypeA == tg::Dtype::MxE2m1) {
      padMultiplierA = 2;
    }
    if (options.mDtypeB == tg::Dtype::MxE2m1) {
      padMultiplierB = 2;
    }
  }
  TLLM_CHECK_ERROR((padMultiplierA * tg::dtypeGetNumBits(options.mDtypeA) * options.mK / 8) % 16 ==
                     0,
                   "K dimension of A must be aligned to 16 bytes.");
  TLLM_CHECK_ERROR((padMultiplierB * tg::dtypeGetNumBits(options.mDtypeB) * options.mK / 8) % 16 ==
                     0,
                   "K dimension of B must be aligned to 16 bytes.");

  // Auto-deduce the dtype of the block scaling factors of C, or validate the user's choice.
  if (tg::dtypeIsBlockFmt(options.mDtypeC)) {
    auto const defaultDtypeSfC = tg::dtypeGetBlockSfType(options.mDtypeC);
    if (options.mDtypeSfC == tg::Dtype::Void) {
      if (updateOptions) {
        options.mDtypeSfC = defaultDtypeSfC;
      } else {
        return false;
      }
    } else {
      // Allow either the default quantized type, or Fp32 (keep the high-precision scaling factors).
      TLLM_CHECK_ERROR(options.mDtypeSfC == defaultDtypeSfC ||
                         options.mDtypeSfC == tg::Dtype::Fp32 ||
                         options.mDtypeSfC == tg::Dtype::Bfloat16,
                       "Unsupported dtypeSfC=",
                       tg::dtypeToString(options.mDtypeSfC),
                       " for dtypeC=",
                       tg::dtypeToString(options.mDtypeC),
                       ". Supported: ",
                       tg::dtypeToString(defaultDtypeSfC),
                       " (default) or Fp32 or Bfloat16.");
    }
  } else if (options.mDtypeSfC != tg::Dtype::Void) {
    TLLM_LOG_WARNING("Got dtypeSfC=",
                     tg::dtypeToString(options.mDtypeSfC),
                     " but dtypeC=",
                     tg::dtypeToString(options.mDtypeC),
                     " does not use block scales");
    GEMM_UPDATE_OR_ERROR(options.mDtypeSfC, tg::Dtype::Void);
  }

  if (tg::dtypeIsBlockFmt(options.mDtypeC)) {
    TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell");

    TLLM_CHECK_ERROR(options.mSfLayoutC == tg::SfLayout::R128c4 ||
                       options.mSfLayoutC == tg::SfLayout::R8c4 ||
                       options.mSfLayoutC == tg::SfLayout::Linear,
                     "Only the 128x4, 8x4 and linear SF layouts are supported for C.");
    if (!options.mTransposeMmaOutput) {
      TLLM_CHECK_ERROR(options.mEpilogueTileN % options.mSfBlockSizeC == 0,
                       "EpilogueTileN must be a multiple of the number of elements per SF for C");
    }
    if (options.mSfLayoutC != tg::SfLayout::Linear) {
      int const numSfTileRowsC = options.mSfLayoutC == tg::SfLayout::R128c4 ? 128 : 8;
      int const tileTokenDim = options.mTransposeMmaOutput ? options.mTileN : options.mTileM;
      TLLM_CHECK_ERROR_FMT(tileTokenDim % numSfTileRowsC == 0,
                           "Tile%s (%d) must be a multiple of %d for C SF layout %s",
                           options.mTransposeMmaOutput ? "N" : "M",
                           tileTokenDim,
                           numSfTileRowsC,
                           tg::sfLayoutToString(options.mSfLayoutC).c_str());
    }

    int numEltsPerSfC = options.mSfBlockSizeC;
    int const hiddenDim = options.mTransposeMmaOutput ? options.mM : options.mN;
    int const hiddenGranularity = 4 * numEltsPerSfC;
    TLLM_CHECK_ERROR(hiddenDim % hiddenGranularity == 0,
                     "Hidden dim (",
                     hiddenDim,
                     ") must be a multiple of ",
                     hiddenGranularity,
                     " for block-scaled outputs.");
    int const validHiddenDim = options.mTransposeMmaOutput ? options.mValidM : options.mValidN;
    TLLM_CHECK_ERROR(validHiddenDim % numEltsPerSfC == 0,
                     "Valid hidden dim (",
                     validHiddenDim,
                     ") must be a multiple of ",
                     numEltsPerSfC,
                     " for block-scaled outputs.");
    TLLM_CHECK_ERROR(!options.mTransposeMmaOutput || options.mUseShuffledMatrix,
                     "Transposing block-scaled outputs requires shuffled matrix.");
  }

  // Set epilogue tile sizes to the output tile sizes, when epilogue tile sizes are incorrect.
  if (options.mTileM % options.mEpilogueTileM != 0) {
    TLLM_LOG_WARNING("TileM (",
                     options.mTileM,
                     ") must be divisible by EpilogueTileM (",
                     options.mEpilogueTileM,
                     "). Setting EpilogueTileM to TileM");
    if (updateOptions) {
      options.mEpilogueTileM = options.mTileM;
    } else {
      return false;
    }
  }

  if (options.mTileN % options.mEpilogueTileN != 0) {
    TLLM_LOG_WARNING("TileN (",
                     options.mTileN,
                     ") must be divisible by EpilogueTileN (",
                     options.mEpilogueTileN,
                     "). Setting EpilogueTileN to TileN");
    if (updateOptions) {
      options.mEpilogueTileN = options.mTileN;
    } else {
      return false;
    }
  }

  // On Hopper, epilogue tile sizes are the same as output tiles.
  if (!isBlackwell &&
      (options.mEpilogueTileM != options.mTileM || options.mEpilogueTileN != options.mTileN)) {
    TLLM_LOG_WARNING(
      "Overwriting epilogueTileM and epilogueTileN to match tileM and tileN respectively");
    if (updateOptions) {
      options.mEpilogueTileM = options.mTileM;
      options.mEpilogueTileN = options.mTileN;
    } else {
      return false;
    }
  }

  // Unsupported epilogue tile size.
  if (options.mMmaM == 128 && options.mEpilogueTileM != options.mTileM) {
    TLLM_LOG_WARNING(
      "When MmaM = 128, EpilogueTileM must be equal to TileM. Setting EpilogueTileM to TileM");
    if (updateOptions) {
      options.mEpilogueTileM = options.mTileM;
    } else {
      return false;
    }
  }

  TLLM_CHECK_ERROR(options.mM > 0 && options.mN > 0 && options.mK > 0,
                   "M, N and K must be larger than 0");
  TLLM_CHECK_ERROR(options.mNumSlicesForSplitK > 0, "Split K must be larger than 0.");

  if (options.mUseShuffledMatrix) {
    auto const shuffleBlockSize = getShuffleBlockSize(options.mEpilogueTileM);
    TLLM_CHECK_ERROR(options.mM % shuffleBlockSize == 0 && options.mValidM % shuffleBlockSize == 0,
                     "M/validM must be a multiple of shuffle block size (",
                     shuffleBlockSize,
                     ") when useShuffledMatrix");
  }

  TLLM_CHECK_ERROR(options.mMmaM / (options.mClusterDimX > 1 ? 2 : 1) <= options.mEpilogueTileM,
                   "EpilogueTileM must be larger or equal than mmaM.");

  if (options.mClusterDimX >= 2) {
    TLLM_CHECK_ERROR(options.mMmaM == 256, "Only mmaM = 256 is supported for 2CTA UTCMMA.");
    TLLM_CHECK_ERROR(options.mMmaN % 16 == 0, "mmaN needs to be multiple of 16 for 2CTA UTCMMA.");
  }

  TLLM_CHECK_ERROR(
    options.mTileM % options.mEpilogueTileM == 0 && options.mTileN % options.mEpilogueTileN == 0,
    "TileM and TileN must be divisible by EpilogueTileM and EpilogueTileN respectively.");
  TLLM_CHECK_ERROR(
    (options.mClusterDimX == 1 || options.mClusterDimX == 2 || options.mClusterDimX == 4) &&
      (options.mClusterDimY == 1 || options.mClusterDimY == 2 || options.mClusterDimY == 4),
    "GEMM only support cluster sizes in X and Y of 1, 2 and 4, but found ",
    options.mClusterDimX,
    " and ",
    options.mClusterDimY);
  TLLM_CHECK_ERROR(options.mClusterDimZ == 1 || options.mNumSlicesForSplitK > 1,
                   "Cluster DimZ is only allowed for split-k.");
  TLLM_CHECK_ERROR(options.mTileM <= 128, "GEMM does not support TileM > 128.");

  if (options.mClusterDimY > 1) {
    TLLM_CHECK_ERROR(options.mClusterDimX >= 2,
                     "When mClusterDimY > 1, options.mClusterDimX has to at least be 2.");
  }

  if (options.mClusterDimX > 2 || options.mClusterDimY > 1) {
    TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps,
                     "Wider CGA sizes requires options.mUseTwoTmaLoadWarps");
    TLLM_CHECK_ERROR(
      options.mClusterDimZ == 1,
      "Only options.mClusterDimZ == 1 is supported when having CGA larger or equal than 2x1x1.");
  }

  if (options.mUseFlexibleClusterDims) {
    TLLM_CHECK_ERROR(options.mClusterDimX >= 2 && options.mFallbackClusterDimX >= 2,
                     "mClusterDimX and mFallbackClusterDimX can only be 2 or 4 for now.");
    TLLM_CHECK_ERROR(options.mFallbackClusterDimX > 0,
                     "options.mFallbackClusterDimX needs to be positive");
    TLLM_CHECK_ERROR(options.mFallbackClusterDimY > 0,
                     "options.mFallbackClusterDimY needs to be positive");
    TLLM_CHECK_ERROR(options.mClusterDimX % options.mFallbackClusterDimX == 0,
                     "mClusterDimX needs to be a multiple of mFallbackClusterDimX");
    TLLM_CHECK_ERROR(options.mClusterDimY % options.mFallbackClusterDimY == 0,
                     "mClusterDimY needs to be a multiple of mFallbackClusterDimY");
  }

  // FIXME: this is a bug in DeepSeek Fp8.
  if (options.mUseDeepSeekFp8) {
    TLLM_CHECK_ERROR(options.mK % (options.mNumSlicesForSplitK * options.mTileK) == 0,
                     "K must be a multiple of TileK * numSlicesForSplitK for DeepSeekFp8");
  }

  // Check all-reduce options.
  if (options.mAllReduceAlgo == AllReduceAlgo::OneShot) {
    // One shot is implemented with PTX cp.reduce.async.bulk.tensor which supports only the
    // following types for reduce add: u32, s32, u64, f32, f16, bf16.
    //
    // See: https://docs.nvidia.com/cuda/parallel-thread-execution/
    //   #data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor
    std::set<tg::Dtype> dtypeSupported{tg::Dtype::UInt32,
                                       tg::Dtype::Int32,
                                       tg::Dtype::UInt64,
                                       tg::Dtype::Fp32,
                                       tg::Dtype::Fp16,
                                       tg::Dtype::Bfloat16};
    TLLM_CHECK_ERROR(dtypeSupported.find(options.mDtypeC) != dtypeSupported.end(),
                     "Unsupported output dtype ",
                     tg::dtypeToString(options.mDtypeC));
  } else if (options.mAllReduceAlgo == AllReduceAlgo::TwoShot) {
    // TODO(anchengc):
    // Input dtype == output dtype -> can perform all-reduce in-place.
    // Input dtype != output dtype -> must perform all-reduce out of place.
    TLLM_CHECK_ERROR_FMT(
      options.mDtypeC == options.mDtypeAcc,
      "Not implemented - mixed dtype (dtypeC (%s) != dtypeAcc (%s)) requires out of place update",
      tg::dtypeToString(options.mDtypeC).c_str(),
      tg::dtypeToString(options.mDtypeAcc).c_str());
  }
  if (options.mAllReduceAlgo != AllReduceAlgo::None) {
    TLLM_CHECK_ERROR(options.mUseTmaStore, "Non-TMA store with all-reduce is not implemented");
  }

  if (updateOptions) {
    if (options.mNumSlicesForSplitK == 1) {
      // No split-k.
      options.mSplitK = SplitK::None;
    } else if (options.mNumSlicesForSplitK > 1 && options.mClusterDimZ == 1) {
      // Split-k with exchange through gmem.
      options.mSplitK = SplitK::Gmem;
    } else {
      // Split-k with exchange through Dsmem.
      options.mSplitK = SplitK::Dsmem;
    }
  }
  // For GMEM-based split-K, we write 4 elements at once.
  if (options.mSplitK == SplitK::Gmem) {
    TLLM_CHECK_ERROR((options.mM * options.mN) % 4 == 0,
                     "M * N must be a multiple of 4 for Split-K");
  }

  if (options.mNumSlicesForSplitK > 1) {
    if ((options.mEpilogueTileM != options.mTileM || options.mEpilogueTileN != options.mTileN) &&
        !options.mUseDeepSeekFp8) {
      TLLM_LOG_WARNING(
        "Overwriting epilogueTileM and epilogueTileN to match tileM and tileN respectively");
      if (updateOptions) {
        options.mEpilogueTileM = options.mTileM;
        options.mEpilogueTileN = options.mTileN;
      } else {
        return false;
      }
    }
  }
  if (options.mSplitK == SplitK::Dsmem) {
    TLLM_CHECK_ERROR(options.mClusterDimZ == options.mNumSlicesForSplitK,
                     "CGA size must be equal to the number of slices in split-k");
  }

  if (options.mUseShuffledMatrix && !options.mTransposeMmaOutput) {
    TLLM_CHECK_ERROR(
      !options.mUseDeepSeekFp8,
      "DeepSeek Fp8 is not supported when using shuffled matrix and non-transposed mma output");
    TLLM_CHECK_ERROR(
      options.mEpilogueLdtmBits == 32,
      "EpilogueLdtmBits must be 32 when using shuffled matrix and non-transposed mma output");
    TLLM_CHECK_ERROR(
      options.mEpilogueLdtmDps == 32,
      "EpilogueLdtmDps must be 32 when using shuffled matrix and non-transposed mma output");
    TLLM_CHECK_ERROR(
      options.mUseTmaStore,
      "TMA store is required when using shuffled matrix and non-transposed mma output");
    // When doing unshuffle in the epilogue, one fragment of epilogue tile must have at least one
    // shuffle block.
    auto minEpilogueTileN = getShuffleBlockSize(options.mEpilogueTileM);
    TLLM_CHECK_ERROR_FMT(options.mEpilogueTileN >= minEpilogueTileN,
                         "EpilogueTileN (%d) must be a larger than the shuffle block size (%d) "
                         "when using shuffled matrix and non-transposed mma output",
                         options.mEpilogueTileN,
                         minEpilogueTileN);
  }

  // Maps numStagesMma to (stagesWithinWorkTile, stagesAcrossWorkTile) if not already set.
  // If (-1, -1) -> (numStagesMma / min(2, numStagesMma), min(2, numStagesMma))
  // If ( m, -1) -> (m, numStagesMma / m)
  // If (-1,  n) -> (numStagesMma / n, n)
  if (options.mNumStagesMmaWithinWorkTile == -1 && options.mNumStagesMmaAcrossWorkTile == -1) {
    if (updateOptions) {
      options.mNumStagesMmaAcrossWorkTile = std::min(2, options.mNumStagesMma);
      options.mNumStagesMmaWithinWorkTile =
        options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile;
    } else {
      return false;
    }
  } else if (options.mNumStagesMmaWithinWorkTile == -1) {
    if (updateOptions) {
      options.mNumStagesMmaWithinWorkTile =
        options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile;
    } else {
      return false;
    }
  } else if (options.mNumStagesMmaAcrossWorkTile == -1) {
    if (updateOptions) {
      options.mNumStagesMmaAcrossWorkTile =
        options.mNumStagesMma / options.mNumStagesMmaWithinWorkTile;
    } else {
      return false;
    }
  }
  // Check mma stages.
  TLLM_CHECK_ERROR_FMT(options.mNumStagesMmaWithinWorkTile * options.mNumStagesMmaAcrossWorkTile ==
                           options.mNumStagesMma &&
                         options.mNumStagesMmaAcrossWorkTile <= 2,
                       "Condition numStagesMmaWithinWorkTile (%d) * numStagesMmaAcrossWorkTile "
                       "(%d) == numStagesMma (%d) && numStagesMmaAcrossWorkTile (%d) <= 2 must be "
                       "satisfied. Check arguments.",
                       options.mNumStagesMmaWithinWorkTile,
                       options.mNumStagesMmaAcrossWorkTile,
                       options.mNumStagesMma,
                       options.mNumStagesMmaAcrossWorkTile);
  // Mma stage must be 1 for pre-Hopper.
  TLLM_CHECK_ERROR(isBlackwell || options.mNumStagesMma == 1,
                   "Mma stage must be 1 for pre-Hopper. Found ",
                   options.mNumStagesMma);
  // DeepSeek Fp8
  if (!options.mUseDeepSeekFp8) {
    TLLM_CHECK_ERROR(options.mNumStagesMmaWithinWorkTile == 1,
                     "Non-DeepSeekFp8 requires numStagesMmaWithinWorkTile == 1");
    if (options.mNumStagesMma > 1) {
      TLLM_CHECK_ERROR(isPersistentScheduler(options.mTileScheduler),
                       "Non-DeepSeekFp8 requires persistent scheduler when using numStagesMma >1");
    }
  }
  if (options.mUseDeepSeekFp8) {
    TLLM_CHECK_ERROR(options.mClusterDimX == 1, "2CTA Gemm is not supported for DeepSeekFp8");
  }
  if (options.mUseDeepSeekFp8) {
    TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3,
                     "A and B dtype must be E4m3 for DeepSeek Fp8. Found dtypeA=",
                     tg::dtypeToString(options.mDtypeA),
                     " dtypeB=",
                     tg::dtypeToString(options.mDtypeB));

    TLLM_CHECK_ERROR(isBlackwell, "DeepSeek Fp8 is not supported for Hopper");
    TLLM_CHECK_ERROR(options.mAllReduceAlgo == AllReduceAlgo::None,
                     "DeepSeek Fp8 does not support AllReduce");

    // Check that TileK = 128 for correct scaling of every 128 channels.
    TLLM_CHECK_ERROR(options.mTileK == 128, "Tile-K must be equal to 128 for DeepSeek Fp8");
    TLLM_CHECK_ERROR(options.mK % options.mTileK == 0, "K must be a multiple of TileK");
    // Tile sizes of the output hidden dimension.
    auto hiddenDimPerOutputTile = options.mTransposeMmaOutput ? options.mTileM : options.mTileN;
    auto hiddenDimPerEpilogueTile =
      options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN;
    auto hiddenDimPerMma = options.mTransposeMmaOutput ? options.mMmaM : options.mMmaN;
    auto hiddenDimName = options.mTransposeMmaOutput ? "M" : "N";
    TLLM_CHECK_WARNING(options.mNumStagesMmaWithinWorkTile > 1,
                       "DeepSeekFp8 recommends setting \"-numStagesMmaWithinWorkTile 2\".");
    // Update the number of stages of the MMA accumulator pipeline. TODO: enable by default for
    // deepseek.
    // options.mNumStagesMma = 2;
    // Use two MMA warps to reduce mbar trywait latency. TODO: enable by default for deepseek.
    // options.mUseTwoMmaWarps = true;

    // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8.
    TLLM_CHECK_ERROR(options.mK % 128 == 0 && options.mValidK % 128 == 0,
                     "GEMM-K and validK must be a multiple of 128 when using DeepSeek Fp8. Found ",
                     options.mK,
                     " and validK=",
                     options.mValidK);

    // Check that the output tile N can be processed with the epilogue tile granularity.
    TLLM_CHECK_ERROR((hiddenDimPerOutputTile / 2) % hiddenDimPerEpilogueTile == 0,
                     "DeepSeek Fp8 requires Tile",
                     hiddenDimName,
                     " / 2 (",
                     hiddenDimPerOutputTile / 2,
                     ") being a multiple of EpilogueTile",
                     hiddenDimName,
                     " (",
                     hiddenDimPerEpilogueTile,
                     ")");
    // Check that the output tile N can be processed with the epilogue tile granularity.
    TLLM_CHECK_ERROR((hiddenDimPerOutputTile / 2) % hiddenDimPerMma == 0,
                     "DeepSeek Fp8 requires Tile",
                     hiddenDimName,
                     " / 2 (",
                     hiddenDimPerOutputTile / 2,
                     ") being a multiple of mma",
                     hiddenDimName,
                     " (",
                     hiddenDimPerMma,
                     ")");
  }

  TLLM_CHECK_ERROR(options.mNumEpilogueWarps == 4 || options.mNumEpilogueWarps == 8,
                   "mNumEpilogueWarps has to be either 4 or 8.");


  // Number of iterations in K dimension after padding.
  // Note the perCtaK in each CTA in the splitK group are padded to the same number of iterations.
  // E.g., K = 512, TileK = 128, numSlicesForSplitK = 3. Then the padded K is
  //
  //   ceil(512 / (128*3)) * (128*3) = 768
  //
  int const paddedK = divUpMul(options.mK, options.mTileK * options.mNumSlicesForSplitK);
  int const perCtaK = paddedK / options.mNumSlicesForSplitK;
  // However, number of iterations is clamped to multiples of tileK within individual CTAs
  // E.g., K = 448, TileK = 64, numSlicesForSplitK = 4.
  //
  //   paddedK                        = 512
  //   perCtaK                        = 128
  //   clampedPerCtaK for CTA 0, 1, 2 = 128
  //   clampedPerCtaK for CTA 3       = 64
  int const paddingForK = paddedK - options.mK;
  int const clampedAndPaddedPerCtaK = divUpMul(perCtaK - paddingForK, options.mTileK);
  if (options.mUseUnrollLoop2xForMma) {
    // Check that the padded K and clamped padded K (K rounded to next multiple of tileK) is a
    // multiple of 2*TileK when UnrollLoop2x is enabled. This is to avoid deadlock when mma runs
    // even-numbered loop while the other warps run odd-numbered loop.
    //
    bool notSupported = (perCtaK % (options.mTileK * 2) != 0) ||
                        (clampedAndPaddedPerCtaK % (options.mTileK * 2) != 0);
    if (notSupported) {
      TLLM_LOG_WARNING("Size K / splitK must be a multiple of TileK * 2. Found TileK=",
                       options.mTileK,
                       " and K=",
                       options.mK,
                       " (paddedK=",
                       paddedK,
                       " clampedAndPaddedPerCtaK=",
                       clampedAndPaddedPerCtaK,
                       ") and numSlicesForSplitK=",
                       options.mNumSlicesForSplitK,
                       ". Disabling unrollLoop2xForMma.");
      if (updateOptions) {
        options.mUseUnrollLoop2xForMma = false;
      } else {
        return false;
      }
    }
  }
  if (options.mNumSlicesForSplitK > 1) {
    TLLM_CHECK_ERROR(
      perCtaK * (options.mNumSlicesForSplitK - 1) < options.mK,
      "K must be greater than perCtaK * (numSlicesForSplitK - 1) to ensure each CTA has work");
  }

  if (options.mSchedHostTask != SchedHostTask::Self) {
    TLLM_CHECK_ERROR(isPersistentScheduler(options.mTileScheduler),
                     "SchedHostTask requires a persistent tile scheduler.");
  }

  if (!isBlackwell && options.mTileScheduler == TileScheduler::Persistent) {
    TLLM_LOG_WARNING("Persistent scheduling is not supported on Hopper. Use StaticPersistent or "
                     "PersistentSm90 instead. Fallback to Static scheduling.");
    if (updateOptions) {
      options.mTileScheduler = TileScheduler::Static;
    } else {
      return false;
    }
  }

  if (isBlackwell && !options.mUseCustomMmaSchedule && !options.mUseDeepSeekFp8 &&
      isPersistentScheduler(options.mTileScheduler)) {
    if (updateOptions) {
      options.mUseCustomMmaSchedule = true;
    } else {
      TLLM_CHECK_ERROR(false,
                       "Persistent scheduler and !UseCustomMmaSchedule is not supported."
      );
    }
  }

  if (options.mEnablesDelayedEarlyExit && options.mEnablesEarlyExit) {
    TLLM_LOG_WARNING("Only one of early exit and delayed early exit should be enabled. Disabling "
                     "delayed early exit");
    if (updateOptions) {
      options.mEnablesDelayedEarlyExit = false;
    } else {
      return false;
    }
  }

  // This check prevents the triggering of the secondary (PREEXIT) from executing before the wait
  // for primary (ACQBULK). This could lead to the following confusing situation, which we want to
  // avoid:
  //
  // Kernel 3 is written with the assumption that it can read the output of
  // kernel 1 *without* ACQBULK and the output of kernel 2 *with* ACQBULK.
  // However, when we allow PREEXIT and ACQBULK to be executed out of order,
  // this is not guaranteed.
  //
  // Time:      ---->
  //
  // Kernel 1:  ----PREEXIT-----------FLUSH
  // Kernel 2:      -------PREEXIT----ACQBULK---FLUSH
  // Kernel 3:  Warp 0:           ---- (!) Output of 1,2 is not yet visible
  // -----------------------
  //            Warp 1:           ---- (!) We normally assume that 1 is visible is not yet
  //            visible- Warp 2:           -------------------ACQBULK-- Kernel 1,2 output visible
  //            ----------
  TLLM_CHECK_ERROR(
    (options.mGridWaitForPrimaryA || !options.mGridTriggerSecondaryA),
    "A: If a task triggers a secondary kernel, it must also wait for primary kernel.");
  TLLM_CHECK_ERROR(
    (options.mGridWaitForPrimaryB || !options.mGridTriggerSecondaryB),
    "B: If a task triggers a secondary kernel, it must also wait for primary kernel.");

  if (options.mPerTokenSfDtype == tg::Dtype::Void) {
    if (options.mUsePerTokenSfA && options.mUsePerTokenSfB) {
      options.mPerTokenSfDtype = tg::Dtype::Fp32;
    } else {
      options.mPerTokenSfDtype = tg::Dtype::Bfloat16;
    }
  }
  if (options.mUsePerTokenSfA || options.mUsePerTokenSfB) {
    // Checks applicable to both MetaFP8 and RoutingScalesOnInput
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8,
                     "DeepSeek FP8 and per-token scaling are not compatible");
    TLLM_CHECK_ERROR(isBlackwell, "Per-token scaling is not supported for Hopper");
    if (options.mUsePerTokenSfA && options.mUsePerTokenSfB) {
      // MetaFP8 case
      TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3,
                       "A and B dtype must be E4m3 for Meta Fp8. Found dtypeA=",
                       tg::dtypeToString(options.mDtypeA),
                       " dtypeB=",
                       tg::dtypeToString(options.mDtypeB));
      TLLM_CHECK_ERROR(options.mDtypeC == tg::Dtype::Fp32 ||
                         options.mDtypeC == tg::Dtype::Bfloat16 ||
                         options.mDtypeC == tg::Dtype::Fp16 || options.mDtypeC == tg::Dtype::E4m3 ||
                         options.mDtypeC == tg::Dtype::E2m1,
                       "Only Fp32, Bfloat16, Fp16, E4m3, E2m1 output dtypes are supported for Meta "
                       "Fp8");
    }
    TLLM_CHECK_ERROR(
      options.mPerTokenSfDtype == tg::Dtype::Fp32 ||
        options.mPerTokenSfDtype == tg::Dtype::Bfloat16 ||
        options.mPerTokenSfDtype == tg::Dtype::Fp16,
      "Only Fp32, Bfloat16, Fp16 output dtypes are supported for the per-token scaling factors");
  }

  // The generation should support non K-major layouts for both A and B; however, it is unclear if
  // there is a use-case
  TLLM_CHECK_ERROR((options.mLayoutA == MatrixLayout::MajorK) ||
                     (options.mLayoutB == MatrixLayout::MajorK),
                   "At least one matrix must be in k-major layout");

  // Some features are currently only support when both matrices are in K-major format
  if (options.mLayoutA != MatrixLayout::MajorK || options.mLayoutB != MatrixLayout::MajorK) {
    TLLM_CHECK_ERROR(isBlackwell, "Non K-major layouts are only supported on Blackwell");
    TLLM_CHECK_ERROR(options.mSplitK == SplitK::None, "Non K-major layouts do not support split K");
  }
  if (options.mLayoutA == MatrixLayout::MajorMn) {
    TLLM_CHECK_ERROR(tg::dtypeGetNumBits(options.mDtypeA) >= 8,
                     "Subbyte types do not support m-major layout");
  }
  if (options.mLayoutB == MatrixLayout::MajorMn) {
    TLLM_CHECK_ERROR(tg::dtypeGetNumBits(options.mDtypeB) >= 8,
                     "Subbyte types do not support n-major layout");
  }

  if ((options.mLayoutA == MatrixLayout::BlockMajorK) ||
      (options.mLayoutB == MatrixLayout::BlockMajorK)) {
    bool const isBlockA = options.mLayoutA == MatrixLayout::BlockMajorK;

    int32_t const padMultiplier = (isBlockA) ? padMultiplierA : padMultiplierB;
    int32_t const elemSizeInBits =
      (isBlockA) ? tg::dtypeGetNumBits(options.mDtypeA) : tg::dtypeGetNumBits(options.mDtypeB);
    int32_t const elemsIn128B = 128 * 8 /* Bits in byte */ / elemSizeInBits;

    // Number of non-zero elements in the k dimension.
    int32_t const nzTileK = options.mTileK >> static_cast<int32_t>(isBlockA && isSparseA);
    // Number of 128B SMEM slices per tile.
    int32_t const smemSlicesPerTile = padMultiplier * nzTileK / elemsIn128B;

    if (smemSlicesPerTile > 2) {
      if (options.mBlockK != elemsIn128B / padMultiplier) {
        // This is to prevent a bug when the TMA box width is truncated to 128B (after padding)
        // and multiple TMA instructions are loading multiple non-contiguous slices each.
        // E.g. TMA #0 loads slices (0,2), TMA #1 loads slices (1,3)
        TLLM_LOG_WARNING("TileK=",
                         options.mTileK,
                         " with ",
                         padMultiplier,
                         "x padding spans across ",
                         smemSlicesPerTile,
                         " 128B SMEM slices. Setting blockK to ",
                         elemsIn128B / padMultiplier);
        GEMM_UPDATE_OR_ERROR(options.mBlockK, elemsIn128B / padMultiplier);
      }
    } else {
      // The larger blockK (128B vs 64B) is generally 1-2% more performant.
      if (options.mBlockK != elemsIn128B && options.mBlockK != elemsIn128B / padMultiplier) {
        TLLM_LOG_WARNING("Setting blockK to ", elemsIn128B);
        GEMM_UPDATE_OR_ERROR(options.mBlockK, elemsIn128B);
      }
    }

    if (options.mBlockK > options.mTileK) {
      TLLM_CHECK_ERROR(
        options.mBlockK % options.mTileK == 0,
        "If block size is greater than tile size, block size must be a multiple of tile size");
    } else if (options.mBlockK < options.mTileK) {
      TLLM_CHECK_ERROR(
        options.mTileK % options.mBlockK == 0,
        "If tile size is greater than block size, tile size must be a multiple of block size");
    }
  }

  if (usesFusedBiasShuffle(options.mFusedBiasShuffleMode)) {
    TLLM_CHECK_ERROR(options.mUseShuffledMatrix,
                     "mFusedBiasShuffleMode requires useShuffledMatrix");
    TLLM_CHECK_ERROR(isBiasTypeMn(options.mBiasType),
                     "mFusedBiasShuffleMode is only supported for BiasType::Mn");
  }

  if (!isBiasTypeNone(options.mBiasType)) {
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8 || isBiasTypeMn(options.mBiasType),
                     "Only BiasType::Mn is supported for DeepSeek Fp8");
    TLLM_CHECK_ERROR(!(options.mUsePerTokenSfA && options.mUsePerTokenSfB),
                     "Bias is not supported for Meta Fp8");
  }

  TLLM_CHECK_ERROR(options.mBiasDtype == tg::Dtype::Fp32 ||
                     options.mBiasDtype == tg::Dtype::Bfloat16 ||
                     options.mBiasDtype == tg::Dtype::Fp16,
                   "Bias dtype must be Fp32, Bfloat16, or Fp16");

  if (options.mUseMaxTmemOverlap) {
    TLLM_CHECK_ERROR(options.mUseTmaStore, "mUseMaxTmemOverlap only works with TMA store");
    TLLM_CHECK_ERROR(options.mNumSlicesForSplitK == 1,
                     "mUseMaxTmemOverlap does not work with splitK");
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8,
                     "mUseMaxTmemOverlap does not work with mUseDeepSeekFp8");
  }

  if (options.mNumEpilogueWarps > 4) {
    TLLM_CHECK_ERROR(options.mUseTmaStore,
                     "Using more than 4 warps for epilogue only works with TMA store");
    TLLM_CHECK_ERROR(options.mNumSlicesForSplitK == 1,
                     "Using more than 4 warps for epilogue does not work with splitK");
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8,
                     "Using more than 4 warps for epilogue does not work with mUseDeepSeekFp8");

    auto const numEpilogueWrpGrps = options.mNumEpilogueWarps / 4;
    TLLM_CHECK_ERROR(options.mTileN % (options.mEpilogueTileN * numEpilogueWrpGrps) == 0,
                     "TileN must be a multiple of EpilogueTileN * numEpilogueWrpGrps");
  }

  if (options.mFineGrainedConsumerA) {
    // TODO This should be easy to enable, just need to modify the consumer arrival count.
    TLLM_CHECK_ERROR(options.mDtypeA == options.mDtypeMmaA,
                     "FineGrained pipeline does not support casting");
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8,
                     "FineGrained Consumer does not support DeepSeek FP8 format");
    TLLM_CHECK_ERROR(padMultiplierA == 1,
                     "FineGrained consumer does not support bit padding for dtypeA=",
                     tg::dtypeToString(options.mDtypeA));
  }
  if (options.mFineGrainedConsumerB) {
    TLLM_CHECK_ERROR(options.mDtypeB == options.mDtypeMmaB,
                     "FineGrained pipeline does not support casting");
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8,
                     "FineGrained Consumer does not support DeepSeek FP8 format");
    TLLM_CHECK_ERROR(padMultiplierB == 1,
                     "FineGrained consumer does not support bit padding for dtypeB=",
                     tg::dtypeToString(options.mDtypeB));
  }
  if (options.mFineGrainedConsumerA || options.mFineGrainedConsumerB) {
    TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps,
                     "Separate load warps must be used when FineGrained pipeline is enabled");

    TLLM_CHECK_ERROR(!options.mGridWaitForPrimaryEarlyExit,
                     "FineGrained consumer should not wait for primary");


    // TODO Update the FineGrained pipeline to use CGA barriers. Note, the SplitK implementation
    // does not use CGA barriers within the DMA pipelines, so that feature is enabled.
    TLLM_CHECK_ERROR(options.mClusterDimX == 1 && options.mClusterDimY == 1,
                     "FineGrained pipeline does not support 2 CTA mode or features requiring CGA "
                     "barriers in DMA pipeline");
  }
  if (options.mFineGrainedForceValid) {
    TLLM_CHECK_ERROR(options.mDtypeC != tg::Dtype::Fp32,
                     "FineGrainedForceValid is not implemented for Fp32 output");
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8,
                     "FineGrainedForceValid is not implemented for DeepSeekFp8");
  }
  if (options.mFineGrainedProducer) {
    TLLM_CHECK_ERROR(options.mUseTmaStore, "FineGrained producer requires Tma store");
    TLLM_CHECK_ERROR(options.mFineGrainedForceValid,
                     "FineGrained producer should always force valid values");
    TLLM_CHECK_ERROR(!options.mGridTriggerSecondaryA,
                     "FineGrained producer cannot trigger secondary in A");
    TLLM_CHECK_ERROR(!options.mGridTriggerSecondaryB,
                     "FineGrained producer cannot trigger secondary in B");
    TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8,
                     "FineGrained producer is not implemented for DeepSeekFp8");
  }

  if (updateOptions) {
    // Init kernel traits.
    int32_t const tileMForTraits = options.mTileM;
    int32_t const tileNForTraits = options.mTileN;
    options.mKernelTraits = KernelTraits(options.mDtypeA,
                                         options.mDtypeB,
                                         options.mDtypeC,
                                         options.mDtypeAcc,
                                         options.mDtypeMmaA,
                                         options.mDtypeMmaB,
                                         options.mMmaKind,
                                         options.mSparsityA,
                                         options.mMmaK,
                                         tileMForTraits,
                                         tileNForTraits,
                                         options.mTileK,
                                         options.mMmaTileK,
                                         options.mEpilogueTileM,
                                         options.mEpilogueTileN,
                                         options.mSfBlockSizeA,
                                         options.mSfBlockSizeB,
                                         options.mNumStagesA,
                                         options.mNumStagesB,
                                         options.mNumStagesMma,
                                         options.mNumStagesTmemSfA,
                                         options.mNumStagesTmemSfB,
                                         options.mNumSlicesForSplitK,
                                         options.mSplitK,
                                         options.mUseTmaStore,
                                         options.mTransposeMmaOutput,
                                         options.mAllReduceAlgo,
                                         options.mFuseUtccpWithUtcmma,
                                         options.mUseMaxTmemOverlap,
                                         options.mUseCustomizedMma3xNvFp4,
                                         options.mNumEpilogueWarps,
                                         isPersistentScheduler(options.mTileScheduler),
                                         options.mUseDeepSeekFp8,
                                         options.mUsePerTokenSfA,
                                         options.mUsePerTokenSfB,
                                         /* useTwoCtas*/ options.mClusterDimX >= 2,
                                         options.mBiasType,
                                         options.mFusedBiasShuffleMode
    );
  }

  return true;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

inline bool getDoesScaleC(tg::Dtype dtypeC) {
  // Need to scale/quantize the output C matrix when the output type is Fp8 or NvFp4.
  return dtypeC == tg::Dtype::E4m3 || dtypeC == tg::Dtype::E2m1;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

inline bool getDoesScaleAb(tg::Dtype dtypeA,
                           tg::Dtype dtypeB,
                           bool useDeepSeekFp8,
                           bool useMetaFp8) {
  // Need to scale/dequantize the input A/B matrices when the input type is Fp8 or NvFp4 and
  // DeepSeekFp8 is not used.
  bool const doesScaleAb{
    dtypeA == tg::Dtype::E2m1 || dtypeB == tg::Dtype::E2m1 ||
    ((dtypeA == tg::Dtype::E4m3 || dtypeB == tg::Dtype::E4m3) && !useDeepSeekFp8 && !useMetaFp8)};
  return doesScaleAb;
}

//////////////////////////////////////////////////////////////////////////////////////////////////

inline bool getDoesScaleAct(tg::Dtype dtypeA,
                            tg::Dtype dtypeB,
                            bool useDeepSeekFp8,
                            bool useMetaFp8,
                            EltwiseActType eltwiseActType) {
  // Only non-linear activations require separate scaleAct.
  bool const isLinearAct = eltwiseActType == EltwiseActType::None;
  return !isLinearAct && getDoesScaleAb(dtypeA, dtypeB, useDeepSeekFp8, useMetaFp8);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

inline bool getKernelDoesScaleC(tg::Dtype dtypeA,
                                tg::Dtype dtypeB,
                                tg::Dtype dtypeC,
                                bool useDeepSeekFp8,
                                bool useMetaFp8) {
  // In the Gemm/BatchedGemm kernels, dequantScaleAb and quantScaleC are combined into one single
  // scaling factor (called scaleC). As a result, we combine the logic for getDoesScaleAb and
  // getDoesScaleC.
  return getDoesScaleC(dtypeC) || getDoesScaleAb(dtypeA, dtypeB, useDeepSeekFp8, useMetaFp8);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename Config> inline CUresult loadCubinData(CUmodule* module, Config const& config) {
  // Trtllm links the cubin into the executable while Flashinfer loads the cubin from storage.
#ifdef TLLM_GEN_EXPORT_FLASHINFER
#ifdef TLLM_GEN_GEMM_CUBIN_PATH
  static const std::string tllm_gen_gemm_cubin_path = std::string(TLLM_GEN_GEMM_CUBIN_PATH);
  const std::string sha256 = config.mHash ? config.mHash : "";
  std::string fileName = config.mFunctionName;
  if (!fileName.empty()) {
    fileName[0] = static_cast<char>(std::toupper(static_cast<unsigned char>(fileName[0])));
  }
  const std::string& data =
    flashinfer::trtllm_cubin_loader::getCubin(tllm_gen_gemm_cubin_path + "/" + fileName + ".cubin",
                                              sha256);
  CUresult result = cuModuleLoadData(module, data.c_str());
#else
  static_assert(false, "TLLM_GEN_GEMM_CUBIN_PATH macro is not defined when compiling");
#endif // TLLM_GEN_GEMM_CUBIN_PATH
#else
  CUresult result = cuModuleLoadData(module, config.mData);
#endif // TLLM_GEN_EXPORT_FLASHINFER
  return result;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

} // namespace gemm

#ifdef TLLM_GEN_EXPORT_INTERFACE

#undef TLLM_CHECK_ERROR
#undef TLLM_CHECK_ERROR_FMT
#undef TLLM_CHECK_WARNING
#undef TLLM_LOG_WARNING
#undef TLLM_LOG_INFO
#undef TLLM_LOG_ERROR

#endif // TLLM_GEN_EXPORT_INTERFACE

} // namespace batchedGemm