willychan21/ParallelKernelBench_Problems

1from typing import List, Optional, Tuple, Union2 3import torch4import torch.distributed as dist5 6 7class _AllToAll(torch.autograd.Function):8    @staticmethod9    def forward(ctx, group, input, output_split_sizes, input_split_sizes):10        ctx.group = group11        ctx.output_split_sizes = output_split_sizes12        ctx.input_split_sizes = input_split_sizes13        if dist.get_world_size(group=group) == 1:14            return input.contiguous()15        input = input.contiguous()16        if output_split_sizes is None:17            output = torch.empty_like(input)18        else:19            output = torch.empty(20                size=(sum(output_split_sizes), input.size(1)),21                dtype=input.dtype,22                device=input.device,23            )24        dist.all_to_all_single(25            output,26            input,27            output_split_sizes=output_split_sizes,28            input_split_sizes=input_split_sizes,29            group=group,30        )31        return output32 33    @staticmethod34    def backward(ctx, grad_output):35        return (36            None,37            _AllToAll.apply(38                ctx.group, grad_output, ctx.input_split_sizes, ctx.output_split_sizes39            ),40            None,41            None,42        )43 44 45def _all_to_all(46    group: dist.ProcessGroup,47    input: torch.Tensor,48    output_split_sizes: Optional[List[int]],49    input_split_sizes: Optional[List[int]],50) -> torch.Tensor:51    return _AllToAll.apply(group, input, output_split_sizes, input_split_sizes)52 53 54def _preprocess(55    expert_mask: torch.Tensor,56    num_experts: int,57    ep_group: dist.ProcessGroup,58) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]:59    ep_size = ep_group.size()60    num_local_experts = num_experts // ep_size61    rank = dist.get_rank(ep_group)62    num_local_tokens_per_expert = expert_mask.sum(dim=(1, 2))63    input_splits = (64        num_local_tokens_per_expert.reshape(ep_size, num_local_experts).sum(dim=1).tolist()65    )66    num_local_tokens_per_expert_flat = num_local_tokens_per_expert.contiguous().view(-1)67    output_size = ep_size * num_local_tokens_per_expert_flat.numel()68    num_global_tokens_per_expert_flat = torch.empty(69        output_size,70        dtype=num_local_tokens_per_expert.dtype,71        device=num_local_tokens_per_expert.device,72    )73    dist.all_gather_into_tensor(74        num_global_tokens_per_expert_flat, num_local_tokens_per_expert_flat, group=ep_group75    )76    num_global_tokens_per_expert = num_global_tokens_per_expert_flat.view(77        ep_size, num_local_tokens_per_expert.size(0)78    )79    start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts80    num_global_tokens_per_local_expert = num_global_tokens_per_expert[81        :, start_idx:end_idx82    ].contiguous()83    output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist()84    num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(85        dim=086    ).to(torch.device("cpu"), non_blocking=True)87    num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(88        -1, num_local_experts89    ).to(torch.device("cpu"), non_blocking=True)90    return (91        input_splits,92        output_splits,93        num_global_tokens_per_local_expert,94        num_global_sum_tokens_per_local_expert,95    )96 97 98def _permute(99    tokens: torch.Tensor, routing_map: torch.Tensor100) -> Tuple[torch.Tensor, torch.Tensor]:101    num_tokens, _ = tokens.shape102    num_experts = routing_map.shape[0]103    routing_map = routing_map.bool()104    token_indices = (105        torch.arange(num_tokens, device=routing_map.device)106        .unsqueeze(0)107        .expand(num_experts, -1)108    )109    sorted_indices = token_indices.masked_select(routing_map)110    permuted_input = tokens.index_select(0, sorted_indices)111    return permuted_input, sorted_indices112 113 114def _sort_chunks_by_idxs(115    input: torch.Tensor,116    split_sizes: Union[torch.Tensor, List[int]],117    sorted_idxs: List[int],118) -> torch.Tensor:119    if isinstance(split_sizes, torch.Tensor):120        split_sizes = split_sizes.tolist()121    chunks = torch.split(input, split_sizes, dim=0)122    return torch.cat([chunks[i] for i in sorted_idxs], dim=0)123 124 125def _generate_weights_idx(126    routing_weights: torch.Tensor,127    selected_experts: torch.Tensor,128    num_experts: int,129) -> torch.Tensor:130    num_tokens, topk = routing_weights.shape131    weights_idx = torch.zeros(132        (num_tokens, num_experts),133        dtype=routing_weights.dtype,134        device=routing_weights.device,135    )136    weights_idx.scatter_add_(1, selected_experts, routing_weights)137    return weights_idx138 139 140def _unpermute(141    tokens: torch.Tensor,142    routing_weights: torch.Tensor,143    hidden_states_shape: torch.Size,144    permutation_mapping: torch.Tensor,145    routing_map: torch.Tensor,146) -> torch.Tensor:147    tokens_weight = routing_weights.T.contiguous().masked_select(routing_map.bool())148    tokens = tokens * tokens_weight.unsqueeze(-1)149    hidden_dim = hidden_states_shape[-1]150    unpermuted_tokens = torch.zeros(151        hidden_states_shape, device=tokens.device, dtype=tokens.dtype152    )153    expanded_mapping = permutation_mapping.unsqueeze(1).expand(-1, hidden_dim)154    unpermuted_tokens.scatter_add_(0, expanded_mapping, tokens)155    return unpermuted_tokens156 157 158def token_pre_all2all(159    hidden_states: torch.Tensor,160    expert_mask: torch.Tensor,161    num_experts: int,162    input_splits: List[int],163    output_splits: List[int],164    num_global_tokens_per_local_expert: torch.Tensor,165    group: Optional[dist.ProcessGroup] = None,166) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Size]:167    group = group or dist.group.WORLD168    hidden_dim = hidden_states.size(-1)169    hidden_states = hidden_states.reshape(-1, hidden_dim)170    org_hidden_states_shape = hidden_states.shape171    routing_map = expert_mask.sum(dim=1)172 173    local_permuted_hidden_states, local_input_permutation_mapping = _permute(174        hidden_states, routing_map175    )176    expected_tokens = sum(input_splits)177    actual_tokens = local_permuted_hidden_states.shape[0]178    if expected_tokens != actual_tokens:179        raise RuntimeError(180            f"EP split mismatch: input_splits sum ({expected_tokens}) != "181            f"permuted tokens ({actual_tokens})"182        )183 184    global_permuted_hidden_states = _all_to_all(185        group, local_permuted_hidden_states, output_splits, input_splits186    )187    num_local_experts = num_experts // dist.get_world_size(group)188    permute_order = (189        torch.arange(num_experts).reshape(-1, num_local_experts).T.ravel().tolist()190    )191    split_sizes = num_global_tokens_per_local_expert.ravel().tolist()192    global_permuted_hidden_states = _sort_chunks_by_idxs(193        global_permuted_hidden_states, split_sizes, permute_order194    )195    return (196        global_permuted_hidden_states,197        routing_map,198        local_input_permutation_mapping,199        org_hidden_states_shape,200    )201 202 203def tokens_post_all2all(204    expert_outputs: torch.Tensor,205    routing_weights: torch.Tensor,206    selected_experts: torch.Tensor,207    num_experts: int,208    input_splits: List[int],209    output_splits: List[int],210    num_global_tokens_per_local_expert: torch.Tensor,211    routing_map: torch.Tensor,212    local_input_permutation_mapping: torch.Tensor,213    org_hidden_states_shape: torch.Size,214    group: Optional[dist.ProcessGroup] = None,215) -> torch.Tensor:216    group = group or dist.group.WORLD217    num_local_experts = num_experts // dist.get_world_size(group)218    unpermute_order = (219        torch.arange(num_experts).reshape(num_local_experts, -1).T.ravel().tolist()220    )221    split_sizes = num_global_tokens_per_local_expert.T.ravel().tolist()222    expert_outputs = _sort_chunks_by_idxs(223        expert_outputs, split_sizes, unpermute_order224    )225    unpermute_outputs = _all_to_all(group, expert_outputs, input_splits, output_splits)226    weights_idx = _generate_weights_idx(routing_weights, selected_experts, num_experts)227    unpermute_outputs = _unpermute(228        unpermute_outputs,229        weights_idx,230        org_hidden_states_shape,231        local_input_permutation_mapping,232        routing_map,233    )234    return unpermute_outputs235 236 237def expert_forward(238    x: torch.Tensor,239    gate_proj: torch.nn.Linear,240    up_proj: torch.nn.Linear,241    down_proj: torch.nn.Linear,242) -> torch.Tensor:243    gate = torch.nn.functional.silu(gate_proj(x))244    up = up_proj(x)245    return down_proj(gate * up)246 247 248def solution(249    hidden_states: torch.Tensor,250    gate_weight: torch.Tensor,251    gate_bias: Optional[torch.Tensor],252    gate_proj: torch.nn.Linear,253    up_proj: torch.nn.Linear,254    down_proj: torch.nn.Linear,255    num_experts: int,256    top_k: int,257    group: Optional[dist.ProcessGroup] = None,258) -> torch.Tensor:259    group = group or dist.group.WORLD260    hidden_dim = hidden_states.size(-1)261    num_tokens = hidden_states.reshape(-1, hidden_dim).size(0)262 263    router_logits = torch.nn.functional.linear(264        hidden_states.reshape(-1, hidden_dim), gate_weight, gate_bias265    )266    routing_weights, selected_experts = torch.topk(267        torch.softmax(router_logits, dim=-1), top_k, dim=-1268    )269    expert_mask = torch.nn.functional.one_hot(270        selected_experts, num_classes=num_experts271    ).permute(2, 1, 0)272 273    input_splits, output_splits, num_global_tokens_per_local_expert, _ = _preprocess(274        expert_mask, num_experts, group275    )276 277    (278        global_permuted_hidden_states,279        routing_map,280        local_input_permutation_mapping,281        org_hidden_states_shape,282    ) = token_pre_all2all(283        hidden_states,284        expert_mask,285        num_experts,286        input_splits,287        output_splits,288        num_global_tokens_per_local_expert,289        group,290    )291 292    expert_outputs = expert_forward(293        global_permuted_hidden_states, gate_proj, up_proj, down_proj294    )295 296    out = tokens_post_all2all(297        expert_outputs,298        routing_weights,299        selected_experts,300        num_experts,301        input_splits,302        output_splits,303        num_global_tokens_per_local_expert,304        routing_map,305        local_input_permutation_mapping,306        org_hidden_states_shape,307        group,308    )309    return out310