/*
* 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 <functional>
#include <numeric>
#include <optional>
#include <unordered_map>

#include "BatchedGemmOptions.h"
#include "KernelParams.h"
#include "trtllm/gen/CudaKernelLauncher.h"

#ifdef TLLM_GEN_EXPORT_INTERFACE
#ifdef TLLM_GEN_EXPORT_FLASHINFER
#include "flashinferMetaInfo.h"
#else
#include "KernelMetaInfo.h"
#endif // TLLM_GEN_EXPORT_FLASHINFER
#endif // TLLM_GEN_EXPORT_INTERFACE

namespace batchedGemm {

namespace batchedGemm {

////////////////////////////////////////////////////////////////////////////////////////////////////
//
// BatchedGemmData
//
////////////////////////////////////////////////////////////////////////////////////////////////////

struct BatchedGemmData {
  struct ProblemDimensions {
    // The number of batches.
    int32_t mNumBatches{0};
    // The number of tokens. Set to 0 if routeAct is false.
    int32_t mNumTokens{0};
    // Whether the batch is on the M dimension.
    bool mBatchM{true};
    // The maximum number of CTAs (CGAs) in the token dimension.
    // Need to be set if mNumTokens > 0 and the token per batch
    // distribution is not known at launch time.
    // In this case, the kernel will launch mMaxNumCtasInTokenDim units in token dim and exit
    // early if the batched-dim index is larger or equal to ptrNumNonExitingCtas.
    // The legacy CTA-based name is kept for interface compatibility even though they are indexed at
    // CGA granularity.
    int32_t mMaxNumCtasInTokenDim{0};

    // Either mBatchedM or mBatchedN must be set when mNumTokens == 0, otherwise not used.
    // The number of tokens in each batch on the M dimension if batchM,
    // otherwise not used.
    // The number of elements in the array is mNumBatches.
    // E.g. to implement a BMM with each batch having M tokens, one needs to set mBatchedM to
    // {M, M, M, .. mNumBatches times ..}
    std::vector<int32_t> mBatchedM{};
    // The number of tokens in each batch on the N dimension if batchN,
    // otherwise not used.
    // The number of elements in the array is mNumBatches.
    // E.g. to implement a BMM with each batch having N tokens, one needs to set mBatchedN to
    // {N, N, N, .. mNumBatches times ..}
    std::vector<int32_t> mBatchedN{};

    // The M dimension.
    // It is the total number of tokens if A is the activation matrix.
    // It is the total number of output channels if A is the weight matrix.
    // ValidM/N/K by default assumes to be full range of M/N/K respectively. If we pad M/N/K due to
    // alignment of other constraints, then we can specify ValidM/N/K to indicate the valid range.
    int32_t mM{0};
    int32_t mValidM{0};
    // The N dimension.
    // It is the total number of tokens if B is the activation matrix.
    // It is the total number of output channels if B is the weight matrix.
    int32_t mN{0};
    int32_t mValidN{0};
    // The K dimension. It is the hidden dimension of the input matrices.
    int32_t mK{0};
    int32_t mValidK{0};
    // The rank id of the current device in the multi-gpu space.
    int32_t mRank{0};
    // The number of devices in tensor-parallel group.
    int32_t mWorldSize{1};
  };

  struct InputBuffers {
    // The matrix A. The data type is controlled by options.mDtypeA.
    //
    // Sparsity is only supported with batchN.
    // Let S be the sparsity ratio (1 for dense, 2 for sparse).
    //
    // If (routeAct == true && batchM), the shape is [M, K]
    // Elseif (batchStrideInTokens > 0)
    //   If batchM:
    //      Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), K].
    //      Logical strides are [K, 1].
    //
    //   If batchN:
    //      If layoutA is MatrixLayout::MajorK
    //         Logical shape is [B, divUpMul(M, tileM), K / S].
    //         Logical strides are [divUpMul(M, tileM) * K / S, K / S, 1].
    //      If layoutA is MatrixLayout::MajorMn (sparsity not supported)
    //         Logical shape is [B, K, divUpMul(M, tileM)].
    //         Logical strides are [K * divUpMul(M, tileM), divUpMul(M, tileM), 1].
    //      If layoutA is MatrixLayout::BlockMajorK
    //         Logical shape is [B, K / S / blockK, divUpMul(M, tileM), blockK].
    //         Logical strides are [K / S * divUpMul(M, tileM), divUpMul(M, tileM) * blockK, blockK,
    //         1]. where blockK is 128B.
    // Else // batchStrideInTokens == 0
    //   If batchM: (sparsity not supported)
    //      Logical shape is [M, K].
    //      Logical strides are [K, 1].
    //
    //   If batchN:
    //      Logical shape is [B, divUpMul(M, tileM), K / S].
    //      Logical strides are [divUpMul(M, tileM) * K / S, K / S, 1].
    void const* mPtrA{nullptr};

    // The block scaling factors to dequantize A.
    //
    // If (routeAct == true && batchM), the shape is [M, K / 16]
    // Else
    //   If DeepSeek FP8 recipe is used:
    //      If transposeMmaOutput is false, shape is [K / 128, M].
    //      Otherwise, shape is [M / 128, K / 128].
    //    The rightmost dimension is contiguous in memory.
    //
    //   If DeepSeek FP8 recipe is not used, but for MxFp{4,8}, MxInt4 and NvFp4 formats:
    //    If the layout is R128c4,
    //      M must be a multiple of 128.
    //      K must be a multiple of 4 * P, where P is the scaling block size.
    //      The "logical" shape is: [paddedM, K / P].
    //      The R128c4 layout is: [paddedM / 128, K / P / 4, 512].
    //      The shape we use for TMA is: [paddedM / 128, K / P / 4, 2, 256].
    //    If the layout is R128c16,
    //      M must be a multiple of 128.
    //      K must be a multiple of 16 * P, where P is the scaling block size.
    //      The "logical" shape is: [paddedM, K / P].
    //      The R128c16 layout is: [paddedM / 128, K / P / 16, 2048].
    //      The shape we use for TMA is: [paddedM / 128, K / P / 16, 8, 256].
    //  Where paddedM is M if (routeAct == true && batchM), or
    //  sum(divUpMul(M[bi], tileM) for bi in B) if batchM,
    //  otherwise divUpMul(M, tileM) * B.
    //  Dtype is Dtype::Fp32 if DeepSeek FP8 recipe is used, otherwise Dtype is Dtype::E4m3 for
    //  NvFp4, Dtype::UE8m0 for MxFp{4,8} formats, Dtype::Bfloat16 for MxInt4.
    //
    // Otherwise should be set to nullptr.
    void const* mPtrSfA{nullptr};

    // The per-token scaling factors from scale A.
    //
    // This is used for either:
    //   * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is
    //   Dtype::Float32
    //   * One-sided per-token scaling. The dtype is Dtype::Bfloat16 by default.
    //
    // if (batchM (A is activations)):
    //     Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B)]
    //
    // if (batchN (A is weights)):
    //     Logical shape is [B, divUpMul(M, tileM)]
    //
    void const* mPtrPerTokenSfA{nullptr};

    // The matrix B. The data type is controlled by options.mDtypeB.
    //
    // If (routeAct == true && batchN), the shape is [N, K]
    //
    // Else if (batchStrideInTokens > 0)
    //   If batchN:
    //      Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), K].
    //      Logical strides are [K, 1].
    //
    //   If batchM:
    //      If layoutB is MatrixLayout::MajorK
    //         Logical shape is [B, divUpMul(N, tileN), K].
    //         Logical strides are [divUpMul(N, tileN) * K, K, 1].
    //      If layoutB is MatrixLayout::MajorMn
    //         Logical shape is [B, K, divUpMul(N, tileN)].
    //         Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN), 1].
    //      If layoutB is MatrixLayout::BlockMajorK
    //         Logical shape is [B, K / blockK, divUpMul(N, tileN), blockK].
    //         Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN) * blockK, blockK, 1].
    //         where blockK is 128B.
    //
    // Else // batchStrideInTokens == 0
    //   If batchN:
    //      Logical shape is [N, K].
    //      Logical strides are [K, 1].
    //
    //   If batchM:
    //      Logical shape is [B, divUpMul(N, tileN), K].
    //      Logical strides are [divUpMul(N, tileN) * K, K, 1].
    void const* mPtrB{nullptr};

    // The scaling factors to dequantize B.
    //
    //
    //
    // Else
    //   If DeepSeek FP8 recipe is used:
    //      If transposeMmaOutput is false, shape is [paddedN / 128, K / 128].
    //      Otherwise, shape is [K / 128, paddedN].
    //      The rightmost dimension is contiguous in memory.
    //
    //   If DeepSeek FP8 recipe is not used, but for MxFp{4,8} and NvFp4 formats:
    //    If the layout is R128c4,
    //       paddedN must be a multiple of 128.
    //       K must be a multiple of 64.
    //       The R128c4 layout is: [paddedN / 128, K / P / 4, 512], where P is the scaling block
    //       size. The shape we use for TMA is: [paddedN / 128, K / P / 4, 2, 256]
    //
    //    If the layout is R8c4,
    //       paddedN must be a multiple of 8.
    //       K must be a multiple of 64.
    //       The R8c4 layout is: [paddedN / 8, K / P / 4, 32], where P is the scaling block size.
    //       The shape we use for TMA is: [paddedN / 8, K / P / 4 / repeats, repeats * 32]
    //       where repeats = min(tileK / P / 4, 8)
    //
    // where paddedN is N if (routeAct == true && batchN),
    // or sum(divUpMul(N[bi], tileN) for bi in B) if batchN,
    // otherwise divUpMul(N, TileN) * B.
    //
    // Dtype is Dtype::Fp32 if DeepSeek FP8 recipe is used, otherwise Dtype::E4m3.
    //
    // Otherwise should be set to nullptr.
    void const* mPtrSfB{nullptr};

    // The per-token scaling factors from scale B.
    //
    // This is used for either:
    //   * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is
    //   Dtype::Float32
    //   * One-sided per-token scaling. The dtype is Dtype::Bfloat16 by default.
    //
    // if (batchM (B is weights)):
    //     Logical shape is [B, divUpMul(N, tileN)]
    //
    // if (batchN (B is activations)):
    //     Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B)]
    void const* mPtrPerTokenSfB{nullptr};

    // The sparsity information of A, if structured sparsity is used.
    // Only supported for batchN (A is weights).
    //
    // When sparsityA is Any_2_4:
    //     2 elements are non-zero in any chunk of 4 elements.
    //     A 4-bit index indicates the position of the non-zero elements.
    //     The shape in Uint8 is: [B, divUpMul(M, tileM), K / 8]
    //     (two 4-bit indices packed into one UInt8)
    //
    // When sparsityA is Pairwise_4_8:
    //     4 elements are non-zero in any chunk of 8 elements.
    //     The zero and non-zero elements are grouped in pairs.
    //     A 4-bit index indicates the position of the non-zero pairs.
    //     The shape in Uint8 is: [B, divUpMul(M, tileM), K / 16]
    //     (two 4-bit indices packed into one UInt8)
    //
    // If sparsityA is Dense, this should be set to nullptr.
    void const* mPtrSparsityInfoA{nullptr};

    // The bias applied after the GEMM and before the activation function.
    // The bias is applied before applying the global scaling factor. I.e.
    // C = act(A * B + bias') * scaleC
    // scaleC = dequantA * dequantB * quantC
    // Thus, the bias' = bias / (dequantA * dequantB), where the bias is the original bias.
    //
    // If batchM, BiasType must be N, and bias shape is [B, N].
    // The bias is broadcasted along the M dimension.
    //
    // If batchN BiasType must be M, and bias shape is [B, M].
    // The bias is broadcasted along the N dimension.
    //
    // If BiasType is Mn, the bias is a full 2D matrix applied element-wise (D = A*B + C):
    //   The row dimension is:
    //     If mPtrPermutedIdxToBiasRowIdx == nullptr:
    //       If batchM: sum(divUpMul(M[bi], tileM) for bi in B).
    //       If batchN: sum(divUpMul(N[bi], tileN) for bi in B).
    //     If mPtrPermutedIdxToBiasRowIdx != nullptr:
    //       The first dimension is indexed by `mPtrPermutedIdxToBiasRowIdx[permuted_padded_index]`.
    //   The hidden dimension is:
    //     If batchM: N.
    //     If batchN: M.
    //   If `mFusedBiasShuffleMode == gemm::FusedBiasShuffleMode::None`, the buffer is row-major
    //   after any fused-act interleave and after the shuffled-matrix reordering expected by the
    //   epilogue.
    //   If `mFusedBiasShuffleMode == gemm::FusedBiasShuffleMode::Shuffle`, the buffer is row-major
    //   [biasRow, hidden] after any fused-act interleave, but before shuffled-matrix reordering.
    //   If `mFusedBiasShuffleMode == gemm::FusedBiasShuffleMode::ReorderAndShuffle`, the buffer is
    //   row-major [biasRow, hidden] before both fused-act bias interleave and shuffled-matrix
    //   reordering.
    //
    // The dtype is float32.
    void const* mPtrBias{nullptr};

    // Optional map from permuted padded row index in the batched dimension to the row index in
    // mPtrBias to load for BiasType::Mn.
    //
    // If this pointer is nullptr, mPtrBias is interpreted directly in the permuted padded row-major
    // layout described above, i.e. biasRow = permutedRow.
    //
    // If this pointer is non-null, the shape is
    // [sum(divUpMul(M[bi], tileM) for bi in B)] for batchM or
    // [sum(divUpMul(N[bi], tileN) for bi in B)] for batchN.
    //
    // The dtype is int32_t.
    int32_t const* mPtrPermutedIdxToBiasRowIdx{nullptr};

    // The output tensor scaling factor for Fp8 (not DeepSeek FP8) and NvFp4 quantization.
    // TensorRT-LLM API requires a scaling factor on the device.
    // scaleC = dequantA * dequantB * quantC,
    // where dequantA is global dequantization scaling factor of A
    //    if dtypeA is FP8, it transforms the range from [-448, 448] to [-amaxA, amaxA]
    //    if dtypeA is NvFp4, it transforms the range from [-448 * 6, 448 * 6] to [-amaxA, amaxA],
    //    otherwise it is 1.
    // dequantB is defined similarly to dequantA.
    // quantC is the quantization scaling factor of C.
    //    if dtypeC is FP8, it transforms the range from [-amaxC, amaxC] to [-448, 448]
    //    if dtypeC is NvFp4, it transforms the range from [-amaxC, amaxC] to [-448 * 6, 448 * 6],
    //    otherwise it is 1.
    // Shape is [B].
    float const* mPtrScaleC{nullptr};

    // The pre-activation scaling factor (typically dequantA * dequantB) for non-gated non-linear
    // activation.
    // Only used when non-linear activation is applied (e.g., GELU, Relu2, Silu).
    // When used, scaleC should be quantScaleC only, and this scale is applied before the
    // activation. Shape is [B].
    float const* mPtrScaleAct{nullptr};

    // The output gate scale for Fp8 (not DeepSeek FP8) and NvFp4 quantization.
    // TensorRT-LLM API requires a scaling factor on the device.
    // scaleGate = dequantA * dequantB,
    // where dequantA is global dequantization scaling factor of A
    //    if dtypeA is FP8, it transforms the range from [-448, 448] to [-amaxA, amaxA]
    //    if dtypeA is NvFp4, it transforms the range from [-448 * 6, 448 * 6] to [-amaxA, amaxA],
    //    otherwise it is 1.
    // dequantB is defined similarly to dequantA.
    // Shape is [B].
    float const* mPtrScaleGate{nullptr};

    // The clamp limit for the accumulator before applying the activation.
    // Shape is [B].
    // Clamp is INF if nullptr.
    // When the input is FP8 or NVFP4, the clamp has to be scaled by limit' = limit / dequantAb.
    // If applied on SwiGlu, it will be:
    //
    //   x_glu    = x_glu.clamp(min=None, max=limit)
    //   x_linear = x_linear.clamp(min=-limit, max=limit)
    //
    // The given clamp limit applies to the dequantized values, so the order of operations would
    // look something like this:
    //
    // x0 = x0 * dqAb
    // x0 = clamp(x0, none, limit)
    // x0 = x0 * sigmoid(alpha * x0)
    // x1 = dqAb * x1
    // x1 = clamp(x1, -limit, limit)
    // out = qC * (x1 + beta) * x0
    //
    // Given that the dqAb and qC are combined into scaleC, we can bring the dqAb into the clamp
    // limit and apply the clamping prior to dequantization:
    //
    // x0 = clamp(x0, none, limit / dqAb)
    // x0 = x0 * dqAb
    // x0 = x0 * sigmoid(alpha * x0)
    // x1 = clamp(x1, -limit / dqAb, limit / dqAb)
    // scaleC = dqAb * qC
    // beta' = beta / dqAb
    // out = scaleC * (x1 + beta') * x0
    //
    // Note this assumes that dequantScaleAb == scaleGate which is true in TRT-LLM MoE use-case
    //
    float const* mPtrClampLimit{nullptr};

    // The alpha and beta for SwiGlu or GeGlu.
    // gatedActivation <- (x0 + beta) * activation(x1, alpha)
    // Shape is [B].
    // Alpha is 1.f if nullptr.
    // Beta is 0.f if nullptr.
    // The formula for SwiGlu (for GeGlu, replace sigmoid with phi):
    //
    //   out_glu  = x_glu * torch.sigmoid(alpha * x_glu) * (x_linear + beta)
    //
    // The beta is added before applying the global scaling factor. I.e.
    // x_linear = (x_linear + beta') * scaleC
    // Thus, the beta' = beta / (dequantA * dequantB), where the beta is the original beta.
    float const* mPtrGatedActAlpha{nullptr};
    float const* mPtrGatedActBeta{nullptr};

    // Param is used when the kernel is configured with -routeAct true.
    // The inputs are not padded, but the outputs are padded to divUpMul(M[bi], tileM) for batchM or
    // divUpMul(N[bi], tileN) for batchN.
    // If -routeAct is false, the params are not used and should be set to zero.

    // The routeMap for the input tokens.
    // Map of expanded token index (counting the previous padded tokens) to the batch index
    // the token belongs to.
    // The shape is
    // [divUpMul(numTokens + numBatches * (tileM/N - 1), tileM/N)]
    // The dtype is int32_t.
    //
    // There are 3 tokens [0, 1, 2] such that [0, 1] belong to batch [B0] and [2] to batch [B1].
    // Let's assume that the padded size is 4.
    //
    // The expanded indices for tokens [0, 1, 2] are:
    // expandedIdx[0] = 0
    // expandedIdx[1] = 1
    // expandedIdx[2] = divUpMul(2, 4) + 0 = 4
    //
    // The route map is [B0, B0, X, X, B1, X, X, X] where X could be any value.
    int32_t const* mPtrRouteMap;

    //////////////////////////////////////////////////////////////////////////////////////////////////
    //
    // Batching information parameters.
    //
    //////////////////////////////////////////////////////////////////////////////////////////////////

    // In some cases, some CTAs must early-exit. E.g. when the grid size is set statically, but the
    // actual workload is decided at runtime. This element on the device contains the number of CTAs
    // that do not early-exit. The number corresponds to the X dim of the grid when the output is
    // not transposed (i.e. batchM). To the Y dim, otherwise. The size is 1 and the dtype is
    // int32_t. Used if isStaticBatch == false, otherwise set to nullptr. The pointer points to a
    // scalar and the dtype is int32_t. The pointed value must be >= 0.
    int32_t const* mPtrNumNonExitingCtas;

    // Pointer to total number of padded tokens.
    // Computed as
    // int32_t totalNumPaddedTokens{0};
    // for (int bi = 0; bi < options.mNumBatches; bi++) {
    //   totalNumPaddedTokens += batchM ? divUpMul(options.mBatchedM[bi], options.mTileM *
    //   options.mClusterDimX)
    //                                  : divUpMul(options.mBatchedN[bi], options.mTileN);
    // }
    // The size is 1 and the dtype is int32_t.
    // If isStaticBatch == true, ptrTotalNumPaddedTokens should be set to nullptr and
    // totalNumPaddedTokens is used.
    int32_t const* mPtrTotalNumPaddedTokens;

    // Pointer to the map from the CTA index (in X/Y dim) to the batch index.
    // The legacy CTA-based name is kept for interface compatibility even though routed MoE paths
    // index this map at CGA granularity.
    // E.g. with listM = 128,255,32 and tileM = 128, should be equal to
    // ctaIdxXyToBatchIdx = [0, 1, 1, 2]
    // If isStaticBatch == true, ptrCtaIdxXyToBatchIdx should be set to nullptr and
    // ctaIdxXyToBatchIdx is used.
    // The shape is
    // [divUp(numTokens + numBatches * (tileM/N - 1), tileM/N)]
    int32_t const* mPtrCtaIdxXyToBatchIdx;

    // Pointer from the CTA index X/Y to the expanded tile index where the expanded tile index is
    // computed as:
    //
    // int expandedIdx = 0;
    // for (int bi = 0; bi < batchIdx-1; ++bi) {
    //   expandIdx = divUpMul(numTokens[bi], TileM/N);
    // }
    // expandIdx += <index in the batch>
    // E.g. with numTokens = [128,255,32] and tileM = 128, should be equal to
    // ptrCtaIdxXyToMnLimit = [128, 256, 383, 416]
    // The legacy CTA-based name is kept for interface compatibility even though routed MoE paths
    // index this map at CGA granularity.
    // The shape is
    // [divUp(numTokens + numBatches * (tileM/N - 1), tileM/N)]
    int32_t const* mPtrCtaIdxXyToMnLimit;

    // Map from permuted token index to expanded token index.
    // Used for MoE finalize direct-store path when mUseCMultiCast is enabled.
    // The dtype is int32_t. Set to nullptr when not using CMulticast.
    int32_t const* mPtrPermutedIdxToExpandedIdx{nullptr};

    //////////////////////////////////////////////////////////////////////////////////////////////////
    //
    // Multicast barrier parameters
    //
    //////////////////////////////////////////////////////////////////////////////////////////////////

    // When mUseCMultiCast is set, we require a barrier in global memory on each
    // GPU to track completion of the multicast stores.
    //
    // The value of the barrier must be set to zero before the kernel is executed on any GPU.
    // The barrier is self-resetting. It does not have to be reset by another kernel.
    // When mUseCMultiCast is false, a nullptr may be passed.
    uint32_t* mPtrMulticastCompletionBarUc;
    uint32_t* mPtrMulticastCompletionBarMc;

    //////////////////////////////////////////////////////////////////////////////////////////////////
    //
    // MoE finalize parameters for direct register-to-gmem store
    //
    //////////////////////////////////////////////////////////////////////////////////////////////////

    // Pointer to expert weights for MoE finalize.
    // The dtype depends on the use case.
    void const* mPtrExpertWeightsPtr{nullptr};
    // Global counter for SW-emulated dynamic tile scheduling. When dynamic scheduling is enabled,
    // must be initialized to gridDim.x * gridDim.y (gridDim.z is reserved for K-dim in CGA splitK)
    // before each kernel launch.
    // Set to nullptr if not using the PersistentSm90 tile scheduler.
    // Shape is [1].
    uint32_t* mPtrDynamicTileCounter{nullptr};
  };

  struct OutputBuffers {
    // The output matrix C. The data type is controlled by options.mDtypeC.
    //
    // If batchM:
    //    Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), N].
    //    Logical strides are [N, 1].
    //
    // If batchN:
    //    Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), M].
    //    Logical strides are [M, 1].
    void* mPtrC{nullptr};

    // The output block scaling factors for C.
    //
    // If MxFp{4,8} and NvFp4 formats are used,
    // The "logical" shape is (P is the scaling block size):
    //    if batchM: [paddedM, N / P]
    //    if batchN: [paddedN, M / P]
    // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B),
    // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B).
    //
    // If the layout is R128c4,
    //    paddedOuter must be a multiple of 128.
    //    inner must be a multiple of 64.
    //    The R128c4 layout is: [paddedOuter / 128, inner / P / 4, 512]
    //    The shape we use for TMA is: [paddedOuter / 128, inner / P / 4, 2, 256]
    //    where inner = N if batchM, otherwise M.
    //    where paddedOuter = paddedM if batchM, otherwise paddedN.
    //
    // If the layout is R8c4,
    //    paddedOuter must be a multiple of 8.
    //    inner must be a multiple of 64.
    //    The R8c4 layout is: [paddedOuter / 8, inner / P / 4, 32]
    //    The shape we use for TMA is: [paddedOuter / 8, inner / P / 4 / repeats, repeats * 32]
    //    where repeats = min(tileInner / P / 4, 8),
    //    where tileInner = tileN if batchM, otherwise tileM,
    //    where paddedOuter = paddedM if batchM, otherwise paddedN.
    //    where inner = N if batchM, otherwise M.
    //
    // The dtype is Dtype::E4m3.
    //
    // If DeepSeek FP8 quantization recipe is used,
    // If batchM:
    //   The shape is [N / 128, paddedM],
    //   where paddedM is sum(divUpMul(M[bi], tileM) for bi in B).
    // If batchN:
    //   The shape is [M / 128, paddedN],
    //   where paddedN is sum(divUpMul(N[bi], tileN) for bi in B).
    // The rightmost dimension is contiguous in memory.
    // The dtype is Dtype::Float32.
    void* mPtrSfC{nullptr};

#ifdef TLLM_RUBIN_FEATURES
#ifdef TLLM_TEST
    // The buffer to invalidate prior to triggering the secondary kernel.
    // The data type which is used to determine the invalid value is controlled by options.mDtypeC.
    //
    // If batchM:
    //    Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), N].
    //    Logical strides are [N, 1].
    //
    // If batchN:
    //    Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), M].
    //    Logical strides are [M, 1].
    void* mPtrInvalidation{nullptr};

    // The scaling factors for testing data invalidation, for MxFp{4,8} and NvFp4 formats.
    //
    // This has the same logical shape and layout as mPtrSfC.
    void* mPtrSfInvalidation{nullptr};
#endif // TLLM_TEST
#endif // TLLM_RUBIN_FEATURES
  };

  ProblemDimensions mProblemDimensions;
  InputBuffers mInputBuffers;
  OutputBuffers mOutputBuffers;
};

////////////////////////////////////////////////////////////////////////////////////////////////////
//
// BatchedGemmInterface
//
////////////////////////////////////////////////////////////////////////////////////////////////////

class BatchedGemmInterface {
public:
  using ModuleCache = std::unordered_map<std::string, std::tuple<CUmodule, CUfunction>>;

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

  BatchedGemmInterface(int32_t rankId = 0,
                       bool const exportsCubin = false,
                       int32_t const numRotations = 1)
    : mRankId(rankId)
    , mExportsCubin(exportsCubin)
    , mNumRotations(numRotations) {}

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

#ifndef TLLM_GEN_EXPORT_INTERFACE
  // Generates and compiles the kernel using either nvcc or nvrtc.
  BatchedGemmConfig generateAndCompileKernel(BatchedGemmConfig const& batchedGemmConfig) const;
#endif

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

  // Launch the cubin from the provided config. It calls all necessary memsets for internal buffers.
  // Provided config must be validated with isValidConfig before the call.
  //
  // pinnedHostBuffer: if provided, it must be pinned host memory buffer of 4 bytes.
  int32_t run(BatchedGemmConfig const& config,
              void* workspace,
              BatchedGemmData const& batchedGemmData,
              void* cudaStream,
#ifdef TLLM_TEST
              void* cudaEvent,
#endif // TLLM_TEST
              int32_t multiProcessorCount,
              bool usePdl = true,
              void* pinnedHostBuffer = nullptr,
              std::optional<std::reference_wrapper<ModuleCache>> moduleCache = std::nullopt) {
    // Get options from config and data.
    auto options = getOptionsFromConfigAndData(config, batchedGemmData);

    bool const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM;
    bool const useDeepSeekFp8 = options.mUseDeepSeekFp8 && options.mDtypeA == tg::Dtype::E4m3 &&
                                options.mDtypeB == tg::Dtype::E4m3;
    // Get the grid dimensions.
    auto [numCtaBatch, numCtaTile, numCtaInner] =
      getGridDim(options, batchedGemmData.mProblemDimensions.mMaxNumCtasInTokenDim);
    auto workspaceSizes = getWorkspaceSizesInBytes(config, batchedGemmData);
    float* dPtrRowMax{nullptr};
    uint32_t* dPtrRowMaxBars{nullptr};

    // Set the completion barriers to 0 if needed.
    if (useDeepSeekFp8 && options.mFusedAct) {
      dPtrRowMax = reinterpret_cast<float*>(alignPtr(reinterpret_cast<char*>(workspace), 1024));
      dPtrRowMaxBars = reinterpret_cast<uint32_t*>(
        alignPtr(reinterpret_cast<char*>(dPtrRowMax) + workspaceSizes[0], 1024));
      auto err = cudaMemsetAsync((void*)dPtrRowMaxBars,
                                 0x00,
                                 workspaceSizes[1],
                                 reinterpret_cast<cudaStream_t>(cudaStream));
      if (err != cudaSuccess) {
        return 1;
      }
    }

    // Initialize the dynamic tile counter if dynamic scheduling is enabled.
    // The counter must be initialized to the fixed grid dim.
    // TODO: we may avoid H2D copy by offloading the operation to some kernel.
    if (batchedGemmData.mInputBuffers.mPtrDynamicTileCounter != nullptr) {
      auto [numFixedCtaBatch, numFixedCtaTile, numFixedCtaInner] =
        getFixedGridDim(options, multiProcessorCount);
      static uint32_t* pinnedCounterValue_ = nullptr;
      if (pinnedHostBuffer) {
        // Use the pre-allocated pinned host memory if provided.
        pinnedCounterValue_ = reinterpret_cast<uint32_t*>(pinnedHostBuffer);
      }
      // Use pinned host memory for the source value to support CUDA graph, or it silently breaks
      // during CUDA graph replay and corrupts the data.
      if (pinnedHostBuffer == nullptr && pinnedCounterValue_ == nullptr) {
        // Allocate pinned host memory if not provided.
        auto err = cudaMallocHost((void**)&pinnedCounterValue_, sizeof(uint32_t));
        if (err != cudaSuccess) {
          return 1;
        }
      }
      *pinnedCounterValue_ = numFixedCtaBatch * numFixedCtaTile;
      auto err = cudaMemcpyAsync((void*)batchedGemmData.mInputBuffers.mPtrDynamicTileCounter,
                                 pinnedCounterValue_,
                                 sizeof(uint32_t),
                                 cudaMemcpyHostToDevice,
                                 reinterpret_cast<cudaStream_t>(cudaStream));
      if (err != cudaSuccess) {
        return 1;
      }
    }

    auto kernelParams =
      KernelParamsSetup::setKernelParams(options,
                                         batchM,
                                         batchedGemmData.mInputBuffers.mPtrA,
                                         batchedGemmData.mInputBuffers.mPtrB,
                                         batchedGemmData.mOutputBuffers.mPtrC,
                                         batchedGemmData.mInputBuffers.mPtrSfA,
                                         batchedGemmData.mInputBuffers.mPtrSfB,
                                         batchedGemmData.mInputBuffers.mPtrPerTokenSfA,
                                         batchedGemmData.mInputBuffers.mPtrPerTokenSfB,
                                         batchedGemmData.mInputBuffers.mPtrSparsityInfoA,
                                         batchedGemmData.mInputBuffers.mPtrBias,
                                         batchedGemmData.mInputBuffers.mPtrPermutedIdxToBiasRowIdx,
                                         batchedGemmData.mOutputBuffers.mPtrSfC,
                                         batchedGemmData.mInputBuffers.mPtrScaleC,
                                         batchedGemmData.mInputBuffers.mPtrScaleAct,
                                         batchedGemmData.mInputBuffers.mPtrScaleGate,
                                         batchedGemmData.mInputBuffers.mPtrClampLimit,
                                         batchedGemmData.mInputBuffers.mPtrGatedActAlpha,
                                         batchedGemmData.mInputBuffers.mPtrGatedActBeta,
#ifdef TLLM_RUBIN_FEATURES
#ifdef TLLM_TEST
                                         batchedGemmData.mOutputBuffers.mPtrInvalidation,
                                         batchedGemmData.mOutputBuffers.mPtrSfInvalidation,
#endif // TLLM_TEST
#endif // TLLM_RUBIN_FEATURES
                                         batchedGemmData.mInputBuffers.mPtrRouteMap,
                                         dPtrRowMax,
                                         dPtrRowMaxBars,
                                         batchedGemmData.mProblemDimensions.mRank,
                                         batchedGemmData.mProblemDimensions.mWorldSize,
                                         batchedGemmData.mInputBuffers.mPtrMulticastCompletionBarUc,
                                         batchedGemmData.mInputBuffers.mPtrMulticastCompletionBarMc,
                                         batchedGemmData.mInputBuffers.mPtrPermutedIdxToExpandedIdx,
                                         batchedGemmData.mInputBuffers.mPtrExpertWeightsPtr,
                                         batchedGemmData.mInputBuffers.mPtrNumNonExitingCtas,
                                         batchedGemmData.mInputBuffers.mPtrTotalNumPaddedTokens,
                                         batchedGemmData.mInputBuffers.mPtrCtaIdxXyToBatchIdx,
                                         batchedGemmData.mInputBuffers.mPtrCtaIdxXyToMnLimit,
                                         numCtaBatch,
                                         batchedGemmData.mInputBuffers.mPtrDynamicTileCounter);

    // The size of the grid.
    auto grid = getLaunchGrid(options,
                              batchedGemmData.mProblemDimensions.mMaxNumCtasInTokenDim,
                              multiProcessorCount);

    BatchedGemmConfig batchedGemmConfig = config;
#ifndef TLLM_GEN_EXPORT_INTERFACE
    // Generate and compile the kernel if data is not provided.
    if (config.mData == nullptr) {
      batchedGemmConfig = generateAndCompileKernel(batchedGemmConfig);
    }
    TLLM_CHECK_ERROR(batchedGemmConfig.mCudaRunner != nullptr, "CudaRunner is not set");
    batchedGemmConfig.mCudaRunner->run((void*)&kernelParams,
                                       (void*)cudaStream,
                                       grid,
                                       /* cluster */ {},
                                       /* instanceId */ batchedGemmConfig.mInstanceIdx
