submission 553122

from task import input_t, output_t

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

# ACF: find best recompute_w_u 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 = _find_acf("recompute_w_u_fwd_*.acf")
_cfg = {"block_sizes": [], "num_warps": 4, "num_stages": 1}
if _acf:
    _cfg["advanced_controls_file"] = _acf

@helion.kernel(
    static_shapes=True,
    dot_precision="ieee",
    config=helion.Config(**_cfg),
)
def project_kv(
    k: torch.Tensor,     # [B, T, H, K] -- key vectors
    v: torch.Tensor,     # [B, T, H, V] -- value vectors
    beta: torch.Tensor,  # [B, T, H]    -- writing strength (scalar per position)
    A: torch.Tensor,     # [B, T, H, BT] -- WY transform matrix (from UT transform)
    g: torch.Tensor,     # [B, T, H]    -- gating/decay values (negative, so exp(g) in (0,1])
) -> tuple[torch.Tensor, torch.Tensor]:
    B, T, H, K = k.shape
    V = v.shape[-1]
    # Specialize chunk size, K, V as compile-time constants
    C = hl.specialize(A.shape[-1])  # 64
    K = hl.specialize(K)
    V = hl.specialize(V)

    w_out = torch.empty_like(k)
    u_out = torch.empty_like(v)

    # Flatten batch*head for parallelization
    BH = B * H
    for flat_bh, rt in hl.tile([BH, T], block_size=[1, C]):
        b_idx = flat_bh.begin // H
        h_idx = flat_bh.begin % H

        # Load A matrix [C, C], scalars, and vectors for this chunk
        a_chunk = A[b_idx, rt, h_idx, :].to(torch.float32)
        beta_chunk = beta[b_idx, rt, h_idx].to(torch.float32)
        g_chunk = g[b_idx, rt, h_idx].to(torch.float32)
        k_chunk = k[b_idx, rt, h_idx, :].to(torch.float32)
        v_chunk = v[b_idx, rt, h_idx, :].to(torch.float32)

        # Scale: v * beta and k * beta * exp(g)
        v_scaled = v_chunk * beta_chunk[:, None]
        k_scaled = k_chunk * (beta_chunk * torch.exp(g_chunk))[:, None]

        # Two matmuls: u = A @ v_scaled, w = A @ k_scaled
        u_out[b_idx, rt, h_idx, :] = torch.matmul(a_chunk, v_scaled).to(v.dtype)
        w_out[b_idx, rt, h_idx, :] = torch.matmul(a_chunk, k_scaled).to(k.dtype)

    return w_out, u_out


def custom_kernel(data: input_t) -> output_t:
    k, v, beta, A, g = data
    return project_kv(k, v, beta, A, g)