perf: drop swiglu warp-rows variant for an M=8 block-size rule

- delete the warp-per-row kernel and the (6912,1536) M=2/4/8 dispatch table; under rotated cold weights the warp path is 2-6% slower than CTA reuse at M=2/4, and the table had been tuned against L2-resident timing
- a single CTA-reuse kernel now serves all M in [1, 8]; block size is 256 threads for M in [1, 7] and 128 for M=8, where the shorter shared-memory reduction tree wins
- document in docs/developer/swiglu_benchmark.md that the earlier operator numbers were L2-resident: the fused kernel sits at the dual-stream cold-read floor (702 vs 699 GB/s at (6912,1536); 369 vs 370 GB/s at (11008,4096)) and wide matrices cap at ~370-400 GB/s even for pure reads, so the reported M=8 -23% regression does not survive the cold regime
- update docs/developer/cuda_kernels.md accordingly

Benchmark: L20 (sm_89), PyTorch 2.11.0+cu128, rotated weight copies >= 240 MB to defeat the 96 MB L2; end-to-end through the built module at (6912,1536) reaches 738-752 GB/s for M in [1, 4] and 702 GB/s at M=8, about +8% at M=2/4 and +6% at M=8 over the removed warp path

csrc/kernels/bf16_swiglu.cu

@@ -1,4 +1,8 @@
 // Fused small-M BF16 SwiGLU primitive for decode-time dense MLP layers.