#ifdef TLLM_TEST
                                       ,
                                       (void*)cudaEvent
#endif // TLLM_TEST
    );
    return 0;
#endif

    CUmodule cuModule;
    CUfunction cuFunction;

    if (moduleCache.has_value()) {
      ModuleCache& moduleCacheRef = moduleCache.value().get();

      // Modules are associated with a specific context, so the context is included in the key
      CUcontext ctx;
      unsigned long long ctxId;
      cuCtxGetCurrent(&ctx);
      cuCtxGetId(ctx, &ctxId);

      // Reinterpret the ctxId as a string to avoid needing a custom hash or converting it to a
      // string in decimal representation.
      std::string const ctxName =
        std::string(reinterpret_cast<char*>(&ctxId), sizeof(unsigned long long) / sizeof(char));
      std::string const funcName = std::string(batchedGemmConfig.mFunctionName);
      auto const moduleKey = ctxName + funcName;
      auto module = moduleCacheRef.find(moduleKey);

      // Use cache if module is found, otherwise load and insert into cache
      if (module != moduleCacheRef.end()) {
        cuFunction = std::get<1>(module->second);
      } else {
        gemm::loadCubinData(&cuModule, batchedGemmConfig);
        cuModuleGetFunction(&cuFunction, cuModule, batchedGemmConfig.mFunctionName);
        moduleCacheRef.insert(std::make_pair(moduleKey, std::make_tuple(cuModule, cuFunction)));
      }
    } else {
      gemm::loadCubinData(&cuModule, batchedGemmConfig);
      cuModuleGetFunction(&cuFunction, cuModule, batchedGemmConfig.mFunctionName);
    }

    // Prepare the grid/block.
    dim3 block3{static_cast<uint32_t>(batchedGemmConfig.mNumThreadsPerCTA),
                static_cast<uint32_t>(1),
                static_cast<uint32_t>(1)};
    dim3 grid3{(grid.size() > 0 ? static_cast<uint32_t>(grid[0]) : 1u),
               (grid.size() > 1 ? static_cast<uint32_t>(grid[1]) : 1u),
               (grid.size() > 2 ? static_cast<uint32_t>(grid[2]) : 1u)};
    // Prepare the cluster size.
    dim3 cluster3{static_cast<uint32_t>(options.mClusterDimX),
                  static_cast<uint32_t>(options.mClusterDimY),
                  static_cast<uint32_t>(options.mClusterDimZ)};

    // Whether PDL can safely be enabled
    const bool pdlSafe = batchedGemmConfig.mOptions.mGridWaitForPrimaryRouting ||
                         batchedGemmConfig.mOptions.mGridWaitForPrimaryEarlyExit ||
                         batchedGemmConfig.mOptions.mGridWaitForPrimaryA ||
                         batchedGemmConfig.mOptions.mGridWaitForPrimaryB
