baseten-admin/wan2.2_production_values

Wan2.2 Production Values — VSA Block-Sparse Attention Inputs + Reference

Real, captured production inputs for the Video Sparse Attention (VSA) fine block-sparse attention stage of Wan2.2 T2V-A14B, recorded from an actual generate.py run at 832×480 / 81 frames, plus the reference Triton kernel and the scoring rule.

This is a fixed, non-gameable benchmark for proposing a faster VSA forward kernel: optimize on these exact tensors, score against the reference at the tolerance below. You do not get to substitute your own input generator — that loophole (sliding-window / head-broadcast q2k_idx, small-magnitude N(0,1) q/k, all-full blocks, loose tolerance) is exactly what these captures close.

Files

Shapes (t2v-A14B @ 832×480 / 81f)

S_padded = 39936 = 624 cubes × 64        (cube-major tiled/padded)
real tokens = 32760 = 21 × 30 × 52       (before cube padding)
Nq = Nkv  = 624 = ceil(21/4)·ceil(30/4)·ceil(52/4) = 6 × 8 × 13
topk = 78  (87.5% sparsity),  D = 128,  sm_scale = 1/sqrt(128)

Operator contract

Non-causal block-sparse attention, BSHD:

out = softmax((q @ k_selected^T) / sqrt(D)) @ v_selected

For each (batch, head, query_cube), topk_idx names the selected KV cubes; vbs[c] masks padded tokens inside cube c (attend only the first vbs[c] of 64).

Scoring / loss (the rule)

A candidate kernel cand(q, k, v, topk_idx, vbs) -> o (bf16, (1,39936,40,128)) passes a capture iff, against the reference output o_ref:

1. 1.finite: every element of o is finite (no NaN/Inf) — enforced on every capture, including call_outlier.pt;
2. 2.per-head cosine ≥ 0.9999: min over (b,h) of cos(o[b,:,h,:], o_ref[b,:,h,:]) ≥ 0.9999 (a global average can hide a few wrong heads — it is checked per head);
3. 3.elementwise: max|o − o_ref| ≤ 3e-2 and mean|o − o_ref| ≤ 1e-4 (bf16-level).

o_ref is triton_block_sparse_attn_forward(...)[0] for the clean captures, and the bundled o_triton for call_outlier.pt.

Win condition: pass all four captures, and have lower warmed, kernel-only latency (CUDA events, B200 / sm_100a) than the deployed baseline. For reference, the deployed CUDA kernel hits cosine ≈ 0.999999 on the clean captures, stays finite on the outlier head, and runs at ~4.7–5.0 ms on these inputs.

Reference scorer

import torch, glob, os
from standalone_kernel import triton_block_sparse_attn_forward

def per_head_cos(a, b):              # a,b: (1, S, H, D)
    a = a.float(); b = b.float()
    a = a.permute(0,2,1,3).reshape(a.shape[2], -1)   # (H, S*D)
    b = b.permute(0,2,1,3).reshape(b.shape[2], -1)
    return torch.nn.functional.cosine_similarity(a, b, dim=1)   # (H,)

def score(cand, path):
    d = torch.load(path, weights_only=True)
    q, k, v = d["q"].cuda(), d["k"].cuda(), d["v"].cuda()
    idx = d.get("topk_idx", d.get("idx")).cuda().contiguous()
    vbs = d["vbs"].cuda().contiguous()
    o_ref = d["o_triton"].cuda() if "o_triton" in d else \
            triton_block_sparse_attn_forward(q, k, v, idx, vbs)[0]
    o = cand(q, k, v, idx, vbs)                       # <-- your kernel
    finite = bool(torch.isfinite(o).all())
    coss   = per_head_cos(o, o_ref)
    diff   = (o.float() - o_ref.float()).abs()
    ok = finite and coss.min() >= 0.9999 and diff.max() <= 3e-2 and diff.mean() <= 1e-4
    print(f"{os.path.basename(path):16s} finite={finite} "
          f"min_head_cos={coss.min():.6f} max|Δ|={diff.max():.2e} -> {'PASS' if ok else 'FAIL'}")
    return ok

# all(score(my_kernel, f) for f in sorted(glob.glob('call_*.pt')))

Load

import torch
d = torch.load("call_0021.pt", weights_only=True)
q, k, v = d["q"], d["k"], d["v"]            # (1, 39936, 40, 128) bf16
topk_idx = d["topk_idx"]                     # (1, 40, 624, 78) int32
vbs = d["vbs"]                               # (624,) int32
# call_outlier.pt uses key "idx" instead of "topk_idx" and ships o_triton/m_triton.

Captured from Wan2.2 T2V-A14B (Apache-2.0 model). Tensors are intermediate attention activations; no prompts or weights are included.