Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
This commit is contained in:
co-authored by
Kangyan-Zhou
parent
f969664172
commit
78a4b446c6
@@ -284,17 +284,9 @@ class ExpertLocationMetadata:
|
|||||||
# -------------------------------- usage ------------------------------------
|
# -------------------------------- usage ------------------------------------
|
||||||
|
|
||||||
def logical_to_all_physical(
|
def logical_to_all_physical(
|
||||||
self,
|
self, layer_id: int, logical_expert_id: int
|
||||||
layer_id: int,
|
|
||||||
logical_expert_id: int,
|
|
||||||
require_global_experts: bool = False,
|
|
||||||
) -> List[int]:
|
) -> List[int]:
|
||||||
# Use CPU copy to avoid GPU→CPU sync on every call, which is expensive in update weights scenario
|
# Use CPU copy to avoid GPU→CPU sync on every call, which is expensive in update weights scenario
|
||||||
if require_global_experts:
|
|
||||||
num_physical_experts = self.logical_to_all_physical_map_cpu[layer_id].shape[
|
|
||||||
-1
|
|
||||||
]
|
|
||||||
return list(torch.arange(0, num_physical_experts))
|
|
||||||
return [
|
return [
|
||||||
physical_expert_id
|
physical_expert_id
|
||||||
for physical_expert_id in self.logical_to_all_physical_map_cpu[
|
for physical_expert_id in self.logical_to_all_physical_map_cpu[
|
||||||
@@ -363,10 +355,14 @@ def _compute_logical_to_all_physical_map(
|
|||||||
)
|
)
|
||||||
|
|
||||||
# Replace by the nearest physical expert
|
# Replace by the nearest physical expert
|
||||||
if nearest_expert != -1:
|
mapped_physical_experts = logical_to_all_physical_map[layer_id][
|
||||||
logical_to_all_physical_map[layer_id][logical_expert_id] = [
|
logical_expert_id
|
||||||
nearest_expert
|
|
||||||
]
|
]
|
||||||
|
if (
|
||||||
|
nearest_expert != -1
|
||||||
|
and nearest_expert not in mapped_physical_experts
|
||||||
|
):
|
||||||
|
mapped_physical_experts[0] = nearest_expert
|
||||||
|
|
||||||
logical_to_all_physical_map = _pad_nested_array(
|
logical_to_all_physical_map = _pad_nested_array(
|
||||||
logical_to_all_physical_map, pad_value=-1
|
logical_to_all_physical_map, pad_value=-1
|
||||||
|
|||||||
@@ -539,12 +539,9 @@ class FusedMoE(torch.nn.Module):
|
|||||||
# This is a shared expert.
|
# This is a shared expert.
|
||||||
physical_expert_ids = [expert_id]
|
physical_expert_ids = [expert_id]
|
||||||
else:
|
else:
|
||||||
require_global_experts = getattr(
|
|
||||||
param, "_sglang_require_global_experts", False
|
|
||||||
)
|
|
||||||
physical_expert_ids = (
|
physical_expert_ids = (
|
||||||
global_expert_location_metadata.logical_to_all_physical(
|
global_expert_location_metadata.logical_to_all_physical(
|
||||||
self.layer_id, expert_id, require_global_experts
|
self.layer_id, expert_id
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -1129,6 +1126,7 @@ class FlashInferFP4MoE(FusedMoE):
|
|||||||
local_expert_offset=self.moe_ep_rank * self.num_local_experts,
|
local_expert_offset=self.moe_ep_rank * self.num_local_experts,
|
||||||
local_num_experts=self.num_local_experts,
|
local_num_experts=self.num_local_experts,
|
||||||
routed_scaling_factor=self.moe_runner_config.routed_scaling_factor,
|
routed_scaling_factor=self.moe_runner_config.routed_scaling_factor,
|
||||||
|
tile_tokens_dim=None,
|
||||||
routing_method_type=RoutingMethodType.DeepSeekV3,
|
routing_method_type=RoutingMethodType.DeepSeekV3,
|
||||||
do_finalize=True,
|
do_finalize=True,
|
||||||
output=symm_output,
|
output=symm_output,
|
||||||
|
|||||||
@@ -1245,6 +1245,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
|||||||
routed_scaling_factor=(
|
routed_scaling_factor=(
|
||||||
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
routed_scaling_factor if routed_scaling_factor is not None else 1.0
|
||||||
),
|
),
|
||||||
|
tile_tokens_dim=None,
|
||||||
routing_method_type=routing_method_type,
|
routing_method_type=routing_method_type,
|
||||||
use_shuffled_weight=False,
|
use_shuffled_weight=False,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -739,6 +739,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):
|
|||||||
else 1.0
|
else 1.0
|
||||||
),
|
),
|
||||||
use_routing_scales_on_input=use_routing_scales_on_input,
|
use_routing_scales_on_input=use_routing_scales_on_input,
|
||||||
|
tile_tokens_dim=None,
|
||||||
routing_method_type=routing_method_type,
|
routing_method_type=routing_method_type,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|||||||
@@ -681,6 +681,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
|
|||||||
layer.moe_ep_rank * layer.num_local_experts, # local_expert_offset
|
layer.moe_ep_rank * layer.num_local_experts, # local_expert_offset
|
||||||
layer.num_local_experts, # local num experts
|
layer.num_local_experts, # local num experts
|
||||||
None,
|
None,
|
||||||
|
None, # tile_tokens_dim
|
||||||
1, # routing_method_type, renormalize
|
1, # routing_method_type, renormalize
|
||||||
True, # do finalize
|
True, # do finalize
|
||||||
output=symm_output,
|
output=symm_output,
|
||||||
|
|||||||
Reference in New Issue
Block a user