From d6f757dc13d7e8653cf1549ef4a2becfdebdea54 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Thu, 3 Sep 2026 06:33:29 +0800 Subject: [PATCH] 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 | 4 +- csrc/kernels/{gemv => }/bf16_gemv.cu | 145 ++----------------------- csrc/kernels/{gemv => }/bf16_swiglu.cu | 0 docs/developer/cuda_kernels.md | 25 +++-- tests/extension/test_gemv.py | 2 +- 5 files changed, 28 insertions(+), 148 deletions(-) rename csrc/kernels/{gemv => }/bf16_gemv.cu (67%) rename csrc/kernels/{gemv => }/bf16_swiglu.cu (100%) diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 7f32c23..91b633a 100644 --- a/csrc/CMakeLists.txt +++ b/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 - gemv/bf16_swiglu.cu + bf16_gemv.cu + bf16_swiglu.cu rotary_emb.cu ) diff --git a/csrc/kernels/gemv/bf16_gemv.cu b/csrc/kernels/bf16_gemv.cu similarity index 67% rename from csrc/kernels/gemv/bf16_gemv.cu rename to csrc/kernels/bf16_gemv.cu index 58278a8..1803d59 100644 --- a/csrc/kernels/gemv/bf16_gemv.cu +++ b/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 -__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(x); - const auto* w4 = reinterpret_cast(weight) + - static_cast(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(&wv_raw); -#pragma unroll - for (int row = 0; row < Rows; ++row) { - const uint4 xv_raw = - x4[static_cast(row) * vectors + vector]; - const auto* xv = reinterpret_cast(&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 -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 -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 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 && - (reinterpret_cast(x) & 15u) == 0u && - (reinterpret_cast(weight) & 15u) == 0u; - if constexpr (Rows == 4) { - if (aligned_rows && use_warp_tiled_kernel(n, k)) { - constexpr int kWarpsPerBlock = kWarpTiledThreads / kWarpSize; - const int blocks = (n + kWarpsPerBlock - 1) / kWarpsPerBlock; - bf16_gemv_aligned_warp_tiled_kernel - <<>>( + // 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(x) & 15u) == 0u && + (reinterpret_cast(weight) & 15u) == 0u && + static_cast(n) * k <= kSmallWeightLimit) { + bf16_gemv_kernel + <<>>( x, weight, bias, output, n, k ); return; } } - if (aligned_rows && use_half_cta_kernel(n, k)) { - bf16_gemv_kernel - <<>>( - x, weight, bias, output, n, k - ); - return; - } bf16_gemv_kernel<<>>( 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(x.data_ptr()); const auto* weight_ptr = reinterpret_cast(weight.data_ptr()); diff --git a/csrc/kernels/gemv/bf16_swiglu.cu b/csrc/kernels/bf16_swiglu.cu similarity index 100% rename from csrc/kernels/gemv/bf16_swiglu.cu rename to csrc/kernels/bf16_swiglu.cu diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index fa13848..0c3c132 100644 --- a/docs/developer/cuda_kernels.md +++ b/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_swiglu` | `gemv/bf16_swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) | +| `bf16_gemv` | `bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (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 -output row and computes all M results together, reusing the weight row across -tokens. For measured aligned M=4 medium projections, a 128-thread CTA instead -assigns one output to each of four warps. That removes the CTA-wide reduction -barrier and exposes four neighboring outputs without changing accumulation. +row-major weights `[N, K]`. One CTA computes an output row for all M tokens +together, reusing the weight row across tokens. CTA size is 256 threads, +except for small weight matrices (`N*K <= 12 MiB`) at `M=8`, where a +128-thread CTA measured 5-9% faster on L20. Variant selection is otherwise +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 -only when both tensors and every row are 16-byte aligned; all other calls keep -the general arbitrary-K path. Accumulation is FP32; the optional BF16 bias is -fused before the BF16 store. The launcher uses the current CUDA stream, is -CUDA Graph capture-safe, and requires sm_80 or newer. +so arbitrary `K` and storage offsets stay correct. Accumulation is FP32; the +optional BF16 bias is 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 diff --git a/tests/extension/test_gemv.py b/tests/extension/test_gemv.py index 1ebecc4..8aa3fd2 100644 --- a/tests/extension/test_gemv.py +++ b/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)