refactor: drop gemv variant shape tables and flatten kernel dir

- delete the warp-tiled kernel and both per-shape (N,K) selector tables; block size is 256 threads everywhere except M=8 with N*K <= 12 MiB, which keeps a 128-thread CTA
- HBM-streaming measurements (weight copies rotated through L2, the real decode regime) show the variants within ~3% on L20 because the kernel is bandwidth-bound; the retired tables were tuned against an L2-resident loop and sometimes picked the slowest variant ((2048,8192) M=8: coop128 6% slower than coop256)
- a shape no longer switches kernels (and accumulation order) with M, removing one shape-dependent nondeterminism source
- remove the stale split-K launcher comment
- move bf16_gemv.cu and bf16_swiglu.cu from csrc/kernels/gemv/ to csrc/kernels/ beside rotary_emb.cu; the family keeps no shared headers
- rename test_bf16_gemv_matches_half_cta_edge_bands to test_bf16_gemv_matches_m8_edge_bands and update docs/developer/cuda_kernels.md

Benchmark: L20 (sm_89), PyTorch 2.11.0+cu128, interleaved CUDA-event timing with rotated weight copies exceeding the 96MB L2; variant spread <=3% across 14 shapes x M in {1,2,4,8}, and the retained rule wins 5-9% at M=8 small weights ((512,3584), (1536,1536), (6912,1536))

csrc/CMakeLists.txt

@@ -70,8 +70,8 @@ set(KERNEL_SRCS
     attention/prefill.cu
     attention/paged_decode.cu
     attention/paged_prefill.cu
-    gemv/bf16_gemv.cu
+    bf16_gemv.cu
-    gemv/bf16_swiglu.cu
+    bf16_swiglu.cu
     rotary_emb.cu
 )
 

csrc/kernels/bf16_gemv.cu

