Skip to content

feat: large-scale EP(part 3: refactor - FusedMoe for redundant expert) #4495

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
juney-nvidia merged 1 commit into from
May 21, 2025
Merged

feat: large-scale EP(part 3: refactor - FusedMoe for redundant expert) #4495

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 64 additions & 42 deletions tensorrt_llm/_torch/modules/fused_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -272,6 +272,10 @@ class FusedMoE(nn.Module):

In min-latency mode, setting `reduce_results=False` disables the AllReduce in the FusedMoE module, so any necessary AllReduce operations must be added explicitly in the model definition.
AttentionDP should be turned off for min-latency mode.

When we have redundant expert, we have more weight slots than `num_experts`, in that case, we separate the concepts of expert and slot.
Expert is the concept from model's perspective while slot is the concept from model engine's perspective.
There should be at lease `num_experts` slots in the model engine. More than that is OK, in that case, some experts may have multiple replicas.
"""

def __init__(
Expand Down Expand Up @@ -326,11 +330,25 @@ def __init__(

self.intermediate_size_per_partition = intermediate_size // self.tp_size

self.expert_size_per_partition = num_experts // self.ep_size
self.expert_start = self.ep_rank * self.expert_size_per_partition
self.expert_end = min(
self.expert_start + self.expert_size_per_partition,
self.num_experts)
# self.expert_slots_per_partition will be replaced with real slots_per_partition to enable redundant expert slots
self.expert_slots_per_partition = num_experts // self.ep_size
assert self.expert_slots_per_partition * self.ep_size >= num_experts, "total slots should be at lease num_experts"
if self.smart_router:
assert self.expert_slots_per_partition == num_experts // self.ep_size,\
"Smart router should not have redundant slots"
self.num_slots = self.expert_slots_per_partition * self.ep_size
# Here the meaning of expert_size_per_partition is the number of expert slots that each rank has.
self.expert_size_per_partition = self.expert_slots_per_partition
self.slot_start = self.ep_rank * self.expert_size_per_partition
self.slot_end = self.slot_start + self.expert_size_per_partition

self.initial_global_assignments = [
(ep_rank * self.num_experts // self.ep_size + local_slot_id) %
self.num_experts for ep_rank in range(self.ep_size)
for local_slot_id in range(self.expert_slots_per_partition)
]
self.initial_local_expert_ids = self.initial_global_assignments[
self.slot_start:self.slot_end]

max_num_tokens = model_config.max_num_tokens
# The maximum number of tokens in MoE are multiplied by DP size when attention DP is enabled
Expand All @@ -354,7 +372,7 @@ def __init__(
# around 16k tokens per expert, which is well into the compute bound domain.
self.tune_max_num_tokens = min(
self.moe_max_num_tokens,
16384 * num_experts // routing_method.get_experts_per_token(),
16384 * self.num_slots // routing_method.get_experts_per_token(),
)
self.has_been_profiled = False
self.has_been_profiled_min_latency = False
Expand Down Expand Up @@ -453,45 +471,45 @@ def is_trtllm(self):
def is_cutlass(self):
return not self.is_trtllm()

def get_quant_scales(self, expert_start, expert_end):
def get_quant_scales(self, slot_start, slot_end):
assert self.smart_router

if self.has_fp8_block_scales:
return FusedMoEQuantScalesFP8BlockScales(
fc_weight_scales=self.w3_w1_weight_scaling_factor.narrow(
0, expert_start, expert_end - expert_start),
0, slot_start, slot_end - slot_start),
proj_weight_scales=self.w2_weight_scaling_factor.narrow(
0, expert_start, expert_end - expert_start),
0, slot_start, slot_end - slot_start),
)
elif self.has_nvfp4:
return FusedMoEQuantScalesNVFP4(
fc1_act_global=self.fc31_input_scale,
fc1_weight_block=self.w3_w1_weight_scale.narrow(
0, expert_start, expert_end - expert_start),
fc1_global=self.fc31_alpha.narrow(0, expert_start,
expert_end - expert_start),
0, slot_start, slot_end - slot_start),
fc1_global=self.fc31_alpha.narrow(0, slot_start,
slot_end - slot_start),
fc2_act_global=self.fc2_input_scale,
fc2_weight_block=self.w2_weight_scale.narrow(
0, expert_start, expert_end - expert_start),
fc2_global=self.fc2_alpha.narrow(0, expert_start,
expert_end - expert_start),
0, slot_start, slot_end - slot_start),
fc2_global=self.fc2_alpha.narrow(0, slot_start,
slot_end - slot_start),
)
elif self.has_w4afp8:
return FusedMoEQuantScalesW4A8(
scale_1_interleaved=self.fc31_weight_scale.narrow(
0, expert_start, expert_end - expert_start),
0, slot_start, slot_end - slot_start),
scale_2_interleaved=self.fc2_weight_scale.narrow(
0, expert_start, expert_end - expert_start),
0, slot_start, slot_end - slot_start),
pre_quant_scale_1=self.fc31_act_scale.narrow(
0, expert_start, expert_end - expert_start),
0, slot_start, slot_end - slot_start),
pre_quant_scale_2=self.fc2_act_scale.narrow(
0, expert_start, expert_end - expert_start),
0, slot_start, slot_end - slot_start),
zero_1=torch.Tensor(),
zero_2=torch.Tensor(),
alpha_1=self.fc31_alpha.narrow(0, expert_start,
expert_end - expert_start),
alpha_2=self.fc2_alpha.narrow(0, expert_start,
expert_end - expert_start),
alpha_1=self.fc31_alpha.narrow(0, slot_start,
slot_end - slot_start),
alpha_2=self.fc2_alpha.narrow(0, slot_start,
slot_end - slot_start),
)
else:
return self.quant_scales
Expand Down Expand Up @@ -1081,13 +1099,13 @@ def forward_trtllmgen(self, x: torch.Tensor,
self.w3_w1_weight_scaling_factor,
self.w2_weight,
self.w2_weight_scaling_factor,
self.num_experts,
self.num_slots,
top_k,
n_group,
topk_group,
self.intermediate_size_per_partition,
self.
expert_start, # local_expert_start; use ep_rank if stride!=1
slot_start, # local_expert_start; use ep_rank if stride!=1
self.expert_size_per_partition, # local_expert_size
routed_scaling_factor,
)
Expand All @@ -1110,13 +1128,13 @@ def forward_trtllmgen(self, x: torch.Tensor,
self.fc31_scale_c.data,
self.fc31_alpha.data,
self.fc2_alpha.data,
self.num_experts,
self.num_slots,
top_k,
n_group,
topk_group,
self.intermediate_size_per_partition,
self.
expert_start, # local_expert_start; use ep_rank if stride!=1
slot_start, # local_expert_start; use ep_rank if stride!=1
self.expert_size_per_partition, # local_expert_size
routed_scaling_factor,
)
Expand Down Expand Up @@ -1282,8 +1300,10 @@ def load_expert_w2_weight(w2_weight,
# CPU memory bandwidth better.
threads = []

for expert_id in range(self.expert_start, self.expert_end):
expert_idx = expert_id - self.expert_start
for local_slot_id, expert_id in enumerate(
self.initial_local_expert_ids):
# expert_idx is the local slot index of current rank
expert_idx = local_slot_id

if self.weight_loading_mode == MoEWeightLoadingMode.VANILLA:
w1_weight = weights[f"{expert_id}.w1.weight"]
Expand Down Expand Up @@ -1342,7 +1362,7 @@ def _load_fp8_block_scales_scales(self, weights: Dict):
load_weight_shard(weights[f"{expert_id}.w2.weight_scale_inv"],
self.tp_size, self.tp_rank,
TensorParallelMode.ROW)
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]

w2_scales = torch.stack(all_w2_scales)
Expand All @@ -1352,14 +1372,14 @@ def _load_fp8_block_scales_scales(self, weights: Dict):
load_weight_shard(weights[f"{expert_id}.w3.weight_scale_inv"],
self.tp_size, self.tp_rank,
TensorParallelMode.COLUMN)
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]

all_w1_scales = [
load_weight_shard(weights[f"{expert_id}.w1.weight_scale_inv"],
self.tp_size, self.tp_rank,
TensorParallelMode.COLUMN)
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]

w3_w1_scales = torch.cat(
Expand Down Expand Up @@ -1458,7 +1478,8 @@ def load_expert_w2_weight_scale_fp8(w2_weight_scale,
dst_w2_weight_scale: torch.Tensor):
dst_w2_weight_scale.copy_(w2_weight_scale[...].reshape([]))

for expert_id in range(self.expert_start, self.expert_end):
for local_slot_id, expert_id in enumerate(
self.initial_local_expert_ids):
if self.weight_loading_mode == MoEWeightLoadingMode.VANILLA:
w1_weight_scale = weights[f"{expert_id}.w1.weight_scale"]
w3_weight_scale = weights[f"{expert_id}.w3.weight_scale"]
Expand All @@ -1472,7 +1493,7 @@ def load_expert_w2_weight_scale_fp8(w2_weight_scale,
f"Unknown weight loading mode in MoE: {self.weight_loading_mode}"
)

expert_idx = expert_id - self.expert_start
expert_idx = local_slot_id

load_expert_w3_w1_weight_scale_fp8_qdq(
w1_weight_scale, w3_weight_scale,
Expand Down Expand Up @@ -1617,7 +1638,8 @@ def load_expert_fc2_alpha_nvfp4(w2_weight_scale_2,
dst_w2_alpha.copy_(1.0 /
(final_fc2_input_scale * w2_weight_scale_2))

for expert_id in range(self.expert_start, self.expert_end):
for local_slot_id, expert_id in enumerate(
self.initial_local_expert_ids):
if self.weight_loading_mode == MoEWeightLoadingMode.VANILLA:
w1_weight_scale = weights[f"{expert_id}.w1.weight_scale"]
w3_weight_scale = weights[f"{expert_id}.w3.weight_scale"]
Expand All @@ -1640,7 +1662,7 @@ def load_expert_fc2_alpha_nvfp4(w2_weight_scale_2,
f"Unknown weight loading mode in MoE: {self.weight_loading_mode}"
)

expert_idx = expert_id - self.expert_start
expert_idx = local_slot_id

load_expert_w3_w1_weight_scale_nvfp4(
w1_weight_scale, w3_weight_scale,
Expand All @@ -1664,11 +1686,11 @@ def _load_int4_groupwise_scales(self, weights: Dict):
assert (len(self.interleave) == 2)
all_w3_input_scales = [
load_weight_shard(weights[f"{expert_id}.w3.input_scale"])
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]
all_w1_input_scales = [
load_weight_shard(weights[f"{expert_id}.w1.input_scale"])
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]
all_w3_w1_input_scales = torch.max(torch.stack(all_w3_input_scales),
torch.stack(all_w1_input_scales))
Expand All @@ -1681,13 +1703,13 @@ def _load_int4_groupwise_scales(self, weights: Dict):
load_weight_shard(weights[f"{expert_id}.w3.weight_scale_inv"],
self.tp_size, self.tp_rank,
TensorParallelMode.COLUMN)
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]
all_w1_scales = [
load_weight_shard(weights[f"{expert_id}.w1.weight_scale_inv"],
self.tp_size, self.tp_rank,
TensorParallelMode.COLUMN)
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]
all_w3_w1_scales = torch.cat(
[torch.stack(all_w3_scales),
Expand All @@ -1709,7 +1731,7 @@ def _load_int4_groupwise_scales(self, weights: Dict):
# fc2 scales
all_w2_input_scales = [
load_weight_shard(weights[f"{expert_id}.w2.input_scale"])
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]
all_w2_input_scales = torch.stack(all_w2_input_scales).to(self.dtype)
all_w2_input_scales = torch.ones_like(
Expand All @@ -1721,7 +1743,7 @@ def _load_int4_groupwise_scales(self, weights: Dict):
load_weight_shard(weights[f"{expert_id}.w2.weight_scale_inv"],
self.tp_size, self.tp_rank,
TensorParallelMode.ROW)
for expert_id in range(self.expert_start, self.expert_end)
for expert_id in self.initial_local_expert_ids
]
if self.sm_version == 89:
w2_scales = torch.stack(all_w2_scales).to(torch.float16).view(
Expand Down