"""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)  # boolean mask (no NaN)
    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)

    # (1) correctness at moderate T
    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'}")

    # (2) determinism
    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'}")

    # (3) scale — flex at big T (eager T x T mask would be H*T*T*2 bytes)
    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()