willychan21/ParallelKernelBench_Problems

1import torch2import torch.distributed as dist3import triton4import triton.language as tl5from typing import Tuple6 7@triton.jit8def block_fp8_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):9    pid = tl.program_id(axis=0)10    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)11 12    x = tl.load(x_ptr + offs).to(tl.float32)13    14    # FP8 E4M3 max value is 448.015    s = tl.max(tl.abs(x)) / 448.016    17    # Prevent division by zero if all elements in the block are 018    s_safe = tl.where(s == 0.0, 1.0, s)19    20    y = (x / s_safe).to(y_ptr.dtype.element_ty)21    22    tl.store(y_ptr + offs, y)23    tl.store(s_ptr + pid, s)24 25def solution(local_tensor: torch.Tensor, block_size: int = 128) -> Tuple[torch.Tensor, torch.Tensor]:26    assert local_tensor.size(-1) % block_size == 0, "Last dimension must be divisible by block_size"27    28    y_local = torch.empty_like(local_tensor, dtype=torch.float8_e4m3fn)29    s_local = local_tensor.new_empty(30        *local_tensor.size()[:-1], local_tensor.size(-1) // block_size, dtype=torch.float3231    )32    33    grid = (triton.cdiv(local_tensor.numel(), block_size),)34    block_fp8_quant_kernel[grid](local_tensor, y_local, s_local, BLOCK_SIZE=block_size)35    36    if dist.is_initialized():37        world_size = dist.get_world_size()38        39        y_local_u8 = y_local.view(torch.uint8)40        y_gather_u8 = [torch.empty_like(y_local_u8) for _ in range(world_size)]41        dist.all_gather(y_gather_u8, y_local_u8)42        43        y_global_u8 = torch.cat(y_gather_u8, dim=0)44        y_global = y_global_u8.view(torch.float8_e4m3fn)45        46        s_gather = [torch.empty_like(s_local) for _ in range(world_size)]47        dist.all_gather(s_gather, s_local)48        49        s_global = torch.cat(s_gather, dim=0)50    else:51        y_global = y_local52        s_global = s_local53        54    return y_global, s_global55