perf(cuda): widen the attention query tile on sm_89 and sm_120

benchmarks/attention-wide/README.md

@@ -0,0 +1,74 @@
+# 128-query attention tile
+
+The full pass (causal or bidirectional) and 128-aligned causal prefix windows of
+at least 640 query rows use a 128-query, 256-thread block instead of 64 queries
+and 128 threads on sm_89 and sm_120. Each query row keeps the same 32-key tile
+walk and epilogue, so the outputs are byte-identical; only occupancy and K/V
+reuse change. Code:
+[`attention.cu`](../../src/backends/cuda/qwen3_5/attention.cu).
+
+## Why
+
+Both tiles stay at one block per SM (67,584 and 101,376 bytes of shared memory,
+the latter the sm_89/sm_120 opt-in limit), so the query tile decides warps per
+SM: four for 64 queries, eight for 128, which also halves K/V reads per query.
+
+| device | tile | regs/thread | smem/block | warps/SM |
+| --- | --- | ---: | ---: | ---: |
+| L20 (sm_89) | 64 x 128 | 244 | 67,584 | 4 |
+| L20 (sm_89) | 128 x 256 | 244 | 101,376 | 8 |
+| RTX 5060 Laptop (sm_120) | 64 x 128 | 228 | 67,584 | 4 |
+| RTX 5060 Laptop (sm_120) | 128 x 256 | 222 | 101,376 | 8 |
+
+## Protocol
+
+Frozen before the runs in [`plan.json`](plan.json): one variable (the launch
+tile), seeded incompressible bfloat16 inputs, Qwen3.5-27B attention shape
+(Hq=24, Hk=4, head dim 256), one excluded warmup pair per length, then
+interleaved A B B A, 10 repetitions, each sample the mean of 5 launches between
+two CUDA events, medians per length, two independent trials per device. Kernel
+and public-entry-point latency only, not HTTP.
+
+## Results
+
+Public `cs1_attention_gated`, old -> new, medians of the two trials:
+
+| tokens | L20 (sm_89) | RTX 5060 Laptop (sm_120) |
+| ---: | --- | --- |
+| 512 | 0.1300 -> 0.1296 ms (-0.3%) | 0.7720 -> 0.7720 ms (0.0%) |
+| 640 | 0.1747 -> 0.1698 (-2.8%) | 1.0478 -> 0.8279 (-21.0%) |
+| 1,024 | 0.3323 -> 0.3144 (-5.4%) | 2.2597 -> 1.8092 (-19.9%) |
+| 2,048 | 0.9859 -> 0.8998 (-8.7%) | 6.9033 -> 5.5222 (-20.0%) |
+| 4,096 | 3.4010 -> 2.9960 (-11.9%) | 23.9185 -> 19.4021 (-18.9%) |
+| 8,192 | 12.0261 -> 10.4388 (-13.2%) | 87.1708 -> 71.5387 (-17.9%) |
+
+Kernel arms alone: L20 -2.2% at 640 rows to -13.8% at 8,192, and +4.9% at 512
+rows, which is why the tile is gated at 640; sm_120 -13.0% to -22.4% from 512
+rows up. Per-trial medians and delta ranges are in [`results.json`](results.json).
+
+## Fidelity
+
+- Byte-identical outputs at every measured length in both trials on both
+  devices, and prefix windows at q_base 64/128/1,024 match the full pass.
+- `compute-sanitizer --tool memcheck`: 0 errors at 640, 1,024, 4,096 and 8,192
+  rows on both devices.
+- Attention kernel tests pass with the wider tile active, including
+  `windowed_attention_matches_full_pass_bit_for_bit` and the float64 reference
+  at 700 and 2,048 rows.
+
+## Limits
+
+- Measured on one L20 and one RTX 5060 Laptop. `H200`/sm_90 is unverified and
+  stays on the 64-query tile.
+- `cs1_attention_gated_shared` keeps 64 queries; `cs1_attention_gated_prefix`
+  uses the wider tile only when `q_base % 128 == 0`. The bidirectional entry
+  point shares `launch<>` and widens the same way (it always has `q_base = 0`).
+- No HTTP or Public231 C1 measurement; a serving-level claim needs the full-C1
+  protocol on the same named hardware as its baseline.
+
+## Reproduction
+
+The per-batch records from both runs are in commit `a3bb8c4` on this branch;
+this tree keeps the frozen protocol and the per-trial results. Build each
+revision for the target (`build.sh <out> 89` on CUDA 12.1, `120` on CUDA 12.9)
+and follow the protocol above.

