[9/N] MoE Refactor: cleanup dispatcher interfaces (#11847)
This commit is contained in:
Cheng Wan
@@ -74,7 +74,6 @@ from sglang.srt.layers.linear import (
|
||||
)
|
||||
from sglang.srt.layers.logits_processor import LogitsProcessor
|
||||
from sglang.srt.layers.moe import (
|
||||
get_deepep_mode,
|
||||
get_moe_a2a_backend,
|
||||
should_use_flashinfer_cutlass_moe_fp4_allgather,
|
||||
should_use_flashinfer_trtllm_moe,
|
||||
@@ -112,10 +111,7 @@ from sglang.srt.model_loader.weight_utils import default_weight_loader
|
||||
from sglang.srt.server_args import get_global_server_args
|
||||
from sglang.srt.single_batch_overlap import SboFlags
|
||||
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
|
||||
from sglang.srt.two_batch_overlap import (
|
||||
MaybeTboDeepEPDispatcher,
|
||||
model_forward_maybe_tbo,
|
||||
)
|
||||
from sglang.srt.two_batch_overlap import model_forward_maybe_tbo
|
||||
from sglang.srt.utils import (
|
||||
BumpAllocator,
|
||||
LazyValue,
|
||||
@@ -649,19 +645,6 @@ class DeepseekV2MoE(nn.Module):
|
||||
else None
|
||||
)
|
||||
|
||||
self.deepep_dispatcher = MaybeTboDeepEPDispatcher(
|
||||
group=parallel_state.get_tp_group().device_group,
|
||||
router_topk=self.top_k,
|
||||
permute_fusion=True,
|
||||
num_experts=self.num_experts,
|
||||
num_local_experts=config.n_routed_experts // self.tp_size,
|
||||
hidden_size=config.hidden_size,
|
||||
params_dtype=config.torch_dtype,
|
||||
deepep_mode=get_deepep_mode(),
|
||||
async_finish=True,
|
||||
return_recv_hook=True,
|
||||
)
|
||||
|
||||
self._enable_a2a_moe = (
|
||||
get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake()
|
||||
)
|
||||
@@ -874,7 +857,7 @@ class DeepseekV2MoE(nn.Module):
|
||||
router_logits = self.gate(hidden_states)
|
||||
if not self._fuse_shared_experts_inside_sbo:
|
||||
shared_output = self._forward_shared_experts(hidden_states)
|
||||
topk_weights, topk_idx, _ = self.topk(
|
||||
topk_output = self.topk(
|
||||
hidden_states,
|
||||
router_logits,
|
||||
num_token_non_padded=forward_batch.num_token_non_padded,
|
||||
@@ -883,9 +866,7 @@ class DeepseekV2MoE(nn.Module):
|
||||
),
|
||||
)
|
||||
else:
|
||||
topk_weights, topk_idx, _ = self.topk.empty_topk_output(
|
||||
hidden_states.device
|
||||
)
|
||||
topk_output = self.topk.empty_topk_output(hidden_states.device)
|
||||
|
||||
if self._fuse_shared_experts_inside_sbo:
|
||||
shared_output = None
|
||||
@@ -896,9 +877,7 @@ class DeepseekV2MoE(nn.Module):
|
||||
|
||||
final_hidden_states = self.experts(
|
||||
hidden_states=hidden_states,
|
||||
topk_idx=topk_idx,
|
||||
topk_weights=topk_weights,
|
||||
forward_batch=forward_batch,
|
||||
topk_output=topk_output,
|
||||
**(
|
||||
dict(
|
||||
forward_shared_experts=_forward_shared_experts_and_put_results,
|
||||
@@ -960,7 +939,7 @@ class DeepseekV2MoE(nn.Module):
|
||||
with get_global_expert_distribution_recorder().with_current_layer(
|
||||
self.layer_id
|
||||
):
|
||||
state.topk_weights_local, state.topk_idx_local, _ = self.topk(
|
||||
state.topk_output = self.topk(
|
||||
hidden_states=hidden_states,
|
||||
router_logits=router_logits,
|
||||
num_token_non_padded=state.forward_batch.num_token_non_padded,
|
||||
@@ -969,21 +948,13 @@ class DeepseekV2MoE(nn.Module):
|
||||
),
|
||||
)
|
||||
else:
|
||||
state.topk_idx_local = torch.full(
|
||||
(0, self.top_k), -1, dtype=torch.int, device=hidden_states.device
|
||||
)
|
||||
state.topk_weights_local = torch.empty(
|
||||
(0, self.top_k), dtype=torch.float32, device=hidden_states.device
|
||||
)
|
||||
state.topk_output = self.topk.empty_topk_output(hidden_states.device)
|
||||
|
||||
def op_dispatch_a(self, state):
|
||||
if self.ep_size > 1:
|
||||
self.experts.deepep_dispatcher.dispatch_a(
|
||||
self.experts.dispatcher.dispatch_a(
|
||||
hidden_states=state.hidden_states_mlp_input,
|
||||
input_global_scale=None,
|
||||
topk_idx=state.pop("topk_idx_local"),
|
||||
topk_weights=state.pop("topk_weights_local"),
|
||||
forward_batch=state.forward_batch,
|
||||
topk_output=state.pop("topk_output"),
|
||||
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
||||
)
|
||||
|
||||
@@ -992,32 +963,29 @@ class DeepseekV2MoE(nn.Module):
|
||||
with get_global_expert_distribution_recorder().with_current_layer(
|
||||
self.layer_id
|
||||
):
|
||||
state.dispatch_output = self.experts.deepep_dispatcher.dispatch_b(
|
||||
state.dispatch_output = self.experts.dispatcher.dispatch_b(
|
||||
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
||||
)
|
||||
|
||||
def op_experts(self, state):
|
||||
state.hidden_states_experts_output = self.experts.moe_impl(
|
||||
state.hidden_states_experts_output = self.experts.run_moe_core(
|
||||
dispatch_output=state.dispatch_output,
|
||||
)
|
||||
|
||||
def op_combine_a(self, state):
|
||||
if self.ep_size > 1:
|
||||
self.experts.deepep_dispatcher.combine_a(
|
||||
self.experts.dispatcher.combine_a(
|
||||
hidden_states=state.pop("hidden_states_experts_output"),
|
||||
topk_idx=state.dispatch_output.topk_idx,
|
||||
topk_ids=state.dispatch_output.topk_ids,
|
||||
topk_weights=state.dispatch_output.topk_weights,
|
||||
forward_batch=state.forward_batch,
|
||||
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
||||
)
|
||||
state.pop("dispatch_output")
|
||||
|
||||
def op_combine_b(self, state):
|
||||
if self.ep_size > 1:
|
||||
state.hidden_states_after_combine = (
|
||||
self.experts.deepep_dispatcher.combine_b(
|
||||
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
||||
)
|
||||
state.hidden_states_after_combine = self.experts.dispatcher.combine_b(
|
||||
tbo_subbatch_index=state.get("tbo_subbatch_index"),
|
||||
)
|
||||
|
||||
def op_output(self, state):
|
||||
|
||||
Reference in New Issue
Block a user