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_lora(238    x: torch.Tensor,239    gate_proj: torch.nn.Linear,240    up_proj: torch.nn.Linear,241    down_proj: torch.nn.Linear,242    lora_gate_A: torch.Tensor,243    lora_gate_B: torch.Tensor,244    lora_up_A: torch.Tensor,245    lora_up_B: torch.Tensor,246    lora_down_A: torch.Tensor,247    lora_down_B: torch.Tensor,248) -> torch.Tensor:249    F = torch.nn.functional250    xa_g = F.linear(x, lora_gate_A)251    gate_x = gate_proj(x) + F.linear(xa_g, lora_gate_B)252    gate = F.silu(gate_x)253    xa_u = F.linear(x, lora_up_A)254    up = up_proj(x) + F.linear(xa_u, lora_up_B)255    y = gate * up256    xa_d = F.linear(y, lora_down_A)257    return down_proj(y) + F.linear(xa_d, lora_down_B)258 259 260def solution(261    hidden_states: torch.Tensor,262    gate_weight: torch.Tensor,263    gate_bias: Optional[torch.Tensor],264    gate_proj: torch.nn.Linear,265    up_proj: torch.nn.Linear,266    down_proj: torch.nn.Linear,267    lora_gate_A: torch.Tensor,268    lora_gate_B: torch.Tensor,269    lora_up_A: torch.Tensor,270    lora_up_B: torch.Tensor,271    lora_down_A: torch.Tensor,272    lora_down_B: torch.Tensor,273    num_experts: int,274    top_k: int,275    group: Optional[dist.ProcessGroup] = None,276) -> torch.Tensor:277    group = group or dist.group.WORLD278    hidden_dim = hidden_states.size(-1)279    num_tokens = hidden_states.reshape(-1, hidden_dim).size(0)280 281    router_logits = torch.nn.functional.linear(282        hidden_states.reshape(-1, hidden_dim), gate_weight, gate_bias283    )284    routing_weights, selected_experts = torch.topk(285        torch.softmax(router_logits, dim=-1), top_k, dim=-1286    )287    expert_mask = torch.nn.functional.one_hot(288        selected_experts, num_classes=num_experts289    ).permute(2, 1, 0)290 291    input_splits, output_splits, num_global_tokens_per_local_expert, _ = _preprocess(292        expert_mask, num_experts, group293    )294 295    (296        global_permuted_hidden_states,297        routing_map,298        local_input_permutation_mapping,299        org_hidden_states_shape,300    ) = token_pre_all2all(301        hidden_states,302        expert_mask,303        num_experts,304        input_splits,305        output_splits,306        num_global_tokens_per_local_expert,307        group,308    )309 310    expert_outputs = expert_forward_lora(311        global_permuted_hidden_states,312        gate_proj,313        up_proj,314        down_proj,315        lora_gate_A,316        lora_gate_B,317        lora_up_A,318        lora_up_B,319        lora_down_A,320        lora_down_B,321    )322 323    out = tokens_post_all2all(324        expert_outputs,325        routing_weights,326        selected_experts,327        num_experts,328        input_splits,329        output_splits,330        num_global_tokens_per_local_expert,331        routing_map,332        local_input_permutation_mapping,333        org_hidden_states_shape,334        group,335    )336    return out337