benchmarks/attention-wide/plan.json

@@ -0,0 +1,29 @@
+{
+  "hypothesis": "A 128-query, 256-thread attention block raises occupancy from four to eight warps per SM (shared memory still allows one block) and reads each K/V tile once per 128 queries instead of 64, while every query row keeps the same 32-key tile walk and bfloat16 rounding, so outputs stay bit-identical.",
+  "independent_variable": "Only the flash-attention launch tile: flash_kernel<64,128> versus flash_kernel<128,256>.",
+  "fixed_controls": {
+    "shape": "Hq=24, Hk=4, head dim 256, ldv=Hk*head_dim (Qwen3.5-27B full attention)",
+    "inputs": "seeded mt19937 20261101, incompressible random bfloat16, same buffers for both arms",
+    "protocol": "1 excluded warmup pair per length, then interleaved A B B A, 10 repetitions, each sample is the mean of 5 launches between two CUDA events, medians per length, two independent trials per device",
+    "cache_stream": "warm cache, no clock changes, one non-blocking stream per process",
+    "quality": "byte comparison of both arms at every length, prefix windows at q_base 64/128/1024 against the full pass, compute-sanitizer memcheck"
+  },
+  "run_budget": "10 A-B-B-A repetitions per length per trial, two trials per device; nothing added after seeing results.",
+  "acceptance": {
+    "fidelity": "Byte-identical outputs at every measured length; 0 sanitizer errors; attention kernel tests pass with the wide block active.",
+    "latency": "Lower median at every enabled length on each measured device; 512 rows stays on the 64-query block, where the wide tile is 4.9% slower on L20.",
+    "scope": "sm_89 and sm_120 only, windows of at least 640 rows; H200/sm_90, the shared-prefix kernel and end-to-end serving are unverified."
+  },
+  "stop": "Stop on any parity mismatch or sanitizer error, or at the end of the budget.",
+  "revisions": {
+    "base_git_sha": "5fa5179d57f9267bb0d20162004d8ce26352aa9d",
+    "devices": [
+      "NVIDIA L20, sm_89, 92 SMs, driver 580.126.20, CUDA 12.1, arch=compute_89,code=sm_89",
+      "NVIDIA GeForce RTX 5060 Laptop GPU, sm_120, 26 SMs, driver 575.57.08, CUDA 12.9, -arch=sm_120"
+    ]
+  },
+  "records": {
+    "per_trial_results": "../results.json",
+    "programs": "in commit a3bb8c4 on the PR branch; this tree keeps the protocol and the per-trial results"
+  }
+}

benchmarks/attention-wide/results.json

