submission 553135

from task import input_t, output_t

import torch
import helion
import helion.language as hl
from pathlib import Path

# ACF: find best chunk_fwd_h ACF on B200
def _find_acf(pattern):
    bp = Path("/opt/booster_pack")
    if not bp.exists():
        return None
    for p in sorted(bp.glob(pattern)):
        return str(p)
    return None

_acf = None  # _find_acf("chunk_fwd_h_*.acf")  # disabled for now
_cfg = {"block_sizes": [8], "num_warps": 1, "num_stages": 1}
if _acf:
    _cfg["advanced_controls_file"] = _acf

@helion.kernel(
    static_shapes=True,
    dot_precision="ieee",
    config=helion.Config(**_cfg),
)
def chunk_state_pass(
    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

    # Outer: parallel over (batch*head, value_dim tiles)
    for flat, tv in hl.tile([BH, V], block_size=[1, None]):  # None = from config block_sizes
        b_idx = flat.begin // H
        h_idx = flat.begin % H

        # Hidden state [K, tv] — sequential across chunks
        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

            # Store state snapshot before update
            h_out[b_idx, chunk_idx, h_idx, :, tv] = state.to(k.dtype)

            # Delta correction: v_new = u - w @ state
            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)

            # Gating: decay each timestep toward chunk end
            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)
            # Apply decay to diff (gated values) instead of keys
            diff_gated = diff * alpha[:, None]

            # Decay state and accumulate new info: state += k^T @ diff_gated
            state = state * torch.exp(g_end)
            upd = hl.dot(k[b_idx, tc, h_idx, :].T, diff_gated, out_dtype=torch.float32)
            state = state + upd

    return h_out, v_out


def custom_kernel(data: input_t) -> output_t:
    k, w, u, g = data
    return chunk_state_pass(k, w, u, g)