@@ -14,7 +14,6 @@ namespace {
 constexpr int kThreads = 256;
 constexpr int kHalfCtaThreads = 128;
 constexpr int kWarpSize = 32;
-constexpr int kWarpTiledThreads = 128;
 
 __device__ __forceinline__ float warp_sum(float value) {
 #pragma unroll
@@ -159,120 +158,6 @@ __global__ void bf16_gemv_kernel(
     }
 }
 
-template <int Rows>
-__global__ void bf16_gemv_aligned_warp_tiled_kernel(
-    const __nv_bfloat16* __restrict__ x,
-    const __nv_bfloat16* __restrict__ weight,
-    const __nv_bfloat16* __restrict__ bias,
-    __nv_bfloat16* __restrict__ output,
-    int n,
-    int k
-) {
-    constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize;
-    const int lane = threadIdx.x & (kWarpSize - 1);
-    const int warp = threadIdx.x / kWarpSize;
-    const int output_index = blockIdx.x * kWarpsPerBlock + warp;
-    if (output_index >= n) {
-        return;
-    }
-
-    // The launcher selects this path only when each row is 16-byte aligned.
-    // Four independent output rows per CTA remove the block-wide reduction
-    // barrier and improve occupancy for the medium LLaMA projection bands.
-    const int vectors = k / 8;
-    const auto* x4 = reinterpret_cast<const uint4*>(x);
-    const auto* w4 = reinterpret_cast<const uint4*>(weight) +
-        static_cast<int64_t>(output_index) * vectors;
-    float sums[Rows] = {};
-    for (int vector = lane; vector < vectors; vector += kWarpSize) {
-        const uint4 wv_raw = w4[vector];
-        const auto* wv = reinterpret_cast<const __nv_bfloat162*>(&wv_raw);
-#pragma unroll
-        for (int row = 0; row < Rows; ++row) {
-            const uint4 xv_raw =
-                x4[static_cast<int64_t>(row) * vectors + vector];
-            const auto* xv = reinterpret_cast<const __nv_bfloat162*>(&xv_raw);
-#pragma unroll
-            for (int pair = 0; pair < 4; ++pair) {
-                sums[row] = fmaf(
-                    __bfloat162float(__low2bfloat16(xv[pair])),
-                    __bfloat162float(__low2bfloat16(wv[pair])),
-                    sums[row]
-                );
-                sums[row] = fmaf(
-                    __bfloat162float(__high2bfloat16(xv[pair])),
-                    __bfloat162float(__high2bfloat16(wv[pair])),
-                    sums[row]
-                );
-            }
-        }
-    }
-
-#pragma unroll
-    for (int row = 0; row < Rows; ++row) {
-        sums[row] = warp_sum(sums[row]);
-        if (lane == 0) {
-            if (bias != nullptr) {
-                sums[row] += __bfloat162float(bias[output_index]);
-            }
-            output[row * n + output_index] = __float2bfloat16_rn(sums[row]);
-        }
-    }
-}
-
-template <int Rows>
-constexpr bool use_warp_tiled_kernel(int n, int k) {
-    // These bands are intentionally narrow and are validated by the common
-    // transformer benchmark. The 256-thread cooperative kernel remains the
-    // fallback for arbitrary K, larger projections, and M=2 (where the
-    // single-warp reduction regresses the current vectorized kernel).
-    if constexpr (Rows == 4) {
-        return (n == 1024 && k == 4096) ||
-            (n == 4096 && k == 4096) ||
-            (n == 11008 && k == 4096) ||
-            (n == 4096 && k == 11008);
-    }
-    return false;
-}
-
-template <int Rows>
-constexpr bool use_half_cta_kernel(int n, int k) {
-    // A 128-thread CTA reduces synchronization and scheduling overhead for
-    // selected medium decode projections. Keep the selector exact: long-K
-    // and bandwidth-saturated shapes regress, and the winning bands differ
-    // materially with the number of reused input rows.
-    if constexpr (Rows == 1) {
-        return n == 8192 && k == 2048;
-    }
-    if constexpr (Rows == 2) {
-        return (n == 4096 && k == 4096) ||
-            (n == 11008 && k == 4096) ||
-            (n == 3584 && k == 3584) ||
-            (n == 2048 && k == 2048) ||
-            (n == 8192 && k == 2048);
-    }
-    if constexpr (Rows == 4) {
-        return (n == 5120 && k == 5120) ||
-            (n == 3584 && k == 3584) ||
-            (n == 2048 && k == 2048) ||
-            (n == 8192 && k == 2048);
-    }
-    if constexpr (Rows == 8) {
-        return (n == 4096 && k == 4096) ||
-            (n == 11008 && k == 4096) ||
-            (n == 4096 && k == 11008) ||
-            (n == 1024 && k == 4096) ||
-            (n == 5120 && k == 5120) ||
-            (n == 512 && k == 3584) ||
-            (n == 3584 && k == 3584) ||
-            (n == 1024 && k == 8192) ||
-            (n == 2048 && k == 2048) ||
-            (n == 8192 && k == 2048) ||
-            (n == 2048 && k == 8192);
-    }
-    return false;
-}
-
 template <int Rows>
 void launch_bf16_gemv(
     const __nv_bfloat16* x,
@@ -283,27 +168,23 @@ void launch_bf16_gemv(
     int k,
     cudaStream_t stream
 ) {
-    const bool aligned_rows = k % 8 == 0 &&
+    // Decode is HBM weight-streaming bound: with weights rotated through L2,
+    // 128/256-thread CTAs measure within noise on L20 except for small
+    // weight matrices at the largest decode batch, where the smaller CTA
+    // wins 5-9% (see docs/developer/decode_linear_benchmark.md).
+    constexpr int64_t kSmallWeightLimit = int64_t{12} << 20;
+    if constexpr (Rows == 8) {
+        if (k % 8 == 0 &&
             (reinterpret_cast<uintptr_t>(x) & 15u) == 0u &&
-        (reinterpret_cast<uintptr_t>(weight) & 15u) == 0u;
+            (reinterpret_cast<uintptr_t>(weight) & 15u) == 0u &&
-    if constexpr (Rows == 4) {
+            static_cast<int64_t>(n) * k <= kSmallWeightLimit) {
-        if (aligned_rows && use_warp_tiled_kernel<Rows>(n, k)) {
-            constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize;
-            const int blocks = (n + kWarpsPerBlock - 1) / kWarpsPerBlock;
-            bf16_gemv_aligned_warp_tiled_kernel<Rows>
-                <<<blocks, kWarpTiledThreads, 0, stream>>>(
-                    x, weight, bias, output, n, k
-                );
-            return;
-        }
-    }
-    if (aligned_rows && use_half_cta_kernel<Rows>(n, k)) {
             bf16_gemv_kernel<Rows, kHalfCtaThreads>
                 <<<n, kHalfCtaThreads, 0, stream>>>(
                     x, weight, bias, output, n, k
                 );
             return;
         }
+    }
     bf16_gemv_kernel<Rows, kThreads><<<n, kThreads, 0, stream>>>(
         x, weight, bias, output, n, k
     );
@@ -366,8 +247,6 @@ torch::Tensor bf16_gemv(
     auto output = x.dim() == 1 ? torch::empty({n}, x.options())
                                : torch::empty({m, n}, x.options());
 
-    // Small-N decode shapes (GQA k/v projections) cannot fill the GPU with
-    // one block per output row; split K across extra blocks and reduce.
     const auto* x_ptr = reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
     const auto* weight_ptr =
         reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr());

docs/developer/cuda_kernels.md

@@ -14,27 +14,28 @@ selected by guarded model dispatchers described below.
 | `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention |
 | `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) |
 | `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
-| `bf16_gemv` | `gemv/bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (sm_80+) |
+| `bf16_gemv` | `bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (sm_80+) |
-| `bf16_swiglu` | `gemv/bf16_swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) |
+| `bf16_swiglu` | `bf16_swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) |
 | `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) |
 
 ### BF16 GEMV primitive
 
 `astrai.extension.bf16_gemv(x, weight, bias=None)` accepts a contiguous BF16
 input shaped `[K]` or `[M, K]`, with `M` in `[1, 8]` and any positive `K`, and
-row-major weights `[N, K]`. The general path assigns one 256-thread CTA to an
+row-major weights `[N, K]`. One CTA computes an output row for all M tokens
-output row and computes all M results together, reusing the weight row across
+together, reusing the weight row across tokens. CTA size is 256 threads,
-tokens. For measured aligned M=4 medium projections, a 128-thread CTA instead
+except for small weight matrices (`N*K <= 12 MiB`) at `M=8`, where a
-assigns one output to each of four warps. That removes the CTA-wide reduction
+128-thread CTA measured 5-9% faster on L20. Variant selection is otherwise
-barrier and exposes four neighboring outputs without changing accumulation.
+intentionally shape-free: under HBM-streaming conditions (weights rotated
+through L2, as in real decode) the kernel is bandwidth-bound and block-size
+choice measures within noise, so earlier per-shape variant tables were
+removed along with the warp-tiled kernel.
 
 The weight stream uses 128-bit vectorized loads anchored at each row's first
 16-byte-aligned address with scalar head/tail sweeps for unaligned remainders,
-so arbitrary `K` and storage offsets stay correct. The warp-tiled path is used
+so arbitrary `K` and storage offsets stay correct. Accumulation is FP32; the
-only when both tensors and every row are 16-byte aligned; all other calls keep
+optional BF16 bias is fused before the BF16 store. The launcher uses the
-the general arbitrary-K path. Accumulation is FP32; the optional BF16 bias is
+current CUDA stream, is CUDA Graph capture-safe, and requires sm_80 or newer.
-fused before the BF16 store. The launcher uses the current CUDA stream, is
-CUDA Graph capture-safe, and requires sm_80 or newer.
 
 Model `Linear` calls route through the lightweight linear backend. Set
 `ASTRAI_GEMV=0` for an unconditional `F.linear` fallback, `1` to force the

tests/extension/test_gemv.py

@@ -93,7 +93,7 @@ def test_bf16_gemv_matches_common_transformer_shapes(m, n, k):
         (8, 2048, 8192),
     ],
 )
-def test_bf16_gemv_matches_half_cta_edge_bands(m, n, k):
+def test_bf16_gemv_matches_m8_edge_bands(m, n, k):
     torch.manual_seed(2026 + m + n + k)
     x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
     weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)