#ifdef TLLM_RUBIN_FEATURES
                         || batchedGemmConfig.mOptions.mLamportConsumerA ||
                         batchedGemmConfig.mOptions.mLamportConsumerB
#endif // TLLM_RUBIN_FEATURES
      ;

    // Run the kernel.
    auto result = trtllm::gen::launchKernel((void*)&kernelParams,
                                            cudaStream,
                                            batchedGemmConfig.mSharedMemSize,
                                            cuFunction,
                                            block3,
                                            grid3,
                                            cluster3,
                                            usePdl && pdlSafe);
    if (result != CUDA_SUCCESS) {
      return result;
    }
    // If a module cache has not been given, unload the module to avoid leaking
    if (!moduleCache.has_value()) {
      cuModuleUnload(cuModule);
    }
    return 0;
  }

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

  // Initializes the buffers before the world sync. Must be called before run.
  int32_t runInitBeforeWorldSync(BatchedGemmConfig const& /* config */,
                                 BatchedGemmData const& /* data */,
                                 void* /* cudaStream */) const {
    return 0;
  }

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

  size_t getWorkspaceSizeInBytes(BatchedGemmConfig const& config,
                                 BatchedGemmData const& data) const {
    auto workspaceSizes = getWorkspaceSizesInBytes(config, data);
    auto size = std::accumulate(workspaceSizes.begin(), workspaceSizes.end(), 0);
    // Additional 1023 bytes to align the pointer to 1024
    return size > 0 ? size + 1023 : 0;
  }

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

  // Returns the list of all available cubin configurations
  BatchedGemmConfig const* getBatchedGemmConfigs() const {
#ifdef TLLM_GEN_EXPORT_INTERFACE
    return tensorrt_llm::kernels::tllmGenBatchedGemmList;
#else
    return nullptr;
#endif
  }

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

  // Returns the number of available cubin configurations
  size_t getNumBatchedGemmConfigs() const {
#ifdef TLLM_GEN_EXPORT_INTERFACE
    return sizeof(tensorrt_llm::kernels::tllmGenBatchedGemmList) /
           sizeof(tensorrt_llm::kernels::tllmGenBatchedGemmList[0]);
#else
    return 0;
#endif
  }

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

  // Returns the grid dimensions of the current kernel when using fixed grid size.
  std::tuple<int32_t, int32_t, int32_t> getFixedGridDim(BatchedGemmOptions const& options,
                                                        int32_t multiProcessorCount,
                                                        int32_t multiProcessorOccupancy = 1) const {

    // For StaticPersistent scheduler, limit grid size based on SM count.
    assert(multiProcessorCount > 0 &&
           "multiProcessorCount must be provided when using StaticPersistent scheduler");
    // The cluster size spanned in the XY dimension.
    auto clusterSizeXy = options.mClusterDimX * options.mClusterDimY;
    // The maximum number of CTAs a GPU can run across the XY dimension.
    auto numCtasXy = multiProcessorCount / options.mNumSlicesForSplitK;
    // Round down to the nearest multiple of the cluster size.
    numCtasXy = (numCtasXy / clusterSizeXy) * clusterSizeXy;
    // Account for available occupancy for the kernel. Defaults to 1 CTA per multiprocessor.
    numCtasXy *= multiProcessorOccupancy;
    return std::make_tuple(numCtasXy, 1, options.mNumSlicesForSplitK);
  }

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

  // Returns the grid dimensions of the current kernel.
  std::tuple<int32_t, int32_t, int32_t> getGridDim(
    BatchedGemmOptions const& options,
    std::optional<int32_t> maxNumCtasInBatchDim = std::nullopt) const {
    bool const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM;
    int32_t const clusterDimInBatchDim = batchM ? options.mClusterDimX : options.mClusterDimY;

    int32_t numCtasBatch{0};
    // For normal BMM, mNumTokens == 0 and the number of CTAs is known to host.
    if (options.mIsStaticBatch) {
      for (int32_t bi = 0; bi < options.mNumBatches; ++bi) {
        numCtasBatch +=
          batchM ? gemm::divUp(options.mBatchedM[bi], options.mTileM * options.mClusterDimX) *
                     options.mClusterDimX
                 : gemm::divUp(options.mBatchedN[bi], options.mTileN * options.mClusterDimY) *
                     options.mClusterDimY;
      }
    }
    // For MoE, mNumTokens != 0 and the number of active batched-dimension entries is known only
    // at runtime. We launch maximally possible number of CTAs and use ptrNumNonExitingCtas to
    // determine the actual number of active entries to run. The legacy CTA-based name is kept for
    // interface compatibility even though routed MoE paths may count these entries at CGA
    // granularity.
    else if ((options.mEnablesEarlyExit || options.mEnablesDelayedEarlyExit) &&
             options.mNumTokens != 0) {
      assert(maxNumCtasInBatchDim.has_value() &&
             "maxNumCtasInBatchDim must be provided when options.mNumTokens != 0");
      numCtasBatch = maxNumCtasInBatchDim.value() * clusterDimInBatchDim;
    } else {
      throw std::invalid_argument("Invalid combination of options");
    }

    int32_t numCtasTile =
      batchM ? gemm::divUp(options.mN, options.mTileN) : gemm::divUp(options.mM, options.mTileM);
    if (batchM) {
      numCtasTile = gemm::divUpMul(numCtasTile, options.mClusterDimY);
    } else {
      numCtasTile = gemm::divUpMul(numCtasTile, options.mClusterDimX);
    }
    int32_t const numCtasInner = options.mNumSlicesForSplitK;
    return std::make_tuple(numCtasBatch, numCtasTile, numCtasInner);
  }

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

  // Returns the number of CTAs of the current kernel.
  std::vector<int32_t> getLaunchGrid(BatchedGemmOptions const& options,
                                     std::optional<int32_t> maxNumCtasInBatchDim = std::nullopt,
                                     int32_t multiProcessorCount = 0) const {
    bool const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM;
    // Determine if the scheduler requires a fixed grid dimension.
    bool const isFixedGridDim = (options.mTileScheduler == gemm::TileScheduler::StaticPersistent ||
                                 options.mTileScheduler == gemm::TileScheduler::PersistentSm90);
    // Get the grid dimensions.
    auto [numCtaBatch, numCtaTile, numCtaInner] = isFixedGridDim
                                                    ? getFixedGridDim(options, multiProcessorCount)
                                                    : getGridDim(options, maxNumCtasInBatchDim);
    std::vector<int32_t> grid = isFixedGridDim || batchM
                                  ? std::vector<int32_t>{numCtaBatch, numCtaTile, numCtaInner}
                                  : std::vector<int32_t>{numCtaTile, numCtaBatch, numCtaInner};
    return grid;
  }

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

  // Returns the number of CTAs of the current kernel.
  int32_t getNumCtas(BatchedGemmOptions const& options,
                     std::optional<int32_t> maxNumCtasInBatchDim = std::nullopt) const {
    auto [dim0, dim1, dim2] = getGridDim(options, maxNumCtasInBatchDim);
    return dim0 * dim1 * dim2;
  }

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

  // Creates GemmOptions from kernel and data.
  BatchedGemmOptions getOptionsFromConfigAndData(BatchedGemmConfig const& config,
                                                 BatchedGemmData const& data) const {
    BatchedGemmOptions options;
    options = config.mOptions;
    options.mM = data.mProblemDimensions.mM;
    options.mN = data.mProblemDimensions.mN;
    options.mK = data.mProblemDimensions.mK;
    options.mValidM = data.mProblemDimensions.mValidM;
    options.mValidN = data.mProblemDimensions.mValidN;
    options.mValidK = data.mProblemDimensions.mValidK;
    options.mBatchedM = data.mProblemDimensions.mBatchedM;
    options.mBatchedN = data.mProblemDimensions.mBatchedN;
    options.mBatchMode = data.mProblemDimensions.mBatchM ? BatchedGemmOptions::BatchMode::BatchM
                                                         : BatchedGemmOptions::BatchMode::BatchN;
    options.mNumBatches = data.mProblemDimensions.mNumBatches;
    options.mNumTokens = data.mProblemDimensions.mNumTokens;
    return options;
  }

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

  // Returns true if the configuration of the cubin can be executed for the given params.
  bool isValidConfig(BatchedGemmConfig const& config, BatchedGemmData const& data) const {
    // Get options from config and data.
    auto options = getOptionsFromConfigAndData(config, data);

    // Check options without modifications.
    return checkAndUpdateBatchedGemmOptions(options,
                                            config.mSm,
                                            /* worldSize */ 1,
                                            /* updateOptions */ false);
  }

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

