willychan21/ParallelKernelBench_Problems

1from typing import Optional2 3import torch4import torch.distributed as dist5import torch.nn.functional as F6 7 8def _all_gather_sequence(x: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:9    world_size = dist.get_world_size(group=group)10    batch, local_seq = x.shape[:2]11    out = torch.empty(12        batch,13        world_size * local_seq,14        *x.shape[2:],15        dtype=x.dtype,16        device=x.device,17    )18    send = x.transpose(0, 1).contiguous()19    recv = torch.empty(20        world_size * local_seq,21        batch,22        *x.shape[2:],23        dtype=x.dtype,24        device=x.device,25    )26    dist.all_gather_into_tensor(recv, send, group=group)27    return recv.transpose(0, 1).contiguous().view_as(out)28 29 30def _kda_forward(31    q: torch.Tensor,32    k: torch.Tensor,33    v: torch.Tensor,34    g: torch.Tensor,35    beta: torch.Tensor,36    a_log: torch.Tensor,37    dt_bias: torch.Tensor,38    lower_bound: float,39) -> torch.Tensor:40    batch, seq_len, heads, key_dim = q.shape41    value_dim = v.shape[-1]42    out_dtype = q.dtype43 44    dt_bias = dt_bias.float().reshape(heads, key_dim)45    a_scale = a_log.float().exp().view(1, 1, heads, 1)46    decay = torch.exp(lower_bound * torch.sigmoid(a_scale * (g.float() + dt_bias)))47    beta = beta.float().sigmoid()48 49    scale = float(key_dim) ** -0.550    q_float = F.normalize(q.float(), p=2, dim=-1) * scale51    k_float = F.normalize(k.float(), p=2, dim=-1)52    v_float = v.float()53 54    q_float = q_float.permute(0, 2, 1, 3).contiguous()55    k_float = k_float.permute(0, 2, 1, 3).contiguous()56    v_float = v_float.permute(0, 2, 1, 3).contiguous()57    decay = decay.permute(0, 2, 1, 3).contiguous()58    beta = beta.permute(0, 2, 1).contiguous()59 60    batch_heads = batch * heads61    q_float = q_float.reshape(batch_heads, seq_len, key_dim)62    k_float = k_float.reshape(batch_heads, seq_len, key_dim)63    v_float = v_float.reshape(batch_heads, seq_len, value_dim)64    decay = decay.reshape(batch_heads, seq_len, key_dim)65    beta = beta.reshape(batch_heads, seq_len)66 67    state = torch.zeros(68        batch_heads, key_dim, value_dim, dtype=torch.float32, device=q.device69    )70    output = torch.empty(71        batch_heads, seq_len, value_dim, dtype=torch.float32, device=q.device72    )73    for step in range(seq_len):74        q_t = q_float[:, step]75        k_t = k_float[:, step]76        v_t = v_float[:, step]77        state = decay[:, step].unsqueeze(-1) * state78        projected = torch.bmm(k_t.unsqueeze(1), state).squeeze(1)79        update = (v_t - projected) * beta[:, step].unsqueeze(-1)80        state = state + k_t.unsqueeze(-1) * update.unsqueeze(1)81        output[:, step] = torch.bmm(q_t.unsqueeze(1), state).squeeze(1)82 83    output = output.reshape(batch, heads, seq_len, value_dim)84    output = output.permute(0, 2, 1, 3).contiguous()85    return output.to(dtype=out_dtype)86 87 88@torch.no_grad()89def solution(90    q: torch.Tensor,91    k: torch.Tensor,92    v: torch.Tensor,93    g: torch.Tensor,94    beta: torch.Tensor,95    a_log: torch.Tensor,96    dt_bias: torch.Tensor,97    cp_group: Optional[dist.ProcessGroup] = None,98    tp_group: Optional[dist.ProcessGroup] = None,99) -> torch.Tensor:100    cp_group = cp_group or dist.group.WORLD101    cp_size = dist.get_world_size(group=cp_group)102    cp_rank = dist.get_rank(group=cp_group)103 104    if cp_size > 1:105        q_full = _all_gather_sequence(q, cp_group)106        k_full = _all_gather_sequence(k, cp_group)107        v_full = _all_gather_sequence(v, cp_group)108        g_full = _all_gather_sequence(g, cp_group)109        beta_full = _all_gather_sequence(beta, cp_group)110    else:111        q_full, k_full, v_full, g_full, beta_full = q, k, v, g, beta112 113    out = _kda_forward(114        q_full,115        k_full,116        v_full,117        g_full,118        beta_full,119        a_log,120        dt_bias,121        lower_bound=-5.0,122    )123    if tp_group is not None and dist.get_world_size(group=tp_group) > 1:124        dist.all_reduce(out, op=dist.ReduceOp.SUM, group=tp_group)125 126    if cp_size == 1:127        return out128    local_seq = q.shape[1]129    start = cp_rank * local_seq130    return out[:, start : start + local_seq].contiguous()