DeepInfer/scripts/h20_attn_check.py

"""Validate the SM90 attention path (TileLang sparse MLA) against a torch reference.

Shapes follow DeepSeek-V4.1-Flash: head_dim 512 with RoPE inside, 64 query heads, MQA (kv_group 1).

    python scripts/h20_attn_check.py            # correctness + perf
"""

import torch
from tilelang.profiler import do_bench

from deepinfer.h20_ops.attn_sm90 import (
    CACHE_ELEMENTS_PER_SF,
    decode_mla_mqa_sparse,
    prefill_attn_varlen_sparse,
)

DEV = "cuda"
HEADS = 64
DIM = 512


def ref_attention(q: torch.Tensor, kv: torch.Tensor, indices: torch.Tensor, scale: float, sink: torch.Tensor | None = None) -> torch.Tensor:
    """q [n, h, d], kv [m, 1, d], indices [n, topk] -> [n, h, d] float32."""
    n, h, d = q.shape
    topk = indices.shape[1]
    valid = indices >= 0
    idx = indices.clamp_min(0)
    k = kv[idx.reshape(-1), 0, :].view(n, topk, d).float()  # [n, topk, d]
    v = k
    scores = torch.einsum("nhd,nkd->nhk", q.float(), k) * scale  # [n, h, topk]
    scores = scores.masked_fill(~valid.unsqueeze(1), float("-inf"))
    probs = scores.softmax(dim=-1)
    out = torch.einsum("nhk,nkd->nhd", probs, v)
    lse = torch.logsumexp(scores, dim=-1)
    if sink is not None:
        out = out * torch.sigmoid(lse - sink.float().unsqueeze(0)).unsqueeze(-1)
    return out, lse


def make_indices(n: int, m: int, topk: int, gen: torch.Generator) -> torch.Tensor:
    """Causal top-k like the indexer produces: random subset of [0, position], -1 padded."""
    indices = torch.full((n, topk), -1, dtype=torch.int32, device=DEV)
    for t in range(n):
        hi = min(t + 1, m)
        take = min(topk, hi)
        picked = torch.randperm(hi, generator=gen)[:take]
        indices[t, :take] = picked.to(torch.int32)
    return indices


