"""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()