#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Checks whether nulling an all-True attn_mask helps SDPA on the production cuDNN path."""

import time
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

B, H, N, D = 1, 16, 50345, 128
dev, dt = "cuda:0", torch.bfloat16


def mk():
    return torch.randn(B, H, N, D, device = dev, dtype = dt)


def timed(fn, iters = 20):
    try:
        torch.cuda.synchronize()
        for _ in range(3):
            fn()
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        for _ in range(iters):
            fn()
        torch.cuda.synchronize()
        return (time.perf_counter() - t0) / iters * 1e3
    except torch.OutOfMemoryError:
        # OOM on the dense NxN mask is a memory limit, not a backend rejecting it; don't mislabel UNSUPPORTED.
        torch.cuda.empty_cache()
        return "OOM"
    except Exception as e:  # noqa: BLE001
        return f"UNSUPPORTED ({type(e).__name__})"


def _identity(run, reference):
    """Bitwise match to default dispatch; a backend that cannot run reports why, not a mismatch."""
    if reference is None:
        return "n/a"
    try:
        out = run()
    except torch.OutOfMemoryError:
        torch.cuda.empty_cache()
        return "OOM"
    except Exception:  # noqa: BLE001
        return "unsupported"
    return "yes" if torch.equal(reference, out) else f"no ({(reference - out).abs().max():.1e})"


q, k, v = mk(), mk(), mk()
dense = torch.ones(B, 1, N, N, dtype = torch.bool, device = dev)

backends = {
    "default(dispatch)": None,
    "MATH": [SDPBackend.MATH],
    "FLASH": [SDPBackend.FLASH_ATTENTION],
    "EFFICIENT": [SDPBackend.EFFICIENT_ATTENTION],
    "CUDNN": [SDPBackend.CUDNN_ATTENTION],
}

# Bitwise identity with the default dispatch tells which backend the dispatcher chose.
try:
    reference = F.scaled_dot_product_attention(q, k, v, attn_mask = dense)
except Exception:  # noqa: BLE001 -- no reference: the identity column just reports n/a
    reference = None

print(f"shape B={B} H={H} N={N} D={D} {dt}\n")
print(f"{'backend':<20}{'mask=dense(ms)':>18}{'mask=None(ms)':>18}{'==default(dense)':>19}")
for name, bk in backends.items():

    def run_dense():
        if bk is None:
            return F.scaled_dot_product_attention(q, k, v, attn_mask = dense)
        with sdpa_kernel(bk):
            return F.scaled_dot_product_attention(q, k, v, attn_mask = dense)

    def run_none():
        if bk is None:
            return F.scaled_dot_product_attention(q, k, v, attn_mask = None)
        with sdpa_kernel(bk):
            return F.scaled_dot_product_attention(q, k, v, attn_mask = None)

    dms = timed(run_dense)
    nms = timed(run_none)
    d_s = f"{dms:.2f}" if isinstance(dms, float) else dms
    n_s = f"{nms:.2f}" if isinstance(nms, float) else nms
    print(f"{name:<20}{d_s:>18}{n_s:>18}{_identity(run_dense, reference):>19}")