# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""
PyTorch custom operator wrapping cuDNN's Scaled Dot-Product Attention (SDPA).

Provides ``scaled_dot_product_attention`` as the public entry point, closely
matching the signature of ``torch.nn.functional.scaled_dot_product_attention``.

**Layout convention**: tensors are expected in **BHSD** layout
``(batch, num_heads, seq_len, head_dim)`` — matching both the cuDNN convention
and PyTorch's ``torch.nn.functional.scaled_dot_product_attention`` layout.

Graph caching ensures cuDNN graphs are built once per unique configuration
and reused across calls.
"""

import logging
import math
from typing import Optional, Tuple, Dict
from enum import IntEnum

import torch
import cudnn

_logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Module-level state
# ---------------------------------------------------------------------------

_cudnn_handles = {}
_fprop_cache: Dict[tuple, tuple] = {}
_bprop_cache: Dict[tuple, tuple] = {}

# Dtype mapping (module-level constant)
_TORCH_DTYPE_TO_CUDNN = {
    torch.float16: cudnn.data_type.HALF,
    torch.bfloat16: cudnn.data_type.BFLOAT16,
    torch.float32: cudnn.data_type.FLOAT,
    torch.int32: cudnn.data_type.INT32,
    torch.int64: cudnn.data_type.INT64,
}


# ---------------------------------------------------------------------------
# UID enum — explicit tensor UIDs for graph caching
# ---------------------------------------------------------------------------


class _UIDs(IntEnum):
    Q = 1
    K = 2
    V = 3
    O = 100
    STATS = 101
    DO = 200
    DQ = 201
    DK = 202
    DV = 203
    SEQ_LEN_Q = 300
    SEQ_LEN_KV = 301
    CUM_SEQ_LEN_Q = 302
    CUM_SEQ_LEN_KV = 303


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _get_handle(device: torch.device):
    """Return a lazily-initialised cuDNN handle with the current CUDA stream."""
    if device not in _cudnn_handles:
        _cudnn_handles[device] = cudnn.create_handle()
    stream = _get_current_stream(device)
    cudnn.set_stream(handle=_cudnn_handles[device], stream=stream)
    return _cudnn_handles[device]


def _get_current_stream(device: torch.device):
    """Return the caller's active CUDA stream for the given device."""
    return torch.cuda.current_stream(device).cuda_stream


def _torch_dtype_to_cudnn(dtype: torch.dtype):
    """Map a PyTorch dtype to a cuDNN data_type enum."""
    return _TORCH_DTYPE_TO_CUDNN[dtype]


def _diagonal_alignment_enum(val: int):
    """Convert int sentinel to cudnn.diagonal_alignment enum."""
    if val == 0:
        return cudnn.diagonal_alignment.TOP_LEFT
    return cudnn.diagonal_alignment.BOTTOM_RIGHT


def _make_fprop_cache_key(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attn_scale: float,
    is_causal: bool,
    diagonal_alignment: int,
    left_bound: int,
    right_bound: int,
    has_seq_len_q: bool,
    has_seq_len_kv: bool,
    has_cum_q: bool,
    has_cum_kv: bool,
):
    q_shape, q_stride = tuple(q.shape), tuple(q.stride())
    k_shape, k_stride = tuple(k.shape), tuple(k.stride())
    v_shape, v_stride = tuple(v.shape), tuple(v.stride())
    return (
        "fprop",
        q_shape,
        q_stride,
        q.dtype,
        k_shape,
        k_stride,
        k.dtype,
        v_shape,
        v_stride,
        v.dtype,
        attn_scale,
        is_causal,
        diagonal_alignment,
        left_bound,
        right_bound,
        has_seq_len_q,
        has_seq_len_kv,
        has_cum_q,
        has_cum_kv,
        q.device,
    )


