total 53
drwxr-xr-x 1 agent agent   130 Oct 18 11:07 .
drwxr-xr-x 1 root  root    220 Oct 18 13:49 ..
-rw-r--r-- 1 agent agent  3034 Oct 18 11:07 check.py
-rw-r--r-- 1 agent agent  4816 Oct 18 11:07 nodus_autosave.py
-rw-r--r-- 1 agent agent  5666 Oct 18 11:07 starting_kernel.py
-rw-r--r-- 1 agent agent 13355 Oct 18 11:07 static_check.py
-rw-r--r-- 1 agent agent 19226 Oct 18 11:07 task.py
-rw-r--r-- 1 agent agent  5143 Oct 18 11:07 task_api.py
"""Public task definition for xpr_hd256_attn_bwd: reference, check, timed shapes, standard normal inputs. The verifier also uses edge shapes and hidden input distributions that are not in this file."""
"""Causal attention backward with head dim 256 (dq, dk, dv) on Blackwell, GQA, sliding windows and packed documents
(expert PR task).

Source: Dao-AILab/flash-attention PRs #2964, #2965 and #2966 "[CuTe, SM100] Unify hd256 backward (1/5-3/5)" (Johnson,
merged 2026-10-06, 10-15 and 10-16), which rebuilt FlashAttention-4's head-dim-256 training backward on Blackwell on
top of #2807, #2891 and #2952. Upstream main (965c921, the #2966 merge, 2026-10-16) has two public hd256 backwards: the
default "dedicated" pair (a dQ kernel and a dKdV kernel that each recompute S and dP, 2-CTA, no sliding window) and
the opt-in general kernel FlashAttentionBackwardSm100 (one kernel, 64-row KV tile per CTA, 2-CTA, any window,
FLASH_ATTENTION_HD256_GENERIC_BWD=1), which the authors measured at parity to 11% faster than the pair on most shapes
and 22-29% slower on 2K GQA (B200). Open PR #2967 (head 90419de, 2026-10-18) packs a KV head's query heads into the
general kernel's Q tile, 1-11% faster than the pair on hd256 GQA per its author. The bar on each shape is the fastest
of the three (vendored in kernel_lab/expert_pr/fa4_v2966 and fa4_v2967; the route per shape is in fa4_v2966/b300.py),
at 1.0x its speed in the same run.

Shapes: training of head-dim-256 models per GPU: Gemma-3-style global and sliding-window-1024 local layers (16 query
and 8 KV heads, or 8 and 4), a Qwen3-Next-style full-attention layer (16 query heads, 2 KV heads) and packed
documents of mixed lengths, 8K-16K tokens per step. o and lse come from an fp32 forward over the same inputs.
"""
from task_api import Case, Task
AGENT_MINUTES = 120
D = 256
SCALE = D ** (-0.5)
GAMMA = 2.0 ** (-17)
OUTS = ((9, 'dq'), (10, 'dk'), (11, 'dv'))

