Module neuralops.core.triton_kernels.low_rank

Triton Low-Rank Projection Kernel (Forward & Backward).

Computes low-rank bottleneck projections: y = (x @ V.T) @ U.T + bias

Used by: - Multi-Head Latent Attention (MLA) - SVD-LLM, ASVD, Direct SVD, DOBI SVD, Frandi - Triton SVD Linear Attention low-rank projections

Functions

def bwd(x: torch.Tensor,
U: torch.Tensor,
V: torch.Tensor,
grad_out: torch.Tensor,
bias: torch.Tensor | None = None,
**kwargs: Any) ‑> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]

Expand source code

def bwd(
    x: torch.Tensor,
    U: torch.Tensor,
    V: torch.Tensor,
    grad_out: torch.Tensor,
    bias: torch.Tensor | None = None,
    **kwargs: Any,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
    """Backward entrypoint computing gradients for x, U, V, and bias."""
    x_req = x.detach().requires_grad_(True)
    U_req = U.detach().requires_grad_(True)
    V_req = V.detach().requires_grad_(True)
    b_req = bias.detach().requires_grad_(True) if bias is not None else None
    with torch.enable_grad():
        out = fwd(x_req, U_req, V_req, bias=b_req, **kwargs)
        diff_vars = [x_req, U_req, V_req]
        if b_req is not None:
            diff_vars.append(b_req)
        grads = torch.autograd.grad(out, diff_vars, grad_out)
    return grads[0], grads[1], grads[2], (grads[3] if b_req is not None else None)

Backward entrypoint computing gradients for x, U, V, and bias.

def fwd(x: torch.Tensor,
U: torch.Tensor,
V: torch.Tensor,
bias: torch.Tensor | None = None,
**kwargs: Any) ‑> torch.Tensor

Expand source code

def fwd(
    x: torch.Tensor,
    U: torch.Tensor,
    V: torch.Tensor,
    bias: torch.Tensor | None = None,
    **kwargs: Any,
) -> torch.Tensor:
    """Forward entrypoint for fused Triton low-rank projection: y = (x @ V.T) @ U.T + bias."""
    if not is_supported(x, U, V, bias=bias, **kwargs):
        return low_rank_pytorch_fallback(x, U, V, bias=bias, **kwargs)
    return _TritonLowRankAutogradFunction.apply(x, U, V, bias)

Forward entrypoint for fused Triton low-rank projection: y = (x @ V.T) @ U.T + bias.

def is_supported(x: torch.Tensor,
U: torch.Tensor,
V: torch.Tensor,
bias: torch.Tensor | None = None,
**kwargs: Any) ‑> bool

Expand source code

def is_supported(
    x: torch.Tensor,
    U: torch.Tensor,
    V: torch.Tensor,
    bias: torch.Tensor | None = None,
    **kwargs: Any,
) -> bool:
    """Check whether low-rank projection can run on the Triton kernel."""
    suitable, _ = check_triton_suitability(x, "low_rank")
    return suitable and U.is_cuda and V.is_cuda

Check whether low-rank projection can run on the Triton kernel.