@@ -0,0 +1,60 @@
+{
+  "protocol": "plan.json",
+  "columns": {
+    "occupancy": ["regs", "smem", "blocks_per_sm", "warps_per_sm"],
+    "kernel": ["tokens", "bm64_ms", "bm128_ms", "delta_pct", "delta_range_pct"],
+    "library": ["tokens", "old_ms", "new_ms", "delta_pct", "delta_range_pct"]
+  },
+  "devices": {
+    "L20 (sm_89)": {
+      "occupancy": {"bm64": [244, 67584, 1, 4], "bm128": [244, 101376, 1, 8]},
+      "kernel": [
+        [512, 0.1295, 0.1359, 4.913, [4.761, 5.066]],
+        [640, 0.1763, 0.1724, -2.24, [-2.283, -2.197]],
+        [768, 0.2245, 0.2196, -2.172, [-2.415, -1.928]],
+        [1024, 0.3352, 0.3192, -4.782, [-4.917, -4.647]],
+        [1536, 0.6198, 0.5767, -6.964, [-7.056, -6.871]],
+        [2048, 0.9911, 0.9064, -8.541, [-8.554, -8.528]],
+        [3072, 2.0667, 1.8306, -11.422, [-11.518, -11.325]],
+        [4096, 3.4041, 2.9979, -11.932, [-12.05, -11.815]],
+        [6144, 7.0877, 6.1537, -13.178, [-13.317, -13.039]],
+        [8192, 12.1139, 10.4371, -13.842, [-13.934, -13.751]]
+      ],
+      "library": [
+        [512, 0.13, 0.1296, -0.322, [-0.592, -0.052]],
+        [640, 0.1747, 0.1698, -2.808, [-2.996, -2.621]],
+        [1024, 0.3323, 0.3144, -5.399, [-5.429, -5.368]],
+        [2048, 0.9859, 0.8998, -8.727, [-9.138, -8.317]],
+        [4096, 3.401, 2.996, -11.909, [-11.974, -11.845]],
+        [8192, 12.0261, 10.4388, -13.199, [-13.223, -13.175]]
+      ],
+      "parity": {"kernel_identical": true, "full_pass_exact": true, "prefix_windows_exact": true},
+      "sanitizer": {"runs": 4, "errors": [0, 0, 0, 0]}
+    },
+    "RTX 5060 Laptop (sm_120)": {
+      "occupancy": {"bm64": [228, 67584, 1, 4], "bm128": [222, 101376, 1, 8]},
+      "kernel": [
+        [512, 0.7644, 0.6376, -16.595, [-16.629, -16.562]],
+        [640, 1.12, 0.9152, -18.288, [-18.669, -17.85]],
+        [768, 1.3805, 1.2009, -13.008, [-15.61, -10.414]],
+        [1024, 2.2778, 1.8418, -19.14, [-19.546, -18.748]],
+        [1536, 4.4761, 3.4748, -22.369, [-25.309, -19.114]],
+        [2048, 7.1696, 5.8404, -18.54, [-19.153, -17.975]],
+        [3072, 14.6444, 11.7271, -19.921, [-20.361, -19.499]],
+        [4096, 24.4066, 19.9933, -18.083, [-18.321, -17.855]],
+        [6144, 51.7392, 42.4056, -18.04, [-18.16, -17.925]],
+        [8192, 87.085, 71.8479, -17.497, [-17.627, -17.366]]
+      ],
+      "library": [
+        [512, 0.772, 0.772, 0.005, [0.0, 0.009]],
+        [640, 1.0478, 0.8279, -20.988, [-21.32, -20.652]],
+        [1024, 2.2597, 1.8092, -19.937, [-20.014, -19.861]],
+        [2048, 6.9033, 5.5222, -20.006, [-20.536, -19.477]],
+        [4096, 23.9185, 19.4021, -18.883, [-18.969, -18.797]],
+        [8192, 87.1708, 71.5387, -17.933, [-18.195, -17.67]]
+      ],
+      "parity": {"kernel_identical": true, "full_pass_exact": true, "prefix_windows_exact": true},
+      "sanitizer": {"runs": 4, "errors": [0, 0, 0, 0]}
+    }
+  }
+}

src/backends/cuda/qwen3_5/README.md

@@ -26,6 +26,18 @@ missing. Rebuild this library before starting that worker. This tree does not
 claim that the causal specialization's SASS is bit-identical to the previous
 build; GPU revalidation of Cua-S1 and Open-Jev was not run with this change.
 
+On sm_89 and sm_120 the full pass (causal or bidirectional) and 128-aligned
+causal prefix windows of at least 640 query rows use a 128-query, 256-thread
+block instead of 64 queries and 128 threads: eight warps per SM instead of four,
+one K/V read per 128 queries, and byte-identical results. Measured through
+`cs1_attention_gated` (24/4 heads, head dim 256), old -> new: L20 0.986 -> 0.900
+ms at 2,048 rows and 12.03 -> 10.44 at 8,192; RTX 5060 Laptop 1.048 -> 0.828 at
+640 and 87.17 -> 71.54 at 8,192. `compute-sanitizer` reported no errors on
+either device. `H200`/sm_90, the shared prefix kernel and end-to-end serving are
+unverified, so the wider block is limited to the two measured capabilities.
+Protocol and results:
+[benchmarks/attention-wide](../../../../benchmarks/attention-wide/README.md).
+
 HF Qwen3.5 folders store `linear_attn.A_log` and `linear_attn.norm.weight` as
 FP32. The shared loader rounds those two suffixes to BF16 at upload because the
 kernels consume BF16. Other non-BF16 language tensors still fail. That rounding

src/backends/cuda/qwen3_5/attention.cu

@@ -83,25 +83,31 @@ constexpr int D = 256, BM = 64, BN = 32, THREADS = 128;
 constexpr int LDS = D + 8;  // shared row stride in elements: 528 bytes keeps ldmatrix conflict-free
 constexpr int SMEM_BYTES = (BM + 2 * BN) * LDS * 2;
 
