ldsa

Fused Local Dense Synthesizer Attention (LDSA) window op as a noarch Triton kernel for the 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

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, code: github.com/mlxu995/multihead-LDSA. Builds on the Synthesizer (Tay et al., arXiv: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; Apache-2.0.