willychan21/ParallelKernelBench_Problems

1from typing import Optional2 3import torch4import torch.distributed as dist5from torch.distributed import ProcessGroup6 7 8def _pad_tensor(x: torch.Tensor, dim: int, padding_size: int, padding_value: int = 0) -> torch.Tensor:9    shape = list(x.shape)10    shape[dim] = padding_size11    pad = torch.full(shape, padding_value, dtype=x.dtype, device=x.device)12    return torch.cat([x, pad], dim=dim)13 14 15def _all_to_all(16    local_input: torch.Tensor,17    scatter_dim: int,18    gather_dim: int,19    group: dist.ProcessGroup,20) -> torch.Tensor:21    seq_world_size = dist.get_world_size(group)22    input_list = [t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)]23    output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]24    dist.all_to_all(output_list, input_list, group=group)25    return torch.cat(output_list, dim=gather_dim).contiguous()26 27 28def _all_to_all_single(29    x: torch.Tensor,30    scatter_dim: int,31    gather_dim: int,32    group: dist.ProcessGroup,33) -> torch.Tensor:34    sp_world_size = dist.get_world_size(group)35    assert scatter_dim <= 1 and gather_dim <= 136    if scatter_dim != 0:37        gather_dim_bef = x.shape[gather_dim]38        scatter_dim_bef = x.shape[scatter_dim]39        x = (40            x.reshape(41                [gather_dim_bef, sp_world_size, scatter_dim_bef // sp_world_size] + list(x.shape[2:])42            )43            .transpose(0, 1)44            .reshape(45                [gather_dim_bef * sp_world_size, scatter_dim_bef // sp_world_size] + list(x.shape[2:])46            )47            .contiguous()48        )49    output = torch.empty_like(x)50    dist.all_to_all_single(output, x.contiguous(), group=group)51    if scatter_dim == 0:52        output = torch.cat(output.split(x.size(0) // sp_world_size), dim=gather_dim)53    return output54 55 56def _all_to_all_tensor(57    x: torch.Tensor,58    scatter_dim: int,59    gather_dim: int,60    group: dist.ProcessGroup,61) -> torch.Tensor:62    if scatter_dim <= 1 and gather_dim <= 1:63        return _all_to_all_single(x, scatter_dim, gather_dim, group)64    return _all_to_all(x, scatter_dim, gather_dim, group)65 66 67def solution(68    x: torch.Tensor,69    seq_dim: int,70    head_dim: int,71    group: Optional[ProcessGroup] = None,72) -> torch.Tensor:73    group = group or dist.group.WORLD74    dim_size = x.size(seq_dim)75    sp_world = dist.get_world_size(group)76    if dim_size % sp_world != 0:77        padding_size = sp_world - (dim_size % sp_world)78        x = _pad_tensor(x, seq_dim, padding_size)79    return _all_to_all_tensor(x, scatter_dim=seq_dim, gather_dim=head_dim, group=group)80