/*
* 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 <stdexcept>
#include "trtllm/gen/SfLayoutDecl.h"
#include "trtllm/gen/SparsityDecl.h"
#include "trtllm/gen/CommonUtils.h"
#include <tuple>

#include "TmaDescriptor.h"
#include "Enums.h"
#include "BatchedGemmEnums.h"

// NOTE: keep this code dependency free. It has to be included by the device code and has to be
// compilable with NVRTC.
#include "KernelParamsDecl.h"

namespace batchedGemm {

namespace batchedGemm {

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

// TODO: Find a better header to put this in, that we can include from here.
template <typename T> inline T ceilDiv(T m, T n) {
  return (m + n - T(1)) / n;
}

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

namespace tg = trtllm::gen;

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

namespace KernelParamsSetup {
#ifdef TLLM_ENABLE_CUDA

enum class MatrixType { MatrixA = 0, MatrixB, MatrixC };

//////////////////////////////////////////////////////////////////////////////////////////////////
//
// Utility functions.
//
//////////////////////////////////////////////////////////////////////////////////////////////////

template <typename BatchedGemmOptions> bool useTmaOobOptA(BatchedGemmOptions const& options) {
  return options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM &&
         doesRouteImplUseNoRoute(options.mRouteImpl) && options.mUseTmaOobOpt;
}

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

template <typename BatchedGemmOptions> bool useTmaOobOptB(BatchedGemmOptions const& options) {
  return options.mBatchMode == BatchedGemmOptions::BatchMode::BatchN &&
         doesRouteImplUseNoRoute(options.mRouteImpl) && options.mUseTmaOobOpt;
}

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

template <typename BatchedGemmOptions> bool useTmaOobOptC(BatchedGemmOptions const& options) {
  // C multicast uses direct pointers instead of TMA store
  return options.mUseTmaStore && !options.mUseCMultiCast && options.mUseTmaOobOpt;
}

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

// Create the TMA shape/stride for A/B/C.
template <class GemmOptions>
static auto makeTmaShapeStrideAbc(GemmOptions const& options,
                                  int sizeM,
                                  int sizeN,
                                  int sizeK,
                                  int tileM,
                                  int tileN,
                                  int tileK,
                                  MatrixType matrixType,
                                  int validM = -1,
                                  int validN = -1,
                                  int validK = -1) {
  // Default to padded dimensions if not provided.
  validM = validM < 0 ? sizeM : validM;
  validN = validN < 0 ? sizeN : validN;
  validK = validK < 0 ? sizeK : validK;
  // Weights matrix is A if we transpose the output of MMA (to have it M-major).
  // Otherwise, it is B, when the output of MMA is K-major.
  bool const isWeights = (matrixType == MatrixType::MatrixA && options.mTransposeMmaOutput) ||
                         (matrixType == MatrixType::MatrixB && !options.mTransposeMmaOutput);

  // Whether to use TMA OOB trick to block out padded dummy tokens and saving BW whenever no routing
  // is involved. It applies to batchM and matrixA, or batchN and matrixB, or any case for matrixC.
  bool const useTmaOobOpt = matrixType == MatrixType::MatrixA   ? useTmaOobOptA(options)
                            : matrixType == MatrixType::MatrixB ? useTmaOobOptB(options)
                            : matrixType == MatrixType::MatrixC ? useTmaOobOptC(options)
                                                                : false;

  // The outer dimension. Uses padded dimensions for strides and valid dimensions for shapes.
  auto numTokens =
    (matrixType == MatrixType::MatrixA || matrixType == MatrixType::MatrixC) ? sizeM : sizeN;
  auto numTokensValid =
    (matrixType == MatrixType::MatrixA || matrixType == MatrixType::MatrixC) ? validM : validN;
  // The outer dimension tile size.
  auto ctaTileNumTokens =
    (matrixType == MatrixType::MatrixA || matrixType == MatrixType::MatrixC) ? tileM : tileN;
  // The outer dimension of TMA box shape.
  auto tileNumTokens =
    (matrixType == MatrixType::MatrixC) ? options.mEpilogueTileM : ctaTileNumTokens;

  // The inner dimension.
  auto hiddenSize = (matrixType == MatrixType::MatrixC) ? sizeN : sizeK;
  auto hiddenSizeValid = (matrixType == MatrixType::MatrixC) ? validN : validK;
  // The inner dimension tile size.
  auto ctaTileHiddenSize = (matrixType == MatrixType::MatrixC) ? tileN : tileK;
  // The inner dimension of TMA box shape.
  auto tileHiddenSize =
    (matrixType == MatrixType::MatrixC) ? options.mEpilogueTileN : ctaTileHiddenSize;

  // Swap matrix C sizes if output is transposed.
  if (matrixType == MatrixType::MatrixC && options.mTransposeMmaOutput) {
    std::swap(numTokens, hiddenSize);
    std::swap(numTokensValid, hiddenSizeValid);
    std::swap(ctaTileNumTokens, ctaTileHiddenSize);
    std::swap(tileNumTokens, tileHiddenSize);
  }

  // For a fused activation kernel, the hidden size of output is halved. TODO: That's true for
  // gated activations but not regular activations.
  if (options.mFusedAct && matrixType == MatrixType::MatrixC) {
    hiddenSize /= 2;
    hiddenSizeValid /= 2;
    tileHiddenSize /= 2;
    ctaTileHiddenSize /= 2;
  }

  // The cute tensor shape for A/B: (numTokens, hiddenSize).
  // Note that TMA descriptor expects the first dimension's stride to be
  // 1, so swap the first two dimension so that the hiddenSize dimension comes first.

  // Activations matrix is 2D (sum(divUpMul(M[bi], tileM) for bi in B), K).
  // Use valid dimensions for shape.
  std::vector<uint64_t> shape = {static_cast<uint64_t>(hiddenSizeValid),
                                 static_cast<uint64_t>(numTokensValid)};
  if (useTmaOobOpt /* also implies input/output activation */) {
    // If TMA OOB optimization is used:
    // Shape [hidden, tokens]                      Stride [1, hidden] becomes
    // Shape [hidden, tileN, TmaDimMax, TmaDimMax] Stride [1, hidden, XLargeN - hidden, hidden]
    shape = {static_cast<uint64_t>(hiddenSizeValid),
             static_cast<uint64_t>(ctaTileNumTokens),
             static_cast<uint64_t>(tg::TmaDimMax),
             static_cast<uint64_t>(tg::TmaDimMax)};
  } else if (isWeights) {
    // If the matrix is a weights matrix, we use 3D logical shape (B, M, K) or (B, N, K).
    shape = {static_cast<uint64_t>(hiddenSizeValid),
             static_cast<uint64_t>(numTokensValid),
             static_cast<uint64_t>(options.mNumBatches)};
  }

  // Assemble the stride (strideTokens, 1).
  // Swap the first two dimension as mentioned before.
  std::vector<uint64_t> stride = {1, static_cast<uint64_t>(hiddenSize)};
  if (useTmaOobOpt) {
    stride = {1,
              static_cast<uint64_t>(hiddenSize),
              static_cast<uint64_t>(tg::XLargeN - hiddenSize),
              static_cast<uint64_t>(hiddenSize)};
  } else if (isWeights) {
    stride = {1,
              static_cast<uint64_t>(hiddenSize),
              static_cast<uint64_t>(hiddenSize) * static_cast<uint64_t>(numTokens)};
  }

  // Assemble the box shape
  std::vector<int32_t> tileShape = {tileHiddenSize, tileNumTokens};

  // Alternate layouts (MajorMn and BlockMajorK) do not apply to matrixC
  if (matrixType != MatrixType::MatrixC) {
    // When using 2CTA MMA, we only need to load half of the tile in each CTA for B.
    if (matrixType == MatrixType::MatrixB && tileShape[1] > 1 && options.mClusterDimX >= 2) {
      tileShape[1] /= 2;
    }
    gemm::MatrixLayout layout =
      (matrixType == MatrixType::MatrixA) ? options.mLayoutA : options.mLayoutB;
    // Note, only the weights support non MajorK layouts
    if (layout == gemm::MatrixLayout::MajorMn) {
      // Apply transpose if necessary
      std::swap(shape[0], shape[1]);
      stride[1] = numTokens;
      std::swap(tileShape[0], tileShape[1]);
    } else if (layout == gemm::MatrixLayout::BlockMajorK) {
      // Set shapes based on blocking layout.
      shape = {static_cast<uint64_t>(options.mBlockK),
               static_cast<uint64_t>(numTokens),
               static_cast<uint64_t>(sizeK / options.mBlockK),
               static_cast<uint64_t>(options.mNumBatches)};
      // Strides use padded dimensions
      stride = {1,
                static_cast<uint64_t>(options.mBlockK),
                static_cast<uint64_t>(numTokens * options.mBlockK),
                static_cast<uint64_t>(hiddenSize * numTokens)};

      // If blockK > tileK, then the inner most box size will be based on the tile
      int32_t const tileBlockK = std::min(options.mBlockK, tileHiddenSize);
      tileShape = {tileBlockK, tileNumTokens, tileHiddenSize / tileBlockK};
      // When using 2CTA MMA, we only need to load half of the tile in each CTA for B.
      if (matrixType == MatrixType::MatrixB && options.mClusterDimX >= 2) {
        tileShape[1] /= 2;
      }
    }
  }

  return std::make_tuple(shape, stride, tileShape);
}

