class BeamSparseMoeBlock(nn.Module):
"""Grouped sigmoid ``noaux_tc`` MoE block."""
def __init__(
self, vllm_config: VllmConfig, prefix: str = "", reduce_results: bool = True
) -> None:
super().__init__()
config = vllm_config.model_config.hf_config
parallel_config = vllm_config.parallel_config
quant_config = vllm_config.quant_config
self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe
self.gather_workspace: BeamAllReduceWorkspace | None = None
self.output_boundary_fused = False
self._compilation_config = vllm_config.compilation_config
router_dtype = STR_DTYPE_TO_TORCH_DTYPE[config.moe_router_dtype or "float32"]
self.gate = GateLinear(
config.hidden_size,
config.num_experts,
bias=False,
params_dtype=router_dtype,
out_dtype=router_dtype,
prefix=f"{prefix}.gate",
)
e_score_correction_bias = None
if config.topk_method == "noaux_tc":
e_score_correction_bias = nn.Parameter(
torch.zeros(config.num_experts, dtype=torch.float32)
)
self.gate.register_parameter("e_score_correction_bias", e_score_correction_bias)
self.experts = FusedMoEFactory(
reduce_results=reduce_results or self.is_sequence_parallel,
is_sequence_parallel=self.is_sequence_parallel,
gate=self.gate,
num_experts=config.num_experts,
top_k=config.num_experts_per_tok,
hidden_size=config.hidden_size,
intermediate_size=config.moe_intermediate_size,
renormalize=config.norm_topk_prob,
quant_config=quant_config,
use_grouped_topk=(
config.n_group is not None and config.topk_group is not None
),
num_expert_group=config.n_group,
topk_group=config.topk_group,
scoring_func=config.scoring_func,
e_score_correction_bias=e_score_correction_bias,
routed_scaling_factor=(
config.routed_scaling_factor
if config.routed_scaling_factor is not None
else 1.0
),
apply_routed_scale_to_output=True,
router_logits_dtype=self.gate.out_dtype,
enable_eplb=parallel_config.enable_eplb,
num_redundant_experts=parallel_config.eplb_config.num_redundant_experts,
prefix=f"{prefix}.experts",
)
self._boundary_layer_name = f"{prefix}.experts.finalize"
def defer_finalize(self, max_tokens: int, *, allow_ep: bool = False) -> None:
from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import (
TrtLlmNvFp4ExpertsModular,
TrtLlmNvFp4ExpertsMonolithic,
)
experts_cls = getattr(
self.experts.routed_experts.quant_method, "experts_cls", None
)
if experts_cls in (TrtLlmNvFp4ExpertsModular, TrtLlmNvFp4ExpertsMonolithic):
self.experts.moe_config.defer_moe_finalize(max_tokens, allow_ep=allow_ep)
def fuse_output_post_norm(
self,
post_norm: RMSNorm | None,
eps: float,
*,
workspace: "BeamAllReduceWorkspace | None" = None,
) -> bool:
from vllm.config.kernel import FLASHINFER_MOE_EP_CUTEDSL
from vllm.model_executor.layers.fused_moe.experts.flashinfer_moe_ep import (
FlashInferMoeEpExperts,
)
from vllm.model_executor.layers.fused_moe.flashinfer_moe_ep import (
supports_flashinfer_moe_ep_deferred_finalize,
)
moe = self.experts.moe_config
experts = self.experts.routed_experts
if (
(not self.is_sequence_parallel and moe.dp_size == 1)
or moe.moe_backend != FLASHINFER_MOE_EP_CUTEDSL
or getattr(experts.quant_method, "experts_cls", None)
is not FlashInferMoeEpExperts
or moe.in_dtype != torch.bfloat16
or moe.hidden_dim != moe.hidden_dim_unpadded
or moe.hidden_dim % 256
or not 256 <= moe.hidden_dim <= 8192
or not 1 <= moe.experts_per_token <= 32
or experts.apply_router_weight_on_input
or self.experts.shared_experts is not None
or self.experts.routed_input_transform is not None
or self.experts.routed_output_transform is not None
or moe.pcp_size != 1
or moe.tp_size != 1
or moe.ep_size <= 1
or not supports_flashinfer_moe_ep_deferred_finalize()
):
return False
fuse_gather = (
workspace is not None
and moe.sp_size == workspace.world_size == moe.ep_size
and moe.dp_size == 1
)
if not self.is_sequence_parallel or fuse_gather:
self.gather_workspace = workspace if fuse_gather else None
self.output_boundary_fused = True
moe.defer_moe_finalize(allow_ep=True)
self._compilation_config.static_forward_context[
self._boundary_layer_name
] = self
return True
return False
def forward_boundary(
self,
hidden_states: torch.Tensor,
residual: torch.Tensor,
post_weight: torch.Tensor | None,
post_eps: float | None,
next_weight: torch.Tensor | None,
next_eps: float | None,
scales: tuple[float, float] | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
hidden, normalized = torch.ops.vllm.beam_moe_forward_boundary(
hidden_states,
residual,
post_weight,
post_eps,
next_weight,
next_eps,
*(scales if scales is not None else (None, None)),
self.experts._encode_layer_name(),
)
return hidden, normalized if next_eps is not None else None
def _forward_boundary(
self,
hidden_states: torch.Tensor,
residual: torch.Tensor,
post_weight: torch.Tensor | None,
post_eps: float | None,
next_weight: torch.Tensor | None,
next_eps: float | None,
scales: tuple[float, float] | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
from .moe_finalize import (
beam_moe_finalize_gather_norm_residual_norm,
beam_moe_finalize_norm_residual_norm,
)
local_states = (
sequence_parallel_chunk(hidden_states)
if self.is_sequence_parallel
else hidden_states
)
output = self.experts._forward_impl(local_states, local_states, None)
assert isinstance(output, UnfinalizedMoEOutput)
# Consume the borrowed expert workspace inside the same opaque op.
if self.gather_workspace is None:
return beam_moe_finalize_norm_residual_norm(
output,
residual,
post_weight,
post_eps,
next_weight,
next_eps,
scales,
routed_scale=self.experts.routed_scaling_factor,
)
return beam_moe_finalize_gather_norm_residual_norm(
output,
residual,
post_weight,
post_eps,
next_weight,
next_eps,
scales,
self.gather_workspace,
routed_scale=self.experts.routed_scaling_factor,
)
def forward(
self, hidden_states: torch.Tensor
) -> torch.Tensor | UnfinalizedMoEOutput:
num_tokens = hidden_states.shape[0]
if self.is_sequence_parallel:
hidden_states = sequence_parallel_chunk(hidden_states)
elif self.experts.moe_config.should_defer_moe_finalize(num_tokens):
output = self.experts._forward_impl(hidden_states, hidden_states, None)
assert isinstance(output, UnfinalizedMoEOutput)
return output
output = self.experts(hidden_states=hidden_states, router_logits=hidden_states)
if self.is_sequence_parallel:
output = tensor_model_parallel_all_gather(output, 0)[:num_tokens]
return output