Module neuralops.core.triton_kernels.hgrn2_scan

Triton-accelerated HGRN2 selective scan.

Element-wise recurrence: h_t = f_t * h_{t-1} + iv_t y_t = h_t

Backward -----= dh_t = d_h_t + f_{t+1} * dh_{t+1} (adjoint recurrence) d_f_t = dh_t * h_{t-1} d_iv_t = dh_t

Note: f, iv, h all have shape (B, L, D). No state-dim mixing.

Functions

def hgrn2_scan_triton(f: torch.Tensor, iv: torch.Tensor, res_clip=None) ‑> torch.Tensor

Expand source code

def hgrn2_scan_triton(f: torch.Tensor, iv: torch.Tensor, res_clip=None) -> torch.Tensor:
    """Triton HGRN2 selective scan. f, iv: (B, L, D). Returns h: (B, L, D)."""
    return _HGRN2ScanFn.apply(f, iv, res_clip)

Triton HGRN2 selective scan. f, iv: (B, L, D). Returns h: (B, L, D).

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,