Moe Marlin#

void trt_edgellm::kernel::moeAwqW4A16MarlinGemm(
rt::Tensor const &input,
rt::Tensor &output,
rt::Tensor const &weights,
rt::Tensor const &scales,
rt::Tensor const &sortedTokenIds,
rt::Tensor const &expertIds,
rt::Tensor const &numTokensPostPadded,
rt::Tensor const &topkWeights,
rt::Tensor &workspace,
int64_t moeBlockSize,
int64_t topK,
bool mulTopkWeights,
cudaStream_t stream
)#

MoE W4A16 GEMM using Marlin kernel with AWQ format.

Performs MoE grouped GEMM with 4-bit AWQ quantized weights and 16-bit activations. Uses kU4 format where zero point = 8 is baked into dequantization.

Dequantization: weight_fp16 = (weight_int4 - 8) * scale

Note: Weights must be transformed into Marlin format using awq_marlin_repack.

Internally uses FP32 reduction for numerical accuracy. The workspace buffer must be sized using getMoeMarlinWorkspaceSize() which includes space for both synchronization locks and the FP32 reduction buffer.

Parameters:

  • input – Input activations [numTokens, hiddenDim] (FP16)

  • output – Output tensor [numTokens * topK, outDim] (FP16)

  • weights – Marlin-repacked INT4 weights [numExperts, K/tile, N*tile/pack]

  • scales – Per-group scales [numExperts, numGroups, outDim] (FP16)

  • sortedTokenIds – Sorted token indices [numTokensPadded] (INT32)

  • expertIds – Expert assignment per block [numBlocks] (INT32)

  • numTokensPostPadded – Total padded token count [1] (INT32)

  • topkWeights – Routing weights [numTokensPadded] (FP32)

  • workspace – Workspace buffer sized by getMoeMarlinWorkspaceSize() (INT32)

  • moeBlockSize – MoE processing block size (8, 16, 32, 48, or 64)

  • topK – Number of experts per token

  • mulTopkWeights – Whether to multiply output by topk weights

  • stream – CUDA stream

void trt_edgellm::kernel::moeNvfp4A16MarlinGemm(
rt::Tensor const &input,
rt::Tensor &output,
rt::Tensor const &weights,
rt::Tensor const &blockScales,
rt::Tensor const &globalScales,
rt::Tensor const &sortedTokenIds,
rt::Tensor const &expertIds,
rt::Tensor const &numTokensPostPadded,
rt::Tensor const &topkWeights,
rt::Tensor &workspace,
int64_t moeBlockSize,
int64_t topK,
bool mulTopkWeights,
cudaStream_t stream
)#

MoE FP16/BF16-A, E2M1-W, FP16/BF16-C GEMM using Marlin.

The block scales are raw E4M3 bytes in Marlin order with a fixed K-group size of 16. The per-expert global scales use the activation data type and must already include its exponent adjustment for Marlin’s skip-flop E2M1 conversion.

Parameters:

  • input – Input activations [numTokens, hiddenDim] (FP16 or BF16)

  • output – Output tensor [numTokens * topK, outDim] (same type as input)

  • weights – Marlin-packed E2M1 codes [numExperts, K/16, 2*outDim] (INT32 view)

  • blockScales – Marlin-permuted E4M3 bytes [numExperts, K/16, outDim] (INT8)

  • globalScales – Per-expert adjusted global scales [numExperts] (same type as input)

  • sortedTokenIds – Sorted slot indices [maxRoutedRows] (INT32)

  • expertIds – Expert assignment per routed block (INT32)

  • numTokensPostPadded – Actual padded routed-row count [1] (INT32)

  • topkWeights – Routing weights indexed by unpadded slot id (FP32)

  • workspace – Workspace buffer sized by getMoeMarlinWorkspaceSize() (INT32)

  • moeBlockSize – MoE processing block size (8, 16, 32, 48, or 64)

  • topK – Number of experts per input row

  • mulTopkWeights – Whether to multiply each output slot by its routing weight

  • stream – CUDA stream

int64_t trt_edgellm::kernel::getMoeMarlinWorkspaceSize(
int64_t numTokensPadded,
int64_t outDim,
int64_t moeBlockSize,
int64_t numSMs
)#

Get required workspace size for MoE Marlin GEMM.

Returns the total workspace size needed, which includes:

  • Synchronization locks for thread blocks

  • FP32 reduction buffer (c_tmp) for numerical accuracy

Parameters:

  • numTokensPadded – Maximum padded token count

  • outDim – Output dimension N

  • moeBlockSize – MoE block size

  • numSMs – Number of SMs on the device

Returns:

Required workspace size in number of int32_t elements