Module neuralops.core.triton_kernels.sparse_attn

Triton Block-Sparse Attention Kernel (Forward & Backward).

Supports -----= - Arbitrary 2D or 4D block sparsity masks: (n_blocks_q, n_blocks_k) - Configurable block tile sizes (e.g. 32, 64, 128) - Optional lower-triangular causal masking combined with block sparsity - Online safe softmax reduction in GPU SRAM - Full autograd integration and fallback handling.

Functions

def bwd(q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
grad_out: torch.Tensor,
block_mask: torch.Tensor | None = None,
block_size: int = 64,
is_causal: bool = False,
eps: float = 1e-06,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Expand source code

def bwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    grad_out: torch.Tensor,
    block_mask: torch.Tensor | None = None,
    block_size: int = 64,
    is_causal: bool = False,
    eps: float = 1e-6,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """Backward entrypoint computing gradients for q, k, v."""
    q_req = q.detach().requires_grad_(True)
    k_req = k.detach().requires_grad_(True)
    v_req = v.detach().requires_grad_(True)
    with torch.enable_grad():
        out = fwd(
            q_req,
            k_req,
            v_req,
            block_mask=block_mask,
            block_size=block_size,
            is_causal=is_causal,
            eps=eps,
            **kwargs,
        )
        g_q, g_k, g_v = torch.autograd.grad(out, [q_req, k_req, v_req], grad_out)
    return g_q, g_k, g_v

Backward entrypoint computing gradients for q, k, v.

def fwd(q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
block_mask: torch.Tensor | None = None,
block_size: int = 64,
is_causal: bool = False,
eps: float = 1e-06,
**kwargs: Any) ‑> torch.Tensor

Expand source code

def fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    block_mask: torch.Tensor | None = None,
    block_size: int = 64,
    is_causal: bool = False,
    eps: float = 1e-6,
    **kwargs: Any,
) -> torch.Tensor:
    """Forward entrypoint for fused Triton block-sparse attention."""
    if not is_supported(q, k, v, **kwargs):
        return sparse_attn_pytorch_fallback(
            q,
            k,
            v,
            block_mask=block_mask,
            block_size=block_size,
            is_causal=is_causal,
            eps=eps,
            **kwargs,
        )
    return _TritonSparseAttnAutogradFunction.apply(q, k, v, block_mask, block_size, is_causal, eps)

Forward entrypoint for fused Triton block-sparse attention.

def is_supported(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: Any) ‑> bool

Expand source code

def is_supported(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: Any) -> bool:
    """Check whether block-sparse attention can run on the Triton kernel."""
    head_dim = q.shape[-1]
    seq_len = q.shape[-2]
    suitable, _ = check_triton_suitability(q, "sparse_attn", head_dim=head_dim, seq_len=seq_len)
    return suitable

Check whether block-sparse attention can run on the Triton kernel.