Moe Per Expert Scale Kernels

void trt_edgellm::kernel::applyPerExpertScale(

float *topkWeights,

int32_t const *topkIds,

float const *perExpertScale,

int32_t numTokens,

int32_t topK,

cudaStream_t stream,

)

Apply per-expert multiplicative scale to top-k routing weights.

After the MoE router selects top-k experts and renormalizes their weights, this kernel applies a learned per-expert scale factor:

topkWeights[i] *= perExpertScale[topkIds[i]] 

for all i in [0, numTokens * topK).

Parameters:

• topkWeights – [inout] Flat array of renormalized routing weights, shape [numTokens * topK]. Modified in-place.

• topkIds – [in] Expert indices for each slot, shape [numTokens * topK]. Each value must be in [0, numExperts).

• perExpertScale – [in] Per-expert scale factors, shape [numExperts].

• numTokens – [in] Number of tokens (batch dimension).

• topK – [in] Number of selected experts per token.

• stream – [in] CUDA stream for asynchronous execution.