+// 128-query tile: eight warps per SM instead of four, and one K/V read per 128
+// queries. Same key tiles and order, so results are bit-identical.
+constexpr int WIDE_BM = 128, WIDE_THREADS = 256;
+constexpr int WIDE_SMEM_BYTES = (WIDE_BM + 2 * BN) * LDS * 2;
+constexpr int WIDE_MIN_ROWS = 640;  // measured crossover
+
 // q_base offsets row 0 of the q/gate/out buffers: this launch covers the window
 // [q_base, T_total) while k/v cover [0, T_total). Used by the cached-prefix
 // continuation; q_base = 0 reduces to the plain full pass.
-template <bool Gated, bool Shared, bool Bidirectional = false>
+template <bool Gated, bool Shared, bool Bidirectional = false, int QROWS = BM, int NTHREADS = THREADS>
 __device__ __forceinline__ void flash_body(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv,
                  const bf16* __restrict__ gate, bf16* __restrict__ out, int T, int Hq, int Hk,
                  float scale_log2, int q_base, const bf16* __restrict__ prefix_k,
                  const bf16* __restrict__ prefix_v) {
     extern __shared__ __align__(16) unsigned char smem[];
     bf16* qs = reinterpret_cast<bf16*>(smem);
-    bf16* ks = qs + BM * LDS;
+    bf16* ks = qs + QROWS * LDS;
     bf16* vs = ks + BN * LDS;
     const int h = blockIdx.y, hk = h / (Hq / Hk);
-    const int q0 = q_base + (gridDim.x - 1 - blockIdx.x) * BM;  // the longest blocks first
+    const int q0 = q_base + (gridDim.x - 1 - blockIdx.x) * QROWS;  // the longest blocks first
     const int tid = threadIdx.x, warp = tid / 32, lane = tid % 32;
     const int g = lane / 4, t = lane % 4;
     const int row0 = q0 + warp * 16;  // this warp's first query
 
-    for (int c = tid; c < BM * (D / 8); c += THREADS) {
+    for (int c = tid; c < QROWS * (D / 8); c += NTHREADS) {
         const int r = c / (D / 8), col = (c % (D / 8)) * 8, row = q0 + r;
         const int source_row = min(row, T - 1) - (Shared ? q_base : 0);
         cp_async16(qs + r * LDS + col, q + ((size_t)source_row * Hq + h) * D + col, row < T);
@@ -114,9 +120,9 @@ __device__ __forceinline__ void flash_body(const bf16* __restrict__ q, const bf1
     float m[2] = {-INFINITY, -INFINITY}, l[2] = {0.f, 0.f};
 
     // Causal blocks stop at their last query. A bidirectional prompt reads every key.
-    const int kv_end = Bidirectional ? T : min(T, q0 + BM);
+    const int kv_end = Bidirectional ? T : min(T, q0 + QROWS);
     for (int k0 = 0; k0 < kv_end; k0 += BN) {
-        for (int c = tid; c < BN * (D / 8); c += THREADS) {
+        for (int c = tid; c < BN * (D / 8); c += NTHREADS) {
             const int r = c / (D / 8), col = (c % (D / 8)) * 8, s = k0 + r;
             if constexpr (Shared) {
                 const int row = min(s, T - 1);
@@ -128,7 +134,7 @@ __device__ __forceinline__ void flash_body(const bf16* __restrict__ q, const bf1
             }
         }
         cp_async_commit();
-        for (int c = tid; c < BN * (D / 8); c += THREADS) {
+        for (int c = tid; c < BN * (D / 8); c += NTHREADS) {
             const int r = c / (D / 8), col = (c % (D / 8)) * 8, s = k0 + r;
             if constexpr (Shared) {
                 const int row = min(s, T - 1);
@@ -266,6 +272,15 @@ flash_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16*
     flash_body<Gated, false, Bidirectional>(q, k, v, ldv, gate, out, T, Hq, Hk, scale_log2, q_base, nullptr, nullptr);
 }
 
+template <bool Gated, bool Bidirectional = false>
+__global__ void __launch_bounds__(WIDE_THREADS)
+flash_wide_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv,
+                  const bf16* __restrict__ gate, bf16* __restrict__ out, int T, int Hq, int Hk,
+                  float scale_log2, int q_base) {
+    flash_body<Gated, false, Bidirectional, WIDE_BM, WIDE_THREADS>(q, k, v, ldv, gate, out, T, Hq, Hk,
+                                                                   scale_log2, q_base, nullptr, nullptr);
+}
+
 __global__ void __launch_bounds__(THREADS)
 flash_shared_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, const bf16* __restrict__ v, int ldv,
                     const bf16* __restrict__ gate, bf16* __restrict__ out, int T, int Hq, int Hk,
@@ -274,6 +289,32 @@ flash_shared_kernel(const bf16* __restrict__ q, const bf16* __restrict__ k, cons
     flash_body<true, true>(q, k, v, ldv, gate, out, T, Hq, Hk, scale_log2, prefix, prefix_k, prefix_v);
 }
 
+// Both the decision and the shared-memory opt-in are per device, and a thread can
+// switch devices between calls, so each thread re-checks when the device changes.
+// Returns whether `device` was measured; `error` is that device's opt-in result.
+template <bool Gated, bool Bidirectional>
+bool wide_tile_ready(int device, cudaError_t& error) {
+    struct State {
+        int device = -1;
+        bool supported = false;
+        cudaError_t configured = cudaSuccess;
+    };
+    static thread_local State state;
+    if (state.device != device) {
+        cudaDeviceProp prop{};
+        const bool known = cudaGetDeviceProperties(&prop, device) == cudaSuccess;
+        state.device = device;
+        state.supported = known &&
+            ((prop.major == 8 && prop.minor == 9) || (prop.major == 12 && prop.minor == 0));
+        state.configured = state.supported
+            ? cudaFuncSetAttribute(flash_wide_kernel<Gated, Bidirectional>,
+                                   cudaFuncAttributeMaxDynamicSharedMemorySize, WIDE_SMEM_BYTES)
+            : cudaSuccess;
+    }
+    error = state.configured;
+    return state.supported;
+}
+
 template <bool Gated, bool Bidirectional = false>
 int launch(const void* q, const void* k, const void* v, int ldv, const void* gate, void* out,
            int T, int Hq, int Hk, int Dh, float scale, int q_base, void* stream) {
@@ -284,12 +325,29 @@ int launch(const void* q, const void* k, const void* v, int ldv, const void* gat
     if (Bidirectional && q_base != 0) return cudaErrorInvalidValue;
     if (T == q_base) return cudaSuccess;
     if (Gated && gate == nullptr) return cudaErrorInvalidValue;
+    constexpr float LOG2E = 1.4426950408889634f;
+    const int rows = T - q_base;
+    if (q_base % WIDE_BM == 0 && rows >= WIDE_MIN_ROWS) {
+        int device = 0;
+        cudaError_t ready = cudaGetDevice(&device);
+        if (ready == cudaSuccess) {
+            if (wide_tile_ready<Gated, Bidirectional>(device, ready) && ready == cudaSuccess) {
+                flash_wide_kernel<Gated, Bidirectional>
+                    <<<dim3((rows + WIDE_BM - 1) / WIDE_BM, Hq), WIDE_THREADS, WIDE_SMEM_BYTES,
+                       static_cast<cudaStream_t>(stream)>>>(
+                        static_cast<const bf16*>(q), static_cast<const bf16*>(k),
+                        static_cast<const bf16*>(v), ldv, static_cast<const bf16*>(gate),
+                        static_cast<bf16*>(out), T, Hq, Hk, scale * LOG2E, q_base);
+                return cudaGetLastError();
+            }
+        }
+        if (ready != cudaSuccess) return ready;
+    }
     // Once per specialization (for the device current at the first call).
     static const cudaError_t configured = cudaFuncSetAttribute(
         flash_kernel<Gated, Bidirectional>, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_BYTES);
     if (configured != cudaSuccess) return configured;
-    constexpr float LOG2E = 1.4426950408889634f;
-    flash_kernel<Gated, Bidirectional><<<dim3((T - q_base + BM - 1) / BM, Hq), THREADS, SMEM_BYTES,
+    flash_kernel<Gated, Bidirectional><<<dim3((rows + BM - 1) / BM, Hq), THREADS, SMEM_BYTES,
                          static_cast<cudaStream_t>(stream)>>>(
         static_cast<const bf16*>(q), static_cast<const bf16*>(k), static_cast<const bf16*>(v), ldv,
         static_cast<const bf16*>(gate), static_cast<bf16*>(out), T, Hq, Hk, scale * LOG2E, q_base);