File size: 2,794 Bytes
9f97b05 d9257c4 9f97b05 ef6f823 9f97b05 ef6f823 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | ---
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.
|