| --- |
| license: apache-2.0 |
| tags: |
| - kernel |
| - triton |
| --- |
| |
| # ldsa |
|
|
| Fused **Local Dense Synthesizer Attention (LDSA)** window op as a noarch |
| [Triton](https://github.com/triton-lang/triton) kernel for the |
| [`kernels`](https://github.com/huggingface/kernels) ecosystem. |
|
|
| Given per-frame **synthesized** window logits `a_logits [B, H, W, T]` and values |
| `v [B, H, Dh, T]`, with a window `(left, right)` and `W = left + right + 1`: |
|
|
| ``` |
| a[b,h,k,t] = softmax over k of a_logits[b,h,:,t] # softmax across the window |
| out[b,h,d,t] = Σ_k a[b,h,k,t] · v[b,h,d, t-left+k] # weighted sum; zero outside [0,T) |
| ``` |
|
|
| The kernel fuses the two memory-bound steps — the window softmax and the windowed weighted |
| sum — into **one launch** (fp32 accumulation, no materialized unfold). It does **not** contain |
| the projections that produce its inputs; the caller passes `a_logits` and `v`. |
|
|
| ## Usage |
|
|
| ```python |
| from kernels import get_kernel |
| |
| k = get_kernel("futo-org/ldsa", version=1) |
| out = k.ldsa_local_attention(a_logits, v, left, right) # [B,H,Dh,T]; contiguous, any float |
| ``` |
|
|
| `version=1` pins the `v1` build; omit to track `main` (latest). |
|
|
| ## What LDSA is |
|
|
| LDSA replaces dot-product self-attention with attention weights **synthesized** per frame by a |
| small MLP (no query–key interaction), restricted to a **local window** — cheaper than full |
| attention and effective for speech. Introduced for ASR by: |
|
|
| > M. Xu, S. Li, X.-L. Zhang, *"Transformer-based End-to-End Speech Recognition with Local Dense |
| > Synthesizer Attention,"* ICASSP 2021 — [arXiv:2010.12155](https://arxiv.org/abs/2010.12155), |
| > code: [github.com/mlxu995/multihead-LDSA](https://github.com/mlxu995/multihead-LDSA). |
| > Builds on the Synthesizer ([Tay et al., arXiv:2005.00743](https://arxiv.org/abs/2005.00743)). |
|
|
| ## Differences from the reference (strict sense) |
|
|
| This kernel is the fused **core operation**, generalized; it is not a drop-in of the paper's |
| module. Concretely, versus `LocalDenseSynthesizerAttention` in the reference repo: |
|
|
| 1. **Scope.** The reference module bundles the weight-synthesis MLP (`w1 → ReLU → w2`), the |
| value projection (`w3`) and the output projection (`w_out`). This kernel is **only** the |
| softmax + windowed weighted-sum; the projections stay in cuBLAS on the caller side. That |
| makes it reusable for any synthesized local-attention that can hand over `a_logits` and `v`. |
| 2. **Window shape.** The reference uses a **symmetric/centered** context (size `c`, `(c−1)/2` |
| frames each side — bidirectional, offline). This kernel takes an **arbitrary `(left, right)`**: |
| set `left = right = (c−1)/2` to reproduce the paper, or `right = 0` for a **causal** window |
| (what a streaming recognizer deploys, e.g. `left=14, right=0`, `W=15`) for low latency. |
| 3. **Fusion / memory.** The reference `chunkwise`-unfolds the windowed values to `[B·T, H, c, d_k]` |
| and does softmax + `matmul`; this kernel streams the window in place (no unfold) in a single |
| launch — lower memory traffic. |
| 4. **Numerics.** fp32 accumulation, so it is at least as accurate as the eager op in low |
| precision. Parity vs an fp32-eager reference: `max|Δ|` ≈ **5e-7 (fp32)**, **2e-3 (bf16)**. |
| 5. **Inference-only.** No backward — grad-enabled / CPU calls fall back to the exported |
| `eager_ldsa` reference (which bit-matches the fused op's math). |
|
|
| ## Performance |
|
|
| Fused kernel vs the eager reference (softmax + pad + `W` shift-mul-add passes), measured with |
| `triton.testing.do_bench` under `torch.no_grad()`, bf16, deploy window `left=14, right=0` |
| (`W=15`), `H=8`, `Dh=64`, on an NVIDIA RTX PRO 6000 Blackwell **(under concurrent training load |
| — the back-to-back speedup ratios are robust; absolute times are inflated)**: |
|
|
| | shape `[B, T]` | speedup | |
| |---|---| |
| | 16 × 256 | ~10× | |
| | 16 × 512 | ~4× | |
| | 32 × 1024 | ~6× | |
| | 16 × 2048 | ~6× | |
| | 32 × 3000 | **~18×** | |
|
|
| Roughly **4–18×**, trending up with sequence length as the eager path's per-window memory |
| passes come to dominate. (Numbers taken on an idle GPU may be higher; these are a floor.) |
|
|
| --- |
| Built with [`kernel-builder`](https://github.com/huggingface/kernel-builder); Apache-2.0. |
|
|
