fast-attn-kernels 0.1.0__tar.gz

fast_attn_kernels-0.1.0/PKG-INFO

@@ -0,0 +1,4 @@
+Metadata-Version: 2.4
+Name: fast-attn-kernels
+Version: 0.1.0
+Requires-Dist: torch

fast_attn_kernels-0.1.0/attns/__init__.py

@@ -0,0 +1 @@
+from .linear_bind import LinearAttention, linear_attention

fast_attn_kernels-0.1.0/attns/linear_attn.c

@@ -0,0 +1,336 @@
+#include <stdlib.h>
+#include <string.h>
+#include <math.h>
+
+#ifdef _OPENMP
+#include <omp.h>
+#endif
+
+static inline float phi(float x) { return x > 0.0f ? x + 1.0f : expf(x); }
+
+static inline float dphi(float x, float phix) { return x > 0.0f ? 1.0f : phix; }
+
+/* ============================================================
+   Forward pass: chunked causal linear attention
+   ============================================================ */
+void linear_attention_forward(
+    const float* q, const float* k, const float* v,
+    float* out, float* kv_states, float* z_states,
+    long B, long H, long N, long D, long E,
+    long chunk_size, long num_chunks)
+{
+    long BH = B * H;
+    long qk_bh_stride = N * D;
+    long v_bh_stride  = N * E;
+    long kv_bh_stride = num_chunks * D * E;
+    long z_bh_stride  = num_chunks * D;
+
+    #pragma omp parallel for schedule(dynamic)
+    for (long bh = 0; bh < BH; ++bh) {
+        const float* Qb = q + bh * qk_bh_stride;
+        const float* Kb = k + bh * qk_bh_stride;
+        const float* Vb = v + bh * v_bh_stride;
+        float* Ob       = out + bh * v_bh_stride;
+        float* KVb      = kv_states + bh * kv_bh_stride;
+        float* Zb       = z_states + bh * z_bh_stride;
+
+        float* S       = calloc((size_t)(D * E), sizeof(float));
+        float* zstate  = calloc((size_t)D, sizeof(float));
+        float* qphi    = malloc((size_t)(chunk_size * D) * sizeof(float));
+        float* kphi    = malloc((size_t)(chunk_size * D) * sizeof(float));
+        float* L       = malloc((size_t)(chunk_size * chunk_size) * sizeof(float));
+        float* out_row = malloc((size_t)E * sizeof(float));
+
+        for (long c = 0; c < num_chunks; ++c) {
+            long c_start = c * chunk_size;
+            long rem = N - c_start;
+            long C = chunk_size < rem ? chunk_size : rem;
+
+            memcpy(KVb + c * D * E, S, (size_t)(D * E) * sizeof(float));
+            memcpy(Zb + c * D, zstate, (size_t)D * sizeof(float));
+
+            for (long i = 0; i < C; ++i) {
+                const float* qi = Qb + (c_start + i) * D;
+                const float* ki = Kb + (c_start + i) * D;
+                for (long d = 0; d < D; ++d) {
+                    qphi[i * D + d] = phi(qi[d]);
+                    kphi[i * D + d] = phi(ki[d]);
+                }
+            }
+
+            for (long i = 0; i < C; ++i) {
+                const float* qi = &qphi[i * D];
+                for (long j = 0; j <= i; ++j) {
+                    const float* kj = &kphi[j * D];
+                    float acc = 0.0f;
+                    for (long d = 0; d < D; ++d) acc += qi[d] * kj[d];
+                    L[i * chunk_size + j] = acc;
+                }
+            }
+
+            for (long i = 0; i < C; ++i) {
+                const float* qi = &qphi[i * D];
+                for (long e = 0; e < E; ++e) out_row[e] = 0.0f;
+
+                float z_intra = 0.0f;
+                for (long j = 0; j <= i; ++j) {
+                    float lij = L[i * chunk_size + j];
+                    z_intra += lij;
+                    const float* vj = Vb + (c_start + j) * E;
+                    for (long e = 0; e < E; ++e) out_row[e] += lij * vj[e];
+                }
+
+                float z_inter = 0.0f;
+                for (long d = 0; d < D; ++d) z_inter += qi[d] * zstate[d];
+                for (long e = 0; e < E; ++e) {
+                    float acc = 0.0f;
+                    for (long d = 0; d < D; ++d) acc += qi[d] * S[d * E + e];
+                    out_row[e] += acc;
+                }
+
+                float z = z_intra + z_inter;
+                if (z < 1e-6f) z = 1e-6f;
+
+                float* orow = Ob + (c_start + i) * E;
+                for (long e = 0; e < E; ++e) orow[e] = out_row[e] / z;
+            }
+
+            for (long i = 0; i < C; ++i) {
+                const float* ki = &kphi[i * D];
+                const float* vi = Vb + (c_start + i) * E;
+                for (long d = 0; d < D; ++d) {
+                    zstate[d] += ki[d];
+                    float kd = ki[d];
+                    float* Srow = &S[d * E];
+                    for (long e = 0; e < E; ++e) Srow[e] += kd * vi[e];
+                }
+            }
+        }
+
+        free(S);
+        free(zstate);
+        free(qphi);
+        free(kphi);
+        free(L);
+        free(out_row);
+    }
+}
+
+/* ============================================================
+   Backward pass: recompute-based, chunked, reverse order
+   ============================================================ */
+void linear_attention_backward(
+    const float* q, const float* k, const float* v, const float* dout,
+    const float* kv_states, const float* z_states,
+    float* dq, float* dk, float* dv,
+    long B, long H, long N, long D, long E,
+    long chunk_size, long num_chunks)
+{
+    long BH = B * H;
+    long qk_bh_stride = N * D;
+    long v_bh_stride  = N * E;
+    long kv_bh_stride = num_chunks * D * E;
+    long z_bh_stride  = num_chunks * D;
+
+    #pragma omp parallel for schedule(dynamic)
+    for (long bh = 0; bh < BH; ++bh) {
+        const float* Qb  = q + bh * qk_bh_stride;
+        const float* Kb  = k + bh * qk_bh_stride;
+        const float* Vb  = v + bh * v_bh_stride;
+        const float* DOb = dout + bh * v_bh_stride;
+        const float* KVb = kv_states + bh * kv_bh_stride;
+        const float* Zb  = z_states + bh * z_bh_stride;
+        float* DQb = dq + bh * qk_bh_stride;
+        float* DKb = dk + bh * qk_bh_stride;
+        float* DVb = dv + bh * v_bh_stride;
+
+        float* dS       = calloc((size_t)(D * E), sizeof(float));
+        float* dzstate  = calloc((size_t)D, sizeof(float));
+        float* qraw     = malloc((size_t)(chunk_size * D) * sizeof(float));
+        float* kraw     = malloc((size_t)(chunk_size * D) * sizeof(float));
+        float* qphi     = malloc((size_t)(chunk_size * D) * sizeof(float));
+        float* kphi     = malloc((size_t)(chunk_size * D) * sizeof(float));
+        float* L        = malloc((size_t)(chunk_size * chunk_size) * sizeof(float));
+        float* dL       = malloc((size_t)(chunk_size * chunk_size) * sizeof(float));
+        float* dnum     = malloc((size_t)(chunk_size * E) * sizeof(float));
+        float* dden     = malloc((size_t)chunk_size * sizeof(float));
+        float* dq_total = malloc((size_t)(chunk_size * D) * sizeof(float));
+        float* dk_total = malloc((size_t)(chunk_size * D) * sizeof(float));
+        float* dv_total = malloc((size_t)(chunk_size * E) * sizeof(float));
+        float* dS_local = malloc((size_t)(D * E) * sizeof(float));
+        float* dz_local = malloc((size_t)D * sizeof(float));
+        float* num_row  = malloc((size_t)E * sizeof(float));
+
+        for (long c = num_chunks - 1; c >= 0; --c) {
+            long c_start = c * chunk_size;
+            long rem = N - c_start;
+            long C = chunk_size < rem ? chunk_size : rem;
+
+            const float* Sb = KVb + c * D * E;
+            const float* zb = Zb + c * D;
+
+            /* load + apply feature map */
+            for (long i = 0; i < C; ++i) {
+                const float* qi = Qb + (c_start + i) * D;
+                const float* ki = Kb + (c_start + i) * D;
+                for (long d = 0; d < D; ++d) {
+                    float qr = qi[d];
+                    float kr = ki[d];
+                    qraw[i * D + d] = qr;
+                    kraw[i * D + d] = kr;
+                    qphi[i * D + d] = phi(qr);
+                    kphi[i * D + d] = phi(kr);
+                }
+            }
+
+            /* recompute L = tril(q @ k^T) */
+            for (long i = 0; i < C; ++i) {
+                for (long j = 0; j <= i; ++j) {
+                    float acc = 0.0f;
+                    for (long d = 0; d < D; ++d) acc += qphi[i * D + d] * kphi[j * D + d];
+                    L[i * chunk_size + j] = acc;
+                }
+            }
+
+            /* recompute forward output -> derive dnum, dden */
+            for (long i = 0; i < C; ++i) {
+                const float* qi = &qphi[i * D];
+                for (long e = 0; e < E; ++e) num_row[e] = 0.0f;
+
+                float z_intra = 0.0f;
+                for (long j = 0; j <= i; ++j) {
+                    float lij = L[i * chunk_size + j];
+                    z_intra += lij;
+                    const float* vj = Vb + (c_start + j) * E;
+                    for (long e = 0; e < E; ++e) num_row[e] += lij * vj[e];
+                }
+
+                float z_inter = 0.0f;
+                for (long d = 0; d < D; ++d) z_inter += qi[d] * zb[d];
+                for (long e = 0; e < E; ++e) {
+                    float acc = 0.0f;
+                    for (long d = 0; d < D; ++d) acc += qi[d] * Sb[d * E + e];
+                    num_row[e] += acc;
+                }
+
+                float den = z_intra + z_inter;
+                if (den < 1e-6f) den = 1e-6f;
+
+                const float* dorow = DOb + (c_start + i) * E;
+                float dot_do_out = 0.0f;
+                for (long e = 0; e < E; ++e) {
+                    float out_e = num_row[e] / den;
+                    dnum[i * E + e] = dorow[e] / den;
+                    dot_do_out += dorow[e] * out_e;
+                }
+                dden[i] = -dot_do_out / den;
+            }
+
+            memset(dq_total, 0, (size_t)(C * D) * sizeof(float));
+            memset(dk_total, 0, (size_t)(C * D) * sizeof(float));
+            memset(dv_total, 0, (size_t)(C * E) * sizeof(float));
+            memset(dS_local, 0, (size_t)(D * E) * sizeof(float));
+            memset(dz_local, 0, (size_t)D * sizeof(float));
+
+            /* grads from the inter-chunk term (q @ S_before, q . z_before) */
+            for (long i = 0; i < C; ++i) {
+                const float* qi = &qphi[i * D];
+
+                for (long d = 0; d < D; ++d) {
+                    float acc = 0.0f;
+                    for (long e = 0; e < E; ++e) acc += dnum[i * E + e] * Sb[d * E + e];
+                    acc += dden[i] * zb[d];
+                    dq_total[i * D + d] += acc;
+                }
+
+                for (long d = 0; d < D; ++d) {
+                    float qd = qi[d];
+                    for (long e = 0; e < E; ++e) dS_local[d * E + e] += qd * dnum[i * E + e];
+                }
+
+                for (long d = 0; d < D; ++d) dz_local[d] += qi[d] * dden[i];
+            }
+
+            /* grads from the intra-chunk term (causal local attention) */
+            for (long i = 0; i < C; ++i) {
+                for (long j = 0; j <= i; ++j) {
+                    float dl = 0.0f;
+                    const float* vj = Vb + (c_start + j) * E;
+                    for (long e = 0; e < E; ++e) dl += dnum[i * E + e] * vj[e];
+                    dl += dden[i];
+                    dL[i * chunk_size + j] = dl;
+                }
+            }
+
+            for (long i = 0; i < C; ++i) {
+                const float* qi = &qphi[i * D];
+                for (long j = 0; j <= i; ++j) {
+                    float dl = dL[i * chunk_size + j];
+                    const float* kj = &kphi[j * D];
+                    for (long d = 0; d < D; ++d) {
+                        dq_total[i * D + d] += dl * kj[d];
+                        dk_total[j * D + d] += dl * qi[d];
+                    }
+                }
+            }
+
+            for (long i = 0; i < C; ++i) {
+                for (long j = 0; j <= i; ++j) {
+                    float lij = L[i * chunk_size + j];
+                    for (long e = 0; e < E; ++e) dv_total[j * E + e] += lij * dnum[i * E + e];
+                }
+            }
+
+            /* grads from the recurrent state update: S += k^T v, zstate += sum(k) */
+            for (long i = 0; i < C; ++i) {
+                const float* vi = Vb + (c_start + i) * E;
+                const float* ki = &kphi[i * D];
+
+                for (long d = 0; d < D; ++d) {
+                    float acc = 0.0f;
+                    for (long e = 0; e < E; ++e) acc += vi[e] * dS[d * E + e];
+                    dk_total[i * D + d] += acc + dzstate[d];
+                }
+
+                for (long e = 0; e < E; ++e) {
+                    float acc = 0.0f;
+                    for (long d = 0; d < D; ++d) acc += ki[d] * dS[d * E + e];
+                    dv_total[i * E + e] += acc;
+                }
+            }
+
+            /* propagate accumulated state-gradient to the previous chunk */
+            for (long idx = 0; idx < D * E; ++idx) dS[idx] += dS_local[idx];
+            for (long d = 0; d < D; ++d) dzstate[d] += dz_local[d];
+
+            /* apply feature-map derivative, write results */
+            for (long i = 0; i < C; ++i) {
+                float* dqrow = DQb + (c_start + i) * D;
+                float* dkrow = DKb + (c_start + i) * D;
+                for (long d = 0; d < D; ++d) {
+                    dqrow[d] = dq_total[i * D + d] * dphi(qraw[i * D + d], qphi[i * D + d]);
+                    dkrow[d] = dk_total[i * D + d] * dphi(kraw[i * D + d], kphi[i * D + d]);
+                }
+                float* dvrow = DVb + (c_start + i) * E;
+                for (long e = 0; e < E; ++e) dvrow[e] = dv_total[i * E + e];
+            }
+        }
+
+        free(dS);
+        free(dzstate);
+        free(qraw);
+        free(kraw);
+        free(qphi);
+        free(kphi);
+        free(L);
+        free(dL);
+        free(dnum);
+        free(dden);
+        free(dq_total);
+        free(dk_total);
+        free(dv_total);
+        free(dS_local);
+        free(dz_local);
+        free(num_row);
+    }
+}

