submission 553239

from task import input_t, output_t

import torch
import triton
import triton.language as tl
import helion
import helion.language as hl


FP8_MAX = 448.0
FP8_MIN = -448.0
FP8_EPS = 1e-10
GROUP128 = 128


@helion.kernel(
    static_shapes=True,
    config=helion.Config(block_sizes=[1], num_warps=1, num_stages=1),
)
def normalize_to_range_small(
    data: torch.Tensor,        # [N, G] input rows
    qout: torch.Tensor,        # [N, G] output quantized rows
    scales_out: torch.Tensor,  # [N] output normalization factors
) -> torch.Tensor:
    nrows = data.size(0)
    ncols = hl.specialize(data.size(1))

    for rr in hl.tile(nrows):
        row = data[rr, :].to(torch.float32)
        amax = torch.amax(torch.abs(row), -1)
        scale = torch.clamp(amax, min=FP8_EPS) / FP8_MAX
        qout[rr, :] = torch.clamp(row / scale[:, None], FP8_MIN, FP8_MAX)
        scales_out[rr] = scale

    return qout


@helion.kernel(
    static_shapes=True,
    config=helion.Config(block_sizes=[8], num_warps=4, num_stages=2),
)
def normalize_to_range_mid(
    data: torch.Tensor,        # [N, G] input rows
    qout: torch.Tensor,        # [N, G] output quantized rows
    scales_out: torch.Tensor,  # [N] output normalization factors
) -> torch.Tensor:
    nrows = data.size(0)
    ncols = hl.specialize(data.size(1))

    for rr in hl.tile(nrows):
        row = data[rr, :].to(torch.float32)
        amax = torch.amax(torch.abs(row), -1)
        scale = torch.clamp(amax, min=FP8_EPS) / FP8_MAX
        qout[rr, :] = torch.clamp(row / scale[:, None], FP8_MIN, FP8_MAX)
        scales_out[rr] = scale

    return qout


@helion.kernel(
    static_shapes=True,
    config=helion.Config(block_sizes=[16], num_warps=8, num_stages=4),
)
def normalize_to_range_large(
    data: torch.Tensor,        # [N, G] input rows
    qout: torch.Tensor,        # [N, G] output quantized rows
    scales_out: torch.Tensor,  # [N] output normalization factors
) -> torch.Tensor:
    nrows = data.size(0)
    ncols = hl.specialize(data.size(1))

    for rr in hl.tile(nrows):
        row = data[rr, :].to(torch.float32)
        amax = torch.amax(torch.abs(row), -1)
        scale = torch.clamp(amax, min=FP8_EPS) / FP8_MAX
        qout[rr, :] = torch.clamp(row / scale[:, None], FP8_MIN, FP8_MAX)
        scales_out[rr] = scale

    return qout


@triton.jit
def normalize_to_range_group128_rows_kernel(
    x_ptr,
    q_ptr,
    s_ptr,
    stride_x_row,
    stride_x_col,
    stride_q_row,
    stride_q_col,
    stride_s_row,
    nrows,
    GROUP_SIZE: tl.constexpr,
    BLOCK_ROWS: tl.constexpr,
):
    pid = tl.program_id(0)
    offs_row = pid * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS)
    offs_col = tl.arange(0, GROUP_SIZE)
    row_mask = offs_row < nrows

    x_ptrs = x_ptr + offs_row[:, None] * stride_x_row + offs_col[None, :] * stride_x_col
    q_ptrs = q_ptr + offs_row[:, None] * stride_q_row + offs_col[None, :] * stride_q_col

    x = tl.load(x_ptrs, mask=row_mask[:, None], other=0.0)
    amax = tl.max(tl.abs(x), axis=1)
    scale = tl.maximum(amax, 1.0e-10) / 448.0
    q = tl.maximum(tl.minimum(x / scale[:, None], 448.0), -448.0)

    tl.store(q_ptrs, q, mask=row_mask[:, None])
    tl.store(s_ptr + offs_row * stride_s_row, scale, mask=row_mask)


def normalize_to_range_group128_triton(
    data: torch.Tensor,
    qout: torch.Tensor,
    scales_out: torch.Tensor,
) -> None:
    nrows = data.size(0)

    if nrows <= 512:
        block_rows = 8
        num_warps = 2
        num_stages = 2
    elif nrows < 32768:
        block_rows = 4
        num_warps = 4
        num_stages = 2
    else:
        block_rows = 8
        num_warps = 4
        num_stages = 3

    grid = (triton.cdiv(nrows, block_rows),)
    normalize_to_range_group128_rows_kernel[grid](
        data,
        qout,
        scales_out,
        data.stride(0),
        data.stride(1),
        qout.stride(0),
        qout.stride(1),
        scales_out.stride(0),
        nrows,
        GROUP_SIZE=GROUP128,
        BLOCK_ROWS=block_rows,
        num_warps=num_warps,
        num_stages=num_stages,
    )


def custom_kernel(data: input_t) -> output_t:
    x, x_q, x_s = data
    num_tokens, hidden_dim = x.shape
    num_groups = x_s.shape[1]
    group_size = hidden_dim // num_groups
    nrows = num_tokens * num_groups

    flat_in = x.reshape(nrows, group_size)
    flat_q = x_q.reshape(nrows, group_size)
    flat_s = x_s.reshape(nrows)

    if group_size == GROUP128:
        normalize_to_range_group128_triton(flat_in, flat_q, flat_s)
    elif nrows <= 512:
        normalize_to_range_small(flat_in, flat_q, flat_s)
    elif nrows < 65536:
        normalize_to_range_mid(flat_in, flat_q, flat_s)
    else:
        normalize_to_range_large(flat_in, flat_q, flat_s)
    return x_q, x_s