willychan21/ParallelKernelBench_Problems huggingface dataset:
from __future__ import annotations

import torch
import torch.distributed as dist
from torch import Tensor

_FP8_E4M3_MAX = 448.0


@torch.no_grad()
def _update_amax_history(amax_history: Tensor, cur_abs_max: Tensor) -> Tensor:
    out = torch.roll(amax_history, shifts=-1, dims=0)
    out[-1] = cur_abs_max.to(dtype=out.dtype)
    return out


@torch.no_grad()
def _fp8_round_trip_bf16(x: Tensor, scale: Tensor) -> Tensor:
    xf = x.float()
    qs = xf / scale
    q = qs.to(torch.float8_e4m3fn)
    return (q.float() * scale).to(dtype=x.dtype)


@torch.no_grad()
def solution(flat_param_shard: Tensor, amax_history: Tensor) -> tuple[Tensor, Tensor]:
    world_size = dist.get_world_size()
    p = flat_param_shard.numel()

    cur_abs_max = flat_param_shard.abs().max().to(torch.float32)
    updated_hist = _update_amax_history(amax_history, cur_abs_max)

    scale = updated_hist.max().clamp(min=1e-12).to(torch.float32) / _FP8_E4M3_MAX
    recon = _fp8_round_trip_bf16(flat_param_shard, scale)

    full = torch.empty(world_size * p, dtype=flat_param_shard.dtype, device=flat_param_shard.device)
    dist.all_gather_into_tensor(full, recon.contiguous())

    return full, updated_hist
from typing import Optional

import torch
import torch.distributed as dist
import torch.nn.functional as F


def _siglip_loss(
    image_features: torch.Tensor,
    text_features: torch.Tensor,
    logit_scale: float,
    logit_bias: float,
    negative_only: bool,
) -> torch.Tensor:
    batch = image_features.size(0)
    logits = logit_scale * image_features @ text_features.T + logit_bias

    if negative_only:
        return -F.logsigmoid(-logits).sum() / batch

    labels = -torch.ones((batch, batch), device=logits.device, dtype=logits.dtype)
    labels.diagonal().fill_(1)
    return -F.logsigmoid(labels * logits).sum() / batch


def _exchange(
    group: dist.ProcessGroup,
    recv_from: int,
    send_to: int,
    tensor: torch.Tensor,
) -> torch.Tensor:
    out = torch.empty_like(tensor)
    ops = [
        dist.P2POp(dist.isend, tensor, dist.get_global_rank(group, send_to), group=group),
        dist.P2POp(dist.irecv, out, dist.get_global_rank(group, recv_from), group=group),
    ]
    for req in dist.batch_isend_irecv(ops):
        req.wait()
    return out