def seq_lens(shape, g, device):
    """n sequences of length L, or with 'ragged' n of mean length L in [L/2, 3L/2], built in pairs (L + d, L - d) with
    the first pair fixed at (3L/2, L/2), so the token count (n * L) and the longest sequence depend on the shape only;
    'long' puts one extra sequence of that length first."""
    import torch
    n, L = (shape['n'], shape['L'])
    if shape.get('ragged'):
        h = L // 2
        d = torch.randint(-h, h + 1, (n // 2,), generator=g, device=device, dtype=torch.int64)
        if n >= 2:
            d[0] = h
        lens = torch.stack([L + d, L - d], 1).reshape(-1)
        if n % 2:
            lens = torch.cat([lens, torch.full((1,), L, device=device, dtype=torch.int64)])
    else:
        lens = torch.full((n,), L, device=device, dtype=torch.int64)
    if shape.get('long'):
        lens = torch.cat([torch.tensor([shape['long']], device=device, dtype=torch.int64), lens])
    return lens

def _visible(t0, t1, k0, k1, W, device):
    """[t1 - t0, k1 - k0] bool: key j is visible to query i of the same sequence when j <= i and, with a window W >= 0,
    i - j <= W."""
    import torch
    i = torch.arange(t0, t1, device=device).view(-1, 1)
    j = torch.arange(k0, k1, device=device).view(1, -1)
    m = j <= i
    return m & (i - j <= W) if W >= 0 else m

def _chunks(cu, W, rows):
    """(t0, t1, k0) per query chunk of every sequence: keys k0 .. t1 - 1 hold every key a query of the chunk sees."""
    bounds = cu.tolist()
    for s0, s1 in zip(bounds[:-1], bounds[1:]):
        for t0 in range(s0, s1, rows):
            t1 = min(s1, t0 + rows)
            yield (t0, t1, s0 if W < 0 else max(s0, t0 - W))

def _rows(H, n_keys):
    """Query tokens per reference chunk: the [H, tokens, keys] fp32 intermediates stay at 64 MiB."""
    return max(16, 2 ** 24 // (H * max(1, n_keys)))

def _mm(eq, a, b):
    """einsum(eq, a, b) at fp32 accuracy for fp32 a and bf16-valued b, also when matmuls run in TF32: a splits into a
    TF32-exact head (11 significant bits) and an exact fp32 tail, and every bf16 x head product is exact."""
    import torch
    hi = (a.view(torch.int32) & -8192).view(torch.float32)
    return torch.einsum(eq, hi, b) + torch.einsum(eq, a - hi, b)

def _tf32(fn, args):
    """Runs fn(args) with TF32 matmuls on CUDA (every product is split or exact, see _mm), restoring the setting."""
    import torch
    prev = torch.get_float32_matmul_precision()
    if args[0].is_cuda:
        torch.set_float32_matmul_precision('high')
    try:
        return fn(args)
    finally:
        torch.set_float32_matmul_precision(prev)

def _forward(args):
    """The training forward in fp32: o = bf16(softmax(q k^T / 16) v) and lse = log(sum_j exp(q . k_j / 16)), written
    into args[3] and args[5]."""
    import torch
    q, k, v, o, _, lse, cu, max_len, W = args[:9]
    T, Hq, _ = q.shape
    Hkv = k.shape[1]
    G = Hq // Hkv
    rows = _rows(Hq, max_len if W < 0 else min(max_len, W + 1))
    for t0, t1, k0 in _chunks(cu, W, rows):
        n = t1 - t0
        qc = q[t0:t1].float().view(n, Hkv, G, D)
        kc, vc = (k[k0:t1].float(), v[k0:t1].float())
        s = torch.einsum('nhgd,mhd->hgnm', qc, kc) * SCALE
        s = s.masked_fill(~_visible(t0, t1, k0, t1, W, q.device), float('-inf'))
        L = torch.logsumexp(s, -1, keepdim=True)
        O = _mm('hgnm,mhd->nhgd', torch.exp(s - L), vc)
        o[t0:t1] = O.reshape(n, Hq, D).to(torch.bfloat16)
        lse[:, t0:t1] = L.reshape(Hq, n)

def make_case(shape, seed, device, dist='normal'):
    """q, k, v and dO from the input distribution, generated as [tokens, heads * 256] so row-wise distributions act
    per token; o and lse from the fp32 forward over the same inputs."""
    import torch
    from task_api import gen, randn
    g = gen(device, seed)
    Hq, Hkv, W = (shape['Hq'], shape['Hkv'], shape['window'])
    lens = seq_lens(shape, g, device)
    T = int(lens.sum().item())
    cu = torch.zeros((lens.numel() + 1,), device=device, dtype=torch.int32)
    cu[1:] = torch.cumsum(lens, 0).to(torch.int32)
    bf = torch.bfloat16
    q = randn((T, Hq * D), g, device, dist).view(T, Hq, D).to(bf)
    k = randn((T, Hkv * D), g, device, dist).view(T, Hkv, D).to(bf)
    v = randn((T, Hkv * D), g, device, dist).view(T, Hkv, D).to(bf)
    do = randn((T, Hq * D), g, device, dist).view(T, Hq, D).to(bf)
    o = torch.empty((T, Hq, D), device=device, dtype=bf)
    lse = torch.empty((Hq, T), device=device, dtype=torch.float32)
    max_len = int(lens.max().item())
    _tf32(_forward, [q, k, v, o, do, lse, cu, max_len, W])
    dq = torch.empty((T, Hq, D), device=device, dtype=bf)
    dk = torch.empty((T, Hkv, D), device=device, dtype=bf)
    dv = torch.empty((T, Hkv, D), device=device, dtype=bf)
    return Case(args=[q, k, v, o, do, lse, cu, max_len, W, dq, dk, dv], outputs=(9, 10, 11), shape=shape, dist=dist)

def _reference(args):
    import torch
    q, k, v, o, do, lse, cu, max_len, W = args[:9]
    T, Hq, _ = q.shape
    Hkv = k.shape[1]
    G = Hq // Hkv
    dev, f32 = (q.device, torch.float32)
    DQ, TQ = (torch.empty((T, Hq, D), device=dev, dtype=f32), torch.empty((T, Hq, D), device=dev, dtype=f32))
    DK, TK = (torch.zeros((T, Hkv, D), device=dev, dtype=f32), torch.zeros((T, Hkv, D), device=dev, dtype=f32))
    DV, TV = (torch.zeros_like(DK), torch.zeros_like(DK))
    kf, vf = (k.float(), v.float())
    nk, nv = (kf.pow(2).sum(-1).sqrt(), vf.pow(2).sum(-1).sqrt())
    rows = _rows(Hq, max_len if W < 0 else min(max_len, W + 1))
    for t0, t1, k0 in _chunks(cu, W, rows):
        n = t1 - t0
        qc = q[t0:t1].float().view(n, Hkv, G, D)
        kc, vc = (kf[k0:t1], vf[k0:t1])
        dO = do[t0:t1].float().view(n, Hkv, G, D)
        L = lse[:, t0:t1].view(Hkv, G, n, 1)
        vis = _visible(t0, t1, k0, t1, W, dev)
        s = torch.einsum('nhgd,mhd->hgnm', qc, kc)
        P = torch.exp(s * SCALE - L).masked_fill_(~vis, 0.0)
        dP = torch.einsum('nhgd,mhd->hgnm', dO, vc)
        Od = o[t0:t1].double().view(n, Hkv, G, D)
        delta = (Od * dO.double()).sum(-1).permute(1, 2, 0).unsqueeze(-1)
        dmd = dP.sub_(delta.float())
        dS = P * dmd
        DQ[t0:t1] = (SCALE * _mm('hgnm,mhd->nhgd', dS, kc)).reshape(n, Hq, D)
        DK[k0:t1] += SCALE * _mm('hgnm,nhgd->mhd', dS, qc)
        DV[k0:t1] += _mm('hgnm,nhgd->mhd', P, dO)
        nq = qc.pow(2).sum(-1).sqrt().permute(1, 2, 0).unsqueeze(-1)
        ndo = dO.pow(2).sum(-1).sqrt().permute(1, 2, 0).unsqueeze(-1)
        d1 = (Od.abs() * dO.double().abs()).sum(-1).permute(1, 2, 0).unsqueeze(-1).float()
        kmax = nk[k0:t1].max(0).values.view(Hkv, 1, 1, 1)
        vmax = nv[k0:t1].max(0).values.view(Hkv, 1, 1, 1)
        e = (2 * GAMMA + 2.0 ** (-20)) * SCALE * nq * kmax + 2.0 ** (-20) * L.abs() + 2.0 ** (-21)
        eta = torch.expm1(e)
        a = 2.0 ** (-7) + eta
        b = (1 + eta) * (2 * GAMMA) * (ndo * vmax + d1)
        A = dS.abs_().mul_(a).add_(b * P)
        TQ[t0:t1] = (SCALE * torch.einsum('hgnm,mhd->nhgd', A, kc.abs())).reshape(n, Hq, D)
        TK[k0:t1] += SCALE * torch.einsum('hgnm,nhgd->mhd', A, qc.abs())
        TV[k0:t1] += torch.einsum('hgnm,nhgd->mhd', P.mul_(a), dO.abs())
    exact = {'dq': DQ, 'dk': DK, 'dv': DV}
    tol = {'dq': TQ, 'dk': TK, 'dv': TV}
    res = {}
    for i, nm in OUTS:
        x = exact[nm]
        res[i] = x.to(torch.bfloat16)
        res[nm] = x
        res['tol_' + nm] = tol[nm] * (1 + 2.0 ** (-8)) + 2.0 ** (-7) * x.abs() + 2.0 ** (-100)
    return res

def reference(args):
    """The contract in fp32 from the bf16 inputs (every product exact or split to fp32 accuracy, delta in float64), in
    chunks of query tokens per sequence: P from the given lse, dS = P (dP - delta) with delta = sum(dO * o), dq per
    token, dk and dv summed over every query (and query head of the KV head's group) that sees the key. Returns the
    outputs in bf16 (int keys) and, for the check, the exact values in fp32 plus per-element tolerances 'tol_<name>'."""
    return _tf32(_reference, args)

def check(got, exp):
    """Per element, every output must lie within its tolerance of the exact value: |got - exact| <= tol. The tolerance
    follows the error a correct kernel can make, term by term, summed over the (query, key) pairs that feed the element:
    - logits: e = (2 GAMMA + 2^-20) |q| max|k| / 16 + 2^-20 |lse| + 2^-21, with GAMMA = 2^-17 for an fp32 dot product
      of exact bf16 products (blocked or tensor-core sums reach about 2^-21), the score, lse and scale roundings and an
      approximate exp2 (max|k| over the keys of the reference's query chunk); P may be off by expm1(e) of itself;
    - dP - delta: eps = 2 GAMMA (|dO| max|v| + sum |dO o|), an fp32 dot product and an fp32 delta;
    - per pair, A = |dS| (2^-7 + expm1(e)) + P (1 + expm1(e)) eps: dS rounded once to bf16 (2^-8, the rest headroom
      for fp32 sums over up to 16K keys) and the P error;
    - dq: sum over keys of A |k| / 16; dk: the sum over (query, query head of the group) of A |q| / 16; dv: the same sum
      of P (2^-7 + expm1(e)) |dO| (P rounded once to bf16 for the dV product); each output also gets 2^-7 of its value
      for the final bf16 rounding.
    Non-finite values fail. The bounds take delta from the given o (bf16): that is the contract (delta = sum(dO * o)),
    as every FlashAttention backward computes it. The FlashAttention-4 numerics (P recomputed in fp32 with exp2, bf16 P
    for dV, dS formed in fp32 and rounded once to bf16, fp32 dq/dk/dv sums, bf16 outputs), modeled on CPU in
    tests/test_hd256_attn_bwd.py, use under half of any tolerance on every loop and holdout distribution. A window off
    by one key, a key of the wrong KV head, keys of another packed sequence, delta left out of dS, dP rounded to bf16
    before subtracting delta or a missing softmax scale on dk fails."""
    import torch
    stats = {}
    ok = True
    for i, nm in OUTS:
        g, e, tol = (got[i], exp[nm], exp['tol_' + nm])
        if tuple(g.shape) != tuple(e.shape):
            return (False, {'error': f'{nm}: shape {tuple(g.shape)} != {tuple(e.shape)}'})
        if g.dtype != exp[i].dtype:
            return (False, {'error': f'{nm}: dtype {g.dtype} is not {exp[i].dtype}'})
        gf = g.float()
        finite = torch.isfinite(gf)
        diff = torch.where(finite, (gf - e).abs(), torch.full_like(gf, float('inf')))
        bad = int((diff > tol).sum().item())
        fin = torch.where(finite, diff, torch.zeros_like(diff))
        over = fin / (tol + 1e-30)
        stats[nm] = {'bad': bad, 'nonfinite': int((~finite).sum().item()), 'max_abs': float(fin.max().item()) if g.numel() else 0.0, 'max_over_tol': float(over.max().item()) if g.numel() else 0.0}
        ok = ok and bad == 0
    return (ok, stats)

def _bytes(s):
    """Lower bound: q, o, dO and dq (query heads), k, v, dk and dv (KV heads) moved once in bf16, lse read once."""
    T = s['n'] * s['L'] + s.get('long', 0)
    return T * D * 2 * 4 * (s['Hq'] + s['Hkv']) + T * s['Hq'] * 4

def vllm():
    """The bar: FlashAttention-4's hd256 backward, on each shape the fastest public route (expert_pr/fa4_v2966/b300.py
    ROUTE): the dedicated pair or the general kernel at upstream 965c921 (the #2966 merge, expert_pr/fa4_v2966), or the
    general kernel with GQA head packing at open PR #2967's head (expert_pr/fa4_v2967); called dense for batches of
    equal-length sequences and varlen otherwise, writing dq, dk and dv into the output buffers."""
    from kernel_lab.expert_pr.fa4_v2966 import b300, utils as utils_2966
    from kernel_lab.expert_pr.fa4_v2967 import utils as utils_2967
    from kernel_lab.expert_pr.fa4_v2966.interface import _flash_attn_bwd as bwd_2966
    from kernel_lab.expert_pr.fa4_v2967.interface import _flash_attn_bwd as bwd_2967

    def run(q, k, v, o, do, lse, cu, max_len, W, dq, dk, dv):
        T, Hq, _ = q.shape
        n = cu.shape[0] - 1
        kernel, mode = b300.route(T, Hq, k.shape[1], n, max_len, W)
        if kernel == 'packed':
            bwd = bwd_2967
            utils_2967.HD256_GENERIC_BWD = utils_2967.BWD_PACK_GQA = True
        else:
            bwd = bwd_2966
            utils_2966.HD256_GENERIC_BWD = kernel == 'generic'
        win = {} if W < 0 else {'window_size_left': W, 'window_size_right': 0}
        if mode == 'dense':
            S = T // n
            x4 = [t.view(n, S, t.shape[1], D) for t in (q, k, v, o, do, dq, dk, dv)]
            bwd(*x4[:5], lse.view(Hq, n, S).transpose(0, 1), softmax_scale=SCALE, causal=True, dq=x4[5], dk=x4[6], dv=x4[7], **win)
        else:
            bwd(q, k, v, o, do, lse, softmax_scale=SCALE, causal=True, cu_seqlens_q=cu, cu_seqlens_k=cu, max_seqlen_q=max_len, max_seqlen_k=max_len, dq=dq, dk=dk, dv=dv, **win)
    return run
TASK = Task(name='xpr_hd256_attn_bwd', title='Causal attention backward (dq, dk, dv), head dim 256, GQA, sliding windows and packed documents, 8K-16K tokens (bar: FlashAttention-4 after PR #2966, with open PR #2967, October 2026)', models='Training of head-dim-256 models on Blackwell, per GPU: Gemma-3-style global layers and sliding-window-1024 local layers (16 query / 8 KV heads, or 8 / 4), Qwen3-Next-style full-attention layers (16 query heads, 2 KV heads) and packed batches of documents of mixed lengths, 8K-16K tokens per step.', semantics="q is [T, Hq, 256] bf16, k and v are [T, Hkv, 256] bf16 (Hq a multiple of Hkv; query head h uses KV head h // (Hq / Hkv)): the tokens of all sequences packed back to back. o is [T, Hq, 256] bf16, the forward output rounded to bf16; do is [T, Hq, 256] bf16, the gradient of the output; lse is [Hq, T] float32, the forward log-sum-exp (natural log, scale included). cu_seqlens is [n + 1] int32 with cu_seqlens[0] = 0 and sequence s occupying tokens cu_seqlens[s] .. cu_seqlens[s + 1] - 1; max_seqlen is the longest sequence (a python int); window is a python int, -1 or W >= 0. All tensors contiguous. Query i of a sequence sees key j of the same sequence when j <= i and, if window >= 0, i - j <= window. For each query i, head h and visible key j: score = q[i,h] . k[j,h'] / 16, P = exp(score - lse[h,i]), dP = do[i,h] . v[j,h'], delta = sum over the 256 columns of do[i,h] * o[i,h], dS = P (dP - delta). Write dq[i,h] = sum_j dS k[j,h'] / 16, dk[j,h'] = sum over every (i, h) that sees j with h // (Hq / Hkv) = h' of dS q[i,h] / 16, and dv[j,h'] = the same sum of P do[i,h], fp32 (or better) math rounded to bf16, into dq [T, Hq, 256], dk and dv [T, Hkv, 256] (contiguous). dq, dk and dv arrive filled with NaN like every output. P and dS may each be rounded once to bf16 as tensor-core operands; scores, P, dP, delta and dS are otherwise fp32 or better from the given values (no extra rounding of the scores or dP); summation order is free.", signature='def kernel_fn(q, k, v, o, do, lse, cu_seqlens, max_seqlen, window, dq, dk, dv) -> None  # writes dq, dk, dv', shapes=[{'model': 'Gemma-3-12B-style global layer (16 query / 8 KV heads), one 8K sequence', 'n': 1, 'L': 8192, 'Hq': 16, 'Hkv': 8, 'window': -1}, {'model': 'Gemma-3-12B-style local layer, sliding window 1024, one 8K sequence', 'n': 1, 'L': 8192, 'Hq': 16, 'Hkv': 8, 'window': 1023}, {'model': 'Qwen3-Next-style full-attention layer (16 query / 2 KV heads), 2 x 4K sequences', 'n': 2, 'L': 4096, 'Hq': 16, 'Hkv': 2, 'window': -1}, {'model': 'packed documents (16 / 8 heads): 8 documents of 1K-3K tokens, 16K tokens', 'n': 8, 'L': 2048, 'Hq': 16, 'Hkv': 8, 'window': -1, 'ragged': True}, {'model': 'Gemma-3-4B-style global layer (8 query / 4 KV heads), 4 x 2K sequences', 'n': 4, 'L': 2048, 'Hq': 8, 'Hkv': 4, 'window': -1}, {'model': 'Gemma-3-4B-style local layer (8 / 4 heads), sliding window 1024, packed documents of 2K-6K tokens, 16K tokens', 'n': 4, 'L': 4096, 'Hq': 8, 'Hkv': 4, 'window': 1023, 'ragged': True}], edge_shapes=[], dists=[], make_case=make_case, reference=reference, check=check, vllm=vllm, bytes_moved=_bytes, notes="Per visible (query, key) pair and query head, 5 x 256 multiply-adds: recompute the score (q . k), dP (dO . v), and the three products dV += P dO, dQ += dS k and dK += dS q, about 1.4 TFLOP for a causal 8K sequence with 16 query heads. Compute bound; skip key blocks the causal mask or the window hides entirely, and the tiles of other packed sequences. Head dim 256 doubles the register, shared-memory and tensor-memory footprint of each tile compared with 128: one 128 x 256 bf16 tile is 64 KiB, and fp32 dK and dV for 128 keys fill a whole SM's 512 tensor-memory columns. With GQA every key gets dK and dV contributions from all Hq / Hkv query heads of its group; summing them inside one program avoids fp32 atomics over heads. The expert runs, per shape, two kernels (dQ, then dK and dV, each recomputing S and dP) or one kernel with a 64-row KV tile and fp32 dQ partial sums, with or without the group's query heads packed into one tile; none is best on every shape. delta = rowsum(dO * o) needs only o and dO. Expert times on B300 (us, cold L2): 1x8K 1518, window 1024 510, 2x4K with 2 KV heads 980, 8 packed documents 1124-1206 (by lengths), 4x2K 302, packed 16K window 1024 520 (the pair before the series, which has no window: 1774, -, 1118, 1336, 356, -). The bar is 1.0x the expert speed (geometric mean over the timed and hidden shapes, each hidden shape at least 0.8x).", weight=0.5, gpus=('B300',))
# Starting kernel: correct and deliberately simple. A preprocess writes delta = rowsum(dO * o) and each token's
# sequence bounds; one program per (16 query tokens, query head) walks every key its rows can see in 16-key steps and
# accumulates dq; one program per (16 keys, KV head) walks every query head of its group and every query row that can
# see its keys and accumulates dk and dv. P is recomputed from lse with tl.dot; tiny tiles, no pipelining, S and dP
# computed twice: much slower than the expert kernels.
import torch
import triton
import triton.language as tl


@triton.jit
def _prep(O, DO, DELTA, CU, START, END, T, H, D: tl.constexpr, BT: tl.constexpr):
    s = tl.program_id(0)
    h = tl.program_id(1)
    st = tl.load(CU + s)
    en = tl.load(CU + s + 1)
    rd = tl.arange(0, D)
    for t0 in range(st, en, BT):
        rt = t0 + tl.arange(0, BT)
        m = rt < en
        row = (rt.to(tl.int64) * H + h) * D
        o = tl.load(O + row[:, None] + rd[None, :], mask=m[:, None], other=0.0).to(tl.float32)
        do = tl.load(DO + row[:, None] + rd[None, :], mask=m[:, None], other=0.0).to(tl.float32)
        tl.store(DELTA + h.to(tl.int64) * T + rt, tl.sum(o * do, 1), mask=m)
        if h == 0:
            tl.store(START + rt, tl.zeros((BT,), tl.int32) + st, mask=m)
            tl.store(END + rt, tl.zeros((BT,), tl.int32) + en, mask=m)


@triton.jit
def _dq(Q, K, V, DO, LSE, DELTA, START, DQ, sm_scale, T, H, HKV, W, D: tl.constexpr, BM: tl.constexpr,
        BN: tl.constexpr):
    i0 = tl.program_id(0) * BM
    h = tl.program_id(1)
    hk = h // (H // HKV)
    ri = i0 + tl.arange(0, BM)
    mi = ri < T
    rd = tl.arange(0, D)
    rn = tl.arange(0, BN)
    qrow = (ri.to(tl.int64) * H + h) * D
    q = tl.load(Q + qrow[:, None] + rd[None, :], mask=mi[:, None], other=0.0)
    do = tl.load(DO + qrow[:, None] + rd[None, :], mask=mi[:, None], other=0.0)
    lse = tl.load(LSE + h.to(tl.int64) * T + ri, mask=mi, other=0.0)
    delta = tl.load(DELTA + h.to(tl.int64) * T + ri, mask=mi, other=0.0)
    st = tl.load(START + ri, mask=mi, other=T)
    lo = st
    if W >= 0:
        lo = tl.maximum(st, ri - W)
    j_lo = tl.min(tl.where(mi, lo, T), 0)
    j_hi = tl.minimum(i0 + BM, T)
    dq = tl.zeros((BM, D), tl.float32)
    for j0 in range(j_lo, j_hi, BN):
        rj = j0 + rn
        mj = rj < j_hi
        krow = (rj.to(tl.int64) * HKV + hk) * D
        k = tl.load(K + krow[:, None] + rd[None, :], mask=mj[:, None], other=0.0)
        v = tl.load(V + krow[:, None] + rd[None, :], mask=mj[:, None], other=0.0)
        vis = (rj[None, :] >= lo[:, None]) & (rj[None, :] <= ri[:, None]) & mi[:, None] & mj[None, :]
        s = tl.dot(q, tl.trans(k))
        p = tl.where(vis, tl.exp(s * sm_scale - lse[:, None]), 0.0)
        dp = tl.dot(do, tl.trans(v))
        ds = (p * (dp - delta[:, None])).to(tl.bfloat16)
        dq = tl.dot(ds, k, dq)
    tl.store(DQ + qrow[:, None] + rd[None, :], (dq * sm_scale).to(tl.bfloat16), mask=mi[:, None])


@triton.jit
def _dkdv(Q, K, V, DO, LSE, DELTA, END, DK, DV, sm_scale, T, H, HKV, W, D: tl.constexpr, BM: tl.constexpr,
          BN: tl.constexpr):
    j0 = tl.program_id(0) * BN
    hk = tl.program_id(1)
    G = H // HKV
    rj = j0 + tl.arange(0, BN)
    mj = rj < T
    rd = tl.arange(0, D)
    rm = tl.arange(0, BM)
    krow = (rj.to(tl.int64) * HKV + hk) * D
    k = tl.load(K + krow[:, None] + rd[None, :], mask=mj[:, None], other=0.0)
    v = tl.load(V + krow[:, None] + rd[None, :], mask=mj[:, None], other=0.0)
    en = tl.load(END + rj, mask=mj, other=0)
    hi = en
    if W >= 0:
        hi = tl.minimum(en, rj + W + 1)
    i_hi = tl.max(tl.where(mj, hi, 0), 0)
    dk = tl.zeros((BN, D), tl.float32)
    dv = tl.zeros((BN, D), tl.float32)
    for g in range(0, G):
        h = hk * G + g
        for i0 in range(j0, i_hi, BM):
            ri = i0 + rm
            mi = ri < i_hi
            qrow = (ri.to(tl.int64) * H + h) * D
            q = tl.load(Q + qrow[:, None] + rd[None, :], mask=mi[:, None], other=0.0)
            do = tl.load(DO + qrow[:, None] + rd[None, :], mask=mi[:, None], other=0.0)
            lse = tl.load(LSE + h.to(tl.int64) * T + ri, mask=mi, other=0.0)
            delta = tl.load(DELTA + h.to(tl.int64) * T + ri, mask=mi, other=0.0)
            vis = (ri[:, None] >= rj[None, :]) & (ri[:, None] < hi[None, :]) & mi[:, None] & mj[None, :]
            s = tl.dot(q, tl.trans(k))
            p = tl.where(vis, tl.exp(s * sm_scale - lse[:, None]), 0.0)
            dp = tl.dot(do, tl.trans(v))
            ds = (p * (dp - delta[:, None])).to(tl.bfloat16)
            dv = tl.dot(tl.trans(p.to(tl.bfloat16)), do, dv)
            dk = tl.dot(tl.trans(ds), q, dk)
    tl.store(DK + krow[:, None] + rd[None, :], (dk * sm_scale).to(tl.bfloat16), mask=mj[:, None])
    tl.store(DV + krow[:, None] + rd[None, :], dv.to(tl.bfloat16), mask=mj[:, None])


def kernel_fn(q, k, v, o, do, lse, cu_seqlens, max_seqlen, window, dq, dk, dv):
    T, H, D = q.shape
    HKV = k.shape[1]
    n = cu_seqlens.shape[0] - 1
    delta = torch.empty((H, T), device=q.device, dtype=torch.float32)
    start = torch.empty((T,), device=q.device, dtype=torch.int32)
    end = torch.empty((T,), device=q.device, dtype=torch.int32)
    _prep[(n, H)](o, do, delta, cu_seqlens, start, end, T, H, D=D, BT=16, num_warps=4)
    _dq[(triton.cdiv(T, 16), H)](q, k, v, do, lse, delta, start, dq, D ** -0.5, T, H, HKV, window, D=D, BM=16,
                                 BN=16, num_warps=4, num_stages=1)
    _dkdv[(triton.cdiv(T, 16), HKV)](q, k, v, do, lse, delta, end, dk, dv, D ** -0.5, T, H, HKV, window, D=D,
                                     BM=16, BN=16, num_warps=4, num_stages=1)
"""Check a kernel on the public shapes before you submit it:  python check.py [solution.py]

Correctness against the reference on the public shapes (standard normal inputs), then median CUDA-event times of
your kernel, the starting kernel and vLLM's kernel. The verifier checks much more: edge shapes, hidden input
distributions and cheats (torch math, cached results, timers, streams). It times you against vLLM in a paired,
L2-flushed run in a separate container.
"""
import importlib.util
import sys

import torch

import static_check
import task as T


def load(path, name):
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


def median_us(fn, args, n=50):
    for _ in range(3):
        fn(*args)
    torch.cuda.synchronize()
    ts = []
    for _ in range(n):
        s, e = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)
        s.record()
        fn(*args)
        e.record()
        torch.cuda.synchronize()
        ts.append(s.elapsed_time(e) * 1000.0)
    return sorted(ts)[n // 2]


def main():
    t = T.TASK
    path = sys.argv[1] if len(sys.argv) > 1 else '/workspace/solution.py'
    with open(path) as f:
        ok, why = static_check.check(f.read(), getattr(t, 'langs', ('triton',)))
    if not ok:                                  # the verifier rejects the file before running it
        print(f'REJECTED by the static check: {why}')
        return
    sol = load(path, 'solution')
    base = load('/workspace/starting_kernel.py', 'starting_kernel')
    try:
        vllm = t.vllm()
    except Exception as e:  # noqa: BLE001
        vllm, why = None, e
    for shape in t.shapes:
        case = t.make_case(shape, 1, 'cuda')
        exp = t.reference([a.clone() if torch.is_tensor(a) else a for a in case.args])
        sol.kernel_fn(*case.args)
        torch.cuda.synchronize()
        ok, stats = t.check({i: case.args[i] for i in case.checked}, exp)
        line = f"{T.shape_key(shape) if hasattr(T, 'shape_key') else shape}: " + ('correct' if ok else f'WRONG {stats}')
        if ok:
            line += f', {median_us(sol.kernel_fn, case.args):.2f} us'
            line += f', starting kernel {median_us(base.kernel_fn, t.make_case(shape, 1, "cuda").args):.2f} us'
            if vllm:
                label = getattr(T, 'BASELINE_SHORT', 'vLLM')
                line += f', {label} {median_us(vllm, t.make_case(shape, 1, "cuda").args):.2f} us'
        print(line[:400])


if __name__ == '__main__':
    import contextlib
    import io

    import nodus_autosave           # best-version autosave (hardened export): /workspace/.best

    class _Tee(io.StringIO):
        def write(self, s):
            sys.__stdout__.write(s)
            return super().write(s)
    _out = _Tee()
    with contextlib.redirect_stdout(_out):
        main()
    sys.stdout.flush()
    nodus_autosave.record(sys.argv[1] if len(sys.argv) > 1 else '/workspace/solution.py', _out.getvalue())
"""Static admission for candidate kernels. Runs on the controller before any GPU time is spent, and again in the
judge. The contract is deliberately narrow: a candidate is a Triton module that defines kernel_fn. torch may be used
to allocate and reshape tensors. All arithmetic must run inside Triton kernels (the judge also enforces this at
runtime by recording every torch op the candidate executes).

A task whose langs include 'cute' also admits CuTe DSL kernels (the judge image's `cutlass` 4.7.1 package): the
modules in CUTE_IMPORTS, and no cutlass name whose path has a private part or a part in CUTE_BLOCKED (loaders of
prebuilt or exported binaries, external linking, the CUDA driver, compiler internals, prebuilt operator packages,
test and framework helpers, pointers made from integers). Internals are entered only through CUTE_ENTRY. Such a
module may also use `with` on cutlass context managers (cute.arch.elect_one) and declare shared memory as @cute.struct
classes of annotated fields.
"""
import ast

MAX_BYTES = 48 * 1024
CUTE_MAX_BYTES = 128 * 1024        # CuTe spells out the layouts, pipelines and warp roles that Triton infers
ALLOWED_IMPORTS = {'torch', 'triton', 'triton.language', 'triton.language.extra', 'triton.language.extra.libdevice',
                   'math', 'typing'}
TORCH_CHAINS = {
    'torch', 'torch.empty', 'torch.empty_like', 'torch.empty_strided', 'torch.zeros', 'torch.zeros_like', 'torch.full',
    'torch.float32', 'torch.float16', 'torch.bfloat16', 'torch.float8_e4m3fn', 'torch.float8_e5m2', 'torch.int8',
    'torch.int16', 'torch.int32', 'torch.int64', 'torch.uint8', 'torch.bool', 'torch.Tensor', 'torch.dtype',
    'torch.device', 'torch.Size', 'torch.finfo', 'torch.iinfo', 'torch.cuda.get_device_properties',
    'torch.cuda.current_device', 'torch.cuda.is_available',
}
TRITON_CHAINS = {'triton.jit', 'triton.autotune', 'triton.Config', 'triton.cdiv', 'triton.next_power_of_2',
                 'triton.heuristics', 'triton.set_allocator'}   # set_allocator enables TMA descriptors
CUTE_IMPORTS = {
    'cutlass', 'cutlass.cute', 'cutlass.cute.arch', 'cutlass.cute.math', 'cutlass.cute.runtime', 'cutlass.cute.nvgpu',
    'cutlass.cute.nvgpu.cpasync', 'cutlass.cute.nvgpu.tcgen05', 'cutlass.cute.nvgpu.warp',
    'cutlass.cute.nvgpu.warpgroup', 'cutlass.pipeline', 'cutlass.utils', 'cutlass.utils.blackwell_helpers',
    'cutlass.utils.hopper_helpers', 'cutlass.cutlass_dsl', 'cutlass._mlir.dialects', 'cutlass._mlir.dialects.nvvm',
    'cuda.bindings.driver',
}
# The only ways into internals: device-only inline PTX (nvvm), and the stream handle that wraps torch's current stream.
# The llvm dialect is not admitted: @cute.jit host functions run as native host code, where llvm ops (inline
# assembly, external calls, raw loads and stores) escape every check here.
CUTE_ENTRY = ('cutlass.cutlass_dsl.T', 'cutlass.cutlass_dsl.dsl_user_op', 'cutlass._mlir.dialects.nvvm',
              'cuda.bindings.driver.CUstream')
# No part of a cutlass name may be one of these. cutlass modules re-export what they import, so the CUDA driver
# (cutlass.cuda: load_cubin_module, launch_kernel, stream_create) is reachable as an attribute of kernel modules.
CUTE_BLOCKED = {
    'export', 'load_module', 'find_runtime_libraries', 'ffi', 'extern', 'base_extern', 'BitCode', 'LinkLibraries',
    'compile_to', 'make_ptr', 'cuda', 'cuda_helpers', 'driver', 'runtime', 'base_dsl', 'cutlass_dsl', 'compiler',
    'ir', 'spec', 'BaseDSL', 'CuTeDSL', 'CuteExperimentalDSL', 'FFI', 'ExecutionArgs', 'JitArgAdapterRegistry',
    'register_jit_arg_adapter', 'operators', 'kernels', 'testing', 'experimental', 'jax', 'torch', 'tvm_ffi',
    'ctypes', 'os', 'sys', 'subprocess', 'importlib', 'builtins', 'inspect', 'json', 'tempfile', 'pathlib', 'shutil',
    'pickle', 'threading', 'weakref',
}
CUTE_TORCH_CHAINS = {'torch.cuda.current_stream'}      # its .cuda_stream goes into cuda.bindings.driver.CUstream
FORBIDDEN_NAMES = {'eval', 'exec', 'compile', 'open', '__import__', 'globals', 'locals', 'vars', 'getattr', 'setattr',
                   'delattr', 'breakpoint', 'input', 'id', 'memoryview', 'help', 'exit', 'quit', 'sys', 'os',
                   'print', 'super', 'type', 'object'}
FORBIDDEN_NODES = (ast.ClassDef, ast.Global, ast.Nonlocal, ast.AsyncFunctionDef, ast.Await, ast.Yield, ast.YieldFrom,
                   ast.Try, ast.With, ast.AsyncWith)


def _chain(node):
    parts = []
    while isinstance(node, ast.Attribute):
        parts.append(node.attr)
        node = node.value
    if isinstance(node, ast.Name):
        parts.append(node.id)
        return list(reversed(parts))
    return None


def _canon(node, alias):
    c = _chain(node)
    return '.'.join(alias[c[0]].split('.') + c[1:]) if c and c[0] in alias else None


def _cute_bad(full):
    """Why a dotted cutlass or cuda name is not admitted, or None. cutlass.cute.runtime is the one runtime module."""
    head = next((h for h in CUTE_ENTRY if full == h or full.startswith(h + '.')), None)
    if full == 'cuda.bindings.driver':          # the module itself; only CUstream may be taken from it
        return None
    if full.startswith('cuda') and not head:
        return f'{full} is not allowed (from cuda.bindings.driver only CUstream)'
    parts = full.split('.')
    for i in range(len(head.split('.')) if head else 1, len(parts)):
        if parts[:i + 1] != ['cutlass', 'cute', 'runtime'] and (parts[i].startswith('_') or parts[i] in CUTE_BLOCKED):
            return f'{full} is not allowed: {parts[i]} is private or not kernel writing (CUTE_BLOCKED, static_check.py)'
    return None


def _launders(node, alias, parents):
    """Whether a CUTE_BLOCKED attribute sits in a chain that does not start at an imported name (a module kept in a
    variable, a list item, a call result): the chain checks cannot see where such a chain leads, so it is refused."""
    while isinstance(parents.get(id(node)), ast.Attribute) and parents[id(node)].value is node:
        node = parents[id(node)]
    return _canon(node, alias) is None


def _cute_exempt(node, alias):
    """A with statement on cutlass context managers, or a @cute.struct class of annotated fields only."""
    if isinstance(node, ast.With):
        return all(isinstance(i.context_expr, ast.Call) and (_canon(i.context_expr.func, alias) or '').startswith(
            'cutlass.') and (i.optional_vars is None or isinstance(i.optional_vars, ast.Name)) for i in node.items)
    if isinstance(node, ast.ClassDef):
        return (not node.bases and not node.keywords and len(node.decorator_list) == 1
                and _canon(node.decorator_list[0], alias) == 'cutlass.cute.struct'
                and all(isinstance(s, ast.AnnAssign) and isinstance(s.target, ast.Name) and s.value is None
                        or isinstance(s, ast.Pass) or isinstance(s, ast.Expr) and isinstance(s.value, ast.Constant)
                        for s in node.body))
    return False


def check(code, langs=('triton',)):
    """Return (ok, reason). The reason names the first violation, in words the proposer can act on. langs is the
    task's tuple of admitted kernel languages: ('triton',) or ('triton', 'cute')."""
    cute = 'cute' in langs
    limit = CUTE_MAX_BYTES if cute else MAX_BYTES
    imports = ALLOWED_IMPORTS | CUTE_IMPORTS if cute else ALLOWED_IMPORTS
    torch_ok = TORCH_CHAINS | CUTE_TORCH_CHAINS if cute else TORCH_CHAINS
    if not isinstance(code, str) or not code.strip():
        return False, 'empty candidate'
    if len(code.encode()) > limit:
        return False, f'candidate exceeds {limit} bytes'
    try:
        tree = ast.parse(code)
    except SyntaxError as e:
        return False, f'syntax error on line {e.lineno}: {e.msg}'
    parents = {}
    for p in ast.walk(tree):
        for ch in ast.iter_child_nodes(p):
            parents[id(ch)] = p
    alias = {}
    for node in ast.walk(tree):
        if isinstance(node, ast.Import):
            for a in node.names:
                if a.name not in imports:
                    return False, f'import of {a.name} is not allowed (allowed: torch, triton, triton.language, math' + (
                        ' and the CuTe DSL modules in CUTE_IMPORTS)' if cute else ')')
                alias[a.asname or a.name.split('.')[0]] = a.name if a.asname else a.name.split('.')[0]
        elif isinstance(node, ast.ImportFrom):
            mod = node.module or ''
            if node.level or mod not in imports:
                return False, f'import from {mod or "."} is not allowed'
            for a in node.names:
                if a.name == '*':
                    return False, 'star imports are not allowed'
                full = f'{mod}.{a.name}'
                if mod == 'torch' and full not in torch_ok:
                    return False, f'from torch import {a.name} is not allowed'
                if mod == 'triton' and full not in TRITON_CHAINS and a.name != 'language':
                    return False, f'from triton import {a.name} is not allowed'
                if cute and mod in CUTE_IMPORTS and _cute_bad(full):
                    return False, _cute_bad(full)
                alias[a.asname or a.name] = full
    has_kernel_fn = False
    has_jit = False
    has_cute_kernel = False
    for node in ast.walk(tree):
        if isinstance(node, FORBIDDEN_NODES) and not (cute and _cute_exempt(node, alias)):
            return False, f'{type(node).__name__} is not allowed in candidates'
        if isinstance(node, ast.FunctionDef):
            if node.name.startswith('__'):
                return False, f'dunder function {node.name} is not allowed'
            if node.name == 'kernel_fn' and node in tree.body:
                has_kernel_fn = True
            for d in node.decorator_list:
                c = _chain(d.func if isinstance(d, ast.Call) else d)
                if c:
                    canon = alias.get(c[0], c[0]).split('.') + c[1:]
                    if canon[:2] == ['triton', 'jit']:
                        has_jit = True
                    if canon[:3] == ['cutlass', 'cute', 'kernel']:
                        has_cute_kernel = True
        if isinstance(node, ast.Name) and node.id in FORBIDDEN_NAMES:
            return False, f'use of {node.id} is not allowed'
        if isinstance(node, ast.Attribute) and node.attr.startswith('__'):
            return False, f'dunder attribute {node.attr} is not allowed'
        if cute and isinstance(node, ast.Attribute) and node.attr.startswith('_') and node.attr != '_mlir':
            return False, f'private attribute {node.attr} is not allowed'
        if cute and isinstance(node, ast.Attribute) and node.attr in CUTE_BLOCKED and _launders(node, alias, parents):
            return False, (f'attribute {node.attr} is not allowed here: in CuTe modules the names in CUTE_BLOCKED may '
                           'only appear in a chain that starts at an import')
        if cute and isinstance(node, ast.keyword) and node.arg is None:
            return False, 'keyword unpacking (**) is not allowed in CuTe modules: pass compile options by name'
        if cute and isinstance(node, ast.keyword) and node.arg == 'options' and not (
                isinstance(node.value, ast.Constant) and isinstance(node.value.value, str)
                and 'link' not in node.value.value.lower()):
            return False, 'compile options must be a string literal and must not link libraries'
        if isinstance(node, (ast.Assign, ast.AugAssign, ast.AnnAssign, ast.Delete)):
            targets = node.targets if isinstance(node, (ast.Assign, ast.Delete)) else [node.target]
            for t in targets:
                for sub in ast.walk(t):
                    if isinstance(sub, ast.Attribute):
                        return False, 'assigning to attributes is not allowed'
    for node in ast.walk(tree):
        if not isinstance(node, (ast.Attribute, ast.Name)):
            continue
        p = parents.get(id(node))
        if isinstance(p, ast.Attribute) and p.value is node:
            continue  # only check the longest chain
        c = _chain(node)
        if not c or c[0] not in alias:
            continue
        canon = alias[c[0]].split('.') + c[1:]
        full = '.'.join(canon)
        if canon[0] == 'torch':
            if full not in torch_ok:
                return False, f'{full} is not allowed: allocate with torch, compute in Triton'
        elif canon[0] == 'triton':
            if canon[:2] == ['triton', 'language']:
                continue
            if len(canon) > 1 and full not in TRITON_CHAINS:
                return False, f'{full} is not allowed'
        elif cute and canon[0] in ('cutlass', 'cuda'):
            if _cute_bad(full):
                return False, _cute_bad(full)
            if full in CUTE_IMPORTS:       # a module object passed around would escape these checks
                return False, f'the module {full} may only be used through its attributes'
    if not has_kernel_fn:
        return False, 'the module must define a top-level function named kernel_fn'
    if not (has_jit or cute and has_cute_kernel):
        return False, 'the module must contain at least one ' + ('@triton.jit or @cute.kernel kernel' if cute else
                                                                 '@triton.jit kernel')
    return True, 'ok'
