Control scan dispatch¶
The four scan functions (parallel_scan_log, linear_scan, matrix_scan, matrix_affine_scan) pick their backend at call time from the MINGRU_SCAN environment variable. This guide shows how to force a backend and what each mode guarantees. It assumes you have completed the getting-started tutorial.
The three modes¶
MINGRU_SCAN is read from os.environ on every scan call (not cached at import), so you can set it before launching a process or from inside one via os.environ.
MINGRU_SCAN |
On CUDA tensors | On CPU tensors | Imports triton_scans? |
|---|---|---|---|
auto (default) |
Triton kernel if available, else eager (warns once) | eager | only when a CUDA input reaches the seam |
eager |
eager | eager | never |
triton |
Triton kernel, or raises | raises | yes |
Force the eager path¶
Set MINGRU_SCAN=eager to guarantee the pure-PyTorch implementation everywhere. This never imports the Triton module, so it is the mode to use below the torch >= 2.8 floor, for deterministic A/B timing against Triton, or when you only use the diagonal mixers (whose elementwise log scan is faster in eager form, see Run the benchmarks).
import os
os.environ["MINGRU_SCAN"] = "eager" # set before the first scan call
import torch
from mingru import parallel_scan_log
log_coeffs = -torch.rand(2, 5, 3)
log_values = -torch.rand(2, 6, 3)
print(parallel_scan_log(log_coeffs, log_values).shape) # torch.Size([2, 5, 3])
Require the Triton path¶
Set MINGRU_SCAN=triton to make the Triton kernel mandatory: if a kernel is unavailable (no CUDA device, triton not importable, inputs on CPU, or no kernel registered for that op), the call raises RuntimeError instead of silently downgrading to eager. Use it in CI or a benchmark harness to prove the accelerated path is actually live.
import os
os.environ["MINGRU_SCAN"] = "triton"
import torch
from mingru import parallel_scan_log
# CPU tensors under triton mode -> hard error, never a silent fallback:
parallel_scan_log(-torch.rand(2, 5, 3), -torch.rand(2, 6, 3))
RuntimeError: MINGRU_SCAN=triton requested for 'parallel_scan_log' but Triton is unavailable: CUDA not available
Understand auto (the default)¶
auto is what you want in production: CUDA-resident inputs with a usable Triton kernel run on Triton; everything else (CPU tensors, a missing Triton module, an out-of-envelope shape, or an op with no kernel yet) falls through to the unchanged eager implementation. A fallback that happens despite CUDA inputs warns exactly once per process, naming the reason, so a silent performance cliff can't hide:
UserWarning: MINGRU_SCAN=auto fell back to the eager scan implementation for 'parallel_scan_log' despite CUDA inputs: <reason>
CPU inputs under auto fall back silently (that is the expected path, not a regression) and never import the Triton module.
Check the current status¶
mingru.available() tells you whether the Triton path can run at all: True, or a reason string like "CUDA not available":
An invalid value (anything other than auto/eager/triton) raises ValueError from the scan call, so a typo fails loudly rather than being treated as auto.
Handle decay-active DeltaMinGRU¶
The three modes above govern the four scan ops. DeltaMinGRU's time decay (decay=) sits on a separate seam with a narrower contract: there is no Triton kernel for decay yet (deferred), so a decay-active forward is eager-only no matter which mixer state size or head count you pick.
MINGRU_SCAN |
Decay-active DeltaMinGRU forward |
|---|---|
auto (default) |
eager; warns once per process on a CUDA input, silent on CPU |
eager |
eager |
triton |
raises RuntimeError naming decay as unsupported |
import os
os.environ["MINGRU_SCAN"] = "triton"
import torch
from mingru import DeltaMinGRU
layer = DeltaMinGRU(input_size=16, hidden_size=32, decay="fixed").cuda()
x = torch.randn(2, 8, 16, device="cuda")
delta_t = torch.rand(2, 8, device="cuda")
layer(x, delta_t=delta_t)
RuntimeError: MINGRU_SCAN=triton requested for DeltaMinGRU's decay-active forward, but decay has no Triton kernel (deferred; torch.compile is the documented CUDA path for decay -- see the project docs)
torch.compile is the recommended CUDA path for mixer="delta", with or without decay active; see Choose a mixer for the full GPU execution guidance.
This is delta-specific, not decay-general. RotationMinGRU and GivensMinGRU also accept decay=, and their fused angle-scan Triton path threads decay through the forward and backward directly (a has_decay flag carried through the kernel), so MINGRU_SCAN=triton still runs on Triton for those two mixers with decay active. Only DeltaMinGRU's decay path lacks a kernel today; nothing here forces eager on the other mixers.
For the float32 large-gap saturation behavior (log1p_delta=) and when to reach for decay= at all, see the README's "Time-aware decay" section and Choose a mixer.
You have now¶
forced the eager or Triton backend per process, made the Triton path mandatory where you need proof it is live, know how auto decides and warns, and know why a decay-active DeltaMinGRU stays eager-only under MINGRU_SCAN=triton while RotationMinGRU/GivensMinGRU do not. For the GPU-side walkthrough, see Triton on GPU.