// Create the TMA shape/stride for A/B block scaling factors.
static auto makeTmaShapeStrideSfAb(int mM,
                                   int mN,
                                   int mK,
                                   MatrixType matrixType,
                                   int tileM,
                                   int tileN,
                                   int tileK,
                                   tg::SfLayout layout,
                                   int sfReshapeFactor,
                                   const int32_t numEltsPerSf) {

  // The outer dimension.
  auto numTokens = matrixType == MatrixType::MatrixA ? mM : mN;
  // The inner dimension.
  auto hiddenSize = mK;
  // The outer tile dimension.
  auto numTokensPerTile = matrixType == MatrixType::MatrixA ? tileM : tileN;
  // The inner tile dimension.
  auto hiddenSizePerTile = tileK;

  switch (layout) {
  case tg::SfLayout::R128c4: {
    auto shape = std::vector<uint64_t>{256,
                                       2,
                                       static_cast<uint64_t>(ceilDiv(hiddenSize, numEltsPerSf * 4)),
                                       static_cast<uint64_t>(ceilDiv(numTokens, 128))};

    std::vector<uint64_t> stride(shape.size());
    stride[0] = 1;
    for (size_t i = 1; i < shape.size(); i++) {
      stride[i] = shape[i - 1] * stride[i - 1];
    }

    auto tileShapes =
      std::vector<uint32_t>{256,
                            2,
                            static_cast<uint32_t>(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4)),
                            static_cast<uint32_t>(ceilDiv(numTokensPerTile, 128))};

    return std::make_tuple(shape, stride, tileShapes);
  }


  case tg::SfLayout::R8c4: {
    // The scaling factor tensor packs 8x4 tiles into contiguous 32B blocks.
    //
    // As the inner dimension (k) is often a multiple of the tile size, we can reshape to use
    // fewer read requests, if the tile dimensions allow. It does not reduce the number of
    // instructions.
    //
    // I.e., let's define r = min(⌈hiddenSizePerTile / (numEltsPerSf * 4)⌉, 8)
    //
    // The "logical" tensor is: [outer,      inner / numEltsPerSf]
    // The 8x4 SF layout is:    [⌈outer / 8⌉, inner / (4 * numEltsPerSf), 32]
    // The TMA tensor shape is: [⌈outer / 8⌉, inner / (4 * numEltsPerSf * r), r * 32]
    //
    // The caveat of NumRepeats>1 is we must pad the hidden dimension of SF to multiples of
    // NumRepeats * numEltsPerSf * 4.

    // Detect if the supplied factor is power of 2. E.g., 0b0100 and (0b0100 - 1) == 0b0000.
    int const r = sfReshapeFactor;
    if (r > 0 && (r & (r - 1)) != 0) {
      throw std::runtime_error("mSfReshapeFactor must be positive and a power of 2. Found " +
                               std::to_string(r));
    }

    // Sanitize number of repeats so it doesn't exceed the dimension.
    int const repeats = std::min(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4), r);

    // Detect if the input hidden size K is a multiple of the repeats.
    if (ceilDiv(hiddenSize, numEltsPerSf * 4) % repeats != 0) {
      throw std::runtime_error("SF hiddenSize K (" +
                               std::to_string(ceilDiv(hiddenSize, numEltsPerSf * 4)) +
                               ") must be a multiple of repeats (" + std::to_string(repeats) + ")");
    }

    auto shape =
      std::vector<uint64_t>{static_cast<uint64_t>(repeats * 32),
                            static_cast<uint64_t>(ceilDiv(hiddenSize, numEltsPerSf * 4 * repeats)),
                            static_cast<uint64_t>(ceilDiv(numTokens, 8))};

    std::vector<uint64_t> stride(shape.size());
    stride[0] = 1;
    for (size_t i = 1; i < shape.size(); i++) {
      stride[i] = shape[i - 1] * stride[i - 1];
    }

    auto tileShapes = std::vector<uint32_t>{
      static_cast<uint32_t>(repeats * 32),
      static_cast<uint32_t>(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4 * repeats)),
      static_cast<uint32_t>(ceilDiv(numTokensPerTile, 8))};

    return std::make_tuple(shape, stride, tileShapes);
  }

  default:
    throw std::runtime_error("Unsupported SF layout");
  }
  return std::make_tuple(std::vector<uint64_t>{}, std::vector<uint64_t>{}, std::vector<uint32_t>{});
}

