willychan21/ParallelKernelBench_Problems

1from typing import Optional2 3import torch4import torch.distributed as dist5import torch.nn.functional as F6 7 8def _a2a_sequence_to_heads(9    x: torch.Tensor,10    group: dist.ProcessGroup,11) -> torch.Tensor:12    world_size = dist.get_world_size(group=group)13    batch, local_seq, heads, dim = x.shape14    local_heads = heads // world_size15    send = (16        x.reshape(batch, local_seq, world_size, local_heads, dim)17        .permute(2, 1, 0, 3, 4)18        .contiguous()19    )20    recv = torch.empty_like(send)21    dist.all_to_all_single(recv, send, group=group)22    return (23        recv.permute(2, 0, 1, 3, 4)24        .reshape(batch, world_size * local_seq, local_heads, dim)25        .contiguous()26    )27 28 29def _a2a_heads_to_sequence(30    x: torch.Tensor,31    group: dist.ProcessGroup,32) -> torch.Tensor:33    world_size = dist.get_world_size(group=group)34    batch, seq_len, local_heads, dim = x.shape35    local_seq = seq_len // world_size36    send = (37        x.reshape(batch, world_size, local_seq, local_heads, dim)38        .permute(1, 2, 0, 3, 4)39        .contiguous()40    )41    recv = torch.empty_like(send)42    dist.all_to_all_single(recv, send, group=group)43    return (44        recv.permute(2, 1, 0, 3, 4)45        .reshape(batch, local_seq, world_size * local_heads, dim)46        .contiguous()47    )48 49 50def _gated_delta_recurrent(51    q: torch.Tensor,52    k: torch.Tensor,53    v: torch.Tensor,54    gate: torch.Tensor,55    beta: torch.Tensor,56    a_log: torch.Tensor,57    dt_bias: torch.Tensor,58) -> torch.Tensor:59    batch, seq_len, query_heads, key_dim = q.shape60    value_heads = v.shape[2]61    value_dim = v.shape[-1]62    out_dtype = q.dtype63    scale = float(key_dim) ** -0.564 65    assert value_heads % query_heads == 066    repeat = value_heads // query_heads67 68    q = F.normalize(q.float(), p=2, dim=-1, eps=1e-6)69    k = F.normalize(k.float(), p=2, dim=-1, eps=1e-6)70    q = q.repeat_interleave(repeat, dim=2) * scale71    k = k.repeat_interleave(repeat, dim=2)72 73    a_scale = a_log.float().exp().view(1, 1, value_heads)74    dt_bias = dt_bias.float().view(1, 1, value_heads)75    decay = -a_scale * F.softplus(gate.float() + dt_bias)76 77    q = q.permute(0, 2, 1, 3).contiguous()78    k = k.permute(0, 2, 1, 3).contiguous()79    v = v.float().permute(0, 2, 1, 3).contiguous()80    decay = decay.exp().permute(0, 2, 1).contiguous()81    beta = beta.float().permute(0, 2, 1).contiguous()82 83    batch_heads = batch * value_heads84    q = q.reshape(batch_heads, seq_len, key_dim)85    k = k.reshape(batch_heads, seq_len, key_dim)86    v = v.reshape(batch_heads, seq_len, value_dim)87    decay = decay.reshape(batch_heads, seq_len)88    beta = beta.reshape(batch_heads, seq_len)89 90    state = torch.zeros(91        batch_heads, key_dim, value_dim, dtype=torch.float32, device=q.device92    )93    output = torch.empty(94        batch_heads, seq_len, value_dim, dtype=torch.float32, device=q.device95    )96    for step in range(seq_len):97        q_t = q[:, step]98        k_t = k[:, step]99        v_t = v[:, step]100        state = state * decay[:, step].view(batch_heads, 1, 1)101        projected = torch.bmm(k_t.unsqueeze(1), state).squeeze(1)102        update = (v_t - projected) * beta[:, step].unsqueeze(-1)103        state = state + k_t.unsqueeze(-1) * update.unsqueeze(1)104        output[:, step] = torch.bmm(q_t.unsqueeze(1), state).squeeze(1)105 106    output = output.reshape(batch, value_heads, seq_len, value_dim)107    output = output.permute(0, 2, 1, 3).contiguous()108    return output.to(dtype=out_dtype)109 110 111@torch.no_grad()112def solution(113    q: torch.Tensor,114    k: torch.Tensor,115    v: torch.Tensor,116    gate: torch.Tensor,117    beta: torch.Tensor,118    a_log: torch.Tensor,119    dt_bias: torch.Tensor,120    group: Optional[dist.ProcessGroup] = None,121) -> torch.Tensor:122    group = group or dist.group.WORLD123    q_head = _a2a_sequence_to_heads(q, group)124    k_head = _a2a_sequence_to_heads(k, group)125    v_head = _a2a_sequence_to_heads(v, group)126    gate_head = _a2a_sequence_to_heads(gate.unsqueeze(-1), group).squeeze(-1)127    beta_head = _a2a_sequence_to_heads(beta.unsqueeze(-1), group).squeeze(-1)128 129    out = _gated_delta_recurrent(130        q_head,131        k_head,132        v_head,133        gate_head,134        beta_head,135        a_log,136        dt_bias,137    )138    return _a2a_heads_to_sequence(out, group)