README.md

---
license: apache-2.0
tags:
- kernel
- triton
---

# causal-trailing-mean

A fused, count-normalized **causal trailing-mean pool** (a causal box-filter running mean) as
a noarch [Triton](https://github.com/triton-lang/triton) kernel for the
[`kernels`](https://github.com/huggingface/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

```python
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:

| shape `[B, D, T]` | eager (ms) | kernel (ms) | speedup |
|---|---|---|---|
| 32 × 512 × 256  | 0.221 | 0.056 | **4.0×** |
| 32 × 512 × 512  | 0.203 | 0.062 | **3.3×** |
| 32 × 512 × 1024 | 0.828 | 0.127 | **6.5×** |
| 64 × 512 × 2048 | 4.510 | 0.402 | **11.2×** |
| 16 × 256 × 3000 | 0.545 | 0.137 | **4.0×** |

**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`](https://github.com/huggingface/kernel-builder); Apache-2.0.