willychan21/ParallelKernelBench_Problems huggingface dataset:
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(tensor: torch.Tensor) -> torch.Tensor:
    out = tensor.clone()
    dist.all_reduce(out, op=dist.ReduceOp.SUM)
    return out
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(tensor: torch.Tensor) -> torch.Tensor:
    world_size = dist.get_world_size()
    out = tensor.new_empty((world_size,) + tensor.shape)
    dist.all_gather_into_tensor(out, tensor)
    return out
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(tensor: torch.Tensor, src: int = 0) -> torch.Tensor:
    out = tensor.clone()
    dist.broadcast(out, src=src)
    return out
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(tensor: torch.Tensor, dst: int = 0) -> torch.Tensor:
    out = tensor.clone()
    dist.reduce(out, dst=dst, op=dist.ReduceOp.SUM)
    return out
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(tensor: torch.Tensor, src: int = 0) -> torch.Tensor:
    rank = dist.get_rank()

    if rank == src:
        world_size = dist.get_world_size()
        scatter_list = [chunk.squeeze(0).contiguous() for chunk in tensor.chunk(world_size, dim=0)]
        out = torch.empty_like(scatter_list[0])
    else:
        scatter_list = None
        out = tensor.clone()

    dist.scatter(out, scatter_list=scatter_list, src=src)
    return out
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(tensor: torch.Tensor, dst: int = 0) -> torch.Tensor:
    rank = dist.get_rank()

    if rank == dst:
        world_size = dist.get_world_size()
        gather_list = [torch.empty_like(tensor) for _ in range(world_size)]
    else:
        gather_list = None

    dist.gather(tensor, gather_list=gather_list, dst=dst)

    if rank == dst:
        return torch.stack(gather_list, dim=0)
    return tensor
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(tensor: torch.Tensor) -> torch.Tensor:
    world_size = dist.get_world_size()
    chunk_size = tensor.shape[0] // world_size
    out = tensor.new_empty((chunk_size,) + tensor.shape[1:])
    dist.reduce_scatter_tensor(out, tensor, op=dist.ReduceOp.SUM)
    return out
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(tensor: torch.Tensor) -> torch.Tensor:
    out = torch.empty_like(tensor)
    dist.all_to_all_single(out, tensor)
    return out
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(X_hat: torch.Tensor, dY: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    d_beta = dY.sum(dim=0)
    d_gamma = (dY * X_hat).sum(dim=0)
    dist.all_reduce(d_beta, op=dist.ReduceOp.SUM)
    dist.all_reduce(d_gamma, op=dist.ReduceOp.SUM)
    return d_gamma, d_beta
import torch
import torch.distributed as dist

@torch.no_grad()
def solution(
    indices: torch.Tensor,
    local_shard: torch.Tensor,
) -> torch.Tensor:
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    shard_size = local_shard.shape[0]
    embed_dim = local_shard.shape[1]
    
    indices = indices.contiguous().to(torch.cuda.current_device())
    
    target_ranks = indices // shard_size
    
    send_indices_list = [indices[target_ranks == r] for r in range(world_size)]
    send_counts = torch.tensor([len(idx) for idx in send_indices_list], dtype=torch.long, device='cuda')
    
    recv_counts = torch.zeros(world_size, dtype=torch.long, device='cuda')
    dist.all_to_all_single(recv_counts, send_counts)
    
    non_empty_lists = [idx_list for idx_list in send_indices_list if len(idx_list) > 0]
    if non_empty_lists:
        flat_send_indices = torch.cat(non_empty_lists)
    else:
        flat_send_indices = torch.empty(0, dtype=torch.long, device='cuda')
    
    total_recv = recv_counts.sum().item()
    total_send = send_counts.sum().item()
    received_indices = torch.empty(total_recv, dtype=torch.long, device='cuda')
    
    if total_recv > 0 or total_send > 0:
        dist.all_to_all_single(
            received_indices, 
            flat_send_indices,
            output_split_sizes=recv_counts.tolist(),
            input_split_sizes=send_counts.tolist()
        )
    
    if total_recv > 0:
        local_lookup_indices = received_indices - (rank * shard_size)
        local_lookup_indices = torch.clamp(local_lookup_indices, 0, shard_size - 1)
        retrieved_vectors = local_shard[local_lookup_indices]
    else:
        retrieved_vectors = torch.empty((0, embed_dim), dtype=local_shard.dtype, device='cuda')
    
    output_vectors = torch.empty((len(indices), embed_dim), dtype=local_shard.dtype, device='cuda')
    
    input_split_sizes = recv_counts.cpu().tolist()
    output_split_sizes = send_counts.cpu().tolist()
    
    if len(indices) > 0 or retrieved_vectors.numel() > 0:
        dist.all_to_all_single(
            output_vectors,
            retrieved_vectors,
            output_split_sizes=output_split_sizes,
            input_split_sizes=input_split_sizes
        )
    
    return output_vectors
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(A_local: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    world_size = dist.get_world_size()
    M, K_local = A_local.shape
    K = world_size * K_local

    A_local_t = A_local.transpose(0, 1).contiguous()
    A_t_buf = A_local_t.new_empty((world_size, K_local, M))
    dist.all_gather_into_tensor(A_t_buf, A_local_t)
    A_global_t = A_t_buf.reshape(K, M)

    C_t = torch.matmul(B.transpose(0, 1), A_global_t)
    return C_t.transpose(0, 1)
import torch
import torch.distributed as dist


@torch.no_grad()
def solution(A_local: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
    world_size = dist.get_world_size()
    A_gathered = [torch.empty_like(A_local) for _ in range(world_size)]
    dist.all_gather(A_gathered, A_local)
    A_global = torch.cat(A_gathered, dim=1)
    return torch.matmul(A_global, B)
import torch
import torch.distributed as dist

@torch.no_grad()
def solution(
    A_local: torch.Tensor,
    B_local: torch.Tensor,
) -> torch.Tensor:
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    
    M, K = A_local.shape
    K_B, N = B_local.shape
    
    A_local = A_local.contiguous()
    B_local = B_local.contiguous()
    C_local = torch.matmul(A_local, B_local)
    
    C = C_local.clone()
    dist.all_reduce(C, op=dist.ReduceOp.SUM)
    
    return C
import torch
import torch.distributed as dist

@torch.no_grad()
def solution(
    A: torch.Tensor,
    B: torch.Tensor,
) -> torch.Tensor:
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    
    M, K = A.shape
    K_B, N_local = B.shape
    
    A = A.contiguous()
    B = B.contiguous()
    C_local = torch.matmul(A, B)
    
    C_gathered = [torch.zeros_like(C_local) for _ in range(world_size)]
    dist.all_gather(C_gathered, C_local)
    
    C = torch.cat(C_gathered, dim=1)
    
    return C
import torch
import torch.distributed as dist
import torch.nn.functional as F


@torch.no_grad()
def solution(
    x_local: torch.Tensor,
    W1: torch.Tensor,
    W2: torch.Tensor,
) -> torch.Tensor:
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    M = x_local.shape[0]
    M_local = M // world_size

    x_local = x_local.contiguous()
    shards = [torch.empty_like(x_local) for _ in range(world_size)]
    dist.all_gather(shards, x_local)
    x_full = torch.cat(shards, dim=1)

    a = F.silu(torch.matmul(x_full, W1))
    a_loc = a[rank * M_local : (rank + 1) * M_local].contiguous()
    block = torch.matmul(a_loc, W2)

    H = block.shape[1]
    buf = block.new_zeros((M, H))
    buf[rank * M_local : (rank + 1) * M_local].copy_(block)

    y_local = block.new_empty((M_local, H))
    dist.reduce_scatter_tensor(y_local, buf, op=dist.ReduceOp.SUM)
    return y_local
import torch
import torch.distributed as dist

@torch.no_grad()
def solution(
    A_local: torch.Tensor,
    B_local: torch.Tensor,
) -> torch.Tensor:
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    
    M, K_local = A_local.shape
    K_B, N = B_local.shape
    
    M_local = M // world_size
    
    A_local = A_local.contiguous()
    B_local = B_local.contiguous()
    C_partial = torch.matmul(A_local, B_local)
    
    C_local = torch.empty((M_local, N), dtype=C_partial.dtype, device=C_partial.device)
    dist.reduce_scatter_tensor(C_local, C_partial, op=dist.ReduceOp.SUM)
    
    return C_local
import torch
import torch.distributed as dist
from typing import Tuple

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    half_dim = x.shape[-1] // 2
    x1, x2 = x[..., :half_dim], x[..., half_dim:]
    return torch.cat((-x2, x1), dim=-1)

def solution(
    q_local: torch.Tensor, 
    k_local: torch.Tensor, 
    cos_local: torch.Tensor, 
    sin_local: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]:
    # Reshape cos and sin to broadcast with q and k's head dimension (dim=2)
    cos = cos_local.unsqueeze(2)
    sin = sin_local.unsqueeze(2)
    
    q_embed_local = (q_local * cos) + (rotate_half(q_local) * sin)
    k_embed_local = (k_local * cos) + (rotate_half(k_local) * sin)
    
    if not dist.is_initialized():
        return q_embed_local, k_embed_local
        
    world_size = dist.get_world_size()
    
    q_gather_list = [torch.empty_like(q_embed_local) for _ in range(world_size)]
    k_gather_list = [torch.empty_like(k_embed_local) for _ in range(world_size)]
    
    dist.all_gather(q_gather_list, q_embed_local.contiguous())
    dist.all_gather(k_gather_list, k_embed_local.contiguous())
    
    q_embed_global = torch.cat(q_gather_list, dim=1)
    k_embed_global = torch.cat(k_gather_list, dim=1)
    
    return q_embed_global, k_embed_global
import torch
import torch.distributed as dist

def solution(local_hidden_states: torch.Tensor, local_weight: torch.Tensor, variance_epsilon: float) -> torch.Tensor:
    input_dtype = local_hidden_states.dtype
    # Upcast to float32 for stable variance calculation
    local_hidden_states = local_hidden_states.to(torch.float32)
    
    local_sum_squares = local_hidden_states.pow(2).sum(dim=-1, keepdim=True)
    
    dist.all_reduce(local_sum_squares, op=dist.ReduceOp.SUM)
    
    world_size = dist.get_world_size()
    global_hidden_size = local_hidden_states.shape[-1] * world_size
    variance = local_sum_squares / global_hidden_size
    
    local_hidden_states = local_hidden_states * torch.rsqrt(variance + variance_epsilon)
    
    return local_weight * local_hidden_states.to(input_dtype)