| """Operator-level validation of block attention under FlexAttention. |
| |
| 1) CORRECTNESS: flex block-attention output == reference SDPA-with-additive-block-mask (small T). |
| 2) STABILITY: flex is deterministic across repeated runs. |
| 3) SCALE: flex runs at long T (O(T) memory) where the eager T x T mask would need terabytes. |
| |
| This validates the kernel/mask logic your serving stack (vLLM) needs, independent of the HF model. |
| """ |
| from __future__ import annotations |
| import argparse, torch |
| import torch.nn.functional as F |
| from torch.nn.attention.flex_attention import create_block_mask, flex_attention |
|
|
| NEG = torch.finfo(torch.float32).min |
|
|
|
|
| def make_block_ids(T, n_blocks, dev): |
| """token 0..s = system(-1); then n_blocks equal blocks; last quarter = query(-2).""" |
| ids = torch.empty(T, dtype=torch.long, device=dev) |
| s = max(1, T // 20) |
| q = T - max(1, T // 4) |
| ids[:s] = -1 |
| ids[q:] = -2 |
| body = q - s |
| per = max(1, body // n_blocks) |
| for b in range(n_blocks): |
| ids[s + b*per : s + (b+1)*per] = b |
| ids[s + n_blocks*per : q] = n_blocks - 1 |
| return ids |
|
|
|
|
| def block_mask_mod(bids): |
| def mm(b, h, qi, ki): |
| causal = ki <= qi |
| bq = bids[qi]; bk = bids[ki] |
| return causal & ((bk == -1) | (bq == bk) | (bq == -2)) |
| return mm |
|
|
|
|
| def ref_sdpa(q, k, v, bids): |
| T = q.shape[-2]; dev = q.device |
| causal = torch.tril(torch.ones(T, T, dtype=torch.bool, device=dev)) |
| bq = bids.view(T, 1); bk = bids.view(1, T) |
| allowed = (((bk == -1) | (bq == bk) | (bq == -2)) & causal).view(1, 1, T, T) |
| return F.scaled_dot_product_attention(q, k, v, attn_mask=allowed) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--H", type=int, default=8); ap.add_argument("--D", type=int, default=128) |
| ap.add_argument("--T", type=int, default=1024); ap.add_argument("--blocks", type=int, default=8) |
| ap.add_argument("--big-T", type=int, default=16384) |
| args = ap.parse_args() |
| dev = "cuda"; torch.manual_seed(0) |
|
|
| |
| T = args.T |
| q = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16) |
| k = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16) |
| v = torch.randn(1, args.H, T, args.D, device=dev, dtype=torch.bfloat16) |
| bids = make_block_ids(T, args.blocks, dev) |
| bm = create_block_mask(block_mask_mod(bids), B=None, H=None, Q_LEN=T, KV_LEN=T, device=dev) |
| flex = torch.compile(flex_attention) |
| o_flex = flex(q, k, v, block_mask=bm) |
| o_ref = ref_sdpa(q, k, v, bids) |
| d = (o_flex.float() - o_ref.float()).abs() |
| print(f"(1) CORRECTNESS T={T} H={args.H} D={args.D} blocks={args.blocks}") |
| print(f" max|Δ|={d.max():.4f} mean|Δ|={d.mean():.6f} (bf16 noise floor ~1e-2) -> " |
| f"{'MATCH' if d.max()<0.05 else 'MISMATCH'}") |
|
|
| |
| o2 = flex(q, k, v, block_mask=bm) |
| print(f"(2) STABILITY flex run twice: max|Δ|={(o_flex.float()-o2.float()).abs().max():.6f} " |
| f"-> {'deterministic' if (o_flex.float()-o2.float()).abs().max()<1e-3 else 'nondeterministic'}") |
|
|
| |
| bt = args.big_T |
| qb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16) |
| kb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16) |
| vb = torch.randn(1, args.H, bt, args.D, device=dev, dtype=torch.bfloat16) |
| bidsb = make_block_ids(bt, args.blocks * 8, dev) |
| bmb = create_block_mask(block_mask_mod(bidsb), B=None, H=None, Q_LEN=bt, KV_LEN=bt, device=dev) |
| torch.cuda.reset_peak_memory_stats() |
| ob = flex(qb, kb, vb, block_mask=bmb); torch.cuda.synchronize() |
| peak = torch.cuda.max_memory_allocated() / 1e9 |
| eager_mask_gb = args.H * bt * bt * 2 / 1e9 |
| print(f"(3) SCALE flex at T={bt}: OK, peak mem {peak:.1f} GB, out finite={bool(torch.isfinite(ob).all())}") |
| print(f" (an eager [1,{args.H},{bt},{bt}] score/mask alone would be ~{eager_mask_gb:.0f} GB -> impossible)") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
