perf: preload V in decode split-kv shared mem and cache partial tensors

- Preload V into shared memory alongside K to eliminate per-element KV address lookups in the inner softmax/accum loop (doubles smem)
- Cache split-KV partial tensors (o_part, ml_part) with static tensors instead of per-call allocation in both decode and paged-decode paths
- Force is_causal=True in CUDA decode backend (decode is always causal)

astrai/extension/attention_backend.py

@@ -440,8 +440,7 @@ class CudaBackend(AttentionBackend):
             kv_cache.req_pool_indices,
             kv_indptr,
             kv_cache.max_len,
-            mask=attn_mask,
-            is_causal=is_causal,
+            is_causal=True,
         )
         return out.unsqueeze(1).flatten(2)
 

csrc/kernels/attn_decode.cu

@@ -22,7 +22,17 @@ torch::Tensor attn_decode(
     auto O_view = (layout == BLHD) ? O.transpose(1, 2) : O;
     p.o = (bf16*)O_view.data_ptr();
 
-    alloc_split_partials(p);
+    {
+        static torch::Tensor s_o_part, s_ml_part;
+        int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
+        auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
+        if (!s_o_part.defined() || s_o_part.numel() < o_needed) {
+            s_o_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
+            s_ml_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
+        }
+        p.o_part = (float*)s_o_part.data_ptr();
+        p.ml_part = (float*)s_ml_part.data_ptr();
+    }
     DISPATCH_HEAD_DIM(p.head_dim, dispatch_decode, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;

csrc/kernels/attn_decode_split_kv.cuh

@@ -35,7 +35,9 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
 
     float m = -FLT_MAX, d = 0.0f, acc_reg[8] = {0.0f};
 
-    extern __shared__ __align__(16) bf16 k_smem[];
+    extern __shared__ __align__(16) bf16 smem[];
+    bf16* k_smem = smem;
+    bf16* v_smem = smem + DC_CHUNK * p.head_dim;
 
     // Split-KV: each split processes a contiguous subset of chunks
     int chunks_total = (seq_len + DC_CHUNK - 1) / DC_CHUNK;
@@ -47,8 +49,8 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
         int chunk_start = ci * DC_CHUNK;
         int this_chunk = min(DC_CHUNK, seq_len - chunk_start);
 
-        // Load K into shared memory (addressing via KV policy; paged guards
-        // empty slots with zero-fill).
+        // Load K and V into shared memory (addressing via KV policy;
+        // paged guards empty slots with zero-fill).
         int total = this_chunk * p.head_dim;
         for (int i = threadIdx.y * 32 + lane; i < total;
              i += blockDim.x * blockDim.y) {
@@ -57,6 +59,7 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
             int kc = chunk_start + s;
             KVAddr a = KV::kv_addr(p, kctx, kc, d_dim, true);
             k_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.k) : (bf16)0.f;
+            v_smem[i] = a.valid ? *reinterpret_cast<const bf16*>(a.v) : (bf16)0.f;
         }
         __syncthreads();
 
@@ -82,13 +85,8 @@ __global__ void attn_decode_split_kv_kernel(AttentionParams<bf16> p) {
             float beta  = __expf(partial - new_m);
             d = d * alpha + beta;
 
-            // V read via KV policy; when masked (beta == 0) or the slot is
-            // empty the term vanishes, so no extra branches are needed.
             for (int i = 0; i < hd_per_thread; i++) {
-                KVAddr a = KV::kv_addr(p, kctx, kv_idx, lane * hd_per_thread + i, true);
-                float vv = a.valid
-                    ? __bfloat162float(*reinterpret_cast<const bf16*>(a.v))
-                    : 0.0f;
+                float vv = __bfloat162float(v_smem[s * p.head_dim + lane * hd_per_thread + i]);
                 acc_reg[i] = fmaf(acc_reg[i], alpha, vv * beta);
             }
             m = new_m;

csrc/kernels/attn_dispatchers.cuh

@@ -157,7 +157,7 @@ struct DecodeLauncherScalar {
         int kv_len = KV::host_kv_len(p);
         int chunks_total = (kv_len + DC_CHUNK - 1) / DC_CHUNK;
         p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
-        size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
+        size_t smem = 2 * DC_CHUNK * p.head_dim * sizeof(bf16);
         int g = min(group_size, 32);  // cap at 32 to respect 1024-thread limit
         dim3 grid(p.batch * p.kv_head, 1, p.num_splits);
         dim3 block(32, g);

csrc/kernels/attn_paged_decode.cu

@@ -24,7 +24,17 @@ torch::Tensor attn_paged_decode(
     auto O = torch::empty({q.size(0), q.size(1), q.size(2)}, q.options());
     p.o = (bf16*)O.data_ptr();
 
-    alloc_split_partials(p);
+    {
+        static torch::Tensor s_o_part, s_ml_part;
+        int64_t o_needed = (int64_t)p.batch * p.q_head * MAX_SPLITS * p.head_dim;
+        auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
+        if (!s_o_part.defined() || s_o_part.numel() < o_needed) {
+            s_o_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, p.head_dim}, fopt);
+            s_ml_part = torch::empty({p.batch, p.q_head, MAX_SPLITS, 2}, fopt);
+        }
+        p.o_part = (float*)s_o_part.data_ptr();
+        p.ml_part = (float*)s_ml_part.data_ptr();
+    }
     DISPATCH_HEAD_DIM(p.head_dim, dispatch_paged_decode, p, stream);
     C10_CUDA_CHECK(cudaGetLastError());
     return O;