Module neuralops.core.triton_kernels.gdn_scan

Triton-accelerated Gated DeltaNet chunked scan.

Recurrence -----= inter_v_t = k_t @ S_{t-1} err_t = v_t - inter_v_t norm_sq_t = sum(k_t^2) step_t = alpha_t / max(1 + norm_sq_t, 1) S_t = beta_t * S_{t-1} + step_t * k_t^T @ err_t y_t = q_t @ S_t

All ops are per-(batch, head). S is a (d_k, d_v) matrix kept in registers. Supports d_k = d_v <= 64.

Functions

def gdn_scan_triton(q, k, v, beta, alpha)

Expand source code

def gdn_scan_triton(q, k, v, beta, alpha):
    """Triton-accelerated Gated DeltaNet scan. head_dim must be <= 64."""
    B, H, L, D = q.shape
    if D > MAX_REGISTER_DIM:
        raise ValueError(f"head_dim {D} > {MAX_REGISTER_DIM}")

    q = q.contiguous()
    k = k.contiguous()
    v = v.contiguous()
    beta = beta.contiguous()
    alpha = alpha.contiguous()

    return _GatedDeltaNetScanFn.apply(q, k, v, beta, alpha)

Triton-accelerated Gated DeltaNet scan. head_dim must be <= 64.

Classes

class triton

Methods

def autotune(*a, **kw)

Expand source code

"autotune": lambda *a, **kw: lambda f: f,

def heuristics(*a, **kw)

Expand source code

"heuristics": lambda *a, **kw: lambda f: f,

def jit(f)

Expand source code

"jit": lambda f: f,