submission 554401

#!POPCORN leaderboard gated_deltanet_chunk_fwd_h
#!POPCORN gpu B200_Nebius

from task import input_t, output_t

import base64
import tempfile
from pathlib import Path

import torch
import helion
import helion.language as hl


# Embedded ACF (base64-encoded /opt/booster_pack/chunk_fwd_h_0.acf)
_ACF_B64 = "dxWiJeZ70zxa8yaIXrS2nR5/YR9acn4hUQuuuQBAtk1NaYlOuRNeZE2ZCJPgpRuUgTkBm9Xup9NUyKFZuxpQHy4NGUpw0lKHWcV8hdhPD4bwug64u+dHVGJ0O1wdbs4aUDmcxnHy3RxJxoEtFfIzKmYi8w3kKjm8Y2AZMLMAEH4cJVgYTO6E07mI32mTktUpVB58dOGBr4ZO/U3hTkwLNIfaQR/MHiD+MNq1rpg1mdzFMNLBwhL73AKIz7nG3Y5Mqg9dt968yZzFOuWzYJmfpzsRy3BnDkcQa9YkYnG7pSKlsgfJcj0XTD4+cLp24kIcOdZnfZGKrYr0qO/0M8lBaL3bqTVAsaBEOfohEQBZUHQsTq6sSQ9qI/aV3lCFkbK0ipM6PmaEku/WtNU2batXFl9jNgp61oUdKSz7zeengbw+cVypqN1g8UI3InvHbFEkvdnu8FUiMvK1A+H6La9r1w3Mz6LmxEKemrxphszkpuNQRaN48swQ2/QB5nNzH0SbmaR03UKyM1DQDd9r2YAzIj772mC2j+J+fO9gSEqTclYE0+LsFb0z/oT4p5KGmR7uPAFTJw4kVwb7oPz3SKvIkPehXIROtufCad+Fix/mpuFevP+hhtp8qNkRkKhGBiygvSCcSCQPJRphf+VcYxuhCZs3QE1neTVnOqA/R6njcVwP5sq/FhwkAmfCbq5VkrVseUEyS7fdQGCXHl1mIR2xgFWr9+jwXw4VNTgMED7f8eLOknaM8/M+3dnYn1xhPfujAZt4ROgFiMNCnM9zDZmYjoP8R6JQnViTouHpfDnHEo3/u4tZKtCHvfoBRKGOSDurka18H6DUDACOlySq7yo9hOvHgQ3Ed5W9pRGPNDExivqcw8lyUVo2TMXMvI4Ld0wvpzuQsbxwQ8+G3ZCI7EECfqDpTSik6z9QkGP5cjyAHgUli/5eunOGldC7iDMRsLxIJmkBchh1xkCmO5pYiqOjeSnYETAvIGWdTyTXksi8+eoPlV6yhQgqhTdCGQM3Oa52Qyvnh4vQYEur7JeiqlMvFmI="


def _get_acf_path():
    if hasattr(_get_acf_path, "_path"):
        return _get_acf_path._path
    local = Path("/opt/booster_pack/chunk_fwd_h_0.acf")
    if local.exists():
        _get_acf_path._path = str(local)
    else:
        d = tempfile.mkdtemp(prefix="fwd_h_acf_")
        p = Path(d) / "chunk_fwd_h_0.acf"
        p.write_bytes(base64.b64decode(_ACF_B64))
        _get_acf_path._path = str(p)
    return _get_acf_path._path


_ACF = _get_acf_path()