def _make_bprop_cache_key(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attn_scale: float,
    is_causal: bool,
    diagonal_alignment: int,
    left_bound: int,
    right_bound: int,
    has_seq_len_q: bool,
    has_seq_len_kv: bool,
    has_cum_q: bool,
    has_cum_kv: bool,
    is_deterministic: bool,
):
    q_shape, q_stride = tuple(q.shape), tuple(q.stride())
    k_shape, k_stride = tuple(k.shape), tuple(k.stride())
    v_shape, v_stride = tuple(v.shape), tuple(v.stride())
    return (
        "bprop",
        q_shape,
        q_stride,
        q.dtype,
        k_shape,
        k_stride,
        k.dtype,
        v_shape,
        v_stride,
        v.dtype,
        attn_scale,
        is_causal,
        diagonal_alignment,
        left_bound,
        right_bound,
        has_seq_len_q,
        has_seq_len_kv,
        has_cum_q,
        has_cum_kv,
        is_deterministic,
        q.device,
    )


# ---------------------------------------------------------------------------
# Forward graph builder
# ---------------------------------------------------------------------------


def _build_fprop_graph(
    handle,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attn_scale: float,
    is_causal: bool,
    diagonal_alignment: int,
    left_bound: int,
    right_bound: int,
    seq_len_q: Optional[torch.Tensor],
    seq_len_kv: Optional[torch.Tensor],
    cumulative_seq_len_q: Optional[torch.Tensor],
    cumulative_seq_len_kv: Optional[torch.Tensor],
):
    """Build, validate, and compile a forward SDPA cuDNN graph."""

    io_dtype = _torch_dtype_to_cudnn(q.dtype)

    _logger.debug(f"Building forward graph for q: {q.shape}, k: {k.shape}, v: {v.shape}")

    q_shape, q_stride = tuple(q.shape), tuple(q.stride())
    k_shape, k_stride = tuple(k.shape), tuple(k.stride())
    v_shape, v_stride = tuple(v.shape), tuple(v.stride())

    B, H_q, S_q, D_qk = q_shape
    _, H_v, S_kv, D_v = v_shape

    graph = cudnn.pygraph(
        handle=handle,
        io_data_type=io_dtype,
        intermediate_data_type=cudnn.data_type.FLOAT,
        compute_data_type=cudnn.data_type.FLOAT,
    )

    # -- Input tensors --
    q_t = graph.tensor(name="q", dim=list(q_shape), stride=list(q_stride), data_type=io_dtype, uid=_UIDs.Q)
    k_t = graph.tensor(name="k", dim=list(k_shape), stride=list(k_stride), data_type=io_dtype, uid=_UIDs.K)
    v_t = graph.tensor(name="v", dim=list(v_shape), stride=list(v_stride), data_type=io_dtype, uid=_UIDs.V)

    # -- Optional tensors --
    seq_len_q_t = None
    seq_len_kv_t = None
    cum_q_t = None
    cum_kv_t = None

    if seq_len_q is not None:
        seq_len_q_t = graph.tensor(
            name="seq_len_q", dim=list(seq_len_q.shape), stride=list(seq_len_q.stride()), data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_Q
        )
    if seq_len_kv is not None:
        seq_len_kv_t = graph.tensor(
            name="seq_len_kv", dim=list(seq_len_kv.shape), stride=list(seq_len_kv.stride()), data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_KV
        )
    if cumulative_seq_len_q is not None:
        cum_q_t = graph.tensor(
            name="cum_seq_len_q",
            dim=list(cumulative_seq_len_q.shape),
            stride=list(cumulative_seq_len_q.stride()),
            data_type=cudnn.data_type.INT32,
            uid=_UIDs.CUM_SEQ_LEN_Q,
        )
    if cumulative_seq_len_kv is not None:
        cum_kv_t = graph.tensor(
            name="cum_seq_len_kv",
            dim=list(cumulative_seq_len_kv.shape),
            stride=list(cumulative_seq_len_kv.stride()),
            data_type=cudnn.data_type.INT32,
            uid=_UIDs.CUM_SEQ_LEN_KV,
        )

    # -- Ragged offsets --
    if cum_q_t is not None:
        q_t.set_ragged_offset(cum_q_t)
    if cum_kv_t is not None:
        k_t.set_ragged_offset(cum_kv_t)
        v_t.set_ragged_offset(cum_kv_t)

    # -- Mask configuration --
    use_padding = seq_len_q is not None or seq_len_kv is not None
    lb = left_bound if left_bound >= 0 else None
    rb = right_bound if right_bound >= 0 else None
    if is_causal and rb is None:
        rb = 0

    # -- SDPA forward --
    o_t, stats_t = graph.sdpa(
        name="sdpa",
        q=q_t,
        k=k_t,
        v=v_t,
        generate_stats=True,
        attn_scale=attn_scale,
        use_padding_mask=use_padding,
        seq_len_q=seq_len_q_t,
        seq_len_kv=seq_len_kv_t,
        diagonal_alignment=_diagonal_alignment_enum(diagonal_alignment),
        diagonal_band_left_bound=lb,
        diagonal_band_right_bound=rb,
        compute_data_type=cudnn.data_type.FLOAT,
    )

    # -- Output shapes (BHSD contiguous) --
    o_shape = (B, H_q, S_q, D_v)
    o_stride = (H_q * S_q * D_v, S_q * D_v, D_v, 1)
    o_t.set_uid(_UIDs.O).set_output(True).set_dim(list(o_shape)).set_stride(list(o_stride))
    o_t.set_data_type(io_dtype)

    if cum_q_t is not None:
        o_t.set_ragged_offset(cum_q_t)

    stats_t.set_uid(_UIDs.STATS).set_output(True).set_data_type(cudnn.data_type.FLOAT)
    if cum_q_t is not None:
        stats_t.set_ragged_offset(cum_q_t)

    # -- Build --
    graph.validate()
    graph.build_operation_graph()
    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
    graph.check_support()
    graph.build_plans()

    workspace_size = graph.get_workspace_size()

    return graph, workspace_size


