diff --git a/benchmarks/attention-wide/README.md b/benchmarks/attention-wide/README.md new file mode 100644 index 00000000..74b8d87c --- /dev/null +++ b/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 89` on CUDA 12.1, `120` on CUDA 12.9) +and follow the protocol above. diff --git a/benchmarks/attention-wide/plan.json b/benchmarks/attention-wide/plan.json new file mode 100644 index 00000000..73f10fe9 --- /dev/null +++ b/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" + } +} diff --git a/benchmarks/attention-wide/results.json b/benchmarks/attention-wide/results.json new file mode 100644 index 00000000..7de4e9f2 --- /dev/null +++ b/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]} + } + } +} diff --git a/src/backends/cuda/qwen3_5/README.md b/src/backends/cuda/qwen3_5/README.md index 090e54fb..3e0c6eec 100644 --- a/src/backends/cuda/qwen3_5/README.md +++ b/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 diff --git a/src/backends/cuda/qwen3_5/attention.cu b/src/backends/cuda/qwen3_5/attention.cu index 5193e05f..a0f60ced 100644 --- a/src/backends/cuda/qwen3_5/attention.cu +++ b/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 +template __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(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(q, k, v, ldv, gate, out, T, Hq, Hk, scale_log2, q_base, nullptr, nullptr); } +template +__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(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(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 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, + cudaFuncAttributeMaxDynamicSharedMemorySize, WIDE_SMEM_BYTES) + : cudaSuccess; + } + error = state.configured; + return state.supported; +} + template 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(device, ready) && ready == cudaSuccess) { + flash_wide_kernel + <<(stream)>>>( + static_cast(q), static_cast(k), + static_cast(v), ldv, static_cast(gate), + static_cast(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, cudaFuncAttributeMaxDynamicSharedMemorySize, SMEM_BYTES); if (configured != cudaSuccess) return configured; - constexpr float LOG2E = 1.4426950408889634f; - flash_kernel<<<<(stream)>>>( static_cast(q), static_cast(k), static_cast(v), ldv, static_cast(gate), static_cast(out), T, Hq, Hk, scale * LOG2E, q_base);