sparse-delta-memory 2.0.0__py3-none-any.whl

sparse_delta_memory/phase2_kernels.py

@@ -0,0 +1,267 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+"""Phase-2 kernels of the SDM recurrence: the dual matmul, and the geometry and quantization helpers the CUDA movers
+share. (2.0.0 removed this module's Triton entry movers and quantizer: the CUDA movers, ``movers``, are the only route;
+the description of the movers below is their shared contract.)
+
+Phase 2 carries the bank across chunks. Step ``ti`` covers chunk ``ti`` of all P streams,
+SC = P * C tokens:
+
+  forward   gather -> dual matmul -> DW
+  backward  GWR -> baddbmm -> Phase A -> GDS
+
+DW (decay and write) updates the bank rows the step writes and quantizes the snapshot. GWR
+(``token_gather`` of the bank gradient) gathers the gradient rows of the write slots. Phase A,
+the backward of DW, writes the bank gradient, restores the bank rows and emits three dots per
+write entry. GDS (read-side backward) emits one dot per read entry and scatters the reads'
+bank gradient. Phase B, the backward through the solve, follows in ``memory_ops``.
+
+Layout. Chunks are time-major (chunk ``ti * P + p``), N = Nt * P of them. Per-entry
+tensors are contiguous [N, C, W] (reads: [N, C, R]) and token rows are [N * C, head_dim].
+Tokens at or past ``valid_tokens`` in a stream are padding: they touch no bank row and
+their outputs are zero.
+
+Each ``make_*`` function checks its operands once and returns ``launch(ti)``. A launch
+passes whole tensors plus the step index, so the step loop slices nothing on the host.
+Buffers that may hold one step or every step (snapshot, delta_v) are told apart by their
+row count.
+
+Movers. A one-warp program owns ENT entries as an [ENT, BLOCK_D] tile and loops over
+NUM_D column tiles (``mover_geom``). Steps without padding whose grid covers the entries
+exactly run an unmasked variant.
+
+Determinism: every bank (bank-gradient) element has one writer per launch, so every output
+is bitwise reproducible. An entry whose slot occurs once in its chunk (is_only, q_only)
+writes its row directly. Entries sharing a slot form a group, a contiguous run of Phase 1's
+sorted kind row; the group's owner (its last entry, or for wide rows its last few entries,
+each owning whole column slices) writes the group's sum (``_group_sums``): the members in
+sorted order, in windows of consecutive entries of the kind row, each window reduced as a
+tree of separately rounded f32 products (see ``_group_slices``), in f32, with one rounding
+at the store. The other members write nothing. GDS adds into the bank gradient, each
+element's single writer (a single-entry slot, or a group's owner) rounding once: a group's
+owner, and a single-entry slot with a bf16 gradient, read the row, add the f32 sum or term
+and store once; with an fp32 gradient a single-entry slot adds with one f32 atomic (its only
+add). The lanes that hold copies of a sum agree bit for bit: no sum that crosses lanes
+starts from a product fused into its first cross-lane add (``_group_slices``, ``_row_dot``).
+The last entry of a slot restores its bank row in Phase A.
+
+"""
+
+from typing import Callable, NamedTuple
+
+import torch
+import triton
+import triton.language as tl
+
+from .cuda._device import tbind
+
+try:  # Triton 3.7 (the tested version) has the programmatic-dependent-launch trigger
+    from triton.language.extra.cuda import gdc_launch_dependents
+except ImportError:  # pragma: no cover - without it the trigger is a no-op (same arithmetic)
+    @triton.jit
+    def gdc_launch_dependents():
+        pass
+
+SNAPSHOT_QUANT_MODES = ("none", "fp8", "int4", "bf16")
+# "bf16" (float32 banks; bfloat16 banks run it as "none"): bf16 compact snapshot rows on the CUDA kernels only (the
+# unquantized forward with bf16 stores, gx_pa2's bf16-snapshot rows); never this module's movers (no fallback).
+
+
+# ============================================================
+# Geometry
+# ============================================================
+
+class MoverGeom(NamedTuple):
+    ent: int        # entries per program
+    block_d: int    # columns per tile
+    num_d: int      # column tiles per row
+
+
+def mover_geom(dtype: torch.dtype, head_dim: int, row_operands: int = 2) -> MoverGeom:
+    """Mover tile for bank rows of ``head_dim`` elements of ``dtype``; one warp per program.
+
+    A row is split into 16-byte vectors over up to 32 lanes, each lane holding at most 16
+    elements of it per tile; wider rows loop over column tiles. Each lane covers the same
+    columns of several rows, about four vectors of each of two row operands, so many bank
+    rows are in flight. Kernels with more ``row_operands`` take proportionally fewer rows,
+    but at least four when available. A warp spans whole rows, so the per-entry dots are
+    warp shuffles.
+    """
+    vec = 16 // dtype.itemsize
+    row_vecs = triton.cdiv(head_dim, vec)
+    lanes = min(32, triton.next_power_of_2(row_vecs))
+    per_lane = min(16 // vec, triton.next_power_of_2(triton.cdiv(row_vecs, lanes)))
+    block_d = lanes * per_lane * vec
+    ent = 32 // lanes * max(1, 4 // per_lane)
+    ent = min(ent, max(4, ent * 2 // row_operands))
+    return MoverGeom(ent, block_d, triton.cdiv(head_dim, block_d))
+
+
+class DualMatmulGeom(NamedTuple):
+    block_m: int
+    block_n: int
+    block_k: int
+    dot: str        # tl.dot input_precision
+
+
+def dual_matmul_geom(chunk_size: int, dtype: torch.dtype) -> DualMatmulGeom:
+    """Tiles of the dual matmul: at least two row blocks per chunk up to 64 rows. bf16 banks
+    use TF32 dots; fp32 banks follow ``torch.backends.cuda.matmul.allow_tf32``. IEEE dots
+    run on the FMA units, where 32-row blocks keep the accumulators in registers."""
+    fp32 = dtype == torch.float32
+    dot = "ieee" if fp32 and not torch.backends.cuda.matmul.allow_tf32 else "tf32"
+    block_m = min(64 if dot == "tf32" else 32, max(16, triton.next_power_of_2(chunk_size) // 2))
+    block_k = 16 if chunk_size <= 32 else 32
+    return DualMatmulGeom(block_m, 64, block_k, dot)
+
+
+# ============================================================
+# Kernels
+# ============================================================
+
+
+@triton.jit(do_not_specialize=["ti", "valid_tokens"])
+def _sdm_dual_matmul(
+    x_ptr, qx_ptr,              # [N, C, C] f32: X = (I + diag(beta) A)^-1, QX = QK @ X
+    v_ptr, beta_ptr,            # [N * C, HD] bank dtype, [N * C] f32
+    ret_ptr, inter_ptr,         # [SC, HD] f32 step rows from the gather
+    dv_ptr, dv_step_rows,       # [*, HD] f32 delta_v
+    out_ptr,                    # [N * C, HD] readings, bank dtype
+    res_ptr,                    # [N * C, HD] f32 residual, or None
+    ti, P, valid_tokens,
+    C: tl.constexpr, HD: tl.constexpr,
+    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, DOT: tl.constexpr,
+    TRIG: tl.constexpr,
+):
+    """delta_v = X @ u and readings = QX @ u + inter with u = beta * (v - retrieved), for
+    chunk p of step ti. The first row block also stores residual = v - retrieved.
+    The epilogue's ``inter`` rows are loaded before the k loop (they overlap its loads). TRIG: issue the
+    programmatic-dependent-launch trigger first, so the step's write mover, launched as a programmatic
+    dependent, can start its setup while this kernel runs (forward._dual_trig_wanted); no effect otherwise.
+
+    Grid: (cdiv(C, BLOCK_M) * cdiv(HD, BLOCK_N) * P,), row blocks first, then column
+    blocks, then chunks (one axis, which holds any number of chunks).
+    """
+    if TRIG:
+        gdc_launch_dependents()
+    NUM_M: tl.constexpr = (C + BLOCK_M - 1) // BLOCK_M
+    NUM_N: tl.constexpr = (HD + BLOCK_N - 1) // BLOCK_N
+    pid = tl.program_id(0)
+    pid_m = pid % NUM_M
+    pid_n = pid // NUM_M % NUM_N
+    p = (pid // (NUM_M * NUM_N)).to(tl.int64)
+    chunk = ti.to(tl.int64) * P + p
+    # 64-bit bases of this chunk's rows; the offsets within a chunk fit 32 bits.
+    x_ptr += chunk * C * C
+    qx_ptr += chunk * C * C
+    v_ptr += chunk * C * HD
+    beta_ptr += chunk * C
+    ret_ptr += p * C * HD
+    inter_ptr += p * C * HD
+    dv_ptr += (ti.to(tl.int64) * dv_step_rows + p * C) * HD
+    out_ptr += chunk * C * HD
+    if res_ptr is not None:
+        res_ptr += chunk * C * HD
+    rows = tl.minimum(valid_tokens - ti * C, C)
+    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    n_ok = rn[None, :] < HD
+    m_ok = (rm < rows)[:, None]
+    in_c = (rm < C)[:, None] & n_ok
+    m_rows = rm[:, None] * HD + rn[None, :]
+    inter = tl.load(inter_ptr + m_rows, mask=m_ok & n_ok, other=0.0)
+    acc_dv = tl.zeros((BLOCK_M, BLOCK_N), tl.float32)
+    acc_out = tl.zeros((BLOCK_M, BLOCK_N), tl.float32)
+    for k0 in range(0, C, BLOCK_K):
+        rk = k0 + tl.arange(0, BLOCK_K)
+        k_ok = rk < rows
+        tile = k_ok[:, None] & n_ok
+        k_rows = rk[:, None] * HD + rn[None, :]
+        v = tl.load(v_ptr + k_rows, mask=tile, other=0.0)
+        r = tl.load(ret_ptr + k_rows, mask=tile, other=0.0)
+        res = v.to(tl.float32) - r
+        if res_ptr is not None:
+            if pid_m == 0:
+                tl.store(res_ptr + k_rows, res, mask=(rk < C)[:, None] & n_ok)
+        u = tl.load(beta_ptr + rk, mask=k_ok, other=0.0)[:, None] * res
+        x_off = rm[:, None] * C + rk[None, :]
+        x_ok = (rm < C)[:, None] & (rk < C)[None, :]
+        x = tl.load(x_ptr + x_off, mask=x_ok, other=0.0)
+        acc_dv = tl.dot(x, u, acc_dv, input_precision=DOT)
+        qx = tl.load(qx_ptr + x_off, mask=x_ok, other=0.0)
+        acc_out = tl.dot(qx, u, acc_out, input_precision=DOT)
+    tl.store(dv_ptr + m_rows, tl.where(m_ok, acc_dv, 0.0), mask=in_c)
+    out = tl.where(m_ok, acc_out + inter, 0.0)
+    tl.store(out_ptr + m_rows, out.to(out_ptr.dtype.element_ty), mask=in_c)
+
+
+# ============================================================
+# Movers
+# ============================================================
+
+
+# ============================================================
+# Launchers
+# ============================================================
+
+def _check(t: torch.Tensor, name: str, numel: int, dtype=None) -> torch.Tensor:
+    if not t.is_contiguous() or t.numel() != numel or (dtype is not None and t.dtype != dtype):
+        want = f"{numel} contiguous elements" + (f" of {dtype}" if dtype is not None else "")
+        raise ValueError(f"{name}: expected {want}, got shape {tuple(t.shape)} "
+                         f"stride {t.stride()} {t.dtype}")
+    return t
+
+
+def quant_row(mode: str, head_dim: int) -> tuple[int, torch.dtype]:
+    """(row width, dtype) of a quantized snapshot row."""
+    if mode == "fp8":
+        return head_dim, torch.float8_e4m3fn
+    if mode == "int4":
+        return (head_dim + 1) // 2, torch.uint8
+    raise ValueError(f"snapshot_quant must be one of {SNAPSHOT_QUANT_MODES}, got {mode!r}")
+
+
+def _step_rows(t: torch.Tensor, name: str, rows: int, steps: int, width: int) -> int:
+    """Row stride between steps of ``t``: ``rows`` if it holds every step, 0 if one."""
+    if t.is_contiguous() and t.numel() == steps * rows * width:
+        return rows
+    if t.is_contiguous() and t.numel() == rows * width:
+        return 0
+    raise ValueError(f"{name}: expected {rows} or {steps * rows} contiguous rows of {width}, "
+                     f"got shape {tuple(t.shape)} stride {t.stride()}")
+
+
+def make_dual_matmul(x, qx, v, beta, retrieved, inter, delta_v, readings, residual=None, *,
+                     num_parallel: int, valid_tokens: int) -> Callable[[int], None]:
+    """Step launcher of delta_v = X @ u and readings = QX @ u + inter, u = beta * (v - retrieved).
+
+    x, qx: [N, C, C] f32; v, readings: [N * C, hd] (bank dtype); beta: [N * C] f32;
+    retrieved, inter: [SC, hd] f32 (this step's gather); delta_v: f32 rows of one step or
+    of every step; residual (v - retrieved, f32 [N * C, hd]) is stored when given.
+    """
+    n, c, _ = x.shape
+    hd = v.shape[-1]
+    sc = num_parallel * c
+    steps = n // num_parallel
+    geom = dual_matmul_geom(c, v.dtype)
+    for t, name in ((x, "x"), (qx, "qx")):
+        _check(t, name, n * c * c, torch.float32)
+    for t, name in ((retrieved, "retrieved"), (inter, "inter")):
+        _check(t, name, sc * hd, torch.float32)
+    _check(v, "v", n * c * hd)
+    _check(readings, "readings", n * c * hd, v.dtype)
+    _check(beta, "beta", n * c, torch.float32)
+    if residual is not None:
+        _check(residual, "residual", n * c * hd, torch.float32)
+    dv_rows = _step_rows(delta_v, "delta_v", sc, steps, hd)
+    kernel = tbind(_sdm_dual_matmul, (triton.cdiv(c, geom.block_m) * triton.cdiv(hd, geom.block_n)
+                                      * num_parallel,), v)
+    head = (x, qx, v, beta, retrieved, inter, delta_v, dv_rows, readings, residual)
+    meta = dict(C=c, HD=hd, BLOCK_M=geom.block_m, BLOCK_N=geom.block_n, BLOCK_K=geom.block_k,
+                DOT=geom.dot, TRIG=0, num_warps=4)
+
+    def launch(ti: int) -> None:
+        kernel(*head, ti, num_parallel, valid_tokens, **meta)
+    return launch
+
+

sparse_delta_memory/product_key_topk.py

@@ -0,0 +1,113 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+"""Exact top-k over all sums of product-key factors, with a sparse backward."""
+
+import math
+import operator
+from collections.abc import Sequence
+
+import torch
+
+from .cuda.product_key_topk_anyf_cuda import MAX_FACTORS, MAX_SLOTS
+from .cuda.product_key_topk_anyf_cuda import _module as _anyf_module
+from .cuda.product_key_topk_cuda import _module
+
+
+def _topk_function(module):
+    """The autograd function of the exact top-k kernels of `module` (a cached extension loader), with a sparse backward."""
+
+    class _ProductKeyTopK(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, raw, factors, k):
+            rows = raw.reshape(-1, raw.shape[-1])
+            if rows.stride(-1) != 1:
+                rows = rows.contiguous()
+            val, lin = module().topk_fwd(rows, factors, k)
+            ctx.set_materialize_grads(False)
+            ctx.save_for_backward(lin)
+            ctx.raw_shape = raw.shape
+            ctx.factors = factors
+            ctx.mark_non_differentiable(lin)
+            out_shape = (*raw.shape[:-1], lin.shape[-1])
+            return val.view(out_shape), lin.view(out_shape)
+
+        @staticmethod
+        def backward(ctx, grad_val, _grad_lin):
+            (lin,) = ctx.saved_tensors
+            grad = module().topk_bwd(lin, grad_val.reshape(lin.shape), ctx.factors)
+            return grad.view(ctx.raw_shape), None, None
+    return _ProductKeyTopK
+
+
+_ProductKeyTopK = _topk_function(_module)          # cuda/product_key_topk_cuda.py
+_ProductKeyTopKAnyF = _topk_function(_anyf_module)  # cuda/product_key_topk_anyf_cuda.py
+
+
+def _bit_packed(factors: tuple[int, ...]) -> bool:
+    """Whether the kernels of ``cuda/product_key_topk_cuda.py`` take ``factors``: two to four
+    factors of 1 to 65535 columns, with at most 32 bits of bit-packed slot index."""
+    return (2 <= len(factors) <= 4 and all(1 <= f <= 65535 for f in factors)
+            and sum((f - 1).bit_length() for f in factors) <= 32)
+
+
+def check_product_key_factors(factors: Sequence[int]) -> tuple[int, ...]:
+    """The factor widths as a tuple of ints, or ``ValueError`` if ``product_key_topk`` cannot
+    select over them: it takes 1 to ``MAX_FACTORS`` factors of at least one column each,
+    addressing at most ``MAX_SLOTS`` slots with a row width of at most ``MAX_SLOTS`` (or
+    any factor list of the two- to four-factor kernels)."""
+    if isinstance(factors, (str, bytes)) or not isinstance(factors, Sequence):
+        raise ValueError(f"product-key factors must be a sequence of ints, got {factors!r}")
+    if any(isinstance(f, bool) for f in factors):
+        raise ValueError(f"product-key factors must be ints, got {list(factors)}")
+    try:
+        dims = tuple(operator.index(f) for f in factors)
+    except TypeError:
+        raise ValueError(f"product-key factors must be ints, got {list(factors)}") from None
+    if not 1 <= len(dims) <= MAX_FACTORS or min(dims) < 1:
+        raise ValueError(f"product-key factors must be 1 to {MAX_FACTORS} positive ints, got "
+                         f"{list(dims)}")
+    if not _bit_packed(dims) and (math.prod(dims) > MAX_SLOTS or sum(dims) > MAX_SLOTS):
+        raise ValueError(f"product-key factors {list(dims)} address {math.prod(dims)} slots "
+                         f"over {sum(dims)} columns; at most {MAX_SLOTS} of each.")
+    return dims
+
+
+def product_key_topk(
+    raw: torch.Tensor, factors: Sequence[int], k: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Exact top-k over all sums of one value from each product-key factor, for any number
+    of factors.
+
+    raw [..., sum(factors)] (CUDA, bfloat16 or float32) holds the F >= 1 factors side by side,
+    each at least one column wide (see ``check_product_key_factors`` for the limits). Slot
+    (c_0, ..., c_{F-1}) scores the float32 sum of its factor values, added in factor order;
+    with bfloat16 input and more than two factors, each partial sum is rounded to bfloat16
+    before the next factor is added. With one factor, slot c_0 scores raw[..., c_0]. Returns
+    (val, lin), each [..., kk] with kk = min(k, prod(factors)) and 1 <= k <= 256:
+    lin = (..(c_0 * factors[1] + c_1) * ..) * factors[F-1] + c_{F-1} (int64), strictly
+    ascending, and val = the slot's score rounded to raw.dtype. val is differentiable with
+    respect to raw; the backward is deterministic.
+
+    Selection rule: scores compare as float32 numbers, with -0 below +0 and NaN counted as
+    -inf. Each factor is ranked by value, larger column first on ties; one factor selects its
+    first kk ranks. Stage f = 1..F-1 keeps the best c = min(k, n * m) of the pairs (i, j) of a
+    rank i < n = min(k, prod(factors[:f])) of the previous stage and a rank j < m =
+    min(k, factors[f]) of factor f with (i + 1) * (j + 1) <= c, exact ties going to the
+    larger i + j, then the larger i. Each pair left out has c pairs of lower or equal ranks
+    scoring at least as much, so the selected scores are an exact top-k. With two factors,
+    bfloat16 input and its float32 upcast select the same slots. The kernels never
+    synchronize and need no warmup before CUDA-graph capture.
+
+    Two to four factors of at most 65535 columns and 32 bits of bit-packed slot index run on
+    the kernels of ``cuda/product_key_topk_cuda.py``; every other factor list runs on those
+    of ``cuda/product_key_topk_anyf_cuda.py``, which apply the same rule.
+    """
+    factors = tuple(operator.index(f) for f in factors)
+    if _bit_packed(factors):
+        return _ProductKeyTopK.apply(raw, factors, operator.index(k))
+    return _ProductKeyTopKAnyF.apply(raw, factors, operator.index(k))
+
+
+def _product_key_topk_anyf(raw, factors, k):
+    """``product_key_topk`` on the any-factor kernels whatever the factor list (tests)."""
+    factors = tuple(operator.index(f) for f in factors)
+    return _ProductKeyTopKAnyF.apply(raw, factors, operator.index(k))

sparse_delta_memory/wy_solve.py

@@ -0,0 +1,273 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# The 16 x 16 tile scheme of the level-1 kernel is adapted from the solve_tril kernel of
+# flash-linear-attention (https://github.com/fla-org/flash-linear-attention), which is
+# distributed under the MIT License:
+#
+#   Copyright (c) 2023-2026 Songlin Yang, Yu Zhang, Zhiyuan Li
+#
+#   Permission is hereby granted, free of charge, to any person obtaining a copy
+#   of this software and associated documentation files (the "Software"), to deal
+#   in the Software without restriction, including without limitation the rights
+#   to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+#   copies of the Software, and to permit persons to whom the Software is
+#   furnished to do so, subject to the following conditions:
+#
+#   The above copyright notice and this permission notice shall be included in all
+#   copies or substantial portions of the Software.
+#
+#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+#   SOFTWARE.
+"""Per-chunk WY solve: X = (I + L)^-1 and QX = QK @ X.
+
+For each chunk of C tokens, ``L = tril(beta[:, None] * A, -1)`` is the gated strictly lower
+key-key matrix and ``QK`` the causal query-key matrix (diagonal included). ``wy_solve``
+returns the unit lower-triangular ``X`` and ``QX``, both dense row-major [N, C, C] float32
+with +0.0 above the diagonal. Only the strict lower triangle of ``A`` and the lower
+triangle of ``QK`` are read, so the rest may hold anything.
+
+Level 1: one program per 64 x 64 diagonal block (16 x 16 or 32 x 32 when C <= 32). Its
+16 x 16 diagonal tiles are inverted by forward substitution and the tiles below them by
+matmuls. The same program emits the diagonal blocks of QX and the zeros of its block row.
+
+Level 2 (C > 64): one program per 64-column panel J walks down the blocks below the
+diagonal, X[i, J] = -X[i, i] @ sum_{k=J}^{i-1} L[i, k] X[k, J], and accumulates
+QX[i, J] = sum_{k=J}^{i} QK[i, k] X[k, J] from the same loads of X. Its matmuls take
+16-wide slices of the reduction dimension from memory (a full-float32 dot keeps whole
+operand rows in registers, so wider operands spill); X[i, J] is staged in its own output
+block so that it can be read back that way.
+
+Rows and columns >= C load as zero, i.e. identity rows of I + L. The leading block of a
+lower-triangular inverse is the inverse of the leading block, so any C is exact.
+
+Every dot is full float32 whatever torch.backends.cuda.matmul.allow_tf32 says: the inverse
+amplifies input rounding. Launch configurations are pinned, not autotuned: the warp count
+changes the rounding, so the table is part of the numerical contract, and a fixed launch
+is capture-safe and gives the same bits for a chunk at any N.
+
+The 16 x 16 tile scheme is adapted from the ``solve_tril`` kernel of flash-linear-attention
+(MIT License, notice above).
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from .cuda._device import tbind
+
+# Pinned (num_warps, num_stages) of the two kernels.
+_DIAG_CONFIG = (1, 2)
+_OFFDIAG_CONFIG = (2, 3)
+
+
+@triton.jit
+def _mm(a, b, acc=None):
+    return tl.dot(a, b, acc, input_precision="ieee")
+
+
+@triton.jit
+def _load(base, rows, cols, LIM, C: tl.constexpr):
+    """base[rows, cols] with row pitch C; zero outside the leading [LIM, LIM] block."""
+    return tl.load(base + rows[:, None] * C + cols[None, :],
+                   mask=(rows < LIM)[:, None] & (cols < LIM)[None, :], other=0.0)
+
+
+@triton.jit
+def _store(base, v, rows, cols, LIM, C: tl.constexpr):
+    tl.store(base + rows[:, None] * C + cols[None, :], v,
+             mask=(rows < LIM)[:, None] & (cols < LIM)[None, :])
+
+
+@triton.jit
+def _tile(base, r0, c0, LIM, C: tl.constexpr):
+    o = tl.arange(0, 16)
+    return _load(base, r0 + o, c0 + o, LIM, C)
+
+
+@triton.jit
+def _put(base, v, r0, c0, LIM, C: tl.constexpr):
+    o = tl.arange(0, 16)
+    _store(base, v, r0 + o, c0 + o, LIM, C)
+
+
+@triton.jit
+def _band_beta(beta, r0, LIM):
+    o = tl.arange(0, 16)
+    return tl.load(beta + r0 + o, mask=r0 + o < LIM, other=0.0)[:, None]
+
+
+@triton.jit
+def _lower(a, beta, r0, c0, LIM, C: tl.constexpr):
+    """beta[:, None] * A on the 16 x 16 tile at (r0, c0)."""
+    return _tile(a, r0, c0, LIM, C) * _band_beta(beta, r0, LIM)
+
+
+@triton.jit
+def _subst_start(a, beta, r0, LIM, C: tl.constexpr):
+    """-L of the diagonal tile at r0: rows 0 and 1 of (I + L)^-1 - I are already final."""
+    o = tl.arange(0, 16)
+    return -tl.where(o[:, None] > o[None, :], _lower(a, beta, r0, r0, LIM, C), 0.0)
+
+
+@triton.jit
+def _subst_row(b, a, beta, r0, i, LIM, C: tl.constexpr):
+    """Forward substitution of row i of the diagonal tile at r0."""
+    o = tl.arange(0, 16)
+    r = r0 + i
+    row = tl.load(a + r * C + r0 + o, mask=(r < LIM) & (r0 + o < LIM), other=0.0)
+    row = -tl.where(r0 + o < r, row * tl.load(beta + r, mask=r < LIM, other=0.0), 0.0)
+    row += tl.sum(row[:, None] * b, 0)
+    return tl.where((o == i)[:, None], row, b)
+
+
+@triton.jit
+def _qk_diag(q, r0, LIM, C: tl.constexpr):
+    o = tl.arange(0, 16)
+    return tl.where(o[:, None] >= o[None, :], _tile(q, r0, r0, LIM, C), 0.0)
+
+
+@triton.jit
+def _wy_diag_kernel(A, Beta, QK, X, QX, C: tl.constexpr, BD: tl.constexpr):
+    n = tl.program_id(0).to(tl.int64)
+    off = tl.program_id(1) * BD
+    LIM = C - off
+    blk = n * C * C + off * C + off
+    a, q, x, qx = A + blk, QK + blk, X + blk, QX + blk
+    beta = Beta + n * C + off
+
+    # Diagonal tiles; the independent substitutions are interleaved for latency hiding.
+    x11 = _subst_start(a, beta, 0, LIM, C)
+    if BD >= 32:
+        x22 = _subst_start(a, beta, 16, LIM, C)
+    if BD == 64:
+        x33 = _subst_start(a, beta, 32, LIM, C)
+        x44 = _subst_start(a, beta, 48, LIM, C)
+    for i in range(2, 16):
+        x11 = _subst_row(x11, a, beta, 0, i, LIM, C)
+        if BD >= 32:
+            x22 = _subst_row(x22, a, beta, 16, i, LIM, C)
+        if BD == 64:
+            x33 = _subst_row(x33, a, beta, 32, i, LIM, C)
+            x44 = _subst_row(x44, a, beta, 48, i, LIM, C)
+    o = tl.arange(0, 16)
+    eye = (o[:, None] == o[None, :]).to(tl.float32)
+
+    x11 += eye
+    _put(x, x11, 0, 0, LIM, C)
+    _put(qx, _mm(_qk_diag(q, 0, LIM, C), x11), 0, 0, LIM, C)
+    if BD >= 32:
+        x22 += eye
+        x21 = -_mm(_mm(x22, _lower(a, beta, 16, 0, LIM, C)), x11)
+        _put(x, x21, 16, 0, LIM, C)
+        _put(x, x22, 16, 16, LIM, C)
+        q22 = _qk_diag(q, 16, LIM, C)
+        _put(qx, _mm(q22, x21, _mm(_tile(q, 16, 0, LIM, C), x11)), 16, 0, LIM, C)
+        _put(qx, _mm(q22, x22), 16, 16, LIM, C)
+    if BD == 64:
+        # The association of each product is fixed: it defines the output bits.
+        x33 += eye
+        x44 += eye
+        a32 = _lower(a, beta, 32, 16, LIM, C)
+        a42 = _lower(a, beta, 48, 16, LIM, C)
+        a43 = _lower(a, beta, 48, 32, LIM, C)
+        x32 = -_mm(_mm(x33, a32), x22)
+        x43 = -_mm(_mm(x44, a43), x33)
+        x31 = -_mm(x33, _mm(_lower(a, beta, 32, 0, LIM, C), x11) + _mm(a32, x21))
+        x42 = -_mm(x44, _mm(a42, x22) + _mm(a43, x32))
+        x41 = -_mm(x44, _mm(_lower(a, beta, 48, 0, LIM, C), x11) + _mm(a42, x21) + _mm(a43, x31))
+        _put(x, x31, 32, 0, LIM, C)
+        _put(x, x32, 32, 16, LIM, C)
+        _put(x, x33, 32, 32, LIM, C)
+        _put(x, x41, 48, 0, LIM, C)
+        _put(x, x42, 48, 16, LIM, C)
+        _put(x, x43, 48, 32, LIM, C)
+        _put(x, x44, 48, 48, LIM, C)
+        q32 = _tile(q, 32, 16, LIM, C)
+        q33 = _qk_diag(q, 32, LIM, C)
+        _put(qx, _mm(q33, x31, _mm(q32, x21, _mm(_tile(q, 32, 0, LIM, C), x11))), 32, 0, LIM, C)
+        _put(qx, _mm(q33, x32, _mm(q32, x22)), 32, 16, LIM, C)
+        _put(qx, _mm(q33, x33), 32, 32, LIM, C)
+        q42 = _tile(q, 48, 16, LIM, C)
+        q43 = _tile(q, 48, 32, LIM, C)
+        q44 = _qk_diag(q, 48, LIM, C)
+        _put(qx, _mm(q44, x41, _mm(q43, x31, _mm(q42, x21, _mm(_tile(q, 48, 0, LIM, C), x11)))),
+             48, 0, LIM, C)
+        _put(qx, _mm(q44, x42, _mm(q43, x32, _mm(q42, x22))), 48, 16, LIM, C)
+        _put(qx, _mm(q44, x43, _mm(q43, x33)), 48, 32, LIM, C)
+        _put(qx, _mm(q44, x44), 48, 48, LIM, C)
+
+    # Zeros: the 16 x 16 tiles above the diagonal inside the block, then the rest of the row.
+    r = tl.arange(0, BD)
+    for c in range(0, LIM, BD):
+        cols = c + r
+        m = (cols[None, :] // 16 > r[:, None] // 16) & (r[:, None] < LIM) & (cols[None, :] < LIM)
+        offs = r[:, None] * C + cols[None, :]
+        tl.store(x + offs, 0.0, mask=m)
+        tl.store(qx + offs, 0.0, mask=m)
+
+
+@triton.jit
+def _wy_offdiag_kernel(A, Beta, QK, X, QX, C: tl.constexpr, NB: tl.constexpr):
+    n = tl.program_id(0).to(tl.int64)
+    j = tl.program_id(1)
+    base = n * C * C
+    a, q, x, qx = A + base, QK + base, X + base, QX + base
+    r = tl.arange(0, 64)
+    s = tl.arange(0, 16)
+    cj = j * 64 + r
+    for i in range(j + 1, NB):
+        ri = i * 64 + r
+        beta = tl.load(Beta + n * C + ri, mask=ri < C, other=0.0)[:, None]
+        lx = tl.zeros((64, 64), dtype=tl.float32)
+        qx_ij = tl.zeros((64, 64), dtype=tl.float32)
+        for k in range(j * 64, i * 64, 16):
+            x_kj = _load(x, k + s, cj, C, C)
+            lx = _mm(_load(a, ri, k + s, C, C) * beta, x_kj, lx)
+            qx_ij = _mm(_load(q, ri, k + s, C, C), x_kj, qx_ij)
+        # X[i, j] = -X[i, i] @ lx, with lx staged in the X[i, j] block.
+        _store(x, lx, ri, cj, C, C)
+        tl.debug_barrier()
+        x_ij = tl.zeros((64, 64), dtype=tl.float32)
+        for k in range(i * 64, i * 64 + 64, 16):
+            x_ij = _mm(_load(x, ri, k + s, C, C), _load(x, k + s, cj, C, C), x_ij)
+        tl.debug_barrier()
+        _store(x, -x_ij, ri, cj, C, C)
+        tl.debug_barrier()
+        for k in range(i * 64, i * 64 + 64, 16):
+            cols = k + s
+            q_ii = tl.where(ri[:, None] >= cols[None, :], _load(q, ri, cols, C, C), 0.0)
+            qx_ij = _mm(q_ii, _load(x, cols, cj, C, C), qx_ij)
+        _store(qx, qx_ij, ri, cj, C, C)
+
+
+def wy_solve(A: torch.Tensor, beta: torch.Tensor, QK: torch.Tensor):
+    """X = (I + tril(beta[:, None] * A, -1))^-1 and QX = QK @ X for each chunk.
+
+    A, QK: [N, C, C] float32, any C >= 1. beta: N * C float32 values, e.g. [N, C].
+    Returns (X, QX), contiguous [N, C, C] float32.
+    """
+    if A.dim() != 3 or A.shape[1] != A.shape[2] or QK.shape != A.shape:
+        raise ValueError(
+            f"wy_solve: A and QK must be [N, C, C], got {tuple(A.shape)} and {tuple(QK.shape)}")
+    N, C, _ = A.shape
+    if beta.numel() != N * C:
+        raise ValueError(f"wy_solve: beta must have N * C = {N * C} elements, got {beta.numel()}")
+    if not A.dtype == QK.dtype == beta.dtype == torch.float32:
+        raise TypeError("wy_solve: A, beta and QK must be float32")
+    A, beta, QK = A.contiguous(), beta.reshape(N * C).contiguous(), QK.contiguous()
+    X = torch.empty_like(A)
+    QX = torch.empty_like(A)
+    bd = 16 if C <= 16 else 32 if C <= 32 else 64
+    nb = triton.cdiv(C, bd)
+    w, s = _DIAG_CONFIG
+    tbind(_wy_diag_kernel, (N, nb), A)(A, beta, QK, X, QX, C=C, BD=bd, num_warps=w, num_stages=s)
+    if nb > 1:
+        w, s = _OFFDIAG_CONFIG
+        tbind(_wy_offdiag_kernel, (N, nb - 1), A)(A, beta, QK, X, QX, C=C, NB=nb,
+                                                  num_warps=w, num_stages=s)
+    return X, QX