From 5e1923372cf23f25cc4f0756626dc17a214679e5 Mon Sep 17 00:00:00 2001 From: jinminxi104 Date: Mon, 13 Jul 2026 11:46:12 +0000 Subject: [PATCH] support goupedmm type012 --- dlinfer/vendor/ascend/moe.py | 99 +++++++++++++++++++++++++++++++----- 1 file changed, 86 insertions(+), 13 deletions(-) diff --git a/dlinfer/vendor/ascend/moe.py b/dlinfer/vendor/ascend/moe.py index a0a66033..187cfd8e 100644 --- a/dlinfer/vendor/ascend/moe.py +++ b/dlinfer/vendor/ascend/moe.py @@ -14,6 +14,12 @@ # True -> catch-all (identical scheme to graph capture) # False -> per-chunk row slicing (weight views, each row computed once) _MOE_PREFILL_USE_CATCHALL = os.environ.get("DLINFER_MOE_PREFILL_CATCHALL", "0") == "1" +_MOE_GROUP_LIST_TYPE = int(os.environ.get("DLINFER_MOE_GROUP_LIST_TYPE", "1")) +if _MOE_GROUP_LIST_TYPE not in (0, 1, 2): + raise ValueError( + "DLINFER_MOE_GROUP_LIST_TYPE must be one of 0, 1, or 2, " + f"got {_MOE_GROUP_LIST_TYPE}" + ) class ChunkedMoeWeightLayout: @@ -82,6 +88,8 @@ def _grouped_mlp( group_list: torch.Tensor, group_list_type: int, ): + _check_group_list_type_supported(hidden_states, gate_up_weights, group_list_type) + # up sample up_proj = torch.ops.npu.npu_grouped_matmul( [hidden_states], @@ -107,6 +115,60 @@ def _grouped_mlp( return down_proj +def _check_group_list_type_supported( + hidden_states: torch.Tensor, + gate_up_weights: torch.Tensor, + group_list_type: int, +): + if group_list_type == 2 and ( + hidden_states.dtype != torch.int8 or gate_up_weights.dtype != torch.int8 + ): + raise RuntimeError( + "npu_grouped_matmul group_list_type=2 is only supported when both " + f"x and weight are int8; got x={hidden_states.dtype}, " + f"weight={gate_up_weights.dtype}. bf16 MetaMoE cannot use " + "DLINFER_MOE_GROUP_LIST_TYPE=2." + ) + + +def _token_counts_from_group_list( + group_list: torch.Tensor, + group_list_type: int, + num_experts: int, +): + if group_list_type == 0: + token_counts = group_list.clone() + token_counts[1:] = group_list[1:] - group_list[:-1] + return token_counts + if group_list_type == 1: + return group_list + if group_list_type == 2: + token_counts = group_list.new_zeros((num_experts,)) + group_ids = group_list[:, 0].to(torch.long) + group_sizes = group_list[:, 1] + return token_counts.scatter(0, group_ids, group_sizes) + raise ValueError(f"Unsupported group_list_type: {group_list_type}") + + +def _build_group_list_from_token_counts( + token_counts: torch.Tensor, + group_list_type: int, +): + if group_list_type == 0: + return torch.cumsum(token_counts, dim=0) + if group_list_type == 1: + return token_counts + if group_list_type == 2: + group_ids = torch.arange( + token_counts.numel(), device=token_counts.device, dtype=token_counts.dtype + ) + group_list = torch.stack([group_ids, token_counts], dim=1) + zero_last_key = (token_counts == 0).to(token_counts.dtype) + sort_idx = torch.argsort(zero_last_key, stable=True) + return group_list[sort_idx] + raise ValueError(f"Unsupported group_list_type: {group_list_type}") + + def _apply_mlp_chunked_eager( hidden_states: torch.Tensor, gate_up_weights: torch.Tensor, @@ -153,8 +215,11 @@ def _apply_mlp_chunked_eager( hidden_states[start_row:end_row], gate_up, down, - token_counts[start_expert:end_expert], - 1, + _build_group_list_from_token_counts( + token_counts[start_expert:end_expert], + _MOE_GROUP_LIST_TYPE, + ), + _MOE_GROUP_LIST_TYPE, ) ) start_expert = end_expert @@ -207,9 +272,13 @@ def _apply_mlp_chunked_capturable( end_expert = min(start_expert + chunk_size, num_experts) leading = token_counts[:start_expert].sum().reshape(1) trailing = token_counts[end_expert:].sum().reshape(1) - group_list = torch.cat( + chunk_token_counts = torch.cat( [leading, token_counts[start_expert:end_expert], trailing] ) + group_list = _build_group_list_from_token_counts( + chunk_token_counts, + _MOE_GROUP_LIST_TYPE, + ) if use_chunked_layout and chunked_moe_layout.packed: offset = start_expert + 2 * chunk_idx chunk_len = end_expert - start_expert + 2 @@ -222,7 +291,13 @@ def _apply_mlp_chunked_capturable( down = torch.cat( [zeros_down, down_weights[start_expert:end_expert], zeros_down], dim=0 ) - chunk_out = _grouped_mlp(hidden_states, gate_up, down, group_list, 1) + chunk_out = _grouped_mlp( + hidden_states, + gate_up, + down, + group_list, + _MOE_GROUP_LIST_TYPE, + ) output = chunk_out if output is None else output + chunk_out start_expert = end_expert chunk_idx += 1 @@ -247,15 +322,13 @@ def apply_mlp( hidden_states, gate_up_weights, down_weights, group_list, group_list_type ) - # More experts than aclnnGroupedMatmulV5 supports: split into chunks of at - # most MAX_GROUP_LIST_SIZE groups. Work in per-expert token counts - # (group_list_type=1) regardless of the incoming layout. - if group_list_type == 0: - # group_list is a cumulative sum -> recover per-expert counts - token_counts = group_list.clone() - token_counts[1:] = group_list[1:] - group_list[:-1] - else: - token_counts = group_list + _check_group_list_type_supported( + hidden_states, gate_up_weights, _MOE_GROUP_LIST_TYPE + ) + + # More experts than aclnnGroupedMatmulV5 supports: split into chunks after + # normalizing the dispatch output to per-expert token counts. + token_counts = _token_counts_from_group_list(group_list, group_list_type, num_experts) # Graph capture cannot tolerate the host sync of the eager row-slicing path, # so capture always uses the catch-all path. Prefill keeps the cheaper eager