Triton kernels¶

The mingru.triton_scans surface, imported lazily on first access (see the Triton scans explanation for kernel design and the dispatch how-to for MINGRU_SCAN). Importing mingru alone does not import this module; touching any name below does. The kernel wrappers require a CUDA build of PyTorch (torch >= 2.8); the availability probe, fallback signal, and impl registry are present on every build.

Availability and registry¶

mingru.triton_scans.available ¶

available() -> bool | str

Whether Triton scan kernels can run in this process.

Returns:

Type	Description
`bool or str`	`True` if a CUDA device is present and `triton` is importable. Otherwise a human-readable reason string naming why not (e.g. `"CUDA not available"` or `"triton not importable: ..."`) -- callers should treat anything other than `True` as unavailable and use the string as the fallback/error reason (see `min_gru._dispatch_scan`).

mingru.triton_scans.ScanFallback ¶

Bases: Exception

Signal that the Triton path declines this call; fall back to eager.

Raised by a SCAN_IMPLS entry when the inputs are outside the kernel envelope (unsupported k/v) or when a kernel launch fails. min_gru._dispatch_scan catches this: under MINGRU_SCAN=auto it warns once and runs the eager reference; under MINGRU_SCAN=triton it re-raises as a RuntimeError (never a silent downgrade). This is distinct from "no kernel registered for this op" -- an op simply absent from SCAN_IMPLS is handled by _dispatch_scan before any call is attempted.

mingru.triton_scans.SCAN_IMPLS `module-attribute` ¶

SCAN_IMPLS: dict = {}

Kernel entry points¶

mingru.triton_scans.angle_scan_impl ¶

angle_scan_impl(theta: Tensor, scale: Tensor, gamma: Tensor, b: Tensor, h0: Tensor, perm: Tensor, sgn: Tensor, p2p: Tensor, *, has_scale: int, has_decay: int) -> torch.Tensor

Envelope-guarded entry point for the angle-fused kernel (module callers).

Mirrors the SCAN_IMPLS-style wrappers (e.g. _matrix_affine_scan_impl): validates every shape-agreement invariant _angle_scan_fwd_kernel's pointer arithmetic assumes BEFORE any launch, then funnels any launch/compile failure into ScanFallback -- so min_gru._dispatch_angle_scan only ever needs to catch that one exception type, exactly the contract _dispatch_scan gets from each SCAN_IMPLS entry. Not itself a SCAN_IMPLS entry (the angle-fused path is a module-forward fast path, not one of the four scan ops), so it is looked up by hasattr(triton_scans, "angle_scan_impl") instead.

