causal-trailing-mean

causal-trailing-mean

A fused, count-normalized causal trailing-mean pool (a causal box-filter running mean) as a noarch Triton kernel for the kernels ecosystem.

For each row (batch × channel) over time t:

h[t] = ( sum_{j = max(0, t-cf+1)}^{t} x[j] ) / min(t+1, cf)

i.e. a past-only mean over the last cf frames, count-normalized so sequence-start frames divide by the real window size rather than a zero-diluted cf. It fuses the eager cumsum → pad → subtract → divide → cast (≈5 full [B, D, T] passes + an fp32 round-trip) into a single launch, using a two-cumsum identity s[t] = cumsum(x)[t] − cumsum(x delayed by cf)[t].

Usage

from kernels import get_kernel

k = get_kernel("futo-org/causal-trailing-mean", version=1)
h = k.causal_trailing_mean(x, cf)   # x: [B, D, T] contiguous, any float dtype -> same shape/dtype

version=1 pins the v1 build; omit it to track main (latest).

Numerics & scope

• fp32 accumulation, count-normalized divide, cast to x.dtype on store. Because it accumulates in fp32 it is at least as accurate as the eager op in low precision. Parity vs a fp32 eager reference: max|Δ| ≈ 1e-7 (fp32), 1e-3 (bf16), 2e-4 (fp16).
• Inference-only (no backward): the Triton path is taken only when grad is disabled and T ≤ 4096; otherwise it falls back to an autograd-safe eager reference (also used on CPU). eager_causal_trailing_mean is exported for that path.
• Variant: torch-cuda (the Triton kernel — the op only accelerates CUDA). The exported eager_causal_trailing_mean covers CPU / grad-enabled / oversized-T in-process.

Performance

Fused kernel vs the eager reference (cumsum → pad → subtract → divide → cast, ~5 full [B,D,T] passes + an fp32 round-trip), measured with triton.testing.do_bench under torch.no_grad(), bf16, cf=100, on an NVIDIA RTX PRO 6000 Blackwell:

3–11× depending on sequence length (the win grows with T as the fused single-pass launch displaces more redundant memory traffic); ~670 GB/s at the top end. Reproduce with triton.testing.do_bench(lambda: causal_trailing_mean(x, cf)) vs the exported eager_causal_trailing_mean.

---

Built with kernel-builder; Apache-2.0.