willychan21/ParallelKernelBench_Problems

1from typing import List, Optional, Tuple, Union2 3import torch4import torch.distributed as dist5 6 7def _permute(tokens: torch.Tensor, routing_map: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:8    num_tokens, _ = tokens.shape9    num_experts = routing_map.shape[0]10    routing_map = routing_map.bool()11    token_indices = torch.arange(num_tokens, device=routing_map.device).unsqueeze(0).expand(num_experts, -1)12    sorted_indices = token_indices.masked_select(routing_map)13    permuted_input = tokens.index_select(0, sorted_indices)14    return permuted_input, sorted_indices15 16 17def _sort_chunks_by_idxs(18    input: torch.Tensor,19    split_sizes: Union[torch.Tensor, List[int]],20    sorted_idxs: List[int],21) -> torch.Tensor:22    if isinstance(split_sizes, torch.Tensor):23        split_sizes = split_sizes.tolist()24    chunks = torch.split(input, split_sizes, dim=0)25    return torch.cat([chunks[i] for i in sorted_idxs], dim=0)26 27 28def _all_to_all_forward(29    group: dist.ProcessGroup,30    input: torch.Tensor,31    output_split_sizes: Optional[List[int]],32    input_split_sizes: Optional[List[int]],33) -> torch.Tensor:34    if dist.get_world_size(group) == 1:35        return input.contiguous()36    input = input.contiguous()37    out_size = sum(output_split_sizes) if output_split_sizes else input.size(0)38    output = torch.empty((out_size, input.size(1)), dtype=input.dtype, device=input.device)39    dist.all_to_all_single(40        output, input,41        output_split_sizes=output_split_sizes,42        input_split_sizes=input_split_sizes,43        group=group,44    )45    return output46 47 48def solution(49    hidden_states: torch.Tensor,50    expert_mask: torch.Tensor,51    num_experts: int,52    input_splits: Union[List[int], torch.Tensor],53    output_splits: Union[List[int], torch.Tensor],54    num_global_tokens_per_local_expert: torch.Tensor,55    group: Optional[dist.ProcessGroup] = None,56) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Size]:57    group = group or dist.group.WORLD58    hidden_dim = hidden_states.size(-1)59    hidden_states = hidden_states.reshape(-1, hidden_dim)60    org_hidden_states_shape = hidden_states.shape61    routing_map = expert_mask.sum(dim=1)62 63    local_permuted_hidden_states, local_input_permutation_mapping = _permute(hidden_states, routing_map)64 65    expected_tokens = sum(input_splits) if isinstance(input_splits, list) else int(input_splits.sum().item())66    actual_tokens = local_permuted_hidden_states.shape[0]67    if expected_tokens != actual_tokens:68        raise RuntimeError(69            f"EP split mismatch: input_splits sum ({expected_tokens}) != permuted tokens ({actual_tokens})"70        )71 72    global_permuted_hidden_states = _all_to_all_forward(73        group, local_permuted_hidden_states, output_splits, input_splits74    )75 76    num_local_experts = num_experts // dist.get_world_size(group)77    permute_order = torch.arange(num_experts).reshape(-1, num_local_experts).T.ravel().tolist()78    global_permuted_hidden_states = _sort_chunks_by_idxs(79        global_permuted_hidden_states,80        num_global_tokens_per_local_expert.ravel(),81        permute_order,82    )83 84    return global_permuted_hidden_states, routing_map, local_input_permutation_mapping, org_hidden_states_shape85