togethercomputer/ParallelKernelBench_Problems huggingface dataset:
# 1_allreduce
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
# 2_allgather
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
# 3_broadcast
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
# 4_reduce
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
# 5_scatter
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
# 6_gather
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
# 7_reducescatter
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
# 8_alltoall
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
# 9_layernorm_backward
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
# 10_embedding_lookup
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
# 11_allgather_gemm_AT
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)
# 12_allgather_gemm
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)
# 13_gemm_allreduce
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
# 14_gemm_allgather
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
# 15_combined_sharded_gemms
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
# 16_gemm_reducescatter
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
# 17_rope_allgather
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
# 18_tp_rms_norm
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)
# 19_blocked_fp8_quantize
import torch
import torch.distributed as dist
import triton
import triton.language as tl
from typing import Tuple

@triton.jit
def block_fp8_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)

    x = tl.load(x_ptr + offs).to(tl.float32)
    
    # FP8 E4M3 max value is 448.0
    s = tl.max(tl.abs(x)) / 448.0
    
    # Prevent division by zero if all elements in the block are 0
    s_safe = tl.where(s == 0.0, 1.0, s)
    
    y = (x / s_safe).to(y_ptr.dtype.element_ty)
    
    tl.store(y_ptr + offs, y)
    tl.store(s_ptr + pid, s)

def solution(local_tensor: torch.Tensor, block_size: int = 128) -> Tuple[torch.Tensor, torch.Tensor]:
    assert local_tensor.size(-1) % block_size == 0, "Last dimension must be divisible by block_size"
    
    y_local = torch.empty_like(local_tensor, dtype=torch.float8_e4m3fn)
    s_local = local_tensor.new_empty(
        *local_tensor.size()[:-1], local_tensor.size(-1) // block_size, dtype=torch.float32
    )
    
    grid = (triton.cdiv(local_tensor.numel(), block_size),)
    block_fp8_quant_kernel[grid](local_tensor, y_local, s_local, BLOCK_SIZE=block_size)
    
    if dist.is_initialized():
        world_size = dist.get_world_size()
        
        y_local_u8 = y_local.view(torch.uint8)
        y_gather_u8 = [torch.empty_like(y_local_u8) for _ in range(world_size)]
        dist.all_gather(y_gather_u8, y_local_u8)
        
        y_global_u8 = torch.cat(y_gather_u8, dim=0)
        y_global = y_global_u8.view(torch.float8_e4m3fn)
        
        s_gather = [torch.empty_like(s_local) for _ in range(world_size)]
        dist.all_gather(s_gather, s_local)
        
        s_global = torch.cat(s_gather, dim=0)
    else:
        y_global = y_local
        s_global = s_local
        
    return y_global, s_global
# 20_blocked_fp8_dequantize
import torch
import torch.distributed as dist
import triton
import triton.language as tl