private:
  //////////////////////////////////////////////////////////////////////////////////////////////////

  template <typename Dtype> inline Dtype* alignPtr(Dtype* ptr, int64_t alignment) const {
    assert((alignment & (alignment - 1)) == 0 && "Alignment must be a power of 2");
    return reinterpret_cast<Dtype*>((reinterpret_cast<uintptr_t>(ptr) + alignment - 1) &
                                    ~(alignment - 1));
  }

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

  // Returns the size of the workspace buffers in bytes
  std::vector<size_t> getWorkspaceSizesInBytes(BatchedGemmConfig const& config,
                                               BatchedGemmData const& data) const {

    std::vector<size_t> workspaceSizes;

    // Get options from config and data.
    auto options = getOptionsFromConfigAndData(config, data);

    if (options.mUseDeepSeekFp8 && options.mFusedAct) {
      int32_t totalNumPaddedTokens = 0;
      auto const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM;
      if (!options.mEnablesEarlyExit || options.mNumTokens == 0) {
        for (int32_t bi = 0; bi < options.mNumBatches; ++bi) {
          totalNumPaddedTokens +=
            batchM ? gemm::divUpMul(options.mBatchedM[bi], options.mTileM * options.mClusterDimX)
                   : gemm::divUpMul(options.mBatchedN[bi], options.mTileN * options.mClusterDimY);
        }
      } else {
        // Get tile in token dim.
        auto tileTokensDim =
          batchM ? options.mTileM * options.mClusterDimX : options.mTileN * options.mClusterDimY;
        totalNumPaddedTokens = data.mProblemDimensions.mMaxNumCtasInTokenDim * tileTokensDim;
      }
      // Get options from config.
      auto& options = config.mOptions;

      int const tokenTile =
        batchM ? options.mTileM * options.mClusterDimX : options.mTileN * options.mClusterDimY;

      auto const numTokens = totalNumPaddedTokens;
      auto const intermediateDim = batchM ? options.mN : options.mM;
      auto const intermediateTile = batchM ? options.mTileN : options.mTileM;

      auto const numBytesRowMax = intermediateDim * totalNumPaddedTokens / 128 * sizeof(float);

      auto const numTilesToken = numTokens / tokenTile;
      auto const numTilesInt = intermediateDim / intermediateTile;
      auto const numBytesRowMaxBars = numTilesToken * numTilesInt / 2 * sizeof(uint32_t);

      // TODO: do we need to pad to 1024?
      workspaceSizes.push_back(getSizePaddedToAlignment(numBytesRowMax, 1024));
      workspaceSizes.push_back(getSizePaddedToAlignment(numBytesRowMaxBars, 1024));
    }

    if (options.mTileScheduler == gemm::TileScheduler::PersistentSm90) {
      workspaceSizes.push_back(sizeof(uint32_t));
    }

    return workspaceSizes;
  }

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

  // Returns the size padded to the alignment
  size_t getSizePaddedToAlignment(size_t size, size_t alignment) const {
    assert((alignment & (alignment - 1)) == 0);
    return (size + alignment - 1) & ~(alignment - 1);
  }
  //////////////////////////////////////////////////////////////////////////////////////////////////

private:
  // The rank id of the current device in the multi-gpu space.
  int32_t mRankId;
  // Whether to export the cubin file.
  bool mExportsCubin;
  // The number of rotations.
  int32_t mNumRotations;
};

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

} // namespace batchedGemm

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

} // namespace batchedGemm