// Create the TMA shape/stride for sparsity information of A.
template <class GemmOptions>
static auto makeTmaShapeStrideSparsityInfoA(GemmOptions const& options) {

  // Tensor dimensions.
  auto outerDim = options.mM;
  auto innerDim = tg::getNumBytesSparsityInfo(options.mSparsityA, options.mK);
  // Tile dimensions.
  auto tileOuterDim = options.mTileM;
  auto tileInnerDim = tg::getNumBytesSparsityInfo(options.mSparsityA, options.mTileK);

  auto shape = std::vector<uint64_t>{static_cast<uint64_t>(innerDim),
                                     static_cast<uint64_t>(outerDim),
                                     static_cast<uint64_t>(options.mNumBatches)};

  std::vector<uint64_t> stride(shape.size());
  stride[0] = 1;
  for (size_t i = 1; i < shape.size(); i++) {
    stride[i] = shape[i - 1] * stride[i - 1];
  }

  auto tileShapes =
    std::vector<int32_t>{static_cast<int32_t>(tileInnerDim), static_cast<int32_t>(tileOuterDim)};

  return std::make_tuple(shape, stride, tileShapes);
}

template <class GemmOptions_>
static KernelParams setKernelParams(GemmOptions_ const& options,
                                    bool const batchM,
                                    void const* ptrA,
                                    void const* ptrB,
                                    void* ptrC,
                                    void const* dSfA,
                                    void const* dSfB,
                                    void const* ptrPerTokenSfA,
                                    void const* ptrPerTokenSfB,
                                    [[maybe_unused]] void const* ptrSparsityInfoA,
                                    void const* ptrBias,
                                    int32_t const* ptrPermutedIdxToBiasRowIdx,
                                    void* dSfC,
                                    float const* ptrScaleC,
                                    float const* ptrScaleAct,
                                    float const* ptrScaleGate,
                                    float const* ptrClampLimit,
                                    float const* ptrGatedActAlpha,
                                    float const* ptrGatedActBeta,
                                    int32_t const* routeMap,
                                    float* rowMax,
                                    uint32_t* rowMaxBars,
                                    int rank,
                                    int worldSize,
                                    uint32_t* ptrMulticastCompletionBarUc,
                                    uint32_t* ptrMulticastCompletionBarMc,
                                    int32_t const* permutedIdxToExpandedIdx,
                                    void const* expertWeightsPtr,
                                    int32_t const* ptrNumNonExitingCtas = nullptr,
                                    int32_t const* ptrTotalNumPaddedTokens = nullptr,
                                    int32_t const* ptrCtaIdxXyToBatchIdx = nullptr,
                                    int32_t const* ptrCtaIdxXyToMnLimit = nullptr,
                                    int32_t const maxNumCtas = KernelParams::MaxNumCtas,
                                    uint32_t* ptrDynamicTileCounter = nullptr) {

  static_assert(sizeof(KernelParams) <= 32 * 1024,
                "sizeof(KernelParams) has to be less or equal than 32KB");

  // Create the return struct.
  KernelParams params;

  params.ptrRouteMap = routeMap;
  params.numTokens = options.mNumTokens;

  params.ptrScaleC = ptrScaleC;
  params.ptrScaleAct = ptrScaleAct;
  params.ptrScaleGate = ptrScaleGate;
  params.ptrClampLimit = ptrClampLimit;
  params.ptrGatedActAlpha = ptrGatedActAlpha;
  params.ptrGatedActBeta = ptrGatedActBeta;
  params.ptrPermutedIdxToBiasRowIdx = ptrPermutedIdxToBiasRowIdx;

  int32_t cgaOffset = 0;

  int32_t const isSparseA = tg::isSparse(options.mSparsityA);

  // Compute totalNumPaddedTokens, ctaIdxXyToBatchIdx and ctaIdxXyToMnLimit if the batch dims are
  // known at kernel launch time. Otherwise, these parameters are defined in the device buffers:
  // ptrTotalNumPaddedTokens, ptrCtaIdxXyToBatchIdx and ptrCtaIdxXyToMnLimit respectively.
  // NOTE: The CTA names are kept for compatibility even though they are indexed at CGA granularity.

  if (options.mIsStaticBatch && !options.mIsUniformNumTokensPerBatch) {
    params.totalNumPaddedTokens = 0;
    for (int b = 0; b < options.mNumBatches; b++) {

      int mM = batchM ? options.mBatchedM[b] : options.mM;
      int mN = batchM ? options.mN : options.mBatchedN[b];

      // Skip Tma descriptor creation if expert isn't used
      if (mM == 0 || mN == 0) {
        continue;
      }

      int32_t const clusterDimInBatchDim = batchM ? options.mClusterDimX : options.mClusterDimY;
      int32_t const ctaTile = batchM ? options.mTileM : options.mTileN;
      int32_t const cgaTile = ctaTile * clusterDimInBatchDim;
      int32_t const numCgas = batchM ? (mM + cgaTile - 1) / cgaTile : (mN + cgaTile - 1) / cgaTile;
      // The problem size.
      int32_t mn = batchM ? mM : mN;
      int32_t tokensPerTile = mn;

      // Make sure we do not exceed the launch limit.
      if (cgaOffset + numCgas > KernelParams::MaxNumCtas) {
        throw std::runtime_error("Too many routing entries");
      }

      for (int32_t cga = 0; cga < numCgas; cga++) {
        params.ctaIdxXyToBatchIdx[cgaOffset + cga] = b;
        // This is now an identity map and it is no longer needed.
        // params.cgaIdxXyToTileIdxMn[cgaOffset + cga] = cgaOffset + cga;
        params.ctaIdxXyToMnLimit[cgaOffset + cga] =
          std::min((cgaOffset + cga + 1) * cgaTile, cgaOffset * cgaTile + tokensPerTile);
      }
      cgaOffset += numCgas;

      params.totalNumPaddedTokens += numCgas * cgaTile;
    }
    params.totalNumOutputPaddedTokens = params.totalNumPaddedTokens;
  } else if (options.mIsStaticBatch && options.mIsUniformNumTokensPerBatch) {
    auto numTokens = batchM ? options.mBatchedM[0] : options.mBatchedN[0];
    auto tileTokensDim =
      batchM ? options.mTileM * options.mClusterDimX : options.mTileN * options.mClusterDimY;
    params.batchStrideInCtas = (options.mBatchStrideInTokens + tileTokensDim - 1) / tileTokensDim;
    params.ctasInTokenDimPerBatch = (numTokens + tileTokensDim - 1) / tileTokensDim;
    params.totalNumOutputPaddedTokens =
      params.ctasInTokenDimPerBatch * tileTokensDim * options.mNumBatches;
    if (params.batchStrideInCtas == 0) {
      params.totalNumPaddedTokens = params.ctasInTokenDimPerBatch * tileTokensDim;
    } else {
      params.totalNumPaddedTokens =
        params.ctasInTokenDimPerBatch * tileTokensDim * options.mNumBatches;
    }
    cgaOffset = maxNumCtas;
  } else {
    params.ptrTotalNumPaddedTokens = ptrTotalNumPaddedTokens;
    params.ptrCtaIdxXyToBatchIdx = ptrCtaIdxXyToBatchIdx;
    params.ptrCtaIdxXyToMnLimit = ptrCtaIdxXyToMnLimit;
    cgaOffset = maxNumCtas;
  }

  if (options.mUseDeepSeekFp8 && options.mDtypeC == tg::Dtype::E4m3) {
    params.ptrDqSfsC = reinterpret_cast<float*>(dSfC);
  }

  params.ptrA = ptrA;
  params.ptrB = ptrB;
  params.strideInBytesA = options.mK * tg::dtypeGetNumBits(options.mDtypeA) / 8;
  params.strideInBytesB = options.mK * tg::dtypeGetNumBits(options.mDtypeB) / 8;

  params.ptrSfA = dSfA;
  params.ptrSfB = dSfB;
  params.ptrSfC = dSfC;


  // Do we pad A or B?
  bool doPadA = tg::dtypeNeedsPadding(options.mDtypeA, options.mMmaKind, options.mMmaK, isSparseA);
  bool doPadB = tg::dtypeNeedsPadding(options.mDtypeB, options.mMmaKind, options.mMmaK, isSparseA);

  // cgaOffset counts CGAs. For TMA descriptor sizing we need the total padded token extent,
  // which is cgaOffset * cgaTile = cgaOffset * ctaTile * clusterDimInBatchDim.
  int32_t const clusterDimInBatchDim = batchM ? options.mClusterDimX : options.mClusterDimY;
  auto const paddedTokensInBatchDim =
    cgaOffset * (batchM ? options.mTileM : options.mTileN) * clusterDimInBatchDim;

  if (!batchM) {
    // A is the expert
    if (0 != options.mM % options.mTileM) {
      throw std::runtime_error("0 == mM %% tileM");
    }
    params.tileStridePerBatch = options.mM / options.mTileM;
    params.nm = options.mM;
    // Shape/stride for gmem tensor A.
    auto [shapeA, strideA, tileShapeA] = makeTmaShapeStrideAbc(options,
                                                               options.mM,
                                                               options.mN,
                                                               options.mK >> isSparseA,
                                                               options.mTileM,
                                                               options.mTileN,
                                                               options.mTileK >> isSparseA,
                                                               MatrixType::MatrixA,
                                                               options.mValidM,
                                                               options.mValidN,
                                                               options.mValidK >> isSparseA);
    // Build tma descriptor for A.
    params.tmaA[0] = gemm::buildNdTmaDescriptor(options.mDtypeA,
                                                shapeA,
                                                strideA,
                                                tileShapeA,
                                                const_cast<void*>(ptrA),
                                                doPadA,
                                                /*doSwizzle=*/true);

    // The input is padded:
    // [act0, padding, padding, ... TileN size .., act1, padding, padding, ...]
    auto const inputNumTokens = paddedTokensInBatchDim;

    if (!batchedGemm::doesRouteImplUseLdgsts(options.mRouteImpl)) {
      bool useRouteAct = batchedGemm::doesRouteImplUseTma(options.mRouteImpl);
      // B is the activation
      // Shape/stride for gmem tensor B.
      auto [shapeB, strideB, tileShapeB] =
        makeTmaShapeStrideAbc(options,
                              options.mM,
                              useRouteAct ? options.mNumTokens : inputNumTokens,
                              options.mK,
                              options.mTileM,
                              (useRouteAct ? 1 : options.mTileN),
                              options.mTileK,
                              MatrixType::MatrixB,
                              options.mValidM,
                              useRouteAct ? options.mNumTokens : inputNumTokens,
                              options.mValidK);
      // Build tma descriptor for B.
      params.tmaB[0] = gemm::buildNdTmaDescriptor(options.mDtypeB,
                                                  shapeB,
                                                  strideB,
                                                  tileShapeB,
                                                  const_cast<void*>(ptrB),
                                                  doPadB,
                                                  /*doSwizzle=*/true);
    }

    if (options.mDtypeA == tg::Dtype::E2m1 || options.mDtypeA == tg::Dtype::MxE4m3 ||
        options.mDtypeA == tg::Dtype::MxE2m1 || options.mDtypeA == tg::Dtype::MxInt4) {
      tg::Dtype dTypeSfA{};
      if (options.mDtypeA == tg::Dtype::E2m1) {
        dTypeSfA = tg::Dtype::E4m3;
      } else if (options.mDtypeA == tg::Dtype::MxInt4) {
        dTypeSfA = tg::Dtype::Bfloat16;
      } else {
        dTypeSfA = tg::Dtype::UE8m0;
      }

      int32_t const numEltsPerSfA = options.mSfBlockSizeA;

      // Build TMA descriptor for gmem A block scaling factors.
      auto [shapeSfA, strideSfA, tileShapesSfA] =
        makeTmaShapeStrideSfAb(options.mM * options.mNumBatches,
                               options.mN,
                               options.mK,
                               MatrixType::MatrixA,
                               options.mTileM,
                               options.mTileN,
                               options.mTileK,
                               options.mSfLayoutA,
                               options.mSfReshapeFactor,
                               numEltsPerSfA);
      params.tmaSfA[0] = gemm::buildSfTmaDescriptor(dTypeSfA,
                                                    shapeSfA,
                                                    strideSfA,
                                                    tileShapesSfA,
                                                    const_cast<void*>(dSfA));
    }

    if (isSparseA) {
      // Build TMA descriptor for gmem A sparsity information.
      auto [shapeSparsityInfoA, strideSparsityInfoA, tileShapesSparsityInfoA] =
        makeTmaShapeStrideSparsityInfoA(options);
      params.tmaSparsityInfoA = gemm::buildNdTmaDescriptor(tg::Dtype::UInt8,
                                                           shapeSparsityInfoA,
                                                           strideSparsityInfoA,
                                                           tileShapesSparsityInfoA,
                                                           const_cast<void*>(ptrSparsityInfoA),
                                                           /*doPad=*/false,
                                                           /*doSwizzle=*/true);
    }

    if (options.mDtypeB == tg::Dtype::E2m1 || options.mDtypeB == tg::Dtype::MxE4m3 ||
        options.mDtypeB == tg::Dtype::MxE2m1) {
      tg::Dtype const dTypeSf =
        (options.mDtypeB == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0;
      int32_t const numEltsPerSfB = options.mSfBlockSizeB;

      if (batchedGemm::doesRouteImplUseTma(options.mRouteSfsImpl.value())) {

        // The input is NOT padded:
        // [act0, act1, act2, ...]

        // Build TMA descriptor for gmem B block scaling factors.
        // Pad number of scaling factors to the nearest multiple of 16 because of the TMA 16B
        // alignment requirement.
        auto numSfsInK = options.mK / numEltsPerSfB;
        numSfsInK = ceilDiv(numSfsInK, 16) * 16;

        auto numSfsInValidK = options.mValidK / numEltsPerSfB;
        numSfsInValidK = ceilDiv(numSfsInValidK, 16) * 16;

        auto [shapeSfB, strideSfB, tileShapesSfB] =
          makeTmaShapeStrideAbc(options,
                                options.mM,
                                options.mNumTokens,
                                numSfsInK,
                                options.mTileM,
                                1 /* tileN */,
                                options.mTileK / numEltsPerSfB,
                                MatrixType::MatrixB,
                                options.mValidM,
                                options.mNumTokens,
                                numSfsInValidK);
        params.tmaSfB[0] = gemm::buildNdTmaDescriptor(dTypeSf,
                                                      shapeSfB,
                                                      strideSfB,
                                                      tileShapesSfB,
                                                      const_cast<void*>(dSfB),
                                                      /*doPad=*/false,
                                                      /*doSwizzle=*/true);
      } else if (batchedGemm::doesRouteImplUseNoRoute(options.mRouteSfsImpl.value())) {

        // The input is padded:
        // [act0, padding, padding, ... TileN size .., act1, padding, padding, ...]

        auto const inputNumTokensSfB = paddedTokensInBatchDim;

        // Build TMA descriptor for gmem B block scaling factors.
        auto [shapeSfB, strideSfB, tileShapesSfB] = makeTmaShapeStrideSfAb(options.mM,
                                                                           inputNumTokensSfB,
                                                                           options.mK,
                                                                           MatrixType::MatrixB,
                                                                           options.mTileM,
                                                                           options.mTileN,
                                                                           options.mTileK,
                                                                           options.mSfLayoutB,
                                                                           options.mSfReshapeFactor,
                                                                           numEltsPerSfB);
        params.tmaSfB[0] = gemm::buildSfTmaDescriptor(dTypeSf,
                                                      shapeSfB,
                                                      strideSfB,
                                                      tileShapesSfB,
                                                      const_cast<void*>(dSfB));
      }
    }

    // C is the output activation
    // Skip TMA descriptor setup when using C multicast - direct pointers are used instead
    if (options.mUseTmaStore && !options.mUseCMultiCast) {
      // Shape/stride for gmem tensor C.
      // NOTE: Output is *always* sanitized across the whole MNK range. This ensures maximum
      // compatibility with the next BMM where unwritten part of the output could be polluted by
      // NaNs.
      auto [shapeC, strideC, tileShapeC] = makeTmaShapeStrideAbc(options,
                                                                 options.mM,
                                                                 paddedTokensInBatchDim,
                                                                 options.mK,
                                                                 options.mTileM,
                                                                 options.mTileN,
                                                                 options.mTileK,
                                                                 MatrixType::MatrixC);
      // Build tma descriptor for C.
      params.tmaC[0] = gemm::buildNdTmaDescriptor(options.mDtypeC,
                                                  shapeC,
                                                  strideC,
                                                  tileShapeC,
                                                  ptrC,
                                                  /*doPad=*/false);


    } else {
      params.ptrC = ptrC;
    }

  } else {
    // B is the expert
    if (0 != options.mN % options.mTileN) {
      throw std::runtime_error("0 == mN %% tileN");
    }
    params.tileStridePerBatch = options.mN / options.mTileN;
    params.nm = options.mN;
    // Shape/stride for gmem tensor B.
    auto [shapeB, strideB, tileShapeB] = makeTmaShapeStrideAbc(options,
                                                               options.mM,
                                                               options.mN,
                                                               options.mK,
                                                               options.mTileM,
                                                               options.mTileN,
                                                               options.mTileK,
                                                               MatrixType::MatrixB,
                                                               options.mValidM,
                                                               options.mValidN,
                                                               options.mValidK);
    // Build tma descriptor for B.
    params.tmaB[0] = gemm::buildNdTmaDescriptor(options.mDtypeB,
                                                shapeB,
                                                strideB,
                                                tileShapeB,
                                                const_cast<void*>(ptrB),
                                                doPadB,
                                                /*doSwizzle=*/true);

    if (!batchedGemm::doesRouteImplUseLdgsts(options.mRouteImpl)) {
      // A is the activation
      // Shape/stride for gmem tensor A.
      // The input is padded:
      // [act0, padding, padding, ... tileM size .., act1, padding, padding, ...]
      auto const inputNumTokens = paddedTokensInBatchDim;
      bool useRouteAct = batchedGemm::doesRouteImplUseTma(options.mRouteImpl);
      auto [shapeA, strideA, tileShapeA] =
        makeTmaShapeStrideAbc(options,
                              useRouteAct ? options.mNumTokens : inputNumTokens,
                              options.mN,
                              options.mK,
                              useRouteAct ? 1 : options.mTileM,
                              options.mTileN,
                              options.mTileK,
                              MatrixType::MatrixA,
                              useRouteAct ? options.mNumTokens : inputNumTokens,
                              options.mValidN,
                              options.mValidK);
      // Build tma descriptor for A.
      params.tmaA[0] = gemm::buildNdTmaDescriptor(options.mDtypeA,
                                                  shapeA,
                                                  strideA,
                                                  tileShapeA,
                                                  const_cast<void*>(ptrA),
                                                  doPadA,
                                                  /*doSwizzle=*/true);
    }

    if (options.mDtypeA == tg::Dtype::E2m1 || options.mDtypeA == tg::Dtype::MxE4m3 ||
        options.mDtypeA == tg::Dtype::MxE2m1) {
      tg::Dtype const dTypeSf =
        (options.mDtypeA == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0;
      int32_t const numEltsPerSfA = options.mSfBlockSizeA;

      if (batchedGemm::doesRouteImplUseTma(options.mRouteSfsImpl.value())) {

        // The input is NOT padded:
        // [act0, act1, act2, ...]

        // Build TMA descriptor for gmem A block scaling factors.
        // Pad number of scaling factors to the nearest multiple of 16 because of the TMA 16B
        // alignment requirement.
        auto numSfsInK = options.mK / numEltsPerSfA;
        numSfsInK = ceilDiv(numSfsInK, 16) * 16;

        auto numSfsInValidK = options.mValidK / numEltsPerSfA;
        numSfsInValidK = ceilDiv(numSfsInValidK, 16) * 16;

        auto [shapeSfA, strideSfA, tileShapesSfA] =
          makeTmaShapeStrideAbc(options,
                                options.mNumTokens,
                                options.mN,
                                numSfsInK,
                                1 /* tileM */,
                                options.mTileN,
                                options.mTileK / numEltsPerSfA,
                                MatrixType::MatrixA,
                                options.mNumTokens,
                                options.mValidN,
                                numSfsInValidK);
        params.tmaSfA[0] = gemm::buildNdTmaDescriptor(dTypeSf,
                                                      shapeSfA,
                                                      strideSfA,
                                                      tileShapesSfA,
                                                      const_cast<void*>(dSfA),
                                                      /*doPad=*/false,
                                                      /*doSwizzle=*/true);
      } else if (options.mRouteSfsImpl.value() == batchedGemm::RouteImpl::NoRoute) {

        // The input is padded:
        // [act0, padding, padding, ... tileM size .., act1, padding, padding, ...]
        auto const inputNumTokensSfA = paddedTokensInBatchDim;

        // Build TMA descriptor for gmem A block scaling factors.
        auto [shapeSfA, strideSfA, tileShapesSfA] = makeTmaShapeStrideSfAb(inputNumTokensSfA,
                                                                           options.mN,
                                                                           options.mK,
                                                                           MatrixType::MatrixA,
                                                                           options.mTileM,
                                                                           options.mTileN,
                                                                           options.mMmaTileK,
                                                                           options.mSfLayoutA,
                                                                           options.mSfReshapeFactor,
                                                                           numEltsPerSfA);
        params.tmaSfA[0] = gemm::buildSfTmaDescriptor(dTypeSf,
                                                      shapeSfA,
                                                      strideSfA,
                                                      tileShapesSfA,
                                                      const_cast<void*>(dSfA));
      }
    }

    if (options.mDtypeB == tg::Dtype::E2m1 || options.mDtypeB == tg::Dtype::MxE4m3 ||
        options.mDtypeB == tg::Dtype::MxE2m1) {
      tg::Dtype const dTypeSf =
        (options.mDtypeB == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0;
      int32_t const numEltsPerSfB = options.mSfBlockSizeB;

      // Build TMA descriptor for gmem B block scaling factors.
      auto [shapeSfB, strideSfB, tileShapesSfB] =
        makeTmaShapeStrideSfAb(options.mM,
                               options.mN * options.mNumBatches,
                               options.mK,
                               MatrixType::MatrixB,
                               options.mTileM,
                               options.mTileN,
                               options.mMmaTileK,
                               options.mSfLayoutB,
                               options.mSfReshapeFactor,
                               numEltsPerSfB);
      params.tmaSfB[0] = gemm::buildSfTmaDescriptor(dTypeSf,
                                                    shapeSfB,
                                                    strideSfB,
                                                    tileShapesSfB,
                                                    const_cast<void*>(dSfB));
    }

    // C is the output activation
    // Skip TMA descriptor setup when using C multicast - direct pointers are used instead
    if (options.mUseTmaStore && !options.mUseCMultiCast) {
      // Shape/stride for gmem tensor C.
      // NOTE: Output is *always* sanitized across the whole MNK range. This ensures maximum
      // compatibility with the next BMM where unwritten part of the output could be polluted by
      // NaNs.
      auto [shapeC, strideC, tileShapeC] = makeTmaShapeStrideAbc(options,
                                                                 paddedTokensInBatchDim,
                                                                 options.mN,
                                                                 options.mK,
                                                                 options.mTileM,
                                                                 options.mTileN,
                                                                 options.mTileK,
                                                                 MatrixType::MatrixC);
      // Build tma descriptor for C.
      params.tmaC[0] = gemm::buildNdTmaDescriptor(options.mDtypeC,
                                                  shapeC,
                                                  strideC,
                                                  tileShapeC,
                                                  ptrC,
                                                  /*doPad=*/false);
    } else {
      params.ptrC = ptrC;
    }
  }

  params.k = options.mK;
  params.numBatches = options.mNumBatches;

  params.rank = rank;
  params.tpGrpSize = worldSize;
  params.ptrMulticastCompletionBarUc = ptrMulticastCompletionBarUc;
  params.ptrMulticastCompletionBarMc = ptrMulticastCompletionBarMc;

  params.ptrPartialRowMax = rowMax;
  params.ptrRowMaxCompletionBars = rowMaxBars;

  params.ptrNumNonExitingCtas = ptrNumNonExitingCtas;

  // Set the per-token scale factors for MetaFP8 or scale inputs
  params.ptrPerTokenSfA = ptrPerTokenSfA;
  params.ptrPerTokenSfB = ptrPerTokenSfB;
  params.ptrBias = ptrBias;

  // Set MoE finalize parameters for direct register-to-gmem store
  params.permutedIdxToExpandedIdx = permutedIdxToExpandedIdx;
  params.expertWeightsPtr = expertWeightsPtr;
  // Set the dynamic tile counter pointer for dynamic scheduling.
  params.ptrDynamicTileCounter = ptrDynamicTileCounter;

  return params;
}
#endif
}; // namespace KernelParamsSetup

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

} // namespace batchedGemm

} // namespace batchedGemm