perf: remove split partials memset and overlap decode tile loads

- alloc_split_partials now uses torch::empty: the split kernel writes every slot it owns, so the per-call zeros/full memset was pure overhead (2 kernels per layer per step)
- decode split-KV MMA kernels now run a true multi-stage cp.async pipeline (wait_group<STAGES-1> instead of wait_group<0>), keeping STAGES-1 tile loads in flight; the old wait_group<0> serialized load and compute so deeper STAGES made no difference
- add a fallback path when ntiles < STAGES to avoid a race on the last tile

astrai/extension/attention_backend.py

@@ -30,7 +30,6 @@ Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
 
 import contextvars
 import enum
-import math
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from typing import Optional, Union

csrc/kernels/attn_decode_split_kv_mma.cuh

@@ -76,26 +76,17 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
         cp_async_commit();
     };
 
-    constexpr int BUF_MASK = (Traits::STAGES > 1) ? (Traits::STAGES - 1) : 0;
+    // ---- Multi-stage cp.async pipeline ----
-
+    // Prologue loads STAGES tiles; each loop iteration waits only for the
-    // Prologue
+    // oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
-    if (ti_begin < ti_end) {
+    // tile loads stay in flight and overlap with the current tile's compute.
-        load_tile(ti_begin, 0);
+    constexpr int STAGES = Traits::STAGES;
-    }
+    const int ntiles = ti_end - ti_begin;
-
-    for (int ti = ti_begin; ti < ti_end; ti++) {
-        int buf = (ti - ti_begin) & BUF_MASK;
-
-        cp_async_wait_group<0>();
-        __syncwarp();
-        if constexpr (Traits::STAGES > 1) {
-            if (ti + 1 < ti_end)
-                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
-        }
 
+    auto process_tile = [&](int it, int buf) {
         const bf16* bK = sK + buf * Traits::BC * Traits::LD;
         const bf16* bV = sV + buf * Traits::BC * Traits::LD;
-        int kv0 = ti * Traits::BC;
+        int kv0 = (ti_begin + it) * Traits::BC;
 
         float Sacc[Traits::NC8][4];
         mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
@@ -115,12 +106,29 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams<bf16> p) {
                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
-        __syncwarp();
+    };
 
-        if constexpr (Traits::STAGES == 1) {
+    if (ntiles >= STAGES) {
-            if (ti + 1 < ti_end)
+        #pragma unroll
-                load_tile(ti + 1, 0);
+        for (int i = 0; i < STAGES; i++)
+            load_tile(ti_begin + i, i);
+
+        for (int it = 0; it < ntiles; it++) {
+            cp_async_wait_group<STAGES - 1>();
+            __syncwarp();
+            process_tile(it, it & (STAGES - 1));
+            __syncwarp();
+            if (it + STAGES < ntiles)
+                load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
         }
+    } else {
+        // Fewer tiles than stages: load all, wait for all, process.
+        for (int i = 0; i < ntiles; i++)
+            load_tile(ti_begin + i, i);
+        cp_async_wait_group<0>();
+        __syncwarp();
+        for (int it = 0; it < ntiles; it++)
+            process_tile(it, it);
     }
 
     // ---- write UN-normalised partials for this split ----

csrc/kernels/attn_entry_utils.cuh

@@ -21,11 +21,15 @@ using bf16 = __nv_bfloat16;
                          " (supported: 32, 64, 128, 256)"); \
     }
 
+// The split kernel unconditionally writes every (batch, q_head, split) slot it
+// owns — including empty split ranges, which store m = -FLT_MAX so the combine
+// skips them. Allocators are therefore left uninitialized (torch::empty); the
+// per-call memset (torch::zeros / torch::full) was pure overhead.
 template<typename P>
 inline void alloc_split_partials(P& p) {
     auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::zeros(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
+    auto o_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
-    auto ml_part = torch::full(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, -FLT_MAX, fopt);
+    auto ml_part = torch::empty(at::IntArrayRef{p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
     p.o_part = (float*)o_part.data_ptr();
     p.ml_part = (float*)ml_part.data_ptr();
 }

csrc/kernels/attn_paged_decode_split_kv_mma.cuh

@@ -91,25 +91,17 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
         cp_async_commit();
     };
 
-    constexpr int BUF_MASK = (Traits::STAGES > 1) ? (Traits::STAGES - 1) : 0;
+    // ---- Multi-stage cp.async pipeline ----
-
+    // Prologue loads STAGES tiles; each loop iteration waits only for the
-    if (ti_begin < ti_end) {
+    // oldest outstanding group (wait_group<STAGES-1>) so the STAGES-1 newer
-        load_tile(ti_begin, 0);
+    // tile loads stay in flight and overlap with the current tile's compute.
-    }
+    constexpr int STAGES = Traits::STAGES;
-
+    const int ntiles = ti_end - ti_begin;
-    for (int ti = ti_begin; ti < ti_end; ti++) {
-        int buf = (ti - ti_begin) & BUF_MASK;
-
-        cp_async_wait_group<0>();
-        __syncwarp();
-        if constexpr (Traits::STAGES > 1) {
-            if (ti + 1 < ti_end)
-                load_tile(ti + 1, (ti + 1 - ti_begin) & BUF_MASK);
-        }
 
+    auto process_tile = [&](int it, int buf) {
         const bf16* bK = sK + buf * Traits::BC * Traits::LD;
         const bf16* bV = sV + buf * Traits::BC * Traits::LD;
-        int kv0 = ti * Traits::BC;
+        int kv0 = (ti_begin + it) * Traits::BC;
 
         float Sacc[Traits::NC8][4];
         mma_compute_scores<Traits>(Qa, bK, lane, Sacc);
@@ -128,12 +120,29 @@ __global__ void paged_attn_decode_split_kv_mma_kernel(PagedAttentionParams<bf16>
                                            Sacc, Oacc, m0, m1, l0, l1, lane);
 
         mma_pv_accumulate<Traits>(Sacc, bV, lane, Oacc);
-        __syncwarp();
+    };
 
-        if constexpr (Traits::STAGES == 1) {
+    if (ntiles >= STAGES) {
-            if (ti + 1 < ti_end)
+        #pragma unroll
-                load_tile(ti + 1, 0);
+        for (int i = 0; i < STAGES; i++)
+            load_tile(ti_begin + i, i);
+
+        for (int it = 0; it < ntiles; it++) {
+            cp_async_wait_group<STAGES - 1>();
+            __syncwarp();
+            process_tile(it, it & (STAGES - 1));
+            __syncwarp();
+            if (it + STAGES < ntiles)
+                load_tile(ti_begin + it + STAGES, (it + STAGES) & (STAGES - 1));
         }
+    } else {
+        // Fewer tiles than stages: load all, wait for all, process.
+        for (int i = 0; i < ntiles; i++)
+            load_tile(ti_begin + i, i);
+        cp_async_wait_group<0>();
+        __syncwarp();
+        for (int it = 0; it < ntiles; it++)
+            process_tile(it, it);
     }
 
     auto split_slot = [&](int h) -> size_t {