Module neuralops.core.triton_kernels.linear_attn

Triton Fused Linear Attention Kernel (Forward & Backward).

Supports -----= - Non-negative feature map phi: "elu+1", "relu", "identity", "softmax_free" - Chunked causal streaming linear attention (O(N) memory, tiled in SRAM) - Bidirectional parallel reduction for encoder backbones - Cosformer-style in-register / pre-applied cosine-sine positional modulation - Full autograd integration for training and recovery fine-tuning.

Functions

def bwd(q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
grad_out: torch.Tensor,
is_causal: bool = False,
phi: str = 'elu+1',
eps: float = 1e-06,
cos_weight: torch.Tensor | None = None,
sin_weight: torch.Tensor | None = None,
**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,
    is_causal: bool = False,
    phi: str = "elu+1",
    eps: float = 1e-6,
    cos_weight: torch.Tensor | None = None,
    sin_weight: torch.Tensor | None = None,
    **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,
            is_causal=is_causal,
            phi=phi,
            eps=eps,
            cos_weight=cos_weight,
            sin_weight=sin_weight,
            **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,
is_causal: bool = False,
phi: str = 'elu+1',
eps: float = 1e-06,
cos_weight: torch.Tensor | None = None,
sin_weight: torch.Tensor | None = None,
**kwargs: Any) ‑> torch.Tensor

Expand source code

def fwd(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    is_causal: bool = False,
    phi: str = "elu+1",
    eps: float = 1e-6,
    cos_weight: torch.Tensor | None = None,
    sin_weight: torch.Tensor | None = None,
    **kwargs: Any,
) -> torch.Tensor:
    """Forward entrypoint for fused Triton linear attention."""
    if not is_supported(q, k, v, **kwargs):
        return linear_attn_pytorch_fallback(
            q,
            k,
            v,
            is_causal=is_causal,
            phi=phi,
            eps=eps,
            cos_weight=cos_weight,
            sin_weight=sin_weight,
            **kwargs,
        )
    return _TritonLinearAttnAutogradFunction.apply(
        q, k, v, is_causal, phi, eps, cos_weight, sin_weight
    )

Forward entrypoint for fused Triton linear 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 linear attention can run on the Triton kernel."""
    head_dim = q.shape[-1]
    seq_len = q.shape[-2]
    suitable, _ = check_triton_suitability(q, "linear_attn", head_dim=head_dim, seq_len=seq_len)
    return suitable

Check whether linear attention can run on the Triton kernel.