fast_attn_kernels-0.1.0/attns/linear_bind.py

@@ -0,0 +1,414 @@
+import os
+import sys
+import ctypes
+import platform
+import torch
+
+_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+
+
+def _find_lib():
+    # possible names depending on OS
+    if platform.system() == "Windows": names = ["liblinear_attn.dll", "linear_attn.dll"]
+    elif platform.system() == "Darwin": names = ["liblinear_attn.dylib"]
+    else: names = ["liblinear_attn.so"]
+
+    # possible locations (depends on CMake generator: single-config vs multi-config)
+    search_dirs = [
+        os.path.join(_THIS_DIR, "..", "build"),
+        os.path.join(_THIS_DIR, "..", "build", "Debug"),
+        os.path.join(_THIS_DIR, "..", "build", "Release"),
+        os.path.join(_THIS_DIR, "build"),
+        _THIS_DIR,
+    ]
+
+    for d in search_dirs:
+        for n in names:
+            p = os.path.join(d, n)
+            if os.path.isfile(p): return p, d
+
+    raise FileNotFoundError(f"linear_attn shared library nahi mili. Search kiya: {search_dirs} me {names}")
+
+
+_LIB_PATH, _LIB_DIR = _find_lib()
+
+# Windows pe MinGW-built DLL ki runtime dependencies (libgomp, libwinpthread)
+# resolve karne ke liye us folder ko DLL search path me add karo
+if platform.system() == "Windows":
+    if hasattr(os, "add_dll_directory"): os.add_dll_directory(_LIB_DIR)
+    os.environ["PATH"] = _LIB_DIR + os.pathsep + os.environ.get("PATH", "")
+
+_lib = ctypes.CDLL(_LIB_PATH)
+
+c_long = ctypes.c_long
+c_float_p = ctypes.POINTER(ctypes.c_float)
+
+_lib.linear_attention_forward.argtypes = [
+    c_float_p, c_float_p, c_float_p,
+    c_float_p, c_float_p, c_float_p,
+    c_long, c_long, c_long, c_long, c_long, c_long, c_long,
+]
+_lib.linear_attention_forward.restype = None
+
+_lib.linear_attention_backward.argtypes = [
+    c_float_p, c_float_p, c_float_p, c_float_p,
+    c_float_p, c_float_p,
+    c_float_p, c_float_p, c_float_p,
+    c_long, c_long, c_long, c_long, c_long, c_long, c_long,
+]
+_lib.linear_attention_backward.restype = None
+
+
+def _ptr(t: torch.Tensor):
+    return ctypes.cast(ctypes.c_void_p(t.data_ptr()), c_float_p)
+
+
+class _LinearAttnCFn(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, k, v, chunk_size):
+        assert not q.is_cuda, "yeh C backend sirf CPU tensors ke liye hai"
+        q = q.contiguous().float()
+        k = k.contiguous().float()
+        v = v.contiguous().float()
+
+        B, H, N, D = q.shape
+        E = v.shape[-1]
+        num_chunks = (N + chunk_size - 1) // chunk_size
+
+        out       = torch.empty((B, H, N, E), dtype=torch.float32)
+        kv_states = torch.zeros((B, H, num_chunks, D, E), dtype=torch.float32)
+        z_states  = torch.zeros((B, H, num_chunks, D), dtype=torch.float32)
+
+        _lib.linear_attention_forward(
+            _ptr(q), _ptr(k), _ptr(v),
+            _ptr(out), _ptr(kv_states), _ptr(z_states),
+            c_long(B), c_long(H), c_long(N), c_long(D), c_long(E),
+            c_long(chunk_size), c_long(num_chunks),
+        )
+
+        ctx.save_for_backward(q, k, v, kv_states, z_states)
+        ctx.chunk_size = chunk_size
+        return out
+
+    @staticmethod
+    def backward(ctx, dout):
+        q, k, v, kv_states, z_states = ctx.saved_tensors
+        dout = dout.contiguous().float()
+
+        B, H, N, D = q.shape
+        E = v.shape[-1]
+        num_chunks = kv_states.shape[2]
+
+        dq = torch.zeros_like(q)
+        dk = torch.zeros_like(k)
+        dv = torch.zeros_like(v)
+
+        _lib.linear_attention_backward(
+            _ptr(q), _ptr(k), _ptr(v), _ptr(dout),
+            _ptr(kv_states), _ptr(z_states),
+            _ptr(dq), _ptr(dk), _ptr(dv),
+            c_long(B), c_long(H), c_long(N), c_long(D), c_long(E),
+            c_long(ctx.chunk_size), c_long(num_chunks),
+        )
+
+        return dq, dk, dv, None
+
+
+import torch.nn as nn
+
+try:
+    import triton
+    import triton.language as tl
+    _HAS_TRITON = True
+
+    # ============================================================
+    # Forward kernel: chunked causal linear attention
+    # ============================================================
+    @triton.jit
+    def _la_fwd_kernel(
+        Q, K, V, Out, KV_STATES, Z_STATES,
+        stride_qb, stride_qh, stride_qm, stride_qd,
+        stride_kb, stride_kh, stride_kn, stride_kd,
+        stride_vb, stride_vh, stride_vn, stride_ve,
+        stride_ob, stride_oh, stride_om, stride_oe,
+        stride_kvb, stride_kvh, stride_kvc, stride_kvd, stride_kve,
+        stride_zb, stride_zh, stride_zc, stride_zd,
+        seq_len, head_dim, v_dim, num_chunks,
+        BLOCK_C: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_E: tl.constexpr,
+    ):
+        pid_bh = tl.program_id(0)
+        Q += pid_bh * stride_qh; K += pid_bh * stride_kh
+        V += pid_bh * stride_vh; Out += pid_bh * stride_oh
+        KV_STATES += pid_bh * stride_kvh
+        Z_STATES += pid_bh * stride_zh
+
+        offs_d = tl.arange(0, BLOCK_D)
+        offs_e = tl.arange(0, BLOCK_E)
+        offs_c = tl.arange(0, BLOCK_C)
+
+        S = tl.zeros((BLOCK_D, BLOCK_E), dtype=tl.float32)
+        z_state = tl.zeros((BLOCK_D,), dtype=tl.float32)
+
+        for c in range(0, num_chunks):
+            c_start = c * BLOCK_C
+            offs_m = c_start + offs_c
+            row_mask = offs_m < seq_len
+
+            q_ptrs = Q + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+            k_ptrs = K + offs_m[:, None] * stride_kn + offs_d[None, :] * stride_kd
+            v_ptrs = V + offs_m[:, None] * stride_vn + offs_e[None, :] * stride_ve
+            qk_mask = row_mask[:, None] & (offs_d[None, :] < head_dim)
+            v_mask = row_mask[:, None] & (offs_e[None, :] < v_dim)
+
+            q = tl.load(q_ptrs, mask=qk_mask, other=0.0).to(tl.float32)
+            k = tl.load(k_ptrs, mask=qk_mask, other=0.0).to(tl.float32)
+            v = tl.load(v_ptrs, mask=v_mask, other=0.0).to(tl.float32)
+
+            q = tl.where(q > 0, q + 1.0, tl.exp(q))
+            k = tl.where(k > 0, k + 1.0, tl.exp(k))
+
+            # store the state as it was BEFORE this chunk (needed for backward)
+            kv_ptrs = KV_STATES + c * stride_kvc + offs_d[:, None] * stride_kvd + offs_e[None, :] * stride_kve
+            z_ptrs = Z_STATES + c * stride_zc + offs_d * stride_zd
+            tl.store(kv_ptrs, S.to(tl.float32))
+            tl.store(z_ptrs, z_state.to(tl.float32))
+
+            local_scores = tl.dot(q, tl.trans(k))
+            causal_mask = offs_c[:, None] >= offs_c[None, :]
+            local_scores = tl.where(causal_mask, local_scores, 0.0)
+
+            out_intra = tl.dot(local_scores.to(tl.float32), v)
+            z_intra = tl.sum(local_scores, axis=1)
+
+            out_inter = tl.dot(q, S)
+            z_inter = tl.sum(q * z_state[None, :], axis=1)
+
+            out = out_intra + out_inter
+            z = z_intra + z_inter
+            z = tl.maximum(z, 1e-6)
+            out = out / z[:, None]
+
+            out_ptrs = Out + offs_m[:, None] * stride_om + offs_e[None, :] * stride_oe
+            tl.store(out_ptrs, out.to(Out.dtype.element_ty), mask=v_mask)
+
+            S += tl.dot(tl.trans(k), v)
+            z_state += tl.sum(k, axis=0)
+
+
+    # ============================================================
+    # Backward kernel: recompute-based (no N x N matrices stored)
+    # ============================================================
+    @triton.jit
+    def _la_bwd_kernel(
+        Q, K, V, DOut, DQ, DK, DV, KV_STATES, Z_STATES,
+        stride_qb, stride_qh, stride_qm, stride_qd,
+        stride_kb, stride_kh, stride_kn, stride_kd,
+        stride_vb, stride_vh, stride_vn, stride_ve,
+        stride_ob, stride_oh, stride_om, stride_oe,
+        stride_kvb, stride_kvh, stride_kvc, stride_kvd, stride_kve,
+        stride_zb, stride_zh, stride_zc, stride_zd,
+        seq_len, head_dim, v_dim, num_chunks,
+        BLOCK_C: tl.constexpr, BLOCK_D: tl.constexpr, BLOCK_E: tl.constexpr,
+    ):
+        pid_bh = tl.program_id(0)
+        Q += pid_bh * stride_qh; K += pid_bh * stride_kh
+        V += pid_bh * stride_vh; DOut += pid_bh * stride_oh
+        DQ += pid_bh * stride_qh; DK += pid_bh * stride_kh; DV += pid_bh * stride_vh
+        KV_STATES += pid_bh * stride_kvh
+        Z_STATES += pid_bh * stride_zh
+
+        offs_d = tl.arange(0, BLOCK_D)
+        offs_e = tl.arange(0, BLOCK_E)
+        offs_c = tl.arange(0, BLOCK_C)
+
+        # running gradient-state accumulators (for chunks processed so far, going backward)
+        dS = tl.zeros((BLOCK_D, BLOCK_E), dtype=tl.float32)
+        dz_state = tl.zeros((BLOCK_D,), dtype=tl.float32)
+
+        for c in range(num_chunks - 1, -1, -1):
+            c_start = c * BLOCK_C
+            offs_m = c_start + offs_c
+            row_mask = offs_m < seq_len
+            qk_mask = row_mask[:, None] & (offs_d[None, :] < head_dim)
+            v_mask = row_mask[:, None] & (offs_e[None, :] < v_dim)
+
+            q_ptrs = Q + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+            k_ptrs = K + offs_m[:, None] * stride_kn + offs_d[None, :] * stride_kd
+            v_ptrs = V + offs_m[:, None] * stride_vn + offs_e[None, :] * stride_ve
+            do_ptrs = DOut + offs_m[:, None] * stride_om + offs_e[None, :] * stride_oe
+
+            q_raw = tl.load(q_ptrs, mask=qk_mask, other=0.0).to(tl.float32)
+            k_raw = tl.load(k_ptrs, mask=qk_mask, other=0.0).to(tl.float32)
+            v = tl.load(v_ptrs, mask=v_mask, other=0.0).to(tl.float32)
+            do = tl.load(do_ptrs, mask=v_mask, other=0.0).to(tl.float32)
+
+            q = tl.where(q_raw > 0, q_raw + 1.0, tl.exp(q_raw))
+            k = tl.where(k_raw > 0, k_raw + 1.0, tl.exp(k_raw))
+            dphi_q = tl.where(q_raw > 0, 1.0, q)   # d(elu+1)/dx = q itself when x<=0 (=exp(x)), else 1
+            dphi_k = tl.where(k_raw > 0, 1.0, k)
+
+            kv_ptrs = KV_STATES + c * stride_kvc + offs_d[:, None] * stride_kvd + offs_e[None, :] * stride_kve
+            z_ptrs = Z_STATES + c * stride_zc + offs_d * stride_zd
+            S_before = tl.load(kv_ptrs)          # state BEFORE this chunk
+            z_before = tl.load(z_ptrs)
+
+            # recompute forward quantities needed
+            local_scores = tl.dot(q, tl.trans(k))
+            causal_mask = offs_c[:, None] >= offs_c[None, :]
+            local_scores = tl.where(causal_mask, local_scores, 0.0)
+
+            out_intra = tl.dot(local_scores, v)
+            z_intra = tl.sum(local_scores, axis=1)
+            out_inter = tl.dot(q, S_before)
+            z_inter = tl.sum(q * z_before[None, :], axis=1)
+            z = tl.maximum(z_intra + z_inter, 1e-6)
+
+            out = (out_intra + out_inter) / z[:, None]
+
+            # dOut/dz and dOut/dnumerator
+            d_num = do / z[:, None]                                   # (C, E)
+            d_z = -tl.sum(do * out, axis=1) / z                       # (C,)
+
+            # ---- grads through inter-chunk term ----
+            # out_inter = q @ S_before ; z_inter = sum(q * z_before)
+            dq_inter = tl.dot(d_num, tl.trans(S_before)) + d_z[:, None] * z_before[None, :]
+            dS_local = tl.dot(tl.trans(q), d_num)          # contribution to dS_before from this chunk's numerator
+            dz_state_local = tl.sum(q * d_z[:, None], axis=0)
+
+            # ---- grads through intra-chunk term ----
+            # local_scores = causal_mask * (q @ k^T)
+            d_local_scores = tl.dot(d_num, tl.trans(v)) + d_z[:, None]
+            d_local_scores = tl.where(causal_mask, d_local_scores, 0.0)
+
+            dq_intra = tl.dot(d_local_scores, k)
+            dk_intra = tl.dot(tl.trans(d_local_scores), q)
+            dv_intra = tl.dot(tl.trans(local_scores), d_num)
+
+            dq_total = dq_intra + dq_inter
+            dk_total = dk_intra
+            dv_total = dv_intra
+
+            # ---- grads from state update S += k^T v, z_state += sum(k) ----
+            # accumulated dS (from later chunks, passed backward) also flows into k, v of THIS chunk's update
+            dk_from_state = tl.dot(dS, tl.trans(v)) if False else tl.dot(v, tl.trans(dS))  # placeholder unused
+            # correct: S_next = S_before + k^T @ v  -> dK += dS_next @ v^T-like term, dV += k @ dS_next
+            dk_total += tl.dot(v, tl.trans(dS))     # (C,D) via (C,E)@(E,D)
+            dv_total += tl.dot(k, dS)               # (C,E) via (C,D)@(D,E)
+            dk_total += dz_state[None, :] * 1.0     # from z_state += sum(k): broadcast grad
+            dk_total = dk_total  # dz_state contributes uniformly across rows: handled below properly
+
+            # z_state_next = z_before_next... actually z_state accumulates sum(k) per chunk:
+            # gradient of z_state w.r.t. k is all-ones over rows, scaled by dz_state
+            dk_total += tl.broadcast_to(dz_state[None, :], (BLOCK_C, BLOCK_D))
+
+            # propagate accumulated grads to previous (earlier) chunk's state
+            dS += dS_local
+            dz_state += dz_state_local
+
+            # apply feature-map derivative
+            dq_final = dq_total * dphi_q
+            dk_final = dk_total * dphi_k
+
+            dq_ptrs = DQ + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
+            dk_ptrs = DK + offs_m[:, None] * stride_kn + offs_d[None, :] * stride_kd
+            dv_ptrs = DV + offs_m[:, None] * stride_vn + offs_e[None, :] * stride_ve
+            tl.store(dq_ptrs, dq_final.to(DQ.dtype.element_ty), mask=qk_mask)
+            tl.store(dk_ptrs, dk_final.to(DK.dtype.element_ty), mask=qk_mask)
+            tl.store(dv_ptrs, dv_total.to(DV.dtype.element_ty), mask=v_mask)
+
+
+    # ============================================================
+    # Autograd wrapper
+    # ============================================================
+    class _LinearAttnFn(torch.autograd.Function):
+        @staticmethod
+        def forward(ctx, q, k, v, chunk_size):
+            B, H, N, D = q.shape
+            E = v.shape[-1]
+            BLOCK_C = chunk_size
+            BLOCK_D = triton.next_power_of_2(D)
+            BLOCK_E = triton.next_power_of_2(E)
+            num_chunks = triton.cdiv(N, BLOCK_C)
+
+            out = torch.empty((B, H, N, E), device=q.device, dtype=q.dtype)
+            kv_states = torch.empty((B, H, num_chunks, D, E), device=q.device, dtype=torch.float32)
+            z_states = torch.empty((B, H, num_chunks, D), device=q.device, dtype=torch.float32)
+
+            grid = (B * H,)
+            _la_fwd_kernel[grid](
+                q, k, v, out, kv_states, z_states,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                out.stride(0), out.stride(1), out.stride(2), out.stride(3),
+                kv_states.stride(0), kv_states.stride(1), kv_states.stride(2), kv_states.stride(3), kv_states.stride(4),
+                z_states.stride(0), z_states.stride(1), z_states.stride(2), z_states.stride(3),
+                N, D, E, num_chunks,
+                BLOCK_C=BLOCK_C, BLOCK_D=BLOCK_D, BLOCK_E=BLOCK_E,
+            )
+            ctx.save_for_backward(q, k, v, kv_states, z_states)
+            ctx.BLOCK_C, ctx.BLOCK_D, ctx.BLOCK_E = BLOCK_C, BLOCK_D, BLOCK_E
+            ctx.num_chunks, ctx.N, ctx.D, ctx.E = num_chunks, N, D, E
+            return out
+
+        @staticmethod
+        def backward(ctx, dout):
+            q, k, v, kv_states, z_states = ctx.saved_tensors
+            B, H, N, D = q.shape
+            E = v.shape[-1]
+            dq = torch.empty_like(q)
+            dk = torch.empty_like(k)
+            dv = torch.empty_like(v)
+
+            grid = (B * H,)
+            _la_bwd_kernel[grid](
+                q, k, v, dout.contiguous(), dq, dk, dv, kv_states, z_states,
+                q.stride(0), q.stride(1), q.stride(2), q.stride(3),
+                k.stride(0), k.stride(1), k.stride(2), k.stride(3),
+                v.stride(0), v.stride(1), v.stride(2), v.stride(3),
+                dout.stride(0), dout.stride(1), dout.stride(2), dout.stride(3),
+                kv_states.stride(0), kv_states.stride(1), kv_states.stride(2), kv_states.stride(3), kv_states.stride(4),
+                z_states.stride(0), z_states.stride(1), z_states.stride(2), z_states.stride(3),
+                N, D, E, ctx.num_chunks,
+                BLOCK_C=ctx.BLOCK_C, BLOCK_D=ctx.BLOCK_D, BLOCK_E=ctx.BLOCK_E,
+            )
+            return dq, dk, dv, None
+
+except ImportError: _HAS_TRITON = False
+
+
+def linear_attention(q, k, v, chunk_size=64):
+    if q.is_cuda:
+        if not _HAS_TRITON: raise RuntimeError("Triton not installed, GPU path unavailable")
+        return _LinearAttnFn.apply(q, k, v, chunk_size)
+    else: return _LinearAttnCFn.apply(q, k, v, chunk_size)
+
+
+# ============================================================
+# nn.Module: drop-in linear attention layer
+# ============================================================
+class LinearAttention(nn.Module):
+    def __init__(self, dim, num_heads, chunk_size=64, qkv_bias=False):
+        super().__init__()
+        assert dim % num_heads == 0
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.chunk_size = chunk_size
+
+        self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
+        self.k_proj = nn.Linear(dim, dim, bias=qkv_bias)
+        self.v_proj = nn.Linear(dim, dim, bias=qkv_bias)
+        self.out_proj = nn.Linear(dim, dim, bias=qkv_bias)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        H, Dh = self.num_heads, self.head_dim
+
+        q = self.q_proj(x).view(B, N, H, Dh).transpose(1, 2).contiguous()
+        k = self.k_proj(x).view(B, N, H, Dh).transpose(1, 2).contiguous()
+        v = self.v_proj(x).view(B, N, H, Dh).transpose(1, 2).contiguous()
+
+        out = linear_attention(q, k, v, chunk_size=self.chunk_size)  # (B, H, N, Dh)
+
+        out = out.transpose(1, 2).contiguous().view(B, N, C)
+        return self.out_proj(out)

