From 800981d85a8663c8ce91f229283187b1bb0eb1fc Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Wed, 2 Sep 2026 14:44:09 +0800 Subject: [PATCH] refactor: accept arbitrary K in bf16 gemv with aligned head-tail sweeps - Drop the K % 2 entry rejection and the per-K if/else load-width branch: the weight stream now anchors uint4 loads at each row's first 16-byte-aligned address, with scalar head/tail sweeps covering at most 14 remainder elements, so any positive K and any storage offset is correct - Keep one pure-uint4 loop (no branching inside the loop) for the production case where every x row base is 16-byte aligned (K % 8 == 0 with allocator-aligned tensors) and a scalar-x pairing loop only for unaligned K, where per-row uint4 loads are not addressable; measured cost of scalar x everywhere was up to 2.5x on multi-row shapes (down M=4 28.4us vs 11.3us) - Remove the now-obsolete k_aligned axis and K divisibility gate from the linear dispatch spec since the primitive no longer rejects any K - Add test coverage for unaligned K (7, 12, 100, 1534) at M=1 and M=3 Benchmark: 8x L20 (sm_89, CUDA 12.8), L2-resident microbench, 300 iters; hot path unchanged within noise vs the pure-uint4 kernel (q M=2 5.8us, down M=4 11.3us, lm M=1 391us); full gate green --- astrai/extension/backend/linear.py | 3 - csrc/kernels/gemv/bf16_gemv.cu | 97 ++++++++++++++++++++---------- docs/developer/cuda_kernels.md | 16 +++-- tests/extension/test_gemv.py | 22 ++++--- 4 files changed, 90 insertions(+), 48 deletions(-) diff --git a/astrai/extension/backend/linear.py b/astrai/extension/backend/linear.py index dd896d2..4c85c82 100644 --- a/astrai/extension/backend/linear.py +++ b/astrai/extension/backend/linear.py @@ -107,7 +107,6 @@ def _axes( x_contiguous=x.is_contiguous(), weight_contiguous=weight.is_contiguous(), bias_supported=bias_supported, - k_even=k is not None and k % 2 == 0, ) @@ -122,7 +121,6 @@ _SPEC_CAPABLE = ( & axis("x_contiguous").truthy() & axis("weight_contiguous").truthy() & axis("bias_supported").truthy() - & axis("k_even").truthy() ) _SPEC_AUTO = _SPEC_CAPABLE & Spec.of( @@ -166,7 +164,6 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool: or x.ndim not in (1, 2) or (x.ndim == 2 and not 1 <= x.shape[0] <= 8) or x.shape[-1] != weight.shape[1] - or weight.shape[1] % 2 != 0 or x.device != weight.device or not x.is_contiguous() or not weight.is_contiguous() diff --git a/csrc/kernels/gemv/bf16_gemv.cu b/csrc/kernels/gemv/bf16_gemv.cu index 8bb171b..70cce50 100644 --- a/csrc/kernels/gemv/bf16_gemv.cu +++ b/csrc/kernels/gemv/bf16_gemv.cu @@ -36,26 +36,40 @@ __global__ void bf16_gemv_kernel( const int warp = threadIdx.x / kWarpSize; float sums[Rows] = {}; - if (k % 8 == 0) { - // 128-bit vectorized loads: eight bf16 elements per access halve the - // per-thread iteration count on bandwidth-bound decode shapes. - const int vecs = k / 8; + __shared__ float warp_sums[Rows][kThreads / kWarpSize]; + // Weight row: scalar head/tail around a 16-byte-aligned uint4 middle so + // any K is accepted while keeping 128-bit weight loads, which dominate + // bandwidth on decode shapes. x pairs with scalar loads: it is a tiny + // L1/L2-resident matrix, consecutive threads still touch contiguous + // addresses, and no per-row alignment case analysis is needed. + const __nv_bfloat16* __restrict__ wrow = + weight + static_cast(output_index) * k; + const unsigned whead_raw = + ((16u - (reinterpret_cast(wrow) & 15u)) & 15u) >> 1; + const int whead = static_cast(min(whead_raw, static_cast(k))); + const int wvecs = (k - whead) / 8; + const int wtail_start = whead + wvecs * 8; + const uint4* __restrict__ w4 = reinterpret_cast(wrow + whead); + + // x chunks pair element-for-element with the aligned weight middle. When + // K % 8 == 0 every x row base shares the weight alignment, so one pure + // uint4 loop covers all rows (the production case: head/tail empty, no + // branching inside the loop). Otherwise per-row uint4 loads are not + // 16-byte addressable, and scalar x pairing keeps the kernel correct for + // any K while the weight stream stays vectorized. + if (k % 8 == 0 && + ((reinterpret_cast(x) + 2u * static_cast(whead)) & 15u) == 0u) { const auto* x4 = reinterpret_cast(x); - const auto* w4 = reinterpret_cast(weight) + - static_cast(output_index) * vecs; - for (int v = threadIdx.x; v < vecs; v += blockDim.x) { + for (int v = threadIdx.x; v < wvecs; v += blockDim.x) { const uint4 wv_raw = w4[v]; const auto* wv = reinterpret_cast(&wv_raw); - uint4 xv_raw[Rows]; -#pragma unroll - for (int row = 0; row < Rows; ++row) { - xv_raw[row] = x4[static_cast(row) * vecs + v]; - } #pragma unroll for (int row = 0; row < Rows; ++row) { + const uint4 xv_raw = + x4[(static_cast(row) * wvecs) + v]; const auto* xv = - reinterpret_cast(&xv_raw[row]); + reinterpret_cast(&xv_raw); #pragma unroll for (int p = 0; p < 4; ++p) { sums[row] = fmaf( @@ -72,30 +86,50 @@ __global__ void bf16_gemv_kernel( } } } else { - const int pairs = k / 2; - const auto* x2 = reinterpret_cast(x); - const auto* w2 = - reinterpret_cast(weight) + output_index * pairs; - for (int pair = threadIdx.x; pair < pairs; pair += blockDim.x) { - const __nv_bfloat162 wv = w2[pair]; + for (int v = threadIdx.x; v < wvecs; v += blockDim.x) { + const uint4 wv_raw = w4[v]; + const __nv_bfloat16* wv_s = + reinterpret_cast(&wv_raw); #pragma unroll for (int row = 0; row < Rows; ++row) { - const __nv_bfloat162 xv = x2[row * pairs + pair]; - sums[row] = fmaf( - __bfloat162float(__low2bfloat16(xv)), - __bfloat162float(__low2bfloat16(wv)), - sums[row] - ); - sums[row] = fmaf( - __bfloat162float(__high2bfloat16(xv)), - __bfloat162float(__high2bfloat16(wv)), - sums[row] - ); + const __nv_bfloat16* xv = + x + static_cast(row) * k + whead + 8 * v; +#pragma unroll + for (int s = 0; s < 8; ++s) { + sums[row] = fmaf( + __bfloat162float(xv[s]), + __bfloat162float(wv_s[s]), + sums[row] + ); + } } } } + // Head and tail remainders: plain scalar pairing, at most 14 elements. + for (int i = threadIdx.x; i < whead; i += blockDim.x) { + const float wv = __bfloat162float(wrow[i]); +#pragma unroll + for (int row = 0; row < Rows; ++row) { + sums[row] = fmaf( + __bfloat162float(x[static_cast(row) * k + i]), + wv, + sums[row] + ); + } + } + for (int i = wtail_start + threadIdx.x; i < k; i += blockDim.x) { + const float wv = __bfloat162float(wrow[i]); +#pragma unroll + for (int row = 0; row < Rows; ++row) { + sums[row] = fmaf( + __bfloat162float(x[static_cast(row) * k + i]), + wv, + sums[row] + ); + } + } - __shared__ float warp_sums[Rows][kThreads / kWarpSize]; + #pragma unroll for (int row = 0; row < Rows; ++row) { sums[row] = warp_sum(sums[row]); @@ -171,7 +205,6 @@ torch::Tensor bf16_gemv( ); TORCH_CHECK(weight.size(1) == k, "weight K must match x K"); TORCH_CHECK(k > 0 && n > 0, "N and K must be positive"); - TORCH_CHECK(k % 2 == 0, "K must be even for vectorized bf16 loads"); TORCH_CHECK( k <= std::numeric_limits::max() && n <= std::numeric_limits::max(), diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 148fa0a..b2dbd20 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -14,22 +14,26 @@ model linear dispatcher 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/2/4/8 BF16 linear with FP32 accumulation (sm_80+) | +| `bf16_gemv` | `gemv/bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (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, 2, 4, 8}`, and row-major -weights `[N, K]`. One CTA reduces each output row and computes all M results -together, reusing the weight row across tokens. It uses vectorized -`__nv_bfloat162` loads and FP32 accumulation; the optional BF16 bias is fused +input shaped `[K]` or `[M, K]`, with `M` in `[1, 8]` and any positive `K`, and +row-major weights `[N, K]`. One CTA reduces each output row and computes all M +results together, reusing the weight row across tokens. 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; x loads are vectorized when every row base is +16-byte aligned (always true for K % 8 == 0 with allocator-aligned tensors) +and scalar otherwise. 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 -kernel for any supported M=1/2/4/8 call, or `auto` (the default) to select only +kernel for any supported M in [1, 8], or `auto` (the default) to select only architecture/shape bands that pass both the per-shape and end-to-end gates. M=1 has no automatic SM89 band because isolated winners did not reach the 3% whole-graph gate. Measured SM89 small-M bands are enabled as follows: diff --git a/tests/extension/test_gemv.py b/tests/extension/test_gemv.py index 0c5359a..c49bf67 100644 --- a/tests/extension/test_gemv.py +++ b/tests/extension/test_gemv.py @@ -96,6 +96,21 @@ def test_bf16_gemv_cuda_graph_replay(): torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) +@skip_no_gemv +@pytest.mark.parametrize("n,k", [(64, 7), (64, 12), (33, 100), (256, 1534)]) +def test_bf16_gemv_handles_unaligned_k(n, k): + torch.manual_seed(29) + x = torch.randn(k, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) + actual = bf16_gemv(x, weight) + expected = F.linear(x, weight) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + x3 = torch.randn(3, k, device="cuda", dtype=torch.bfloat16) + actual3 = bf16_gemv(x3, weight) + torch.testing.assert_close(actual3, F.linear(x3, weight), rtol=0.02, atol=0.5) + + @skip_no_gemv def test_bf16_gemv_small_batch_cuda_graph_replay(): torch.manual_seed(31) @@ -125,13 +140,6 @@ def test_bf16_gemv_small_batch_cuda_graph_replay(): ), "M must", ), - ( - lambda: ( - torch.randn(15, device="cuda", dtype=torch.bfloat16), - torch.randn(8, 15, device="cuda", dtype=torch.bfloat16), - ), - "even", - ), ( lambda: ( torch.randn(16, device="cuda", dtype=torch.float16),