Moe Activation Kernels

void trt_edgellm::kernel::moeActivation(

rt::Tensor const &input,

rt::Tensor &output,

int64_t numTokens,

int64_t intermediateDim,

int32_t activationType,

cudaStream_t stream

)

Apply an FP16 or BF16 MoE activation.

Supported activation types:

• kMoeSwiGlu: SwiGLU. The input is [numTokens, 2 * intermediateDim], with all gate values followed by all up values.

• kMoeRelu2: ReLU2. The input is [numTokens, intermediateDim].

• kMoeGeGlu: GeGLU. The input is [numTokens, 2 * intermediateDim], with all gate values followed by all up values.

FP16 and BF16 support all activation types. The input and output data types must match.

Parameters:

• input – Input tensor (FP16 or BF16, GPU).

• output – Output tensor [numTokens, intermediateDim] with the same data type as input.

• numTokens – Number of routed token slots.

• intermediateDim – Intermediate dimension.

• activationType – Activation type from MoeActivationType enum.

• stream – CUDA stream.

Throws:

std::runtime_error – If the tensor shapes, data types, devices, alignment, or activation type are invalid.