submission 554363

from task import input_t, output_t

import torch
import helion
import helion.language as hl

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

# Per-shape configs autotuned with HELION_AUTOTUNE_EFFORT=full on GB200
SHAPE_CONFIGS: dict[tuple[int, int, int], helion.Config] = {
    # Test shapes
    (1, 256, 64): helion.Config(block_sizes=[4], indexing=['pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer'], load_eviction_policies=['', '', ''], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None], reduction_loops=[None]),
    (4, 512, 128): helion.Config(block_sizes=[16], indexing=['pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer'], load_eviction_policies=['', '', ''], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None], reduction_loops=[None]),
    (16, 1024, 64): helion.Config(block_sizes=[32], indexing=['pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer'], load_eviction_policies=['', '', ''], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None], reduction_loops=[None]),
    (1, 4096, 128): helion.Config(block_sizes=[2], indexing=['tensor_descriptor', 'pointer', 'tensor_descriptor', 'pointer', 'tensor_descriptor', 'pointer'], load_eviction_policies=['last', '', 'first'], maxnreg=32, num_sm_multiplier=2, num_stages=7, num_warps=4, pid_type='persistent_interleaved', range_flattens=[False], range_multi_buffers=[True], range_unroll_factors=[4], range_warp_specializes=[False], reduction_loops=[64]),
    (8, 4096, 128): helion.Config(block_sizes=[1], indexing=['tensor_descriptor', 'pointer', 'tensor_descriptor', 'pointer', 'pointer', 'tensor_descriptor'], load_eviction_policies=['', 'first', ''], num_stages=7, num_warps=8, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None], reduction_loops=[16]),
    # Benchmark shapes
    (256, 4096, 128): helion.Config(block_sizes=[8], indexing=['pointer', 'pointer', 'pointer', 'tensor_descriptor', 'tensor_descriptor', 'pointer'], load_eviction_policies=['', 'first', ''], num_stages=1, num_warps=8, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None], reduction_loops=[None]),
    (256, 8192, 128): helion.Config(block_sizes=[16], indexing=['pointer', 'tensor_descriptor', 'pointer', 'pointer', 'pointer', 'pointer'], load_eviction_policies=['first', 'last', 'first'], num_stages=5, num_warps=16, pid_type='flat', range_flattens=[None], range_multi_buffers=[None], range_num_stages=[], range_unroll_factors=[0], range_warp_specializes=[None], reduction_loops=[64]),
    (4096, 7168, 128): helion.Config(block_sizes=[32], indexing=['pointer', 'pointer', 'pointer', 'tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], load_eviction_policies=['first', 'last', 'last'], maxnreg=128, num_sm_multiplier=32, num_stages=1, num_warps=8, pid_type='persistent_interleaved', range_flattens=[None], range_multi_buffers=[True], range_unroll_factors=[0], range_warp_specializes=[None], reduction_loops=[None]),
}


def _make_kernel(config: helion.Config):
    @helion.kernel(static_shapes=True, config=config)
    def fp8_group_quant_kernel(
        grouped_x: torch.Tensor,       # [N, group_size], float32
        grouped_q: torch.Tensor,       # [N, group_size], float32 (output)
        grouped_scales: torch.Tensor,  # [N], float32 (output)
    ) -> None:
        nrows = grouped_x.size(0)
        group_size = hl.specialize(grouped_x.size(1))

        for tile_n in hl.tile(nrows):
            x_row = grouped_x[tile_n, :]
            absmax = torch.amax(torch.abs(x_row), dim=-1)
            denom = torch.clamp(absmax, min=FP8_EPS)
            grouped_scales[tile_n] = denom / FP8_MAX
            grouped_q[tile_n, :] = torch.clamp(x_row * (FP8_MAX / denom[:, None]), FP8_MIN, FP8_MAX)

    return fp8_group_quant_kernel


_KERNELS: dict[tuple[int, int, int], object] = {}
for _shape, _cfg in SHAPE_CONFIGS.items():
    _KERNELS[_shape] = _make_kernel(_cfg)


def custom_kernel(data: input_t) -> output_t:
    x, x_q, x_s = data

    t, h = x.shape
    num_groups = x_s.size(1)
    group_size = h // num_groups

    key = (t, h, group_size)
    kernel = _KERNELS[key]

    grouped_x = x.reshape(t * num_groups, group_size)
    grouped_q = x_q.reshape(t * num_groups, group_size)
    grouped_scales = x_s.reshape(t * num_groups)
    kernel(grouped_x, grouped_q, grouped_scales)

    return x_q, x_s