+// One CTA per output column; each weight pair is read once and reused across
+// all decode rows. Bandwidth-bound in the cold-HBM decode regime, so variant
+// selection beyond the M=8 block-size rule is noise (see
+// docs/developer/swiglu_benchmark.md).
 
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
@@ -11,9 +15,7 @@
 
 namespace {
 
-constexpr int kThreads = 256;
 constexpr int kWarpSize = 32;
-constexpr int kWarps = kThreads / kWarpSize;
 
 __device__ __forceinline__ float warp_sum(float value) {
 #pragma unroll
@@ -27,7 +29,7 @@ __device__ __forceinline__ float round_bf16(float value) {
     return __bfloat162float(__float2bfloat16_rn(value));
 }
 
-template <int Rows>
+template <int Threads, int Rows>
 __global__ void bf16_swiglu_kernel(
     const __nv_bfloat16* __restrict__ x,
     const __nv_bfloat16* __restrict__ up_weight,
@@ -36,6 +38,7 @@ __global__ void bf16_swiglu_kernel(
     int n,
     int k
 ) {
+    constexpr int kWarps = Threads / kWarpSize;
     const int output_index = blockIdx.x;
     const int lane = threadIdx.x & (kWarpSize - 1);
     const int warp = threadIdx.x / kWarpSize;
@@ -120,7 +123,7 @@ __global__ void bf16_swiglu_kernel(
     }
 }
 
-template <int Rows>
+template <int Threads, int Rows>
 void launch_bf16_swiglu(
     const __nv_bfloat16* x,
     const __nv_bfloat16* up_weight,
@@ -130,85 +133,7 @@ void launch_bf16_swiglu(
     int k,
     cudaStream_t stream
 ) {
-    bf16_swiglu_kernel<Rows><<<n, kThreads, 0, stream>>>(
+    bf16_swiglu_kernel<Threads, Rows><<<n, Threads, 0, stream>>>(
-        x, up_weight, gate_weight, output, n, k
-    );
-}
-
-template <int Rows>
-__global__ void bf16_swiglu_warp_rows_kernel(
-    const __nv_bfloat16* __restrict__ x,
-    const __nv_bfloat16* __restrict__ up_weight,
-    const __nv_bfloat16* __restrict__ gate_weight,
-    __nv_bfloat16* __restrict__ output,
-    int n,
-    int k
-) {
-    const int output_index = blockIdx.x;
-    const int row = threadIdx.x / kWarpSize;
-    const int lane = threadIdx.x & (kWarpSize - 1);
-    const int vector_count = k / 8;
-
-    float up_sum = 0.0f;
-    float gate_sum = 0.0f;
-    const auto* x4 = reinterpret_cast<const uint4*>(
-        x + static_cast<int64_t>(row) * k
-    );
-    const auto* up4 = reinterpret_cast<const uint4*>(
-        up_weight + static_cast<int64_t>(output_index) * k
-    );
-    const auto* gate4 = reinterpret_cast<const uint4*>(
-        gate_weight + static_cast<int64_t>(output_index) * k
-    );
-
-    // A warp owns one decode row. Same-address weight reads from sibling
-    // warps are served through the read-only/L1 path, while each row avoids
-    // CTA-wide shared-memory reductions and synchronization.
-    for (int vector_index = lane;
-         vector_index < vector_count;
-         vector_index += kWarpSize) {
-        const uint4 x_raw = x4[vector_index];
-        const uint4 up_raw = up4[vector_index];
-        const uint4 gate_raw = gate4[vector_index];
-        const auto* x_values = reinterpret_cast<const __nv_bfloat162*>(&x_raw);
-        const auto* up_values =
-            reinterpret_cast<const __nv_bfloat162*>(&up_raw);
-        const auto* gate_values =
-            reinterpret_cast<const __nv_bfloat162*>(&gate_raw);
-#pragma unroll
-        for (int pair = 0; pair < 4; ++pair) {
-            const float2 xv = __bfloat1622float2(x_values[pair]);
-            const float2 uv = __bfloat1622float2(up_values[pair]);
-            const float2 gv = __bfloat1622float2(gate_values[pair]);
-            up_sum = fmaf(xv.x, uv.x, up_sum);
-            up_sum = fmaf(xv.y, uv.y, up_sum);
-            gate_sum = fmaf(xv.x, gv.x, gate_sum);
-            gate_sum = fmaf(xv.y, gv.y, gate_sum);
-        }
-    }
-    up_sum = warp_sum(up_sum);
-    gate_sum = warp_sum(gate_sum);
-    if (lane == 0) {
-        up_sum = round_bf16(up_sum);
-        gate_sum = round_bf16(gate_sum);
-        const float silu =
-            round_bf16(gate_sum / (1.0f + expf(-gate_sum)));
-        output[static_cast<int64_t>(row) * n + output_index] =
-            __float2bfloat16_rn(up_sum * silu);
-    }
-}
-
-template <int Rows>
-void launch_bf16_swiglu_warp_rows(
-    const __nv_bfloat16* x,
-    const __nv_bfloat16* up_weight,
-    const __nv_bfloat16* gate_weight,
-    __nv_bfloat16* output,
-    int n,
-    int k,
-    cudaStream_t stream
-) {
-    bf16_swiglu_warp_rows_kernel<Rows><<<n, Rows * kWarpSize, 0, stream>>>(
         x, up_weight, gate_weight, output, n, k
     );
 }
@@ -288,66 +213,51 @@ torch::Tensor bf16_swiglu(
         reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
     const int n_int = static_cast<int>(n);
     const int k_int = static_cast<int>(k);
-    const bool use_warp_rows = n_int == 6912 && k_int == 1536;
 
+    // Block size 256 keeps the weight streams at the HBM bandwidth floor for
+    // M in [1, 7]; M=8 halves the CTA so each thread owns more of the row
+    // and the shared-memory reduction tree shrinks (measured on L20 with
+    // rotated cold weights; larger CTAs only add idle warps).
     switch (m) {
         case 1:
-            launch_bf16_swiglu<1>(
+            launch_bf16_swiglu<256, 1>(
                 x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
             );
             break;
         case 2:
-            if (use_warp_rows) {
+            launch_bf16_swiglu<256, 2>(
-                launch_bf16_swiglu_warp_rows<2>(
                 x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
             );
-            } else {
-                launch_bf16_swiglu<2>(
-                    x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-                );
-            }
             break;
         case 3:
-            launch_bf16_swiglu<3>(
+            launch_bf16_swiglu<256, 3>(
                 x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
             );
             break;
         case 4:
-            if (use_warp_rows) {
+            launch_bf16_swiglu<256, 4>(
-                launch_bf16_swiglu_warp_rows<4>(
                 x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
             );
-            } else {
-                launch_bf16_swiglu<4>(
-                    x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-                );
-            }
             break;
         case 5:
-            launch_bf16_swiglu<5>(
+            launch_bf16_swiglu<256, 5>(
                 x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
             );
             break;
         case 6:
-            launch_bf16_swiglu<6>(
+            launch_bf16_swiglu<256, 6>(
                 x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
             );
             break;
         case 7:
-            launch_bf16_swiglu<7>(
+            launch_bf16_swiglu<256, 7>(
                 x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
             );
             break;
         case 8:
-            if (use_warp_rows) {
+            launch_bf16_swiglu<128, 8>(
-                launch_bf16_swiglu_warp_rows<8>(
                 x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
             );
-            } else {
-                launch_bf16_swiglu<8>(
-                    x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
-                );
-            }
             break;
     }
     C10_CUDA_CHECK(cudaGetLastError());

docs/developer/cuda_kernels.md

@@ -125,12 +125,14 @@ BF16 inputs with `M` in `[1, 8]` and K divisible by 8. It preserves the BF16
 rounding boundaries of the two projection outputs, SiLU output, and final
 product while accumulating dot products in FP32.
 
-The kernel contains two output-row tilings. A CTA-reuse path reads each up/gate
+The kernel is a single CTA-reuse tiling: one CTA per output column reads each
-weight chunk once and applies it to all M rows. The native AstrAI 1B shape
+up/gate weight chunk once and applies it to all M rows. Block size is 256
-`(N,K)=(6912,1536)` uses one warp per decode row for M=2/4/8; on L20 this
+threads for M in [1, 7] and 128 for M=8, where the shorter shared-memory
-removes the shared reductions and barrier and reduces M=4 CUDA-Graph latency
+reduction tree wins under cold-HBM decode traffic. An earlier per-shape
-from 0.0324 ms to 0.0181 ms. Wider LLaMA/GPT-NeoX matrices keep CTA reuse,
+`(6912,1536)` warp-per-row variant and its dispatch table were removed: HBM
-because duplicating their weight reads across row warps regressed 1.3-4.2%.
+measurements with rotated weights showed the table was tuned against L2-cache
+regime timing and was up to 6% slower than CTA reuse at M=2/4; the kernel is
+bandwidth-bound, so finer variant selection is noise.
 
 Dense `MLP` modules route through the SwiGLU backend. `ASTRAI_SWIGLU=0` keeps
 the unfused linear backend, and `1` explicitly forces the fused primitive.

docs/developer/swiglu_benchmark.md

@@ -56,3 +56,28 @@ the existing unfused linear backend, including any independently qualified
 GEMV dispatch. `ASTRAI_SWIGLU=1` remains an explicit benchmark/experimentation
 switch for callers that accept normal BF16 reduction-order variation. A future
 automatic band must repeat both the performance and checkpoint-output gates.
+
+## HBM re-measurement and kernel simplification
+
+The operator numbers above are L2-resident: the AstrAI pair is 40.5 MB,
+smaller than the 96 MB L2, so a tight timing loop re-reads warm weights
+(13.75 us implies ~3.1 TB/s, far above the 864 GB/s spec). Real decode rotates
+~1 GB of per-layer weights through L2 every step, so every call is cold.
+
+Re-measuring with rotated weight copies (>= 240 MB working set) on the same
+L20 showed:
+
+- The fused CTA-reuse kernel sits at the dual-stream cold-read floor
+  (702 vs 699 GB/s at (6912,1536); 369 vs 370 GB/s at (11008,4096)). Wide
+  LLaMA matrices cap at ~370-400 GB/s regardless of kernel, even for a
+  pure-read loop, so the old per-variant gaps there were noise.
+- The `(6912,1536)` warp-per-row variant (formerly M=2/4/8) is 2-6% slower
+  than CTA reuse at M=2/4 under cold weights and no longer wins at M=8 once
+  the CTA drops to 128 threads. It and its dispatch table were deleted.
+- New rule: 256 threads for M in [1, 7], 128 threads for M=8. End-to-end
+  through the built module at (6912,1536): 738-752 GB/s for M in [1, 4] and
+  702 GB/s at M=8 (+6% over the removed warp path).
+
+The M=8 CUDA-Graph regression reported above (`-23.24%`) does not survive the
+cold-weight regime: cuBLAS reaches L2 bandwidth in the warm loop while both
+fused paths converge to the same HBM floor.