submission 553021

#!POPCORN leaderboard fp8_quant
#!POPCORN gpu B200_Nebius
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: map (num_tokens, hidden_dim, group_size) to optimized helion.Config objects.
# Run `python autotune.py` to generate tuned configs, then paste them here.
SHAPE_CONFIGS: dict[tuple, helion.Config] = {
    # Test shapes
    (1, 256, 64): helion.Config(block_sizes=[1], indexing=['pointer', 'pointer', 'pointer'], load_eviction_policies=[''], num_stages=1, num_warps=1, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[True]),
    (4, 512, 128): helion.Config(block_sizes=[2], indexing=['pointer', 'pointer', 'pointer'], load_eviction_policies=[''], num_stages=1, num_warps=2, pid_type='flat', range_flattens=[None, True], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[False]),
    (16, 1024, 64): helion.Config(block_sizes=[1], indexing=['pointer', 'pointer', 'pointer'], load_eviction_policies=['last'], num_stages=1, num_warps=1, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
    (1, 4096, 128): helion.Config(block_sizes=[1], indexing=['pointer', 'pointer', 'pointer'], load_eviction_policies=[''], num_sm_multiplier=1, num_stages=2, num_warps=1, pid_type='persistent_interleaved', range_flattens=[None, None], range_multi_buffers=[None, None], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
    (8, 4096, 128): helion.Config(block_sizes=[1], indexing=['pointer', 'tensor_descriptor', 'pointer'], load_eviction_policies=['last'], num_sm_multiplier=1, num_stages=1, num_warps=1, pid_type='persistent_interleaved', range_flattens=[None, None], range_multi_buffers=[None, False], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
    # Benchmark shapes
    (16, 4096, 128): helion.Config(block_sizes=[1], indexing=['pointer', 'tensor_descriptor', 'pointer'], load_eviction_policies=[''], num_stages=1, num_warps=1, pid_type='flat', range_flattens=[None, True], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 1], range_warp_specializes=[None, None]),
    (256, 4096, 128): helion.Config(block_sizes=[1], indexing=['pointer', 'pointer', 'pointer'], load_eviction_policies=['last'], num_stages=1, num_warps=2, pid_type='flat', range_flattens=[None, True], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, False]),
    (256, 8192, 128): helion.Config(block_sizes=[1], indexing=['pointer', 'pointer', 'tensor_descriptor'], load_eviction_policies=['last'], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, False], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, False]),
    (4096, 7168, 128): helion.Config(block_sizes=[2], indexing=['tensor_descriptor', 'tensor_descriptor', 'tensor_descriptor'], load_eviction_policies=['first'], num_stages=1, num_warps=2, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
}


@helion.kernel(autotune_effort="none")
def fp8_quantize_kernel(
    x: torch.Tensor,
    x_q: torch.Tensor,
    x_s: torch.Tensor,
    group_size: hl.constexpr,
) -> tuple[torch.Tensor, torch.Tensor]:
    num_tokens, hidden_dim = x.shape
    num_groups = hidden_dim // group_size

    for tile_m in hl.tile(num_tokens):
        for g in range(num_groups):
            col_start = g * group_size
            x_group = x[tile_m, col_start:col_start + group_size].to(torch.float32)

            absmax = x_group.abs().amax(dim=1).clamp(min=FP8_EPS)
            scale = absmax / FP8_MAX

            x_q[tile_m, col_start:col_start + group_size] = (
                x_group / scale[:, None]
            ).clamp(FP8_MIN, FP8_MAX)
            x_s[tile_m, g:g + 1] = scale[:, None]

    return x_q, x_s


# Pre-compile and warm up a runner for each shape using bind + compile_config
_RUNNERS: dict[tuple, object] = {}
for (_T, _H, _gsz), _cfg in SHAPE_CONFIGS.items():
    _G = _H // _gsz
    _example_x = torch.empty(_T, _H, dtype=torch.float32, device="cuda")
    _example_x_q = torch.empty(_T, _H, dtype=torch.float32, device="cuda")
    _example_x_s = torch.empty(_T, _G, dtype=torch.float32, device="cuda")
    _bound = fp8_quantize_kernel.bind((_example_x, _example_x_q, _example_x_s, _gsz))
    _runner = _bound.compile_config(_cfg)
    # Warm up to trigger CUDA PTX compilation eagerly
    _runner(_example_x, _example_x_q, _example_x_s, _gsz)
    _RUNNERS[(_T, _H, _gsz)] = _runner
torch.cuda.synchronize()


def custom_kernel(data: input_t) -> output_t:
    x, x_q, x_s = data
    T, H = x.shape
    G = x_s.shape[1]
    gsz = H // G

    runner = _RUNNERS[(T, H, gsz)]
    return runner(x, x_q, x_s, gsz)