Moe Marlin Indices Kernels#

void trt_edgellm::kernel::launchBuildMarlinIndicesKernel(
int32_t const *slotsByExpertWorkspace,
int32_t const *slotsPerExpertWorkspace,
int32_t const *paddedCounts,
int32_t const *paddedOffsets,
float const *topkWeights,
int32_t *sortedTokenIds,
float *topkWeightsFlat,
int32_t *expertIds,
int32_t numTokens,
int32_t topK,
int32_t numExperts,
int32_t moeBlockSize,
cudaStream_t stream
)#

Launch kernel to build Marlin indices from slot lists (per-expert).

void trt_edgellm::kernel::launchBuildDenseMarlinIndicesKernel(
int32_t *sortedTokenIds,
int32_t *expertIds,
int32_t *numTokensPostPadded,
float *topkWeights,
int32_t numTokens,
int32_t paddedRows,
int32_t moeBlockSize,
cudaStream_t stream
)#

Build the trivial Marlin routing arrays for a dense (single-expert, topK=1) GEMM.

The dense case degenerates the MoE routing: every one of the numTokens rows maps to expert 0 in order, so sortedTokenIds is the identity 0..numTokens-1 with padded tail slots set to the out-of-range sentinel numTokens (== numTokens*topK, matching marlin_template.h’s idx < prob_m*top_k masking), expertIds is all zeros, and numTokensPostPadded is the block-aligned row count. topkWeights is filled with 1.0f purely to keep the shared Marlin call signature valid; it is unused because the dense wrapper passes mulTopkWeights=false.

Parameters:

  • sortedTokenIds – [paddedRows] (INT32) output slot ids

  • expertIds – [paddedRows/moeBlockSize] (INT32) output, all zeros

  • numTokensPostPadded – [1] (INT32) output, set to paddedRows

  • topkWeights – [paddedRows] (FP32) output, all 1.0f (unused by dense path)

  • numTokens – Number of real rows M

  • paddedRows – Block-aligned row count ceilDiv(M, moeBlockSize)*moeBlockSize

  • moeBlockSize – Marlin block size (8 for decode, 32 for prefill)

void trt_edgellm::kernel::launchAggregateSlotOutputsKernel(
void const *slotOutputs,
void *aggregatedOutput,
int32_t numTokens,
int32_t topK,
int32_t outDim,
cudaStream_t stream
)#

Launch kernel to aggregate slot outputs back to tokens (sum over topK in slot order).

void trt_edgellm::kernel::launchAggregateSlotOutputsBf16Kernel(
void const *slotOutputs,
void *aggregatedOutput,
int32_t numTokens,
int32_t topK,
int32_t outDim,
cudaStream_t stream
)#

Launch the BF16 slot-output aggregation kernel.

Each output element is accumulated over topK slots in FP32 and converted to BF16 once at the final store.