# Per-shape ACF-optimized configs from autotuning on B200.
# Some shapes autotuned best WITHOUT ACF (advanced_controls_file='').
SHAPE_CONFIGS: dict[tuple, helion.Config] = {
    # Test shapes
    (1, 64, 2, 64, 64): helion.Config(advanced_controls_file=_ACF, block_sizes=[8], indexing=['tensor_descriptor', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', '', '', '', ''], loop_orders=[[0, 1]], num_sm_multiplier=1, num_stages=1, num_warps=2, pid_type='persistent_blocked', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None], static_ranges=[True]),
    (2, 128, 4, 64, 64): helion.Config(advanced_controls_file=_ACF, block_sizes=[4], indexing=['pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['', '', '', '', ''], loop_orders=[[0, 1]], num_stages=1, num_warps=2, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 1], range_warp_specializes=[None, False], static_ranges=[False]),
    (1, 256, 4, 64, 128): helion.Config(advanced_controls_file='', block_sizes=[8], indexing=['pointer', 'pointer', 'pointer', 'pointer', 'tensor_descriptor', 'tensor_descriptor', 'pointer'], l2_groupings=[1], load_eviction_policies=['last', 'last', 'last', '', ''], loop_orders=[[0, 1]], num_stages=1, num_warps=16, pid_type='flat', range_flattens=[None, True], range_multi_buffers=[None, False], range_num_stages=[0, 1], range_unroll_factors=[0, 0], range_warp_specializes=[None, False], static_ranges=[False]),
    # Benchmark shapes
    (1, 64, 1, 64, 64): helion.Config(advanced_controls_file=_ACF, block_sizes=[4], indexing=['pointer', 'pointer', 'tensor_descriptor', 'pointer', 'pointer', 'pointer', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', '', '', 'last'], loop_orders=[[0, 1]], num_sm_multiplier=1, num_stages=1, num_warps=8, pid_type='persistent_interleaved', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
    (2, 512, 3, 64, 64): helion.Config(advanced_controls_file='', block_sizes=[8], indexing=['pointer', 'tensor_descriptor', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', '', '', '', ''], loop_orders=[[0, 1]], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
    (2, 1024, 3, 64, 64): helion.Config(advanced_controls_file=_ACF, block_sizes=[4], indexing=['pointer', 'pointer', 'tensor_descriptor', 'pointer', 'pointer', 'pointer', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', '', '', 'last'], loop_orders=[[0, 1]], num_sm_multiplier=1, num_stages=1, num_warps=8, pid_type='persistent_interleaved', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
    (3, 1024, 4, 100, 100): helion.Config(advanced_controls_file='', block_sizes=[8], indexing=['pointer', 'tensor_descriptor', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', '', '', '', ''], loop_orders=[[0, 1]], num_stages=1, num_warps=4, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
    (4, 1024, 4, 128, 128): helion.Config(advanced_controls_file=_ACF, block_sizes=[16], indexing=['pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['first', '', '', '', ''], loop_orders=[[0, 1]], num_sm_multiplier=1, num_stages=1, num_warps=8, pid_type='persistent_interleaved', range_flattens=[None, True], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
    (2, 1536, 4, 128, 128): helion.Config(advanced_controls_file=_ACF, block_sizes=[8], indexing=['pointer', 'pointer', 'pointer', 'tensor_descriptor', 'pointer', 'pointer', 'tensor_descriptor'], l2_groupings=[1], load_eviction_policies=['last', '', '', '', ''], loop_orders=[[0, 1]], num_stages=1, num_warps=8, pid_type='flat', range_flattens=[None, None], range_multi_buffers=[None, None], range_num_stages=[0, 0], range_unroll_factors=[0, 1], range_warp_specializes=[None, None]),
    (4, 2048, 8, 64, 64): helion.Config(advanced_controls_file=_ACF, block_sizes=[16], indexing=['pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer', 'pointer'], l2_groupings=[1], load_eviction_policies=['', '', '', '', ''], loop_orders=[[0, 1]], num_sm_multiplier=1, num_stages=1, num_warps=8, pid_type='persistent_interleaved', range_flattens=[None, None], range_multi_buffers=[None, False], range_num_stages=[0, 0], range_unroll_factors=[0, 0], range_warp_specializes=[None, None]),
}


def _make_kernel(config: helion.Config):
    @helion.kernel(static_shapes=True, dot_precision="ieee", config=config)
    def kernel(
        k: torch.Tensor,   # [B, T, H, K]
        w: torch.Tensor,   # [B, T, H, K]
        u: torch.Tensor,   # [B, T, H, V]
        g: torch.Tensor,   # [B, T, H]
    ) -> tuple[torch.Tensor, torch.Tensor]:
        B, T, H, K = k.shape
        V = u.shape[-1]
        C = 64
        K = hl.specialize(K)
        V = hl.specialize(V)

        NT = (T + C - 1) // C
        h_out = torch.empty(B, NT, H, K, V, dtype=k.dtype, device=k.device)
        v_out = torch.empty_like(u)

        BH = B * H

        for flat, tv in hl.tile([BH, V], block_size=[1, None]):
            b_idx = flat.begin // H
            h_idx = flat.begin % H
            state = hl.zeros([K, tv], dtype=torch.float32)

            for tc in hl.tile(T, block_size=C):
                chunk_idx = tc.begin // C
                t_end = min(tc.begin + C, T) - 1

                h_out[b_idx, chunk_idx, h_idx, :, tv] = state.to(k.dtype)

                proj = hl.dot(
                    w[b_idx, tc, h_idx, :], state, out_dtype=torch.float32
                )
                diff = u[b_idx, tc, h_idx, tv].to(torch.float32) - proj
                v_out[b_idx, tc, h_idx, tv] = diff.to(u.dtype)

                g_end = g[b_idx, t_end, h_idx].to(torch.float32)
                g_t = g[b_idx, tc, h_idx].to(torch.float32)
                valid = tc.index < T
                alpha = torch.where(valid, torch.exp(g_end - g_t), 0.0)
                k_adj = k[b_idx, tc, h_idx, :] * alpha[:, None]

                state = state * torch.exp(g_end)
                state = state + hl.dot(k_adj.T, diff, out_dtype=torch.float32)

        return h_out, v_out

    return kernel


_KERNELS = {shape: _make_kernel(cfg) for shape, cfg in SHAPE_CONFIGS.items()}


def custom_kernel(data: input_t) -> output_t:
    k, w, u, g = data
    B, T, H, K = k.shape
    V = u.shape[-1]
    kernel = _KERNELS[(B, T, H, K, V)]
    return kernel(k, w, u, g)