F16 Moe Support Kernels

struct F16MoeRoutingBuffers

Device buffers used to sort token-major routes into expert-major order.

Public Members

int32_t *expertCounts = {}

INT32 [E].

int32_t *expertOffsets = {}

INT32 [E + 1].

int32_t *expertWriteOffsets = {}

INT32 [E].

int32_t *sortedToExpanded = {}

INT32 [R].

int32_t *expandedToSorted = {}

INT32 [R].

struct F16MoeGemmMetadata

Device-resident grouped-GEMM descriptors shared by every CuTeDSL variant.

Public Members

int32_t *problemShapes = {}

INT32 [E, 4], storing M, N, K, L.

int32_t *strides = {}

INT32 [E, 3, 2], storing A, B, D strides.

int64_t *addresses = {}

INT64 [E, 3], storing A, B, D addresses.

struct F16MoeGemmSetup

Internal setup for one grouped-GEMM stage.

Public Members

F16MoeGemmMetadata metadata = {}

Device-resident descriptors for this stage.

void const *input = {}

FP16 expert-contiguous input.

void const *weights = {}

FP16 expert weights.

void *output = {}

FP16 expert-contiguous output.

int32_t n = {}

Output columns.

int32_t k = {}

Reduction columns.

cudaError_t trt_edgellm::kernel::buildF16MoeRoutingAndGemmMetadata(

F16MoeRoutingBuffers const &buffers,

F16MoeGemmSetup const &fc1Setup,

F16MoeGemmSetup const &fc2Setup,

int32_t const *topkIds,

int32_t numTokens,

int32_t topK,

int32_t numExperts,

cudaStream_t stream,

) noexcept

Sort token-major routes by expert and populate both grouped-GEMM stages. Uses a single-CTA fast path for up to 256 tokens and a multi-kernel fallback otherwise.

cudaError_t trt_edgellm::kernel::gatherF16MoeHiddenRows(

void const *hiddenStates,

void *gatheredInput,

int32_t const *sortedToExpanded,

int32_t routedRows,

int32_t topK,

int32_t hiddenSize,

cudaStream_t stream,

) noexcept

Gather hidden rows into the expert-major order described by sortedToExpanded.

cudaError_t trt_edgellm::kernel::activateF16Moe(

void const *rawFc1,

void *activatedFc1,

int32_t routedRows,

int32_t moeInterSize,

int32_t activationType,

cudaStream_t stream,

) noexcept

Apply SwiGLU or ReLU-squared in FP32 and store the activated FP16 intermediate.

cudaError_t trt_edgellm::kernel::scatterF16MoeOutput(

void const *routedOutput,

int32_t const *expandedToSorted,

float const *topkWeights,

void *output,

int32_t numTokens,

int32_t topK,

int32_t hiddenSize,

cudaStream_t stream,

) noexcept

Weighted token-major route reduction with FP32 accumulation and one FP16 cast.