fast_attn_kernels-0.1.0/fast_attn_kernels.egg-info/PKG-INFO

@@ -0,0 +1,4 @@
+Metadata-Version: 2.4
+Name: fast-attn-kernels
+Version: 0.1.0
+Requires-Dist: torch

fast_attn_kernels-0.1.0/fast_attn_kernels.egg-info/SOURCES.txt

@@ -0,0 +1,10 @@
+pyproject.toml
+setup.py
+attns/__init__.py
+attns/linear_attn.c
+attns/linear_bind.py
+fast_attn_kernels.egg-info/PKG-INFO
+fast_attn_kernels.egg-info/SOURCES.txt
+fast_attn_kernels.egg-info/dependency_links.txt
+fast_attn_kernels.egg-info/requires.txt
+fast_attn_kernels.egg-info/top_level.txt

fast_attn_kernels-0.1.0/fast_attn_kernels.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

fast_attn_kernels-0.1.0/fast_attn_kernels.egg-info/requires.txt

@@ -0,0 +1 @@
+torch

fast_attn_kernels-0.1.0/fast_attn_kernels.egg-info/top_level.txt

@@ -0,0 +1 @@
+attns

fast_attn_kernels-0.1.0/pyproject.toml

@@ -0,0 +1,8 @@
+[build-system]
+requires = ["setuptools", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "fast-attn-kernels"
+version = "0.1.0"
+dependencies = ["torch"]

fast_attn_kernels-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

fast_attn_kernels-0.1.0/setup.py

@@ -0,0 +1,16 @@
+from setuptools import setup, find_packages, Extension
+
+setup(
+    name="fast-attn-kernels",
+    version="0.1.0",
+    packages=find_packages(),
+    ext_modules=[
+        Extension(
+            "attns.liblinear_attn",
+            sources=["attns/linear_attn.c"],
+            extra_compile_args=["-O3", "-fopenmp"],
+            extra_link_args=["-fopenmp"],
+        )
+    ],
+    install_requires=["torch"],
+)