willychan21/ParallelKernelBench_Problems

1from typing import Any, Optional, Tuple2 3import torch4import torch.nn.functional as F5import torch.distributed as dist6from torch import Tensor7from torch.distributed import ProcessGroup8 9 10def _pad_tensor(x: Tensor, dim: int, padding_size: int, padding_value: int = 0) -> Tensor:11    shape = list(x.shape)12    shape[dim] = padding_size13    pad = torch.full(shape, padding_value, dtype=x.dtype, device=x.device)14    return torch.cat([x, pad], dim=dim)15 16 17def _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor:18    slc = [slice(None)] * len(x.shape)19    slc[dim] = slice(0, -padding_size)20    return x[tuple(slc)]21 22 23def _all_to_all_single(24    x: Tensor,25    scatter_dim: int,26    gather_dim: int,27    group: Optional[dist.ProcessGroup] = None,28    async_op: bool = False,29):30    group = group or dist.group.WORLD31    sp_world_size = dist.get_world_size(group)32    assert scatter_dim <= 1, "scatter_dim must be 0 or 1 when using all_to_all_single!"33    assert gather_dim <= 1, "gather_dim must be 0 or 1 when using all_to_all_single!"34    if scatter_dim != 0:35        gather_dim_bef = x.shape[gather_dim]36        scatter_dim_bef = x.shape[scatter_dim]37        x = (38            x.reshape(39                [gather_dim_bef, sp_world_size, scatter_dim_bef // sp_world_size]40                + list(x.shape[2:])41            )42            .transpose(0, 1)43            .reshape(44                [gather_dim_bef * sp_world_size, scatter_dim_bef // sp_world_size]45                + list(x.shape[2:])46            )47            .contiguous()48        )49 50    output = torch.empty_like(x)51    comm = dist.all_to_all_single(output, x.contiguous(), group=group, async_op=async_op)52 53    if async_op:54 55        def wait():56            comm.wait()57            if scatter_dim == 0:58                return torch.cat(output.split(x.size(0) // sp_world_size), dim=gather_dim)59            else:60                return output61 62        return wait63 64    if scatter_dim == 0:65        output = torch.cat(output.split(x.size(0) // sp_world_size), dim=gather_dim)66    return output67 68 69def _all_to_all(70    local_input: Tensor,71    scatter_dim: int,72    gather_dim: int,73    group: Optional[dist.ProcessGroup] = None,74    async_op: bool = False,75):76    group = group or dist.group.WORLD77    seq_world_size = dist.get_world_size(group)78    input_list = [79        t.contiguous()80        for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)81    ]82    output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]83    comm = dist.all_to_all(output_list, input_list, group=group, async_op=async_op)84    if async_op:85 86        def wait():87            comm.wait()88            return torch.cat(output_list, dim=gather_dim).contiguous()89 90        return wait91    return torch.cat(output_list, dim=gather_dim).contiguous()92 93 94def _all_to_all_tensor(95    x: Tensor,96    scatter_dim: int,97    gather_dim: int,98    group: dist.ProcessGroup,99    async_op: bool = False,100):101    if scatter_dim <= 1 and gather_dim <= 1:102        return _all_to_all_single(x, scatter_dim, gather_dim, group, async_op)103    return _all_to_all(x, scatter_dim, gather_dim, group, async_op)104 105 106class _SeqAllToAll(torch.autograd.Function):107    @staticmethod108    def forward(109        ctx: Any,110        group: dist.ProcessGroup,111        local_input: Tensor,112        scatter_dim: int,113        gather_dim: int,114        async_op: bool,115    ) -> Tensor:116        ctx.group = group117        ctx.scatter_dim = scatter_dim118        ctx.gather_dim = gather_dim119        ctx.async_op = async_op120        return _all_to_all_tensor(local_input, scatter_dim, gather_dim, group, async_op)121 122    @staticmethod123    def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None, None]:124        if ctx.async_op:125            input_t = torch.cat(grad_output[1:], dim=ctx.gather_dim).contiguous()126        else:127            input_t = grad_output[0]128        return (129            None,130            _all_to_all_tensor(131                input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False132            ),133            None,134            None,135            None,136        )137 138 139def gather_heads_scatter_seq(140    x: Tensor, head_dim: int, seq_dim: int, group: Optional[ProcessGroup] = None141) -> Tensor:142    group = group or dist.group.WORLD143    if not group:144        return x145    dim_size = x.size(seq_dim)146    sp_world = dist.get_world_size(group)147    if dim_size % sp_world != 0:148        padding_size = sp_world - (dim_size % sp_world)149        x = _pad_tensor(x, seq_dim, padding_size)150    return _SeqAllToAll.apply(group, x, seq_dim, head_dim, False)151 152 153def gather_seq_scatter_heads(154    x: Tensor,155    seq_dim: int,156    head_dim: int,157    unpadded_dim_size: int = 0,158    async_op: bool = False,159    group: Optional[ProcessGroup] = None,160) -> Tensor:161    group = group or dist.group.WORLD162    if not group:163        return x164    sp_world = dist.get_world_size(group)165    if async_op:166        return _SeqAllToAll.apply(group, x, head_dim, seq_dim, async_op)167    x = _SeqAllToAll.apply(group, x, head_dim, seq_dim, async_op)168    if unpadded_dim_size and unpadded_dim_size % sp_world != 0:169        padding_size = x.size(seq_dim) - unpadded_dim_size170        x = _unpad_tensor(x, seq_dim, padding_size)171    return x172 173 174def _local_attention(175    q: torch.Tensor,176    k: torch.Tensor,177    v: torch.Tensor,178    scale: float,179    causal: bool = False,180) -> torch.Tensor:181    scores = torch.matmul(q, k.transpose(-2, -1)) * scale182    if causal and q.size(1) > 1:183        S = scores.size(-1)184        causal_mask = torch.triu(185            torch.ones(S, S, device=scores.device, dtype=torch.bool),186            diagonal=1,187        )188        scores = scores.masked_fill(causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))189    attn = F.softmax(scores, dim=-1)190    return torch.matmul(attn, v)191 192 193def solution(194    hidden_states: torch.Tensor,195    w_qkv: torch.Tensor,196    w_o: torch.Tensor,197    group: Optional[dist.ProcessGroup] = None,198    num_heads: int = 8,199    causal: bool = False,200) -> torch.Tensor:201    group = group or dist.group.WORLD202    world_size = dist.get_world_size(group)203    if world_size == 1:204        B, S_local, H = hidden_states.shape205        head_dim = H // num_heads206        qkv = F.linear(hidden_states, w_qkv)207        qkv = qkv.view(B, S_local, 3, num_heads, head_dim)208        q, k, v = qkv.unbind(2)209        scale = head_dim**-0.5210        attn_out = _local_attention(q, k, v, scale, causal=causal)211        out = attn_out.reshape(B, S_local, -1)212        return F.linear(out, w_o)213 214    B, S_local, H = hidden_states.shape215    head_dim = (w_qkv.shape[0] // 3) // num_heads216    assert (w_qkv.shape[0] // 3) == num_heads * head_dim217    assert num_heads % world_size == 0, "num_heads must be divisible by world_size"218 219    qkv = F.linear(hidden_states, w_qkv)220    qkv = qkv.view(B, S_local, 3, num_heads, head_dim)221    q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2]222 223    q = gather_seq_scatter_heads(q, seq_dim=1, head_dim=2, group=group)224    kv = torch.stack([k, v], dim=3)225    kv = kv.reshape(B, S_local, 2 * num_heads, head_dim)226    kv = gather_seq_scatter_heads(kv, seq_dim=1, head_dim=2, group=group)227    kv = kv.reshape(B, kv.size(1), num_heads // world_size, 2, head_dim)228    k = kv[:, :, :, 0, :]229    v = kv[:, :, :, 1, :]230 231    scale = head_dim**-0.5232    attn_out = _local_attention(q, k, v, scale, causal=causal)233 234    attn_out = gather_heads_scatter_seq(attn_out, seq_dim=1, head_dim=2, group=group)235 236    out = attn_out.reshape(B, attn_out.size(1), -1)237    return F.linear(out, w_o)238