has_scale/has_decay are keyword-only here (this function's public-ish entry point, called from min_gru._dispatch_angle_scan) -- self-documenting 0/1 flags at the call site instead of a bare trailing 1, 0. The underlying angle_scan_fwd triton_op this function calls keeps its flat positional-int schema unchanged (a torch.library op signature, not a Python call-site convenience).

Parameters mirror angle_scan_fwd's: theta (B, T, n, R, half); scale/gamma (B, T, n); b (B, T, n, k); h0 (B, n, k); perm/sgn/p2p (R, k). k must be in the module's block-size envelope.

mingru.triton_scans.affine_scan_fwd ¶

affine_scan_fwd(A: Tensor, Bm: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Generic k x k affine prefix scan (Kernel 1).

A is (B, T, n, k, k), Bm is (B, T, n, k, v). Returns (Abar, Bbar) with the same shapes, where Abar_t is the running matrix product of A (composition A_current @ A_earlier) and Bbar_t the H_0 = 0 solution of H_t = A_t @ H_{t-1} + B_t.

The functional body (.contiguous + empty_like + a single wrap_triton launch) is what triton_op traces for the fake-tensor meta, so torch.compile gets correct output shapes with no graph break.

mingru.triton_scans.affine_scan_bwd ¶

affine_scan_bwd(A: Tensor, Abar: Tensor, Bbar: Tensor, gAbar: Tensor, gBbar: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Generic k x k affine-scan adjoint (Kernel 3).

Reads the forward input A (B, T, n, k, k) and the forward outputs Abar (B, T, n, k, k) / Bbar (B, T, n, k, v), plus the incoming output grads gAbar / gBbar (same shapes as Abar / Bbar). Returns (dA, dB) -- the grads of the loss w.r.t. the forward inputs A and Bm -- with shapes matching A and Bbar. Same traced-body fake-meta contract as the forward ops (.contiguous + empty_like + one wrap_triton launch), so torch.compile sees correct grad shapes with no graph break.

mingru.triton_scans.linear_scan_fwd ¶

linear_scan_fwd(a: Tensor, b: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Channel-tiled k=1 affine prefix scan (Kernel 1).

a and b are (B, T, D). Returns (A, Bc) of the same shape, where A_t is the running product of a and Bc_t the h_0 = 0 solution of h_t = a_t * h_{t-1} + b_t. Same traced-body fake-meta contract as affine_scan_fwd.

mingru.triton_scans.linear_scan_bwd ¶

linear_scan_bwd(a: Tensor, A: Tensor, Bc: Tensor, gA: Tensor, gBc: Tensor) -> tuple[torch.Tensor, torch.Tensor]

Channel-tiled k=1 affine-scan adjoint.

Reads the forward input a (B, T, D) and outputs A / Bc (B, T, D), plus incoming grads gA / gBc. Returns (da, db), the grads w.r.t. the forward inputs a and b. Same traced-body fake-meta contract as linear_scan_fwd.

mingru.triton_scans.parallel_scan_log_fwd ¶

parallel_scan_log_fwd(log_coeffs: Tensor, log_values: Tensor) -> torch.Tensor

Fused log-space forward scan (Kernel 2).

log_coeffs is (B, T, D) and log_values is (B, T+1, D) (slot 0 = log(h_0), slots 1..T = log(b_t)). Returns the states h_1..h_T as (B, T, D) -- a single tensor, matching min_gru.parallel_scan_log (unlike the affine ops, which return a pair). Same traced-body fake-meta contract as linear_scan_fwd: .contiguous + empty_like + one wrap_triton launch, so torch.compile sees the (B, T, D) output shape with no graph break.

mingru.triton_scans.angle_scan_fwd ¶

angle_scan_fwd(theta: Tensor, scale: Tensor, gamma: Tensor, b: Tensor, h0: Tensor, perm: Tensor, sgn: Tensor, p2p: Tensor, has_scale: int, has_decay: int) -> torch.Tensor

Angle-fused forward (Kernel 4). Returns states h (B,T,n,k).

Same traced-body fake-meta contract as the scan ops (.contiguous + empty_like + one wrap_triton launch), so torch.compile sees the output shape with no graph break.

mingru.triton_scans.angle_scan_bwd ¶

angle_scan_bwd(theta: Tensor, scale: Tensor, gamma: Tensor, b: Tensor, h0: Tensor, out: Tensor, grad_out: Tensor, perm: Tensor, sgn: Tensor, p2p: Tensor, has_scale: int, has_decay: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]

Angle-fused exact stored-state reversible backward.

Reads the forward inputs, the forward output out (the source of every h_{t-1}), and the incoming state grads grad_out. Returns (dtheta, dscale, dgamma, db, dh0) shaped like theta, scale, gamma, b, h0. The module's own autograd then backprops dscale/dgamma through tanh(u) / exp(-lambda dt) to the real parameters, so this op handles only the recurrence.

dscale/dgamma use zeros_like rather than empty_like: the kernel only stores into them when HAS_SCALE/HAS_DECAY is set, so the disabled-feature output would otherwise be uninitialized memory. The disabled-feature grads are never consumed (the module never differentiates through the "ones" placeholder it passed in that case), but zero-init makes the returned tensor correct on its own terms rather than relying on that caller invariant.