willychan21/ParallelKernelBench_Problems huggingface dataset:
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
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
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
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
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
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
import torch
import torch.nn.functional as F
import torch.distributed as dist
from typing import Tuple, Any

def solution(
    hidden_states: torch.Tensor,
    weight: torch.Tensor,
    labels: torch.Tensor,
    old_logprobs: torch.Tensor,
    advantages: torch.Tensor,
    ignore_index: int = -100,
) -> Tuple[torch.Tensor, Any, torch.Tensor, torch.Tensor, torch.Tensor]:
    logits = F.linear(hidden_states, weight)
    logits_flat = logits.view(-1, logits.size(-1))
    labels_flat = labels.view(-1)
    
    per_token_ce = F.cross_entropy(logits_flat, labels_flat, ignore_index=ignore_index, reduction='none')
    
    new_logprobs_flat = -per_token_ce.detach()
    old_logprobs_flat = old_logprobs.view(-1)
    advantages_flat = advantages.view(-1)
    
    valid_mask = (labels_flat != ignore_index)
    n_valid_local = valid_mask.sum().float()
    
    n_valid_global = n_valid_local.clone()
    dist.all_reduce(n_valid_global, op=dist.ReduceOp.SUM)
    n_valid_global_clamped = n_valid_global.clamp(min=1.0)
    
    delta = (new_logprobs_flat - old_logprobs_flat).masked_fill(~valid_mask, 0.0).clamp(min=-20.0, max=20.0)
    ratio = torch.exp(delta)
    
    per_token_pg = -(ratio * advantages_flat).masked_fill(~valid_mask, 0.0)
    local_pg_sum = per_token_pg.sum()
    
    global_pg_sum = local_pg_sum.clone()
    dist.all_reduce(global_pg_sum, op=dist.ReduceOp.SUM)
    true_pg = global_pg_sum / n_valid_global_clamped
    
    w = (ratio.detach() * advantages_flat).masked_fill(~valid_mask, 0.0)
    local_surrogate_sum = (w * per_token_ce).sum()
    
    surrogate = local_surrogate_sum / n_valid_global_clamped
    
    loss = true_pg.detach() + surrogate - surrogate.detach()
    
    ratio_valid = ratio.masked_fill(~valid_mask, 0.0)
    sum_ratio_local = ratio_valid.sum()
    dist.all_reduce(sum_ratio_local, op=dist.ReduceOp.SUM)
    ratio_mean = sum_ratio_local / n_valid_global_clamped
    
    ratio_for_min = ratio.masked_fill(~valid_mask, float('inf'))
    min_ratio_local = ratio_for_min.min() if n_valid_local > 0 else torch.tensor(float('inf'), device=ratio.device)
    dist.all_reduce(min_ratio_local, op=dist.ReduceOp.MIN)
    
    ratio_for_max = ratio.masked_fill(~valid_mask, float('-inf'))
    max_ratio_local = ratio_for_max.max() if n_valid_local > 0 else torch.tensor(float('-inf'), device=ratio.device)
    dist.all_reduce(max_ratio_local, op=dist.ReduceOp.MAX)
    
    k3_local = (ratio - delta - 1.0).masked_fill(~valid_mask, 0.0).sum()
    dist.all_reduce(k3_local, op=dist.ReduceOp.SUM)
    k3_mean = k3_local / n_valid_global_clamped
    
    entropy_local = per_token_ce.detach().masked_fill(~valid_mask, 0.0).sum()
    dist.all_reduce(entropy_local, op=dist.ReduceOp.SUM)
    entropy_mean = entropy_local / n_valid_global_clamped
    
    metrics = torch.stack([ratio_mean, min_ratio_local, max_ratio_local, k3_mean, entropy_mean])
    
    per_token_logprobs = new_logprobs_flat.view_as(labels)
    per_token_loss = per_token_pg.view_as(labels)
    
    return loss, None, per_token_logprobs, per_token_loss, metrics
from typing import List, Optional, Tuple

import torch
import torch.distributed as dist


def _preprocess_impl(
    expert_mask: torch.Tensor,
    num_experts: int,
    ep_group: dist.ProcessGroup,
) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]:
    ep_size = ep_group.size()
    num_local_experts = num_experts // ep_size
    rank = dist.get_rank(ep_group)
    num_local_tokens_per_expert = expert_mask.sum(dim=(1, 2))

    input_splits = num_local_tokens_per_expert.reshape(ep_size, num_local_experts).sum(dim=1).tolist()

    num_global_tokens_per_expert = torch.empty(
        ep_size,
        num_local_tokens_per_expert.size(0),
        dtype=num_local_tokens_per_expert.dtype,
        device=num_local_tokens_per_expert.device,
    )
    dist.all_gather_into_tensor(num_global_tokens_per_expert, num_local_tokens_per_expert, group=ep_group)

    start_idx, end_idx = rank * num_local_experts, (rank + 1) * num_local_experts
    num_global_tokens_per_local_expert = num_global_tokens_per_expert[:, start_idx:end_idx].contiguous()

    output_splits = num_global_tokens_per_local_expert.sum(dim=1).tolist()

    num_global_sum_tokens_per_local_expert = num_global_tokens_per_local_expert.sum(dim=0).to(
        torch.device("cpu"), non_blocking=True
    )
    num_global_tokens_per_local_expert = num_global_tokens_per_local_expert.view(-1, num_local_experts).to(
        torch.device("cpu"), non_blocking=True
    )

    return input_splits, output_splits, num_global_tokens_per_local_expert, num_global_sum_tokens_per_local_expert


def solution(
    expert_mask: torch.Tensor,
    num_experts: int,
    group: Optional[dist.ProcessGroup] = None,
) -> Tuple[List[int], List[int], torch.Tensor, torch.Tensor]:
    group = group or dist.group.WORLD
    return _preprocess_impl(expert_mask=expert_mask, num_experts=num_experts, ep_group=group)
from typing import List, Optional, Union

import torch
import torch.distributed as dist


def solution(
    local_tensor: torch.Tensor,
    input_split_sizes: Optional[Union[List[int], torch.Tensor]] = None,
    output_split_sizes: Optional[Union[List[int], torch.Tensor]] = None,
    group: Optional[dist.ProcessGroup] = None,
) -> torch.Tensor:
    group = group or dist.group.WORLD
    world_size = dist.get_world_size(group)
    if world_size == 1:
        return local_tensor.contiguous()

    local_tensor = local_tensor.contiguous()
    if output_split_sizes is None:
        output = torch.empty_like(local_tensor)
    else:
        out_size = sum(output_split_sizes) if isinstance(output_split_sizes, list) else int(output_split_sizes.sum().item())
        output = torch.empty(
            (out_size, local_tensor.size(1)),
            dtype=local_tensor.dtype,
            device=local_tensor.device,
        )
    dist.all_to_all_single(
        output,
        local_tensor,
        output_split_sizes=output_split_sizes,
        input_split_sizes=input_split_sizes,
        group=group,
    )
    return output