diff --git a/astrai/extension/backend/linear.py b/astrai/extension/backend/linear.py index ec181ba..dd896d2 100644 --- a/astrai/extension/backend/linear.py +++ b/astrai/extension/backend/linear.py @@ -73,7 +73,7 @@ def _axes( x_shape = tuple(x.shape) weight_shape = tuple(weight.shape) m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None) - supported_m = m in (1, 2, 4, 8) + supported_m = m is not None and 1 <= m <= 8 shape_matches = ( weight.ndim == 2 and x.ndim in (1, 2) @@ -164,7 +164,7 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool: or weight.dtype != torch.bfloat16 or weight.ndim != 2 or x.ndim not in (1, 2) - or (x.ndim == 2 and x.shape[0] not in (1, 2, 4, 8)) + 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 @@ -239,7 +239,7 @@ def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor: """Apply a linear projection with safe inference-only GEMV dispatch. ``ASTRAI_GEMV=0`` always uses PyTorch, ``1`` forces GEMV whenever the - primitive can safely handle an M in ``{1, 2, 4, 8}``, and ``auto`` (the + primitive can safely handle any M in ``{1, ..., 8}``, and ``auto`` (the default) uses only architecture/shape bands backed by benchmark evidence. """ # Preserve the shared dispatcher for explicit/context selection and diff --git a/astrai/extension/ops/gemv.py b/astrai/extension/ops/gemv.py index 238685e..54a32b0 100644 --- a/astrai/extension/ops/gemv.py +++ b/astrai/extension/ops/gemv.py @@ -14,7 +14,7 @@ def bf16_gemv( ) -> torch.Tensor: """Compute ``F.linear(x, weight, bias)`` for up to eight BF16 rows. - ``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``{1, 2, 4, 8}``, + ``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``[1, 8]``, and ``weight`` must be a contiguous row-major ``[N, K]`` tensor. The CUDA kernel reuses each weight row across M, accumulates in FP32, and returns BF16. This primitive is inference-only and intentionally performs no diff --git a/csrc/kernels/gemv/bf16_gemv.cu b/csrc/kernels/gemv/bf16_gemv.cu index 3c0358f..8bb171b 100644 --- a/csrc/kernels/gemv/bf16_gemv.cu +++ b/csrc/kernels/gemv/bf16_gemv.cu @@ -34,27 +34,64 @@ __global__ void bf16_gemv_kernel( const int output_index = blockIdx.x; const int lane = threadIdx.x & (kWarpSize - 1); const int warp = threadIdx.x / kWarpSize; - const int pairs = k / 2; - const auto* x2 = reinterpret_cast(x); - const auto* w2 = - reinterpret_cast(weight) + output_index * pairs; float sums[Rows] = {}; - for (int pair = threadIdx.x; pair < pairs; pair += blockDim.x) { - const __nv_bfloat162 wv = w2[pair]; + 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; + 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) { + 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) { - 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] - ); + 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 auto* xv = + reinterpret_cast(&xv_raw[row]); +#pragma unroll + for (int p = 0; p < 4; ++p) { + sums[row] = fmaf( + __bfloat162float(__low2bfloat16(xv[p])), + __bfloat162float(__low2bfloat16(wv[p])), + sums[row] + ); + sums[row] = fmaf( + __bfloat162float(__high2bfloat16(xv[p])), + __bfloat162float(__high2bfloat16(wv[p])), + sums[row] + ); + } + } + } + } 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]; +#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] + ); + } } } @@ -129,8 +166,8 @@ torch::Tensor bf16_gemv( const int64_t k = x.size(-1); const int64_t n = weight.size(0); TORCH_CHECK( - m == 1 || m == 2 || m == 4 || m == 8, - "M must be one of 1, 2, 4, or 8" + m >= 1 && m <= 8, + "M must be in [1, 8]" ); TORCH_CHECK(weight.size(1) == k, "weight K must match x K"); TORCH_CHECK(k > 0 && n > 0, "N and K must be positive"); @@ -160,6 +197,8 @@ 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()); @@ -177,11 +216,31 @@ torch::Tensor bf16_gemv( x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() ); break; + case 3: + launch_bf16_gemv<3>( + x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; case 4: launch_bf16_gemv<4>( x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() ); break; + case 5: + launch_bf16_gemv<5>( + x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + case 6: + launch_bf16_gemv<6>( + x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + case 7: + launch_bf16_gemv<7>( + x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; case 8: launch_bf16_gemv<8>( x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() @@ -201,6 +260,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { py::arg("x"), py::arg("weight"), py::arg("bias") = py::none(), - "M in {1,2,4,8} BF16 GEMV with FP32 accumulation and optional fused bias" + "M in [1, 8] BF16 GEMV with FP32 accumulation and optional fused bias" ); } diff --git a/tests/extension/test_gemv.py b/tests/extension/test_gemv.py index ff2881e..0c5359a 100644 --- a/tests/extension/test_gemv.py +++ b/tests/extension/test_gemv.py @@ -31,7 +31,7 @@ def test_bf16_gemv_matches_linear_shape_families(n, k): @skip_no_gemv -@pytest.mark.parametrize("m", [2, 4, 8]) +@pytest.mark.parametrize("m", [2, 3, 4, 5, 6, 7, 8]) @pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (1536, 6912)]) def test_bf16_gemv_matches_small_decode_batches(m, n, k): torch.manual_seed(19 + m) @@ -120,7 +120,7 @@ def test_bf16_gemv_small_batch_cuda_graph_replay(): [ ( lambda: ( - torch.randn(3, 16, device="cuda", dtype=torch.bfloat16), + torch.randn(9, 16, device="cuda", dtype=torch.bfloat16), torch.randn(8, 16, device="cuda", dtype=torch.bfloat16), ), "M must", diff --git a/tests/extension/test_linear_dispatch.py b/tests/extension/test_linear_dispatch.py index 2b6274b..ae9cf01 100644 --- a/tests/extension/test_linear_dispatch.py +++ b/tests/extension/test_linear_dispatch.py @@ -152,8 +152,8 @@ def test_grad_enabled_and_unsupported_multirow_always_fall_back(monkeypatch): ) assert "=> torch" in explain("linear", x, weight) with torch.no_grad(): - multirow = x.expand(3, -1).contiguous() - assert "=> torch" in explain("linear", multirow, weight) + oversized = torch.randn(9, 1536, device="cuda", dtype=torch.bfloat16) + assert "=> torch" in explain("linear", oversized, weight) @skip_no_gemv