willychan21/ParallelKernelBench_Problems

1import torch2import torch.distributed as dist3import triton4import triton.language as tl5 6@triton.jit7def block_fp8_dequant_kernel(y_ptr, s_ptr, x_ptr, num_elements, BLOCK_SIZE: tl.constexpr):8    pid = tl.program_id(axis=0)9    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)10    mask = offs < num_elements11 12    s = tl.load(s_ptr + pid)13 14    y = tl.load(y_ptr + offs, mask=mask).to(tl.float32)15 16    tl.store(x_ptr + offs, y * s, mask=mask)17 18 19def solution(20    local_y: torch.Tensor,21    local_s: torch.Tensor,22    block_size: int = 128,23) -> torch.Tensor:24    world_size = dist.get_world_size()25 26    chunk_shape = local_y.shape[1:]27    chunk_numel = local_y.numel() // world_size28    num_elements = local_y.numel()29    assert chunk_numel % block_size == 0, (30        f"Chunk size {chunk_numel} must be divisible by block_size ({block_size})"31    )32 33    y_flat = local_y.view(-1)34    s_flat = local_s.view(-1)35    x_flat = torch.empty(num_elements, device=local_y.device, dtype=torch.float32)36 37    if num_elements > 0:38        grid = (triton.cdiv(num_elements, block_size),)39        block_fp8_dequant_kernel[grid](40            y_flat, s_flat, x_flat, num_elements, BLOCK_SIZE=block_size41        )42 43    x = x_flat.view(world_size, *chunk_shape)44    out = torch.empty_like(x)45    dist.all_to_all_single(out, x)46 47    return out48