Implementation:InternLM Lmdeploy Gemm MoeUtils

Knowledge Sources	InternLM_Lmdeploy
Domains	GPU_Kernels, GEMM
Last Updated	2026-02-07 15:00 GMT

Overview

Implements Mixture-of-Experts (MoE) routing utilities including gating, token dispatch, scale-aware dispatch, expert combination, softmax masking, and sampling functions for expert assignment.

Description

This header provides the MoE routing pipeline functions that prepare data for grouped GEMM execution:

• invokeMoeGate_V2: Computes expert routing from gating logits -- produces forward-to-native mappings (f2n, f2E), native-to-forward mappings (en2f), per-expert offsets, routing scales, accumulator masks, with optional softmax normalization and top-K normalization

• invokeMoeDispatch: Dispatches input tokens to expert-specific buffers based on routing decisions

• invokeMoeDispatchScales: Dispatches tokens with per-token scaling factors

• invokeMoeCombine: Combines expert outputs back into the original token order, applying routing scales and optional bias

• invokeMoeSoftmaxMaskTopKGroups: Applies softmax with group-wise top-K masking for grouped expert selection

• SampleUniform / SampleBalanced: CPU-side sampling functions for expert assignment (uniform random and load-balanced strategies)

Constants: kMoeGateMaxTiles = 16, kMoeGateVecSize = 4.

Usage

Called before and after grouped GEMM operations in MoE transformer layers to route tokens to experts and combine results.

Code Reference

Source Location

• Repository: InternLM_Lmdeploy
• File: src/turbomind/kernels/gemm/moe_utils_v2.h

Signature

void invokeMoeGate_V2(int* f2n, int* f2E, int* en2f, int* offsets, float* scales,
                       void* masks, int* accum, const float* logits,
                       int tokens, int tokens_padded, int experts, int exp_per_tok,
                       bool softmax, bool norm_topk, float routed_scale, cudaStream_t st);

void invokeMoeDispatch(Ref<Tensor> out_, const Tensor& src, const int* f2n,
                        int expert_per_token, cudaStream_t st);

void invokeMoeCombine(Ref<Tensor> out_, const Tensor& src, const Tensor& bias,
                       const float* scales, const int* en2f, const int* f2E,
                       const float* dst_scales, int experts_per_token,
                       float bscale, float dst_scale, cudaStream_t st);

Import

#include "src/turbomind/kernels/gemm/moe_utils_v2.h"

I/O Contract

Inputs

Name	Type	Required	Description
logits	const float*	Yes	Gating logits (tokens x experts)
tokens	int	Yes	Number of input tokens
experts	int	Yes	Total number of experts
exp_per_tok	int	Yes	Experts selected per token (top-K)

Outputs

Name	Type	Description
f2n, f2E, en2f	int*	Token routing index mappings
offsets	int*	Per-expert token count offsets
scales	float*	Routing scale factors per token-expert pair

Usage Examples

invokeMoeGate_V2(f2n, f2E, en2f, offsets, scales, masks, accum,
    logits, num_tokens, padded_tokens, num_experts, top_k,
    /*softmax=*/true, /*norm_topk=*/true, 1.0f, stream);
invokeMoeDispatch(expert_input, token_input, f2n, top_k, stream);
// ... grouped GEMM ...
invokeMoeCombine(output, expert_output, bias, scales, en2f, f2E,
    dst_scales, top_k, bscale, dst_scale, stream);

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime