refactor: stride-based attn interface with layout and causal mask

- Replace is_causal + causal_offset with unified causal_offset (-1 = off, >=0 = first Q pos)
- Causal and mask can now coexist (was mutually exclusive)
- Add stride-based addressing for Q/KV/O (layout-agnostic, zero-copy)
- Add layout param ("bhld"/"blhd") parsed in Python, passed as int to C++
- Support 2D [batch, kv_len] and 3D [batch, q_len, kv_len] mask
- Vectorize paged KV gather in Python fallback (was per-token Python loop)
- Extract shared helpers: compute_num_splits, alloc_split_partials, dispatch_head_dim
- Unify paged_decode entry via attn_pack_paged_params
- Update mma_softmax_tile for 3D mask with per-row qrow indexing

csrc/kernels/attn_decode.cu

@@ -5,24 +5,12 @@
 #include "attn_decode_split_kv_mma.cuh"
 #endif
 
-static int decode_num_splits(int base_blocks, int tiles_total) {
-    int sm_count = 0;
-    cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, 0);
-    int n = (2 * sm_count + base_blocks - 1) / base_blocks;
-    return std::max(1, std::min(n, std::min(tiles_total, 32)));
-}
-
 // Scalar fallback: one warp per query head, split-KV across grid.z.
 static void launch_scalar_decode(AttentionParams<bf16>& p) {
     int group_size = p.q_head / p.kv_head;
     int chunks_total = (p.kv_len + DC_CHUNK - 1) / DC_CHUNK;
-    p.num_splits = decode_num_splits(p.batch * p.kv_head, chunks_total);
-
-    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
-    p.o_part = o_part.data_ptr<float>();
-    p.ml_part = ml_part.data_ptr<float>();
+    p.num_splits = compute_num_splits(p.batch * p.kv_head, chunks_total);
+    alloc_split_partials(p);
 
     size_t smem = DC_CHUNK * p.head_dim * sizeof(bf16);
     attn_decode_split_kv_kernel<<<dim3(p.batch * p.kv_head, 1, p.num_splits), dim3(32, group_size), smem>>>(p);
@@ -38,13 +26,8 @@ static void launch_scalar_decode(AttentionParams<bf16>& p) {
 template <int HEAD_DIM, int BC, int STAGES = (HEAD_DIM <= 128) ? 2 : 1>
 static void launch_mma_decode(AttentionParams<bf16>& p) {
     int tiles_total = (p.kv_len + BC - 1) / BC;
-    p.num_splits = decode_num_splits(p.batch * p.kv_head, tiles_total);
-
-    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA);
-    auto o_part = torch::empty({p.batch, p.q_head, p.num_splits, p.head_dim}, fopt);
-    auto ml_part = torch::empty({p.batch, p.q_head, p.num_splits, 2}, fopt);
-    p.o_part = o_part.data_ptr<float>();
-    p.ml_part = ml_part.data_ptr<float>();
+    p.num_splits = compute_num_splits(p.batch * p.kv_head, tiles_total);
+    alloc_split_partials(p);
 
     attn_decode_split_kv_mma_kernel<HEAD_DIM, BC, STAGES><<<dim3(p.kv_head, p.batch, p.num_splits), 32>>>(p);
     attn_decode_combine_kernel<<<p.batch * p.q_head, p.head_dim>>>(p);
@@ -68,35 +51,21 @@ torch::Tensor attn_decode(
     torch::Tensor k,
     torch::Tensor v,
     c10::optional<torch::Tensor> mask,
-    bool is_causal = false,
-    int64_t causal_offset = 0,
-    c10::optional<double> scale = c10::nullopt
+    int64_t causal_offset,
+    double scale,
+    int64_t layout
 ) {
     AttentionParams<bf16> p;
-    attn_pack_params(q, k, v, mask, is_causal, causal_offset, scale, p);
+    attn_pack_params(q, k, v, mask, causal_offset, scale, layout, p);
     TORCH_CHECK(p.q_len == 1, "Q seq_len must be 1");
     TORCH_CHECK(p.head_dim % 32 == 0, "head_dim must be multiple of 32");
 
-    auto O = torch::empty_like(q);
-    p.o = (bf16*)O.data_ptr();
+    // O matches Q's original layout
+    auto O = torch::empty_strided(q.sizes(), q.strides(), q.options());
+    auto O_view = (layout == 1) ? O.transpose(1, 2) : O;
+    p.o = (bf16*)O_view.data_ptr();
 
-    switch (p.head_dim) {
-        case 32:
-            dispatch_decode<32>(p);
-            break;
-        case 64:
-            dispatch_decode<64>(p);
-            break;
-        case 128:
-            dispatch_decode<128>(p);
-            break;
-        case 256:
-            dispatch_decode<256>(p);
-            break;
-        default:
-            TORCH_CHECK(false, "decode: unsupported head_dim ", p.head_dim,
-                        " (supported: 32, 64, 128, 256)");
-    }
+    dispatch_head_dim(p.head_dim, [&]<int D>() { dispatch_decode<D>(p); });
     return O;
 }
 
@@ -106,8 +75,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         py::arg("k"),
         py::arg("v"),
         py::arg("mask") = py::none(),
-        py::arg("is_causal") = false,
-        py::arg("causal_offset") = 0,
-        py::arg("scale") = py::none(),
+        py::arg("causal_offset") = -1,
+        py::arg("scale") = 0.0,
+        py::arg("layout") = 0,
         "GQA decode (tensor-core head-packing on sm_80+, scalar fallback)");
 }