@triton.jit
def block_fp8_dequant_kernel(y_ptr, s_ptr, x_ptr, num_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(axis=0)
    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offs < num_elements

    s = tl.load(s_ptr + pid)

    y = tl.load(y_ptr + offs, mask=mask).to(tl.float32)

    tl.store(x_ptr + offs, y * s, mask=mask)


def solution(
    local_y: torch.Tensor,
    local_s: torch.Tensor,
    block_size: int = 128,
) -> torch.Tensor:
    world_size = dist.get_world_size()

    chunk_shape = local_y.shape[1:]
    chunk_numel = local_y.numel() // world_size
    num_elements = local_y.numel()
    assert chunk_numel % block_size == 0, (
        f"Chunk size {chunk_numel} must be divisible by block_size ({block_size})"
    )

    y_flat = local_y.view(-1)
    s_flat = local_s.view(-1)
    x_flat = torch.empty(num_elements, device=local_y.device, dtype=torch.float32)

    if num_elements > 0:
        grid = (triton.cdiv(num_elements, block_size),)
        block_fp8_dequant_kernel[grid](
            y_flat, s_flat, x_flat, num_elements, BLOCK_SIZE=block_size
        )

    x = x_flat.view(world_size, *chunk_shape)
    out = torch.empty_like(x)
    dist.all_to_all_single(out, x)

    return out
# 21_clip_grad_norm_no_ep
from typing import List, Optional

import torch
import torch.distributed as dist


def _local_pth_sum(grad_tensors: List[torch.Tensor], p: float) -> torch.Tensor:
    dev = None
    acc = None
    for g in grad_tensors:
        if g is None:
            continue
        g_local = g
        if dev is None:
            dev = g_local.device
            acc = torch.tensor(0.0, device=dev, dtype=torch.float32)
        gn = torch.norm(g_local.detach().to(torch.float32), p=p)
        acc = acc + (gn ** p)
    if acc is None:
        acc = torch.tensor(0.0, device=next((t.device for t in grad_tensors if t is not None), torch.device("cuda", 0)), dtype=torch.float32)
    return acc


def _fsdp2_reduce_group(
    grad_tensors: List[torch.Tensor],
    norm_type: float,
    reduce_groups: List[tuple],
) -> torch.Tensor:
    p = float(norm_type)
    val = _local_pth_sum(grad_tensors, p)
    for _, group in reduce_groups:
        if group is not None:
            dist.all_reduce(val, op=dist.ReduceOp.SUM, group=group)
    return val


def solution(
    grad_tensors: List[torch.Tensor],
    max_norm: float,
    norm_type: float = 2.0,
    fsdp_group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    reduce_groups = [("fsdp", fsdp_group)]
    total_p = _fsdp2_reduce_group(grad_tensors, norm_type, reduce_groups)
    total_norm = total_p ** (1.0 / float(norm_type))

    if total_norm > max_norm:
        coef = (max_norm / total_norm)
        for t in grad_tensors:
            if t is not None:
                t.mul_(coef.to(t.device))

    return total_norm
# 22_clip_grad_norm_ep
from typing import List, Optional

import torch
import torch.distributed as dist


def _local_pth_sum(grad_tensors: List[torch.Tensor], p: float) -> torch.Tensor:
    dev = None
    acc = None
    for g in grad_tensors:
        if g is None:
            continue
        g_local = g
        if dev is None:
            dev = g_local.device
            acc = torch.tensor(0.0, device=dev, dtype=torch.float32)
        gn = torch.norm(g_local.detach().to(torch.float32), p=p)
        acc = acc + (gn ** p)
    if acc is None:
        acc = torch.tensor(
            0.0,
            device=next((t.device for t in grad_tensors if t is not None), torch.device("cuda", 0)),
            dtype=torch.float32,
        )
    return acc


def _fsdp2_reduce_group(
    grad_tensors: List[torch.Tensor],
    norm_type: float,
    reduce_groups: List[tuple],
) -> torch.Tensor:
    p = float(norm_type)
    val = _local_pth_sum(grad_tensors, p)
    for _, group in reduce_groups:
        if group is not None:
            dist.all_reduce(val, op=dist.ReduceOp.SUM, group=group)
    return val


def solution(
    non_ep_grad_tensors: List[torch.Tensor],
    ep_grad_tensors: List[torch.Tensor],
    max_norm: float,
    norm_type: float = 2.0,
    ep_size: int = 1,
    fsdp_group: Optional[dist.ProcessGroup] = None,
    ep_fsdp_group: Optional[dist.ProcessGroup] = None,
    ep_group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    if ep_size > 1 and ep_grad_tensors:
        scale = 1.0 / float(ep_size)
        for t in ep_grad_tensors:
            if t is not None:
                t.detach().mul_(scale)

    non_ep_total = _fsdp2_reduce_group(
        non_ep_grad_tensors,
        norm_type=norm_type,
        reduce_groups=[("fsdp", fsdp_group)],
    )

    ep_total = _fsdp2_reduce_group(
        ep_grad_tensors,
        norm_type=norm_type,
        reduce_groups=[("ep_fsdp", ep_fsdp_group), ("ep", ep_group)],
    )

    total_norm = (non_ep_total + ep_total) ** (1.0 / float(norm_type))

    if total_norm > max_norm:
        coef = (max_norm / total_norm)
        for t in non_ep_grad_tensors:
            if t is not None:
                t.mul_(coef.to(t.device))
        for t in ep_grad_tensors:
            if t is not None:
                t.mul_(coef.to(t.device))

    return total_norm
# 23_grad_acc_loss
import torch
import torch.distributed as dist
from typing import Tuple, Optional


def forward(
    loss: torch.Tensor,
    local_valid_tokens: torch.Tensor,
    global_valid_tokens: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
    if local_valid_tokens.item() == 0:
        loss = torch.nan_to_num(loss)

    loss_sum = loss * local_valid_tokens
    dist.all_reduce(loss_sum, op=dist.ReduceOp.SUM)

    normalized_loss = loss_sum / global_valid_tokens
    return normalized_loss, loss_sum


def backward(
    local_valid_tokens: torch.Tensor,
    global_valid_tokens: torch.Tensor,
    grad_normalized_loss: torch.Tensor,
    grad_loss_sum: Optional[torch.Tensor],
) -> torch.Tensor:
    grad_from_normalized = grad_normalized_loss * local_valid_tokens / global_valid_tokens

    if grad_loss_sum is not None:
        grad_from_sum = grad_loss_sum * local_valid_tokens
    else:
        grad_from_sum = torch.zeros_like(grad_normalized_loss, device=grad_normalized_loss.device)

    return grad_from_normalized + grad_from_sum


def solution(
    loss: torch.Tensor,
    local_valid_tokens: torch.Tensor,
    global_valid_tokens: torch.Tensor,
    grad_normalized_loss: torch.Tensor,
    grad_loss_sum: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    normalized_loss, loss_sum = forward(loss, local_valid_tokens, global_valid_tokens)

    grad_loss = backward(
        local_valid_tokens,
        global_valid_tokens,
        grad_normalized_loss,
        grad_loss_sum,
    )

    return normalized_loss, loss_sum, grad_loss
# 24_load_balancing_loss_fn
import torch
import torch.distributed as dist
from typing import Union, Tuple, Optional

def solution(
    gate_logits: Union[torch.Tensor, Tuple[torch.Tensor, ...]],
    num_experts: int,
    top_k: int = 2,
    attention_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    if isinstance(gate_logits, (tuple, list)):
        compute_device = gate_logits[0].device
        concatenated_gate_logits = torch.cat(
            [layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0
        )
    else:
        compute_device = gate_logits.device
        concatenated_gate_logits = gate_logits

    routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
    _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
    
    expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
    
    if attention_mask is None:
        tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
        router_prob_per_expert = torch.mean(routing_weights, dim=0)
    else:
        batch_size, sequence_length = attention_mask.shape
        num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
        
        expert_attention_mask = (
            attention_mask[None, :, :, None, None]
            .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
            .reshape(-1, top_k, num_experts)
            .to(compute_device)
        )
        tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(expert_attention_mask, dim=0)
        
        router_per_expert_attention_mask = (
            attention_mask[None, :, :, None]
            .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
            .reshape(-1, num_experts)
            .to(compute_device)
        )
        router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(router_per_expert_attention_mask, dim=0)

    overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
    overall_loss = overall_loss * num_experts

    if dist.is_available() and dist.is_initialized():
        dist.all_reduce(overall_loss, op=dist.ReduceOp.SUM)
        overall_loss = overall_loss / dist.get_world_size()

    return overall_loss