# ---------------------------------------------------------------------------
# Backward graph builder
# ---------------------------------------------------------------------------


def _build_bprop_graph(
    handle,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    o: torch.Tensor,
    dO: torch.Tensor,
    stats: torch.Tensor,
    attn_scale: float,
    is_causal: bool,
    diagonal_alignment: int,
    left_bound: int,
    right_bound: int,
    seq_len_q: Optional[torch.Tensor],
    seq_len_kv: Optional[torch.Tensor],
    cumulative_seq_len_q: Optional[torch.Tensor],
    cumulative_seq_len_kv: Optional[torch.Tensor],
    is_deterministic: bool,
):
    """Build, validate, and compile a backward SDPA cuDNN graph."""

    io_dtype = _torch_dtype_to_cudnn(q.dtype)

    _logger.debug(f"Building backward graph for q: {q.shape}, k: {k.shape}, v: {v.shape}, o: {o.shape}, dO: {dO.shape}")

    q_shape, q_stride = tuple(q.shape), tuple(q.stride())
    k_shape, k_stride = tuple(k.shape), tuple(k.stride())
    v_shape, v_stride = tuple(v.shape), tuple(v.stride())
    o_shape, o_stride = tuple(o.shape), tuple(o.stride())
    dO_shape, dO_stride = tuple(dO.shape), tuple(dO.stride())

    B, H_q, S_q, D_qk = q_shape
    _, H_k, S_kv, _ = k_shape
    _, H_v, _, D_v = v_shape

    graph = cudnn.pygraph(
        handle=handle,
        io_data_type=io_dtype,
        intermediate_data_type=cudnn.data_type.FLOAT,
        compute_data_type=cudnn.data_type.FLOAT,
    )

    # -- Input tensors --
    q_t = graph.tensor(name="q", dim=list(q_shape), stride=list(q_stride), data_type=io_dtype, uid=_UIDs.Q)
    k_t = graph.tensor(name="k", dim=list(k_shape), stride=list(k_stride), data_type=io_dtype, uid=_UIDs.K)
    v_t = graph.tensor(name="v", dim=list(v_shape), stride=list(v_stride), data_type=io_dtype, uid=_UIDs.V)
    o_t = graph.tensor(name="o", dim=list(o_shape), stride=list(o_stride), data_type=io_dtype, uid=_UIDs.O)
    dO_t = graph.tensor(name="dO", dim=list(dO_shape), stride=list(dO_stride), data_type=io_dtype, uid=_UIDs.DO)
    stats_t = graph.tensor(name="stats", dim=list(stats.shape), stride=list(stats.stride()), data_type=cudnn.data_type.FLOAT, uid=_UIDs.STATS)

    # -- Optional tensors --
    seq_len_q_t = None
    seq_len_kv_t = None
    cum_q_t = None
    cum_kv_t = None

    if seq_len_q is not None:
        seq_len_q_t = graph.tensor(
            name="seq_len_q", dim=list(seq_len_q.shape), stride=list(seq_len_q.stride()), data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_Q
        )
    if seq_len_kv is not None:
        seq_len_kv_t = graph.tensor(
            name="seq_len_kv", dim=list(seq_len_kv.shape), stride=list(seq_len_kv.stride()), data_type=cudnn.data_type.INT32, uid=_UIDs.SEQ_LEN_KV
        )
    if cumulative_seq_len_q is not None:
        cum_q_t = graph.tensor(
            name="cum_seq_len_q",
            dim=list(cumulative_seq_len_q.shape),
            stride=list(cumulative_seq_len_q.stride()),
            data_type=cudnn.data_type.INT32,
            uid=_UIDs.CUM_SEQ_LEN_Q,
        )
    if cumulative_seq_len_kv is not None:
        cum_kv_t = graph.tensor(
            name="cum_seq_len_kv",
            dim=list(cumulative_seq_len_kv.shape),
            stride=list(cumulative_seq_len_kv.stride()),
            data_type=cudnn.data_type.INT32,
            uid=_UIDs.CUM_SEQ_LEN_KV,
        )

    # -- Ragged offsets --
    if cum_q_t is not None:
        q_t.set_ragged_offset(cum_q_t)
        o_t.set_ragged_offset(cum_q_t)
        dO_t.set_ragged_offset(cum_q_t)
    if cum_kv_t is not None:
        k_t.set_ragged_offset(cum_kv_t)
        v_t.set_ragged_offset(cum_kv_t)

    # -- Mask configuration --
    use_padding = seq_len_q is not None or seq_len_kv is not None
    lb = left_bound if left_bound >= 0 else None
    rb = right_bound if right_bound >= 0 else None
    if is_causal and rb is None:
        rb = 0

    # Compute max_total_seq_len for ragged backward
    max_total_seq_len_q = None
    max_total_seq_len_kv = None
    if cumulative_seq_len_q is not None and seq_len_q is not None:
        total = torch.sum(seq_len_q).item()
        max_total_seq_len_q = ((total + 63) // 64) * 64
    if cumulative_seq_len_kv is not None and seq_len_kv is not None:
        total = torch.sum(seq_len_kv).item()
        max_total_seq_len_kv = ((total + 63) // 64) * 64

    # -- SDPA backward --
    dQ_t, dK_t, dV_t = graph.sdpa_backward(
        name="sdpa_backward",
        q=q_t,
        k=k_t,
        v=v_t,
        o=o_t,
        dO=dO_t,
        stats=stats_t,
        attn_scale=attn_scale,
        use_padding_mask=use_padding,
        seq_len_q=seq_len_q_t,
        seq_len_kv=seq_len_kv_t,
        max_total_seq_len_q=max_total_seq_len_q,
        max_total_seq_len_kv=max_total_seq_len_kv,
        diagonal_alignment=_diagonal_alignment_enum(diagonal_alignment),
        diagonal_band_left_bound=lb,
        diagonal_band_right_bound=rb,
        use_deterministic_algorithm=is_deterministic,
    )

    # -- Output shapes (BHSD order) --
    dq_shape, dq_stride = q_shape, q_stride
    dk_shape, dk_stride = k_shape, k_stride
    dv_shape, dv_stride = v_shape, v_stride

    dQ_t.set_uid(_UIDs.DQ).set_output(True).set_dim(list(dq_shape)).set_stride(list(dq_stride))
    dK_t.set_uid(_UIDs.DK).set_output(True).set_dim(list(dk_shape)).set_stride(list(dk_stride))
    dV_t.set_uid(_UIDs.DV).set_output(True).set_dim(list(dv_shape)).set_stride(list(dv_stride))

    if cum_q_t is not None:
        dQ_t.set_ragged_offset(cum_q_t)
    if cum_kv_t is not None:
        dK_t.set_ragged_offset(cum_kv_t)
        dV_t.set_ragged_offset(cum_kv_t)

    # -- Build --
    graph.validate()
    graph.build_operation_graph()
    graph.create_execution_plans([cudnn.heur_mode.A, cudnn.heur_mode.FALLBACK])
    graph.check_support()
    graph.build_plans()

    workspace_size = graph.get_workspace_size()

    return graph, workspace_size


# ---------------------------------------------------------------------------
# Forward custom op
# ---------------------------------------------------------------------------


_lib = torch.library.Library("cudnn", "DEF")

_lib.define(
    "sdpa(Tensor q, Tensor k, Tensor v, float attn_scale, "
    "bool is_causal=False, int diagonal_alignment=0, "
    "int left_bound=-1, int right_bound=-1, "
    "Tensor? seq_len_q=None, Tensor? seq_len_kv=None, "
    "Tensor? cumulative_seq_len_q=None, Tensor? cumulative_seq_len_kv=None"
    ") -> (Tensor, Tensor)"
)

_lib.define(
    "sdpa_bwd(Tensor dO, Tensor q, Tensor k, Tensor v, Tensor o, Tensor stats, "
    "float attn_scale, bool is_causal=False, int diagonal_alignment=0, "
    "int left_bound=-1, int right_bound=-1, "
    "Tensor? seq_len_q=None, Tensor? seq_len_kv=None, "
    "Tensor? cumulative_seq_len_q=None, Tensor? cumulative_seq_len_kv=None, "
    "bool is_deterministic=False"
    ") -> (Tensor, Tensor, Tensor)"
)


def _sdpa_impl(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attn_scale: float,
    is_causal: bool = False,
    diagonal_alignment: int = 0,
    left_bound: int = -1,
    right_bound: int = -1,
    seq_len_q: Optional[torch.Tensor] = None,
    seq_len_kv: Optional[torch.Tensor] = None,
    cumulative_seq_len_q: Optional[torch.Tensor] = None,
    cumulative_seq_len_kv: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    cuDNN SDPA forward (internal). BHSD layout.

    Args:
        q: Query tensor (B, H_q, S_q, D_qk)
        k: Key tensor   (B, H_k, S_kv, D_qk)
        v: Value tensor (B, H_v, S_kv, D_v)
        attn_scale: Attention scale factor (typically 1/sqrt(D_qk))
        is_causal: If True, apply causal mask (right_bound defaults to 0)
        diagonal_alignment: 0 = TOP_LEFT, 1 = BOTTOM_RIGHT
        left_bound: Left sliding window bound (>= 0 to enable, -1 = disabled)
        right_bound: Right sliding window bound (>= 0 to enable, -1 = disabled)
        seq_len_q: Actual query sequence lengths (B, 1, 1, 1) INT32
        seq_len_kv: Actual key/value sequence lengths (B, 1, 1, 1) INT32
        cumulative_seq_len_q: Ragged offset for Q (B+1, 1, 1, 1) INT32
        cumulative_seq_len_kv: Ragged offset for KV (B+1, 1, 1, 1) INT32

    Returns:
        (O, Stats): Output tensor (B, H_q, S_q, D_v) and softmax stats (B, H_q, S_q, 1)
    """

    handle = _get_handle(q.device)

    cache_key = _make_fprop_cache_key(
        q,
        k,
        v,
        attn_scale,
        is_causal,
        diagonal_alignment,
        left_bound,
        right_bound,
        seq_len_q is not None,
        seq_len_kv is not None,
        cumulative_seq_len_q is not None,
        cumulative_seq_len_kv is not None,
    )

    if cache_key not in _fprop_cache:
        graph, workspace_size = _build_fprop_graph(
            handle,
            q,
            k,
            v,
            attn_scale,
            is_causal,
            diagonal_alignment,
            left_bound,
            right_bound,
            seq_len_q,
            seq_len_kv,
            cumulative_seq_len_q,
            cumulative_seq_len_kv,
        )
        _fprop_cache[cache_key] = (graph, workspace_size)

    graph, workspace_size = _fprop_cache[cache_key]

    # Allocate outputs and workspace (BHSD layout)
    # Workspace is per-call — PyTorch's caching allocator recycles the allocation.
    B, H_q, S_q, D_qk = q.shape
    _, H_v, S_kv, D_v = v.shape
    o_gpu = torch.empty(B, H_q, S_q, D_v, dtype=q.dtype, device=q.device)
    stats_gpu = torch.empty(B, H_q, S_q, 1, dtype=torch.float32, device=q.device)
    workspace = torch.empty(max(workspace_size, 1), device=q.device, dtype=torch.uint8)

    # UID → tensor map for sorted pointer extraction
    uid_to_tensor = {
        int(_UIDs.Q): q,
        int(_UIDs.K): k,
        int(_UIDs.V): v,
        int(_UIDs.O): o_gpu,
        int(_UIDs.STATS): stats_gpu,
    }
    if seq_len_q is not None:
        uid_to_tensor[int(_UIDs.SEQ_LEN_Q)] = seq_len_q
    if seq_len_kv is not None:
        uid_to_tensor[int(_UIDs.SEQ_LEN_KV)] = seq_len_kv
    if cumulative_seq_len_q is not None:
        uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_Q)] = cumulative_seq_len_q
    if cumulative_seq_len_kv is not None:
        uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_KV)] = cumulative_seq_len_kv

    graph.execute(uid_to_tensor, workspace, handle=handle)

    return o_gpu, stats_gpu


_lib.impl("sdpa", _sdpa_impl, "CUDA")


@torch.library.register_fake("cudnn::sdpa")
def _sdpa_fake(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    attn_scale: float,
    is_causal: bool = False,
    diagonal_alignment: int = 0,
    left_bound: int = -1,
    right_bound: int = -1,
    seq_len_q: Optional[torch.Tensor] = None,
    seq_len_kv: Optional[torch.Tensor] = None,
    cumulative_seq_len_q: Optional[torch.Tensor] = None,
    cumulative_seq_len_kv: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
    B, H_q, S_q, D_qk = q.shape
    _, H_v, S_kv, D_v = v.shape
    O = torch.empty(B, H_q, S_q, D_v, dtype=q.dtype, device=q.device)
    Stats = torch.empty(B, H_q, S_q, 1, dtype=torch.float32, device=q.device)
    return O, Stats


def _sdpa_bwd_impl(
    dO: torch.Tensor,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    o: torch.Tensor,
    stats: torch.Tensor,
    attn_scale: float,
    is_causal: bool = False,
    diagonal_alignment: int = 0,
    left_bound: int = -1,
    right_bound: int = -1,
    seq_len_q: Optional[torch.Tensor] = None,
    seq_len_kv: Optional[torch.Tensor] = None,
    cumulative_seq_len_q: Optional[torch.Tensor] = None,
    cumulative_seq_len_kv: Optional[torch.Tensor] = None,
    is_deterministic: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    cuDNN SDPA backward (graph-based).

    Returns:
        (dQ, dK, dV) in BHSD layout matching Q, K, V shapes.
    """
    handle = _get_handle(dO.device)

    cache_key = _make_bprop_cache_key(
        q,
        k,
        v,
        attn_scale,
        is_causal,
        diagonal_alignment,
        left_bound,
        right_bound,
        seq_len_q is not None,
        seq_len_kv is not None,
        cumulative_seq_len_q is not None,
        cumulative_seq_len_kv is not None,
        is_deterministic,
    )

    if cache_key not in _bprop_cache:
        graph, workspace_size = _build_bprop_graph(
            handle,
            q,
            k,
            v,
            o,
            dO,
            stats,
            attn_scale,
            is_causal,
            diagonal_alignment,
            left_bound,
            right_bound,
            seq_len_q,
            seq_len_kv,
            cumulative_seq_len_q,
            cumulative_seq_len_kv,
            is_deterministic,
        )
        _bprop_cache[cache_key] = (graph, workspace_size)

    graph, workspace_size = _bprop_cache[cache_key]

    # Allocate gradient outputs and workspace (same shapes as Q, K, V)
    dQ_gpu = torch.empty_like(q)
    dK_gpu = torch.empty_like(k)
    dV_gpu = torch.empty_like(v)
    workspace = torch.empty(max(workspace_size, 1), device=dO.device, dtype=torch.uint8)

    # UID → tensor map for sorted pointer extraction
    uid_to_tensor = {
        int(_UIDs.Q): q,
        int(_UIDs.K): k,
        int(_UIDs.V): v,
        int(_UIDs.O): o,
        int(_UIDs.DO): dO,
        int(_UIDs.STATS): stats,
        int(_UIDs.DQ): dQ_gpu,
        int(_UIDs.DK): dK_gpu,
        int(_UIDs.DV): dV_gpu,
    }
    if seq_len_q is not None:
        uid_to_tensor[int(_UIDs.SEQ_LEN_Q)] = seq_len_q
    if seq_len_kv is not None:
        uid_to_tensor[int(_UIDs.SEQ_LEN_KV)] = seq_len_kv
    if cumulative_seq_len_q is not None:
        uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_Q)] = cumulative_seq_len_q
    if cumulative_seq_len_kv is not None:
        uid_to_tensor[int(_UIDs.CUM_SEQ_LEN_KV)] = cumulative_seq_len_kv

    graph.execute(uid_to_tensor, workspace, handle=handle)

    return dQ_gpu, dK_gpu, dV_gpu


_lib.impl("sdpa_bwd", _sdpa_bwd_impl, "CUDA")


@torch.library.register_fake("cudnn::sdpa_bwd")
def _sdpa_bwd_fake(
    dO: torch.Tensor,
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    o: torch.Tensor,
    stats: torch.Tensor,
    attn_scale: float,
    is_causal: bool = False,
    diagonal_alignment: int = 0,
    left_bound: int = -1,
    right_bound: int = -1,
    seq_len_q: Optional[torch.Tensor] = None,
    seq_len_kv: Optional[torch.Tensor] = None,
    cumulative_seq_len_q: Optional[torch.Tensor] = None,
    cumulative_seq_len_kv: Optional[torch.Tensor] = None,
    is_deterministic: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    return torch.empty_like(q), torch.empty_like(k), torch.empty_like(v)


# ---------------------------------------------------------------------------
# Autograd registration
# ---------------------------------------------------------------------------


def _sdpa_setup_context(ctx, inputs, output):
    q, k, v, attn_scale, is_causal, diagonal_alignment, left_bound, right_bound, seq_len_q, seq_len_kv, cumulative_seq_len_q, cumulative_seq_len_kv = inputs
    o, stats = output

    tensors_to_save = [q, k, v, o, stats]
    ctx.has_seq_len_q = seq_len_q is not None
    ctx.has_seq_len_kv = seq_len_kv is not None
    ctx.has_cum_q = cumulative_seq_len_q is not None
    ctx.has_cum_kv = cumulative_seq_len_kv is not None

    if ctx.has_seq_len_q:
        tensors_to_save.append(seq_len_q)
    if ctx.has_seq_len_kv:
        tensors_to_save.append(seq_len_kv)
    if ctx.has_cum_q:
        tensors_to_save.append(cumulative_seq_len_q)
    if ctx.has_cum_kv:
        tensors_to_save.append(cumulative_seq_len_kv)

    ctx.save_for_backward(*tensors_to_save)

    ctx.attn_scale = attn_scale
    ctx.is_causal = is_causal
    ctx.diagonal_alignment = diagonal_alignment
    ctx.left_bound = left_bound
    ctx.right_bound = right_bound


def _sdpa_backward(ctx, dO, dStats):
    # dO from autograd may have zero strides (e.g. from o.sum().backward()),
    # which cuDNN cannot handle. Make it contiguous.
    dO = dO.contiguous()

    saved = list(ctx.saved_tensors)
    q, k, v, o, stats = saved[:5]
    idx = 5
    seq_len_q = saved[idx] if ctx.has_seq_len_q else None
    if ctx.has_seq_len_q:
        idx += 1
    seq_len_kv = saved[idx] if ctx.has_seq_len_kv else None
    if ctx.has_seq_len_kv:
        idx += 1
    cum_q = saved[idx] if ctx.has_cum_q else None
    if ctx.has_cum_q:
        idx += 1
    cum_kv = saved[idx] if ctx.has_cum_kv else None

    dQ, dK, dV = torch.ops.cudnn.sdpa_bwd(
        dO,
        q,
        k,
        v,
        o,
        stats,
        ctx.attn_scale,
        ctx.is_causal,
        ctx.diagonal_alignment,
        ctx.left_bound,
        ctx.right_bound,
        seq_len_q,
        seq_len_kv,
        cum_q,
        cum_kv,
    )

    # Return gradients for: q, k, v, attn_scale, is_causal, diagonal_alignment,
    # left_bound, right_bound, seq_len_q, seq_len_kv, cum_q, cum_kv
    return dQ, dK, dV, None, None, None, None, None, None, None, None, None


torch.library.register_autograd("cudnn::sdpa", _sdpa_backward, setup_context=_sdpa_setup_context)


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def scaled_dot_product_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attn_mask: Optional[torch.Tensor] = None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale: Optional[float] = None,
    enable_gqa: bool = False,
    *,
    diagonal_alignment: int = 0,
    left_bound: int = -1,
    right_bound: int = -1,
    seq_len_q: Optional[torch.Tensor] = None,
    seq_len_kv: Optional[torch.Tensor] = None,
    cumulative_seq_len_q: Optional[torch.Tensor] = None,
    cumulative_seq_len_kv: Optional[torch.Tensor] = None,
) -> torch.Tensor:
    """cuDNN-accelerated Scaled Dot-Product Attention.

    API closely mirrors ``torch.nn.functional.scaled_dot_product_attention``.

    **Layout**: tensors use **BHSD** layout ``(batch, num_heads, seq_len, head_dim)``,
    matching PyTorch's convention.

    Args:
        query: Query tensor ``(B, H_q, S_q, D)``.
        key: Key tensor ``(B, H_k, S_kv, D)``.
        value: Value tensor ``(B, H_v, S_kv, D_v)``.
        attn_mask: Not yet supported. Must be ``None``.
        dropout_p: Not yet supported. Must be ``0.0``.
        is_causal: If ``True``, applies a causal (upper-triangular) mask.
        scale: Attention scale factor. Defaults to ``1 / sqrt(D)`` when ``None``.
        enable_gqa: When ``False``, raises ``ValueError`` if ``H_q != H_k``.
            When ``True``, grouped-query attention is enabled (cuDNN handles
            this automatically via head dimension broadcast).

        diagonal_alignment: cuDNN extension. ``0`` = TOP_LEFT, ``1`` = BOTTOM_RIGHT.
        left_bound: cuDNN extension. Left sliding-window bound (``-1`` = disabled).
        right_bound: cuDNN extension. Right sliding-window bound (``-1`` = disabled).
        seq_len_q: cuDNN extension. Actual query sequence lengths ``(B, 1, 1, 1)`` INT32.
        seq_len_kv: cuDNN extension. Actual key/value sequence lengths ``(B, 1, 1, 1)`` INT32.
        cumulative_seq_len_q: cuDNN extension. Ragged offset for Q ``(B+1, 1, 1, 1)`` INT32.
        cumulative_seq_len_kv: cuDNN extension. Ragged offset for KV ``(B+1, 1, 1, 1)`` INT32.

    Note:
        For head dimension ``256``, the specialized backward path currently supports
        only plain BHSD tensors and does not support ``seq_len_q``, ``seq_len_kv``,
        ``cumulative_seq_len_q``, or ``cumulative_seq_len_kv``.

    Returns:
        Output tensor ``(B, H_q, S_q, D_v)``.
    """
    if attn_mask is not None:
        raise NotImplementedError("attn_mask is not yet supported by cuDNN SDPA")
    if dropout_p != 0.0:
        raise NotImplementedError("dropout is not yet supported by cuDNN SDPA")
    if not enable_gqa and query.shape[1] != key.shape[1]:
        raise ValueError(f"query has {query.shape[1]} heads but key has {key.shape[1]} heads. " f"Set enable_gqa=True for grouped-query attention.")

    d = query.shape[-1]
    attn_scale = scale if scale is not None else (1.0 / math.sqrt(d))

    o, _stats = torch.ops.cudnn.sdpa(
        query,
        key,
        value,
        attn_scale,
        is_causal=is_causal,
        diagonal_alignment=diagonal_alignment,
        left_bound=left_bound,
        right_bound=right_bound,
        seq_len_q=seq_len_q,
        seq_len_kv=seq_len_kv,
        cumulative_seq_len_q=cumulative_seq_len_q,
        cumulative_seq_len_kv=cumulative_seq_len_kv,
    )
    return o