def _exchange_bidir(
    group: dist.ProcessGroup,
    left: int,
    right: int,
    tensor_to_left: torch.Tensor,
    tensor_to_right: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    from_left = torch.empty_like(tensor_to_right)
    from_right = torch.empty_like(tensor_to_left)
    left_global = dist.get_global_rank(group, left)
    right_global = dist.get_global_rank(group, right)
    ops = [
        dist.P2POp(dist.isend, tensor_to_right, right_global, group=group),
        dist.P2POp(dist.isend, tensor_to_left, left_global, group=group),
        dist.P2POp(dist.irecv, from_right, right_global, group=group),
        dist.P2POp(dist.irecv, from_left, left_global, group=group),
    ]
    for req in dist.batch_isend_irecv(ops):
        req.wait()
    return from_right, from_left


@torch.no_grad()
def solution(
    image_features: torch.Tensor,
    text_features: torch.Tensor,
    logit_scale: float,
    logit_bias: float = 0.0,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)

    loss = _siglip_loss(image_features, text_features, logit_scale, logit_bias, False)

    left = (rank - 1) % world_size
    right = (rank + 1) % world_size
    text_to_left = text_features
    text_to_right = text_features
    num_bidir, remainder = divmod(world_size - 1, 2)

    for _ in range(num_bidir):
        from_right, from_left = _exchange_bidir(group, left, right, text_to_left, text_to_right)
        loss = loss + _siglip_loss(image_features, from_right, logit_scale, logit_bias, True)
        loss = loss + _siglip_loss(image_features, from_left, logit_scale, logit_bias, True)
        text_to_left, text_to_right = from_right, from_left

    if remainder:
        text_recv = _exchange(group, left, right, text_to_right)
        loss = loss + _siglip_loss(image_features, text_recv, logit_scale, logit_bias, True)

    return loss
from typing import Optional, Sequence

import torch
import torch.distributed as dist


def _all_to_all_transpose(
    tensor: torch.Tensor,
    split_dim: int,
    group: dist.ProcessGroup,
) -> list[torch.Tensor]:
    world_size = dist.get_world_size(group)
    chunk = tensor.shape[split_dim] // world_size
    send = [x.contiguous() for x in torch.split(tensor, chunk, dim=split_dim)]
    recv = [torch.empty_like(send[0]) for _ in range(world_size)]
    dist.all_to_all(recv, send, group=group)
    return recv


def _truncate(tensor: torch.Tensor, dim: int, size: int) -> torch.Tensor:
    slices = [slice(None)] * tensor.ndim
    slices[dim % tensor.ndim] = slice(0, size)
    return tensor[tuple(slices)].contiguous()


@torch.no_grad()
def solution(
    x: torch.Tensor,
    s: Sequence[int],
    dim: Sequence[int],
    norm: str = "ortho",
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    dim0, dim1 = int(dim[0]), int(dim[1])

    n0 = int(s[0]) if s[0] is not None else None
    n1 = int(s[1]) if s[1] is not None else None

    x1 = torch.fft.fft(x, n=n0, dim=dim0, norm=norm)

    x1_recv = _all_to_all_transpose(x1, split_dim=dim0, group=group)
    x1_tran = torch.cat(x1_recv, dim=dim1)

    x2 = torch.fft.fft(x1_tran, n=n1, dim=dim1, norm=norm)

    return _truncate(x2, dim1, x2.shape[dim1] // 2 + 1)
from typing import List, Tuple

import torch
import torch.distributed as dist


@torch.no_grad()
def solution(
    H: List[torch.Tensor],
    error: torch.Tensor,
    weights: List[torch.Tensor],
    P: List[torch.Tensor],
    kalman_lambda: float,
    kalman_nue: float = 0.9987,
) -> Tuple[List[torch.Tensor], List[torch.Tensor], torch.Tensor]:
    weights_num = len(weights)
    lam = torch.as_tensor(kalman_lambda, dtype=weights[0].dtype, device=weights[0].device)
    err = error.to(device=weights[0].device, dtype=weights[0].dtype)

    tmp = 0
    for i in range(weights_num):
        tmp = tmp + (lam + torch.matmul(torch.matmul(H[i].T, P[i]), H[i]))

    if dist.is_initialized():
        dist.all_reduce(tmp, op=dist.ReduceOp.SUM)

    A = 1 / tmp

    for i in range(weights_num):
        K = torch.matmul(P[i], H[i])
        weights[i] = weights[i] + A * err * K
        P[i] = (1 / lam) * (P[i] - A * torch.matmul(K, K.T))

    if dist.is_initialized():
        device = weights[0].device
        local_shape = [tensor.shape[0] for tensor in weights]

        shape_list = [
            torch.zeros_like(torch.empty(1), dtype=torch.float64, device=device)
            for _ in range(dist.get_world_size())
        ]
        dist.all_gather_object(shape_list, local_shape)

        weight_tensor = torch.cat([w.reshape(-1) for w in weights], dim=0).to(torch.float64)
        world_shape = [sum(inner_list) for inner_list in shape_list]

        weight_list = [
            torch.zeros(world_shape[i], dtype=torch.float64, device=device)
            for i in range(len(world_shape))
        ]
        dist.all_gather(weight_list, weight_tensor)

        result: List[torch.Tensor] = []
        for i in range(dist.get_world_size()):
            result = result + [
                t.reshape(-1, 1).to(dtype=weights[0].dtype)
                for t in torch.split(weight_list[i], shape_list[i])
            ]
        weights = result

    kalman_lambda_next = (
        torch.as_tensor(kalman_nue, dtype=lam.dtype, device=lam.device) * lam
        + 1
        - torch.as_tensor(kalman_nue, dtype=lam.dtype, device=lam.device)
    )

    return weights, P, kalman_lambda_next
from typing import List, Optional

import torch
import torch.distributed as dist


def _shift(chunks: List[torch.Tensor], group: dist.ProcessGroup) -> List[torch.Tensor]:
    cutoff = len(chunks) - dist.get_rank(group)
    return chunks[cutoff:] + chunks[:cutoff]


def _all_to_all(
    outputs: List[torch.Tensor],
    inputs: List[torch.Tensor],
    group: dist.ProcessGroup,
) -> None:
    outputs = _shift(list(outputs), group)
    inputs = _shift(list(inputs), group)
    if outputs and outputs[0].is_cuda:
        dist.all_to_all(outputs, inputs, group=group)
        return

    output_splits = [out.size(0) for out in outputs]
    input_splits = [inp.size(0) for inp in inputs]
    flat_out = torch.cat(outputs) if outputs else torch.empty(0)
    flat_in = torch.cat(inputs) if inputs else torch.empty(0)
    dist.all_to_all_single(
        flat_out,
        flat_in,
        output_split_sizes=output_splits,
        input_split_sizes=input_splits,
        group=group,
    )
    for out, temp in zip(outputs, flat_out.split(output_splits)):
        out.copy_(temp)


@torch.no_grad()
def solution(
    local_features: torch.Tensor,
    seed_inverse_ids: torch.Tensor,
    counts_sent: List[int],
    counts_received: List[int],
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    gathered = local_features[seed_inverse_ids]
    out = local_features.new_empty((sum(counts_sent),) + local_features.shape[1:])
    _all_to_all(
        list(torch.split(out, counts_sent)),
        list(torch.split(gathered, counts_received)),
        group,
    )
    return out
from typing import List, Optional

import torch
import torch.distributed as dist


def _shift(chunks: List[torch.Tensor], group: dist.ProcessGroup) -> List[torch.Tensor]:
    cutoff = len(chunks) - dist.get_rank(group)
    return chunks[cutoff:] + chunks[:cutoff]


def _all_to_all(
    outputs: List[torch.Tensor],
    inputs: List[torch.Tensor],
    group: dist.ProcessGroup,
) -> None:
    outputs = _shift(list(outputs), group)
    inputs = _shift(list(inputs), group)
    if outputs and outputs[0].is_cuda:
        dist.all_to_all(outputs, inputs, group=group)
        return

    output_splits = [out.size(0) for out in outputs]
    input_splits = [inp.size(0) for inp in inputs]
    flat_out = torch.cat(outputs) if outputs else torch.empty(0)
    flat_in = torch.cat(inputs) if inputs else torch.empty(0)
    dist.all_to_all_single(
        flat_out,
        flat_in,
        output_split_sizes=output_splits,
        input_split_sizes=input_splits,
        group=group,
    )
    for out, temp in zip(outputs, flat_out.split(output_splits)):
        out.copy_(temp)


@torch.no_grad()
def solution(
    grad_output: torch.Tensor,
    seed_inverse_ids: torch.Tensor,
    seed_size: int,
    counts_sent: List[int],
    counts_received: List[int],
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    out = grad_output.new_empty((sum(counts_received),) + grad_output.shape[1:])
    _all_to_all(
        list(torch.split(out, counts_received)),
        list(torch.split(grad_output, counts_sent)),
        group,
    )

    idx = torch.empty((2, out.shape[0]), dtype=torch.int64, device=grad_output.device)
    idx[0] = seed_inverse_ids
    idx[1] = torch.arange(out.shape[0], device=grad_output.device)
    coo = torch.sparse_coo_tensor(
        idx,
        torch.ones(idx.shape[1], dtype=grad_output.dtype, device=idx.device),
        size=(seed_size, idx.shape[1]),
    )
    return torch.sparse.mm(coo, out)
from typing import Optional, Tuple

import torch
import torch.distributed as dist


def _generate_permutation_remainder(
    idx: torch.Tensor,
    world_size: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    owner = (idx % world_size).long()
    send_splits = torch.bincount(owner, minlength=world_size)
    perm = torch.argsort(owner, stable=True).long()
    return perm, send_splits


@torch.no_grad()
def solution(
    idx: torch.Tensor,
    value: torch.Tensor,
    num_nodes: int,
    group: Optional[dist.ProcessGroup] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    if world_size == 1:
        return idx, value

    perm, send_splits = _generate_permutation_remainder(idx, world_size)

    recv_splits = torch.empty_like(send_splits)
    dist.all_to_all_single(recv_splits, send_splits, group=group)

    recv_splits = recv_splits.to("cpu", non_blocking=True)
    send_splits = send_splits.to("cpu", non_blocking=True)
    send_idx = idx[perm]
    send_value = value[perm]
    if idx.is_cuda:
        torch.cuda.current_stream().synchronize()

    recv_count = int(recv_splits.sum().item())
    recv_splits_list = recv_splits.tolist()
    send_splits_list = send_splits.tolist()

    recv_idx = torch.empty((recv_count,), dtype=idx.dtype, device=idx.device)
    dist.all_to_all_single(
        recv_idx,
        send_idx,
        output_split_sizes=recv_splits_list,
        input_split_sizes=send_splits_list,
        group=group,
    )

    recv_value = torch.empty(
        (recv_count, *value.shape[1:]),
        dtype=value.dtype,
        device=value.device,
    )
    dist.all_to_all_single(
        recv_value,
        send_value,
        output_split_sizes=recv_splits_list,
        input_split_sizes=send_splits_list,
        group=group,
    )

    return recv_idx, recv_value
from typing import Optional

import torch
import torch.distributed as dist


@torch.no_grad()
def solution(
    local_embedding_shard: torch.Tensor,
    input_node_ids: torch.Tensor,
    proj_matrix: torch.Tensor,
    num_total_nodes: int,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    rank = dist.get_rank(group)
    shard_size = (num_total_nodes + world_size - 1) // world_size
    embed_dim = local_embedding_shard.shape[1]
    num_queries = input_node_ids.shape[0]

    owner = (input_node_ids // shard_size).clamp(max=world_size - 1)
    sort_idx = torch.argsort(owner, stable=True)
    sorted_ids = input_node_ids[sort_idx]
    sorted_owner = owner[sort_idx]

    send_counts = torch.zeros(
        world_size, dtype=torch.long, device=input_node_ids.device
    )
    send_counts.scatter_add_(0, sorted_owner, torch.ones_like(sorted_owner))
    recv_counts = torch.empty_like(send_counts)
    dist.all_to_all_single(recv_counts, send_counts, group=group)

    send_splits = send_counts.to("cpu").tolist()
    recv_splits = recv_counts.to("cpu").tolist()
    recv_ids = torch.empty(
        int(recv_counts.sum().item()),
        dtype=input_node_ids.dtype,
        device=input_node_ids.device,
    )
    dist.all_to_all_single(
        recv_ids,
        sorted_ids,
        output_split_sizes=recv_splits,
        input_split_sizes=send_splits,
        group=group,
    )

    local_ids = (recv_ids - rank * shard_size).long()
    fetched = local_embedding_shard[local_ids]

    gathered = torch.empty(
        (num_queries, embed_dim), dtype=fetched.dtype, device=fetched.device
    )
    dist.all_to_all_single(
        gathered,
        fetched,
        output_split_sizes=send_splits,
        input_split_sizes=recv_splits,
        group=group,
    )
    emb = gathered[torch.argsort(sort_idx, stable=True)]

    return emb @ proj_matrix
from typing import Optional

import torch
import torch.distributed as dist


def _broadcast_data(
    rank: int,
    world_size: int,
    data: torch.Tensor,
    group: dist.ProcessGroup,
) -> torch.Tensor:
    if world_size == 1:
        return data

    sizes = torch.zeros(world_size, dtype=torch.long, device=data.device)
    sizes[rank] = data.shape[0]
    dist.all_reduce(sizes, op=dist.ReduceOp.SUM, group=group)

    send_splits = [data.shape[0] for _ in range(world_size)]
    recv_splits = sizes.to("cpu").tolist()
    send = data.repeat(*([world_size] + [1] * (data.ndim - 1))).contiguous()
    recv = torch.empty(
        (int(sizes.sum().item()), *data.shape[1:]),
        dtype=data.dtype,
        device=data.device,
    )
    dist.all_to_all_single(
        recv,
        send,
        output_split_sizes=recv_splits,
        input_split_sizes=send_splits,
        group=group,
    )
    return recv


def _calc_ranking(pos_score: torch.Tensor, neg_score: torch.Tensor) -> torch.Tensor:
    scores = torch.cat([pos_score.view(-1, 1), neg_score], dim=1)
    _, indices = torch.sort(torch.sigmoid(scores), dim=1, descending=True)
    return torch.nonzero(indices == 0)[:, 1].view(-1).detach() + 1


@torch.no_grad()
def solution(
    local_pos_scores: torch.Tensor,
    local_neg_scores: torch.Tensor,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    rank = dist.get_rank(group)
    world_size = dist.get_world_size(group)

    pos_scores = _broadcast_data(rank, world_size, local_pos_scores, group)
    neg_scores = _broadcast_data(rank, world_size, local_neg_scores, group)
    return _calc_ranking(pos_scores, neg_scores)
from typing import Optional, Tuple

import torch
import torch.distributed as dist
import torch.nn.functional as F


def _zigzag_get_overlapping_patches(
    data: torch.Tensor,
    seq_dim: int,
    overlap_size: int,
) -> Tuple[torch.Tensor, torch.Tensor]:
    shape = list(data.shape)
    shape[seq_dim : seq_dim + 1] = [2, data.shape[seq_dim] // 2]
    chunks = data.reshape(shape)

    order = list(range(chunks.dim()))
    order.insert(0, order.pop(seq_dim))
    chunks = chunks.permute(order)

    chunk_len = chunks.shape[seq_dim + 1]
    overlaps = chunks.narrow(seq_dim + 1, chunk_len - overlap_size, overlap_size)
    return overlaps[0], overlaps[1]


@torch.no_grad()
def solution(
    x: torch.Tensor,
    weight: torch.Tensor,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    group_ranks = dist.get_process_group_ranks(group)
    group_rank = dist.get_rank(group)
    group_world_size = len(group_ranks)

    batch, hidden, local_seq = x.shape
    chunk_len = local_seq // 2
    pad_size = weight.shape[-1] - 1
    chunk_a, chunk_b = _zigzag_get_overlapping_patches(x, 2, pad_size)

    ops = []
    recv_prev_a = None
    recv_next_b = None

    if group_rank > 0:
        recv_prev_a = torch.empty_like(chunk_a)
        ops.append(
            dist.P2POp(dist.irecv, recv_prev_a, group_ranks[group_rank - 1], group)
        )
    if group_rank < group_world_size - 1:
        ops.append(
            dist.P2POp(
                dist.isend,
                chunk_a.contiguous(),
                group_ranks[group_rank + 1],
                group,
            )
        )

    if group_rank < group_world_size - 1:
        recv_next_b = torch.empty_like(chunk_b)
        ops.append(
            dist.P2POp(dist.irecv, recv_next_b, group_ranks[group_rank + 1], group)
        )
    if group_rank > 0:
        ops.append(
            dist.P2POp(
                dist.isend,
                chunk_b.contiguous(),
                group_ranks[group_rank - 1],
                group,
            )
        )

    for request in dist.batch_isend_irecv(ops):
        request.wait()

    if recv_prev_a is None:
        recv_prev_a = torch.zeros_like(chunk_a)
    if recv_next_b is None:
        recv_next_b = chunk_a.clone().contiguous()

    # Move the two zigzag chunks into batch so both use the same grouped conv.
    x_chunks = x.reshape(batch, hidden, 2, chunk_len).permute(2, 0, 1, 3)
    x_chunks = x_chunks.reshape(2 * batch, hidden, chunk_len)
    padding = torch.cat([recv_prev_a, recv_next_b], dim=0)
    x_padded = torch.cat([padding, x_chunks], dim=-1)

    y = F.conv1d(x_padded, weight, bias=None, stride=1, padding=0, groups=hidden)
    return (
        y.reshape(2, batch, hidden, chunk_len)
        .permute(1, 2, 0, 3)
        .reshape(batch, hidden, local_seq)
    )
from typing import Optional, Tuple

import torch
import torch.distributed as dist


def _vocab_range(partition_vocab_size: int, rank: int) -> Tuple[int, int]:
    start = rank * partition_vocab_size
    return start, start + partition_vocab_size


@torch.no_grad()
def solution(
    vocab_parallel_logits: torch.Tensor,
    target: torch.Tensor,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    rank = dist.get_rank(group=group)

    logits_max = torch.max(vocab_parallel_logits, dim=-1).values
    dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=group)
    vocab_parallel_logits.sub_(logits_max.unsqueeze(dim=-1))

    partition_vocab_size = vocab_parallel_logits.shape[-1]
    vocab_start, vocab_end = _vocab_range(partition_vocab_size, rank)
    target_mask = (target < vocab_start) | (target >= vocab_end)
    masked_target = target - vocab_start
    masked_target = masked_target.masked_fill(target_mask, 0)

    logits_2d = vocab_parallel_logits.reshape(-1, partition_vocab_size)
    target_1d = masked_target.reshape(-1)
    row_ids = torch.arange(logits_2d.shape[0], device=logits_2d.device)
    predicted_logits = logits_2d[row_ids, target_1d].clone().reshape_as(target)
    predicted_logits = predicted_logits.masked_fill(target_mask, 0.0)
    dist.all_reduce(predicted_logits, op=dist.ReduceOp.SUM, group=group)

    exp_logits = torch.exp(vocab_parallel_logits)
    sum_exp_logits = exp_logits.sum(dim=-1)
    dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=group)

    return torch.log(sum_exp_logits) - predicted_logits
from typing import Optional

import torch
import torch.distributed as dist


@torch.no_grad()
def solution(
    teacher_output: torch.Tensor,
    teacher_temp: float,
    n_masked_patches_tensor: torch.Tensor,
    n_iterations: int = 3,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    q = torch.exp(teacher_output.float() / teacher_temp).T
    total_batch = n_masked_patches_tensor.to(device=q.device, dtype=q.dtype).clone()
    dist.all_reduce(total_batch, group=group)

    num_prototypes = q.shape[0]
    total_mass = q.sum()
    dist.all_reduce(total_mass, group=group)
    q /= total_mass

    for _ in range(n_iterations):
        row_sum = q.sum(dim=1, keepdim=True)
        dist.all_reduce(row_sum, group=group)
        q /= row_sum
        q /= num_prototypes

        q /= q.sum(dim=0, keepdim=True)
        q /= total_batch

    q *= total_batch
    return q.T.contiguous()