def quantize_cache(kv: torch.Tensor, block_size: int = 64) -> torch.Tensor:
    """Build a V41-format paged cache ([rows, 528] uint8) from bf16 kv [m, 1, d]."""
    m = kv.shape[0]
    values = kv[:, 0, :].float()
    groups = values.view(m, DIM // CACHE_ELEMENTS_PER_SF, CACHE_ELEMENTS_PER_SF)
    amax = groups.abs().amax(dim=-1).clamp_min(torch.finfo(torch.float32).tiny)
    # UE8M0: round the per-32 scale up to a power of two so the byte encoding is exact.
    exp = torch.ceil(torch.log2(amax / 448.0))
    exp = exp.clamp_min(-126)
    scale = torch.pow(2.0, exp)
    q = (groups / scale.unsqueeze(-1)).to(torch.float8_e4m3fn)
    packed = torch.cat([q.reshape(m, DIM).view(torch.uint8), (exp + 127).to(torch.uint8)], dim=1)  # [m, 528]
    deq = (q.float() * scale.unsqueeze(-1)).reshape(m, DIM).to(torch.bfloat16)
    rows = packed.contiguous()
    n_blocks = (m + block_size - 1) // block_size
    pad = n_blocks * block_size - m
    if pad:
        rows = torch.cat([rows, rows.new_zeros(pad, rows.shape[1])], dim=0)
        deq = torch.cat([deq, deq.new_zeros(pad, deq.shape[1])], dim=0)
    return rows.view(n_blocks, block_size, 1, -1), deq


def main() -> None:
    torch.manual_seed(0)
    gen = torch.Generator(device="cpu").manual_seed(0)
    scale = DIM**-0.5

    # ---------------- prefill: n query tokens packed in one sequence ----------------
    n, m, topk = 512, 1024, 512
    q = torch.randn(n, HEADS, DIM, device=DEV, dtype=torch.bfloat16)
    kv = torch.randn(m, 1, DIM, device=DEV, dtype=torch.bfloat16)
    indices = make_indices(n, m, topk, gen)
    sink = torch.randn(HEADS, device=DEV, dtype=torch.float32) * 2

    out, lse = prefill_attn_varlen_sparse(q, kv, indices, scale, attn_sink=sink)
    ref, _ = ref_attention(q, kv, indices, scale, sink=sink)
    err = (out.float() - ref).abs().max().item()
    print(f"prefill (n={n}, m={m}, topk={topk}, heads={HEADS}): max_abs_err={err:.4f}")

    # ---------------- absolute positions: row index must not act as a causal bound --------
    # The engine's KV buffer is `[cached swa | current swa | compressed]`, so indices are
    # absolute positions that can exceed the query's row number. Deriving causality from the
    # row index (the original TileLang example did) silently drops every index here.
    n_abs, m_abs, topk_abs = 16, 4096, 256
    pad = 3000  # stand-in for `total_k_swa_cached + total_k_swa_cur`
    q_abs = torch.randn(n_abs, HEADS, DIM, device=DEV, dtype=torch.bfloat16)
    kv_abs = torch.randn(m_abs, 1, DIM, device=DEV, dtype=torch.bfloat16)
    gen_abs = torch.Generator(device="cpu").manual_seed(7)
    idx_abs = torch.full((n_abs, topk_abs), -1, dtype=torch.int32, device=DEV)
    for t in range(n_abs):
        lo = pad + t
        picked = lo + torch.randperm(min(64, m_abs - lo - 1), generator=gen_abs)[:64]
        idx_abs[t, : picked.numel()] = picked.to(torch.int32)
    assert (idx_abs >= 0).sum() > 0 and int(idx_abs.max()) > n_abs
    out_abs, _ = prefill_attn_varlen_sparse(q_abs, kv_abs, idx_abs, scale)
    ref_abs, _ = ref_attention(q_abs, kv_abs, idx_abs, scale)
    err_abs = (out_abs.float() - ref_abs).abs().max().item()
    print(f"absolute positions (indices up to {int(idx_abs.max())} vs rows {n_abs}): max_abs_err={err_abs:.4f}")
    assert err_abs < 0.2, "row-index causal bound regression"

    # ---------------- empty window: all-`-1` indices must yield zeros, not NaN ----------------
    # A padded batch looks exactly like this (no valid top-k for any row); a 0/0 softmax here
    # used to turn the entire hidden state into NaN on the next layer.
    q_empty = torch.randn(8, HEADS, DIM, device=DEV, dtype=torch.bfloat16)
    kv_empty = torch.randn(64, 1, DIM, device=DEV, dtype=torch.bfloat16)
    idx_empty = torch.full((8, 128), -1, dtype=torch.int32, device=DEV)
    out_e, lse_e = prefill_attn_varlen_sparse(q_empty, kv_empty, idx_empty, scale)
    print(f"empty window: finite={bool(torch.isfinite(out_e).all())} zero={bool((out_e == 0).all())} lse_max={lse_e.max().item():.3e}")
    assert torch.isfinite(out_e).all() and (out_e == 0).all()

    # ---------------- decode: paged fp8 cache + gather ----------------
    n_dec, m_dec, topk_dec = 8, 4096, 1024
    q2 = torch.randn(n_dec, HEADS, DIM, device=DEV, dtype=torch.bfloat16)
    kv2 = torch.randn(m_dec, 1, DIM, device=DEV, dtype=torch.bfloat16)
    cache, kv_deq = quantize_cache(kv2)

    idx2 = torch.randint(0, m_dec, (n_dec, topk_dec), device=DEV, dtype=torch.int32)
    dec, _ = decode_mla_mqa_sparse(q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink)
    ref2, _ = ref_attention(q2, kv_deq.unsqueeze(1), idx2, scale, sink=sink)
    err2 = (dec.float() - ref2).abs().max().item()
    corr2 = torch.cosine_similarity(dec.float().flatten(), ref2.flatten(), dim=0).item()
    print(f"decode (n={n_dec}, cache={m_dec}, topk={topk_dec}): max_abs_err={err2:.4f} corr={corr2:.6f}")

    ms2 = do_bench(
        lambda: decode_mla_mqa_sparse(q2, cache, DIM, idx2, sched_meta=None, softmax_scale=scale, attn_sink=sink), warmup=5, rep=20
    )
    print(f"decode perf: batch={n_dec} topk={topk_dec} -> {ms2:.2f} ms/step")

    # ---------------- perf: full-size prefill ----------------
    n_big, m_big, topk_big = 4096, 8192, 2048
    qb = torch.randn(n_big, HEADS, DIM, device=DEV, dtype=torch.bfloat16)
    kvb = torch.randn(m_big, 1, DIM, device=DEV, dtype=torch.bfloat16)
    ib = make_indices(n_big, m_big, topk_big, gen)
    ms = do_bench(lambda: prefill_attn_varlen_sparse(qb, kvb, ib, scale), warmup=10, rep=50)
    tflops = n_big * 2 * 2 * topk_big * DIM * HEADS / (ms * 1e-3) / 1e12
    print(f"prefill perf: {ms:.2f} ms  -> {tflops:.1f} TFLOP/s (H20 bf16 peak ~148)")

    print("ATTN_CHECK_DONE")


if __name__ == "__main__":
    main()