clean redundant code in previous PR (#12957)
This commit is contained in:
@@ -61,7 +61,6 @@ from sglang.srt.layers.moe import (
|
|||||||
)
|
)
|
||||||
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
|
from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
|
||||||
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
|
from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
|
||||||
from sglang.srt.layers.moe.kt_ep_wrapper import KTEPWrapperMethod
|
|
||||||
from sglang.srt.layers.moe.topk import TopK
|
from sglang.srt.layers.moe.topk import TopK
|
||||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||||
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
|
from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
|
||||||
@@ -455,42 +454,6 @@ class Glm4MoeSparseMoeBlock(nn.Module):
|
|||||||
else:
|
else:
|
||||||
return self.forward_deepep(hidden_states, forward_batch)
|
return self.forward_deepep(hidden_states, forward_batch)
|
||||||
|
|
||||||
def forward_normal_dual_stream(
|
|
||||||
self,
|
|
||||||
hidden_states: torch.Tensor,
|
|
||||||
should_allreduce_fusion: bool = False,
|
|
||||||
use_reduce_scatter: bool = False,
|
|
||||||
) -> torch.Tensor:
|
|
||||||
|
|
||||||
current_stream = torch.cuda.current_stream()
|
|
||||||
self.alt_stream.wait_stream(current_stream)
|
|
||||||
shared_output = self._forward_shared_experts(hidden_states)
|
|
||||||
|
|
||||||
with torch.cuda.stream(self.alt_stream):
|
|
||||||
# router_logits: (num_tokens, n_experts)
|
|
||||||
router_logits = self.gate(hidden_states)
|
|
||||||
topk_output = self.topk(hidden_states, router_logits)
|
|
||||||
final_hidden_states = self.experts(hidden_states, topk_output)
|
|
||||||
if not _is_cuda or isinstance(self.experts.quant_method, KTEPWrapperMethod):
|
|
||||||
final_hidden_states *= self.routed_scaling_factor
|
|
||||||
|
|
||||||
current_stream.wait_stream(self.alt_stream)
|
|
||||||
with use_symmetric_memory(
|
|
||||||
parallel_state.get_tp_group(), disabled=not is_allocation_symmetric()
|
|
||||||
):
|
|
||||||
final_hidden_states_out = torch.empty_like(final_hidden_states)
|
|
||||||
|
|
||||||
torch.add(final_hidden_states, shared_output, out=final_hidden_states_out)
|
|
||||||
final_hidden_states = final_hidden_states_out
|
|
||||||
if (
|
|
||||||
self.tp_size > 1
|
|
||||||
and not should_allreduce_fusion
|
|
||||||
and not use_reduce_scatter
|
|
||||||
and not should_use_flashinfer_cutlass_moe_fp4_allgather()
|
|
||||||
):
|
|
||||||
final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states)
|
|
||||||
return final_hidden_states
|
|
||||||
|
|
||||||
def forward_normal(
|
def forward_normal(
|
||||||
self,
|
self,
|
||||||
hidden_states: torch.Tensor,
|
hidden_states: torch.Tensor,
|
||||||
|
|||||||
Reference in New Issue
Block a user