class Step5MegaMoE(nn.Module):
"""Router plus EP routed experts."""
def __init__(
self,
vllm_config,
*,
num_experts,
top_k,
hidden_size,
intermediate_size,
quant_config,
activation,
activation_clamp,
prefix,
renormalize,
scoring_func,
e_score_correction_bias,
routed_scaling_factor,
sequence_parallel=False,
):
super().__init__()
self.sequence_parallel = sequence_parallel
self.layer_id = extract_layer_index(prefix)
self.router = create_fused_moe_router(
top_k=top_k,
global_num_experts=num_experts,
renormalize=renormalize,
scoring_func=scoring_func,
e_score_correction_bias=e_score_correction_bias,
routed_scaling_factor=routed_scaling_factor,
)
max_tokens = vllm_config.scheduler_config.max_num_batched_tokens
if sequence_parallel:
tp_size = vllm_config.parallel_config.tensor_parallel_size
max_tokens = (max_tokens + tp_size - 1) // tp_size
self.routed_experts = Step5MegaMoEExperts(
vllm_config,
num_experts=num_experts,
top_k=top_k,
hidden_size=hidden_size,
intermediate_size=intermediate_size,
quant_config=quant_config,
activation=activation,
activation_clamp=activation_clamp,
prefix=f"{prefix}.routed_experts",
max_num_tokens=max_tokens,
)
def forward(self, hidden_states, router_logits, padding=None):
weights, ids = self.router.select_experts(
hidden_states, router_logits, topk_indices_dtype=torch.int64
)
num_tokens = hidden_states.shape[0]
if (
padding is None
and not self.sequence_parallel
and envs.VLLM_MOE_SKIP_PADDING
and is_forward_context_available()
):
mask = get_forward_context().is_padding
padding = mask[:num_tokens] if mask is not None else None
return self.routed_experts(hidden_states, weights, ids, padding)