From a144d7f30647f92981e01048f5135641ad7f0db5 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 2 Sep 2026 13:08:08 +0800 Subject: [PATCH] perf: accelerate decode linear with bf16 gemv - add decode-shape benchmark harness - add bf16 GEMV CUDA primitive with head-dim generic kernel - dispatch decode-time linear layers to gemv for M=1 - extend gemv coverage to small decode batches --- astrai/extension/__init__.py | 4 + astrai/extension/backend/__init__.py | 2 + astrai/extension/backend/linear.py | 264 +++++++++++++++++ astrai/extension/ops/__init__.py | 2 + astrai/extension/ops/gemv.py | 23 ++ astrai/model/components/linear.py | 5 +- csrc/CMakeLists.txt | 2 + csrc/kernels/gemv/bf16_gemv.cu | 206 +++++++++++++ docs/developer/cuda_kernels.md | 41 ++- docs/developer/decode_linear_benchmark.md | 22 ++ scripts/tools/benchmark_gemv.py | 333 ++++++++++++++++++++++ setup.py | 1 + tests/extension/test_gemv.py | 155 ++++++++++ tests/extension/test_linear_dispatch.py | 185 ++++++++++++ 14 files changed, 1242 insertions(+), 3 deletions(-) create mode 100644 astrai/extension/backend/linear.py create mode 100644 astrai/extension/ops/gemv.py create mode 100644 csrc/kernels/gemv/bf16_gemv.cu create mode 100644 docs/developer/decode_linear_benchmark.md create mode 100644 scripts/tools/benchmark_gemv.py create mode 100644 tests/extension/test_gemv.py create mode 100644 tests/extension/test_linear_dispatch.py diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index 32e9425..6e0ba08 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -26,6 +26,7 @@ from astrai.extension.backend import ( attention, attn_backend, get_backend, + linear, ) from astrai.extension.dispatch import ( Axes, @@ -49,6 +50,7 @@ from astrai.extension.ops import ( attn_decode, attn_paged_decode, attn_prefill, + bf16_gemv, ) __all__ = [ @@ -62,9 +64,11 @@ __all__ = [ "attention", "attn_backend", "get_backend", + "linear", "attn_decode", "attn_paged_decode", "attn_prefill", + "bf16_gemv", "is_available", "KERNEL_NAMES", "apply_rotary_emb", diff --git a/astrai/extension/backend/__init__.py b/astrai/extension/backend/__init__.py index 189f8e8..7b43a21 100644 --- a/astrai/extension/backend/__init__.py +++ b/astrai/extension/backend/__init__.py @@ -11,6 +11,7 @@ from astrai.extension.backend.attention import ( attn_backend, get_backend, ) +from astrai.extension.backend.linear import linear from astrai.extension.backend.rotary import apply_rotary_emb __all__ = [ @@ -24,4 +25,5 @@ __all__ = [ "attention", "attn_backend", "get_backend", + "linear", ] diff --git a/astrai/extension/backend/linear.py b/astrai/extension/backend/linear.py new file mode 100644 index 0000000..ec181ba --- /dev/null +++ b/astrai/extension/backend/linear.py @@ -0,0 +1,264 @@ +"""Inference-only dispatch for AstrAI linear layers. + +The CUDA GEMV path is deliberately narrow: automatic selection is enabled +only for single-row BF16 shapes measured to beat ``F.linear`` on a supported +architecture. Every training, prefill, unsupported-layout, and unmeasured +call falls back to PyTorch. +""" + +import logging +import os +from functools import lru_cache +from typing import Optional + +import torch +import torch.nn.functional as F +from torch import Tensor + +from astrai.extension.dispatch import ( + ImplRecord, + Spec, + axis, + get_override, + register_family, + resolve, + tensor_axes, +) +from astrai.extension.loader import is_available +from astrai.extension.ops.gemv import bf16_gemv + +logger = logging.getLogger(__name__) + +# Shape keys are (N, K) for Y[M, N] = X[M, K] @ W[N, K].T. A band is +# automatic only after both the per-shape >=5% and end-to-end decode >=3% +# gates pass and checkpoint greedy output remains stable. M=1 and M=8 remain +# empty on SM89; the safe M=2/4 bands improve real-engine throughput by +# 11.8-14.0%. +_AUTO_GEMV_SHAPES: dict[tuple[int, int], dict[int, frozenset[tuple[int, int]]]] = { + (8, 9): { + 2: frozenset( + { + (256, 1536), + (1536, 1536), + (100000, 1536), + } + ), + 4: frozenset({(256, 1536), (1536, 1536)}), + } +} +_AUTO_GEMV_M = frozenset( + m for architecture in _AUTO_GEMV_SHAPES.values() for m in architecture +) + +_VALID_MODES = {"0", "1", "auto"} +_WARNED_MODES: set[str] = set() + + +def _gemv_mode() -> str: + mode = os.environ.get("ASTRAI_GEMV", "auto").strip().lower() + if mode in _VALID_MODES: + return mode + if mode not in _WARNED_MODES: + _WARNED_MODES.add(mode) + logger.warning( + "ASTRAI_GEMV=%r is invalid; expected 0, 1, or auto; using auto", + mode, + ) + return "auto" + + +def _axes( + x: Tensor, weight: Tensor, bias: Optional[Tensor] = None +) -> dict[str, object]: + 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) + shape_matches = ( + weight.ndim == 2 + and x.ndim in (1, 2) + and bool(x_shape) + and x_shape[-1] == weight_shape[-1] + ) + same_device = x.device == weight.device and ( + bias is None or bias.device == x.device + ) + bias_supported = bias is None or ( + bias.ndim == 1 + and weight.ndim == 2 + and bias.shape[0] == weight.shape[0] + and bias.dtype == torch.bfloat16 + and bias.is_contiguous() + ) + capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None + n = weight_shape[0] if weight.ndim == 2 else None + k = weight_shape[1] if weight.ndim == 2 else None + return tensor_axes( + x, + mode=_gemv_mode(), + capability=capability, + n=n, + k=k, + m=m, + supported_m=supported_m, + shape_matches=shape_matches, + same_device=same_device, + weight_dtype=weight.dtype, + x_contiguous=x.is_contiguous(), + weight_contiguous=weight.is_contiguous(), + bias_supported=bias_supported, + k_even=k is not None and k % 2 == 0, + ) + + +_SPEC_CAPABLE = ( + axis("device_cuda").truthy() + & axis("dtype").in_(torch.bfloat16) + & axis("weight_dtype").in_(torch.bfloat16) + & axis("grad_enabled").eq(False) + & axis("supported_m").truthy() + & axis("shape_matches").truthy() + & axis("same_device").truthy() + & axis("x_contiguous").truthy() + & axis("weight_contiguous").truthy() + & axis("bias_supported").truthy() + & axis("k_even").truthy() +) + +_SPEC_AUTO = _SPEC_CAPABLE & Spec.of( + lambda ax: ( + (ax.get("n"), ax.get("k")) + in _AUTO_GEMV_SHAPES.get(ax.get("capability"), {}).get(ax.get("m"), ()) + ), + "shape is a measured winner for this architecture", +) + + +def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor: + return F.linear(x, weight, bias) + + +def _inference_bf16_gemv( + x: Tensor, weight: Tensor, bias: Optional[Tensor] = None +) -> Tensor: + # Model parameters retain requires_grad=True after eval(). Dispatch is + # already restricted to no-grad, so detached views preserve storage and + # layout while satisfying the primitive's explicit autograd guard. + return bf16_gemv( + x.detach(), + weight.detach(), + bias.detach() if bias is not None else None, + ) + + +@lru_cache(maxsize=None) +def _device_capability(device_index: int) -> tuple[int, int]: + return torch.cuda.get_device_capability(device_index) + + +def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool: + if ( + torch.is_grad_enabled() + or not x.is_cuda + or x.dtype != torch.bfloat16 + 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.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() + or not is_available("bf16_gemv") + ): + return False + return bias is None or ( + bias.device == x.device + and bias.dtype == torch.bfloat16 + and bias.ndim == 1 + and bias.shape[0] == weight.shape[0] + and bias.is_contiguous() + ) + + +def _auto_gemv_shape(x: Tensor, weight: Tensor) -> bool: + capability = _device_capability(x.get_device()) + m = 1 if x.ndim == 1 else x.shape[0] + return (weight.shape[0], weight.shape[1]) in _AUTO_GEMV_SHAPES.get( + capability, {} + ).get(m, ()) + + +def _linear_records() -> list[ImplRecord]: + mode = _gemv_mode() + gemv_priority = 0 if mode == "1" else 100 + auto_priority = 0 if mode == "auto" else 90 + torch_priority = 0 if mode == "0" else 50 + return [ + ImplRecord( + family="linear", + name="gemv", + obj=_inference_bf16_gemv, + spec=_SPEC_CAPABLE, + available=lambda: is_available("bf16_gemv"), + priority=gemv_priority, + ), + ImplRecord( + family="linear", + name="auto_gemv", + obj=_inference_bf16_gemv, + spec=_SPEC_AUTO, + available=lambda: is_available("bf16_gemv"), + priority=auto_priority, + ), + ImplRecord( + family="linear", + name="torch", + obj=_torch_linear, + spec=Spec.always(), + priority=torch_priority, + ), + ] + + +def _fallback_record() -> ImplRecord: + return ImplRecord( + family="linear", + name="torch", + obj=_torch_linear, + spec=Spec.always(), + priority=999, + ) + + +register_family("linear", _axes, _linear_records, _fallback_record) + + +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 + default) uses only architecture/shape bands backed by benchmark evidence. + """ + # Preserve the shared dispatcher for explicit/context selection and + # ASTR_OPS diagnostics, while keeping the default per-layer hot path free + # of axes dictionaries, record sorting, and repeated capability queries. + if get_override("linear") is not None or "linear" in os.environ.get("ASTR_OPS", ""): + return resolve("linear", x, weight, bias).record.obj(x, weight, bias) + + mode = _gemv_mode() + if mode == "0" or (mode == "auto" and not _AUTO_GEMV_SHAPES): + return _torch_linear(x, weight, bias) + if mode == "auto": + m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None) + if m not in _AUTO_GEMV_M: + return _torch_linear(x, weight, bias) + if mode != "0" and _gemv_capable(x, weight, bias): + if mode == "1" or _auto_gemv_shape(x, weight): + return _inference_bf16_gemv(x, weight, bias) + return _torch_linear(x, weight, bias) + + +__all__ = ["linear"] diff --git a/astrai/extension/ops/__init__.py b/astrai/extension/ops/__init__.py index ac0009a..ba06b2f 100644 --- a/astrai/extension/ops/__init__.py +++ b/astrai/extension/ops/__init__.py @@ -7,6 +7,7 @@ from astrai.extension.ops.attention import ( attn_paged_prefill, attn_prefill, ) +from astrai.extension.ops.gemv import bf16_gemv from astrai.extension.ops.rotary import rotary_emb __all__ = [ @@ -15,5 +16,6 @@ __all__ = [ "attn_paged_decode", "attn_paged_prefill", "attn_prefill", + "bf16_gemv", "rotary_emb", ] diff --git a/astrai/extension/ops/gemv.py b/astrai/extension/ops/gemv.py new file mode 100644 index 0000000..238685e --- /dev/null +++ b/astrai/extension/ops/gemv.py @@ -0,0 +1,23 @@ +"""Stateless wrapper for the directly callable BF16 GEMV primitive.""" + +from typing import Optional + +import torch + +from astrai.extension.loader import get_module + + +def bf16_gemv( + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor] = None, +) -> 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}``, + 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 + fallback or model-level dispatch. + """ + return get_module("bf16_gemv").bf16_gemv(x, weight, bias) diff --git a/astrai/model/components/linear.py b/astrai/model/components/linear.py index f3a9aeb..3e50e62 100644 --- a/astrai/model/components/linear.py +++ b/astrai/model/components/linear.py @@ -1,8 +1,9 @@ import torch import torch.nn as nn -import torch.nn.functional as F from torch import Tensor +from astrai.extension.backend.linear import linear + class Linear(nn.Module): def __init__( @@ -21,4 +22,4 @@ class Linear(nn.Module): nn.init.uniform_(self.bias, -bound, bound) def forward(self, x: Tensor) -> Tensor: - return F.linear(x, self.weight, self.bias) + return linear(x, self.weight, self.bias) diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index b08b91f..bdd6abb 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -61,6 +61,7 @@ set(KERNEL_NAMES attn_prefill attn_paged_decode attn_paged_prefill + bf16_gemv rotary_emb ) set(KERNEL_SRCS @@ -68,6 +69,7 @@ set(KERNEL_SRCS attention/prefill.cu attention/paged_decode.cu attention/paged_prefill.cu + gemv/bf16_gemv.cu rotary_emb.cu ) diff --git a/csrc/kernels/gemv/bf16_gemv.cu b/csrc/kernels/gemv/bf16_gemv.cu new file mode 100644 index 0000000..3c0358f --- /dev/null +++ b/csrc/kernels/gemv/bf16_gemv.cu @@ -0,0 +1,206 @@ +// Directly callable small-M BF16 GEMV primitive for decode-time linear layers. + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kWarpSize = 32; + +__device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + return value; +} + +template +__global__ void bf16_gemv_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ weight, + const __nv_bfloat16* __restrict__ bias, + __nv_bfloat16* __restrict__ output, + int n, + int k +) { + 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]; +#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] + ); + } + } + + __shared__ float warp_sums[Rows][kThreads / kWarpSize]; +#pragma unroll + for (int row = 0; row < Rows; ++row) { + sums[row] = warp_sum(sums[row]); + } + if (lane == 0) { +#pragma unroll + for (int row = 0; row < Rows; ++row) { + warp_sums[row][warp] = sums[row]; + } + } + __syncthreads(); + + if (warp == 0) { +#pragma unroll + for (int row = 0; row < Rows; ++row) { + float sum = + lane < (kThreads / kWarpSize) ? warp_sums[row][lane] : 0.0f; + sum = warp_sum(sum); + if (lane == 0) { + if (bias != nullptr) { + sum += __bfloat162float(bias[output_index]); + } + output[row * n + output_index] = __float2bfloat16_rn(sum); + } + } + } +} + +template +void launch_bf16_gemv( + const __nv_bfloat16* x, + const __nv_bfloat16* weight, + const __nv_bfloat16* bias, + __nv_bfloat16* output, + int n, + int k, + cudaStream_t stream +) { + bf16_gemv_kernel<<>>( + x, weight, bias, output, n, k + ); +} + +torch::Tensor bf16_gemv( + torch::Tensor x, + torch::Tensor weight, + py::object bias_object +) { + TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "x and weight must be CUDA tensors"); + TORCH_CHECK(x.device() == weight.device(), "x and weight must share device"); + TORCH_CHECK( + x.scalar_type() == torch::kBFloat16 && + weight.scalar_type() == torch::kBFloat16, + "x and weight must be bf16" + ); + TORCH_CHECK( + x.dim() == 1 || x.dim() == 2, + "x must have shape [K] or [M, K]" + ); + TORCH_CHECK(weight.dim() == 2, "weight must have shape [N, K]"); + TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(), "x and weight must be contiguous"); + TORCH_CHECK( + !x.requires_grad() && !weight.requires_grad(), + "bf16_gemv is inference-only and does not support autograd" + ); + + const int64_t m = x.dim() == 1 ? 1 : x.size(0); + 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" + ); + 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(), + "N or K exceeds the CUDA launcher limit" + ); + + torch::Tensor bias; + const __nv_bfloat16* bias_ptr = nullptr; + if (!bias_object.is_none()) { + bias = bias_object.cast(); + TORCH_CHECK(bias.is_cuda() && bias.device() == x.device(), "bias must share the CUDA device"); + TORCH_CHECK(bias.scalar_type() == torch::kBFloat16, "bias must be bf16"); + TORCH_CHECK(bias.dim() == 1 && bias.size(0) == n, "bias must have shape [N]"); + TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous"); + TORCH_CHECK(!bias.requires_grad(), "bf16_gemv bias does not support autograd"); + bias_ptr = reinterpret_cast(bias.data_ptr()); + } + + const at::cuda::OptionalCUDAGuard guard(x.device()); + const auto* properties = at::cuda::getDeviceProperties(x.device().index()); + TORCH_CHECK(properties->major >= 8, "bf16_gemv requires compute capability 8.0+"); + auto stream = at::cuda::getCurrentCUDAStream(); + auto output = x.dim() == 1 ? torch::empty({n}, x.options()) + : torch::empty({m, n}, x.options()); + + const auto* x_ptr = reinterpret_cast(x.data_ptr()); + const auto* weight_ptr = + reinterpret_cast(weight.data_ptr()); + auto* output_ptr = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + const int n_int = static_cast(n); + const int k_int = static_cast(k); + switch (m) { + case 1: + launch_bf16_gemv<1>( + x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + case 2: + launch_bf16_gemv<2>( + 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 8: + launch_bf16_gemv<8>( + x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + } + C10_CUDA_CHECK(cudaGetLastError()); + return output; +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def( + "bf16_gemv", + &bf16_gemv, + 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" + ); +} diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index b28259b..148fa0a 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -1,6 +1,9 @@ # CUDA Kernels -AstrAI includes optional custom CUDA kernels for attention, rotary embedding, and FP8 GEMM. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend, auto-dispatched for rotary, or invoked through the FP8 linear primitives. +AstrAI includes optional custom CUDA kernels for attention, rotary embedding, +BF16 GEMV, and FP8 GEMM. These are built when `nvcc` is available and CUDA is +detected. BF16 GEMV is directly callable and can be selected by the guarded +model linear dispatcher described below. ## Overview @@ -11,8 +14,43 @@ AstrAI includes optional custom CUDA kernels for attention, rotary embedding, an | `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+) | | `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 +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 +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: + +| M | Automatic `(N, K)` bands | Engine throughput | +|---:|---|---:| +| 2 | `(256,1536)`, `(1536,1536)`, `(100000,1536)` | +14.0% | +| 4 | `(256,1536)`, `(1536,1536)` | +11.8% | + +These A→B→B→A results use the real `InferenceEngine`, including scheduler, +sampling, and CUDA Graph. M=8 stays on PyTorch because its remaining +greedy-stable winners missed the 3% end-to-end gate. Long-K MLP-down bands are +also excluded because their valid BF16 error changed a checkpoint greedy +argmax; the enabled M=2/4 bands matched the baseline greedy output exactly. + +Training, prefill, unmeasured architectures, and losing shape bands always +remain on PyTorch. Use mode `1` only for explicit A/B runs outside this table. + +The primitive remains directly callable and deliberately has no internal +`F.linear` fallback. The model-level backend owns fallback and dispatch policy. + Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist: | Variant | File | Optimization | @@ -202,6 +240,7 @@ astrai/extension/ ├── ops/ │ ├── attention.py # Stateless attention kernel wrappers │ ├── rotary.py # Stateless rotary kernel wrapper +│ ├── gemv.py # Stateless BF16 GEMV primitive │ └── fp8.py # Stateless FP8 primitives (custom_op) ├── fp8.py # FP8 strategy layer (fp8_autocast, recipes) └── backend/ diff --git a/docs/developer/decode_linear_benchmark.md b/docs/developer/decode_linear_benchmark.md new file mode 100644 index 0000000..02bf113 --- /dev/null +++ b/docs/developer/decode_linear_benchmark.md @@ -0,0 +1,22 @@ +# Decode linear shape benchmark + +`scripts/tools/benchmark_gemv.py` records the `F.linear` baseline used to decide +whether a BF16 GEMV or small-M kernel should enter automatic inference dispatch. +It does not change model execution or select a custom kernel. + +The default matrix covers the AstrAI 1B q/k/v/out projections, MLP up/gate/down, +and LM head for `M=1,2,4,8,16,32`. Each shape runs in eager and CUDA Graph replay +modes. Results include device-event latency samples, p50/p90/p99, estimated +effective IO bandwidth, and CUDA kernel launches per call. + +```bash +CUDA_VISIBLE_DEVICES=0 python scripts/tools/benchmark_gemv.py \ + --output results/decode_linear.json \ + --markdown-output results/decode_linear.md +``` + +Use `--shape NAME:N:K` repeatedly to override the preset and `--m-values` to +change the decode batch sizes. Compare each GPU architecture only with its own +baseline; do not use absolute A100-versus-L20 numbers as a dispatch criterion. +Keep the raw JSON as the source of truth and generate tables with +`--markdown-output` rather than transcribing measurements by hand. diff --git a/scripts/tools/benchmark_gemv.py b/scripts/tools/benchmark_gemv.py new file mode 100644 index 0000000..83743ca --- /dev/null +++ b/scripts/tools/benchmark_gemv.py @@ -0,0 +1,333 @@ +"""Benchmark decode-time linear shapes before enabling custom GEMV dispatch. + +The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It +establishes the per-architecture cuBLAS baseline that later GEMV primitives and +dispatch decisions must beat. +""" + +from __future__ import annotations + +import json +import math +import statistics +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Callable, Iterable + +import click +import torch +import torch.nn.functional as F + + +@dataclass(frozen=True) +class LinearShape: + name: str + n: int + k: int + + +DEFAULT_SHAPES = ( + LinearShape("q_proj", 1536, 1536), + LinearShape("k_proj", 256, 1536), + LinearShape("v_proj", 256, 1536), + LinearShape("attn_out", 1536, 1536), + LinearShape("mlp_up", 6912, 1536), + LinearShape("mlp_gate", 6912, 1536), + LinearShape("mlp_down", 1536, 6912), + LinearShape("lm_head", 100000, 1536), +) +DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16} + + +def parse_positive_ints(value: str) -> tuple[int, ...]: + """Parse a comma-separated, duplicate-free list of positive integers.""" + try: + values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(","))) + except ValueError as exc: + raise click.BadParameter("expected comma-separated integers") from exc + if not values or any(item <= 0 for item in values): + raise click.BadParameter("values must be positive integers") + return values + + +def parse_shape(value: str) -> LinearShape: + """Parse NAME:N:K into a benchmark shape.""" + parts = value.split(":") + if len(parts) != 3 or not parts[0]: + raise click.BadParameter("shape must use NAME:N:K") + try: + n, k = (int(item) for item in parts[1:]) + except ValueError as exc: + raise click.BadParameter("N and K must be integers") from exc + if n <= 0 or k <= 0: + raise click.BadParameter("N and K must be positive") + return LinearShape(parts[0], n, k) + + +def estimate_io_bytes( + m: int, n: int, k: int, element_size: int, *, has_bias: bool +) -> int: + """Estimate bytes touched once by Y[M,N] = X[M,K] @ W[N,K].T.""" + elements = m * k + n * k + m * n + if has_bias: + elements += n + return elements * element_size + + +def percentile(values: Iterable[float], quantile: float) -> float: + ordered = sorted(values) + if not ordered: + raise ValueError("percentile requires at least one sample") + rank = (len(ordered) - 1) * quantile + lower = math.floor(rank) + upper = math.ceil(rank) + if lower == upper: + return ordered[lower] + fraction = rank - lower + return ordered[lower] * (1 - fraction) + ordered[upper] * fraction + + +def summarize_latency(samples_ms: list[float]) -> dict[str, float]: + return { + "median_ms": statistics.median(samples_ms), + "p90_ms": percentile(samples_ms, 0.90), + "p99_ms": percentile(samples_ms, 0.99), + "min_ms": min(samples_ms), + "max_ms": max(samples_ms), + } + + +def measure_cuda_ms( + operation: Callable[[], torch.Tensor], *, warmup: int, iterations: int, trials: int +) -> list[float]: + for _ in range(warmup): + operation() + torch.cuda.synchronize() + + samples = [] + for _ in range(trials): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iterations): + operation() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) / iterations) + return samples + + +def count_cuda_kernels( + operation: Callable[[], torch.Tensor], repeats: int = 5 +) -> float: + """Profile a few calls and return the average device events per call.""" + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + acc_events=True, + ) as profile: + for _ in range(repeats): + operation() + torch.cuda.synchronize() + + device_type = torch.autograd.DeviceType.CUDA + events = [event for event in profile.events() if event.device_type == device_type] + return len(events) / repeats + + +def capture_linear( + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None +) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]: + for _ in range(3): + F.linear(x, weight, bias) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = F.linear(x, weight, bias) + return graph, output + + +def benchmark_case( + shape: LinearShape, + m: int, + *, + dtype: torch.dtype, + mode: str, + bias_enabled: bool, + warmup: int, + iterations: int, + trials: int, +) -> dict[str, object]: + x = torch.randn((m, shape.k), device="cuda", dtype=dtype) + weight = torch.randn((shape.n, shape.k), device="cuda", dtype=dtype) + bias = torch.randn(shape.n, device="cuda", dtype=dtype) if bias_enabled else None + + graph = None + graph_output = None + if mode == "graph": + graph, graph_output = capture_linear(x, weight, bias) + + def operation() -> torch.Tensor: + graph.replay() + return graph_output + + else: + + def operation() -> torch.Tensor: + return F.linear(x, weight, bias) + + samples_ms = measure_cuda_ms( + operation, warmup=warmup, iterations=iterations, trials=trials + ) + latency = summarize_latency(samples_ms) + io_bytes = estimate_io_bytes( + m, shape.n, shape.k, x.element_size(), has_bias=bias is not None + ) + median_seconds = latency["median_ms"] / 1000 + + result: dict[str, object] = { + "name": shape.name, + "m": m, + "n": shape.n, + "k": shape.k, + "mode": mode, + "bias": bias is not None, + "estimated_io_bytes": io_bytes, + "effective_bandwidth_gbps": io_bytes / median_seconds / 1e9, + "cuda_kernel_launches_per_call": count_cuda_kernels(operation), + **latency, + "samples_ms": samples_ms, + } + return result + + +def render_markdown(payload: dict[str, object]) -> str: + metadata = payload["metadata"] + assert isinstance(metadata, dict) + results = payload["results"] + assert isinstance(results, list) + + lines = [ + "# Decode linear baseline", + "", + f"- GPU: {metadata['gpu_name']}", + f"- Compute capability: {metadata['compute_capability']}", + f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}", + f"- Dtype: {metadata['dtype']}", + "", + "| Layer | M | N | K | Mode | Median (ms) | p99 (ms) | GB/s | CUDA kernels/call |", + "|---|---:|---:|---:|---|---:|---:|---:|---:|", + ] + for item in results: + assert isinstance(item, dict) + lines.append( + "| {name} | {m} | {n} | {k} | {mode} | {median_ms:.4f} | " + "{p99_ms:.4f} | {effective_bandwidth_gbps:.1f} | " + "{cuda_kernel_launches_per_call:.2f} |".format(**item) + ) + lines.append("") + return "\n".join(lines) + + +def device_metadata(dtype_name: str) -> dict[str, object]: + props = torch.cuda.get_device_properties(0) + return { + "timestamp_utc": datetime.now(timezone.utc).isoformat(), + "gpu_name": props.name, + "compute_capability": f"{props.major}.{props.minor}", + "total_memory_bytes": props.total_memory, + "torch_version": torch.__version__, + "cuda_version": torch.version.cuda, + "dtype": dtype_name, + } + + +@click.command(help=__doc__) +@click.option("--output", type=click.Path(path_type=Path), required=True) +@click.option("--markdown-output", type=click.Path(path_type=Path)) +@click.option("--m-values", default="1,2,4,8,16,32", show_default=True) +@click.option( + "--shape", + "shape_values", + multiple=True, + help="Override defaults with repeatable NAME:N:K shapes.", +) +@click.option("--dtype", type=click.Choice(tuple(DTYPES)), default="bfloat16") +@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both") +@click.option("--bias/--no-bias", default=False) +@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True) +@click.option( + "--iterations", type=click.IntRange(min=1), default=100, show_default=True +) +@click.option("--trials", type=click.IntRange(min=1), default=20, show_default=True) +@click.option("--seed", type=int, default=0, show_default=True) +def benchmark_command( + output: Path, + markdown_output: Path | None, + m_values: str, + shape_values: tuple[str, ...], + dtype: str, + mode: str, + bias: bool, + warmup: int, + iterations: int, + trials: int, + seed: int, +) -> None: + if not torch.cuda.is_available(): + raise click.ClickException("CUDA is required") + + parsed_m = parse_positive_ints(m_values) + shapes = tuple(parse_shape(item) for item in shape_values) or DEFAULT_SHAPES + modes = ("eager", "graph") if mode == "both" else (mode,) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + results = [] + for shape in shapes: + for m in parsed_m: + for current_mode in modes: + click.echo( + f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}" + ) + results.append( + benchmark_case( + shape, + m, + dtype=DTYPES[dtype], + mode=current_mode, + bias_enabled=bias, + warmup=warmup, + iterations=iterations, + trials=trials, + ) + ) + + payload: dict[str, object] = { + "schema_version": 1, + "metadata": device_metadata(dtype), + "parameters": { + "m_values": list(parsed_m), + "shapes": [asdict(shape) for shape in shapes], + "modes": list(modes), + "bias": bias, + "warmup": warmup, + "iterations": iterations, + "trials": trials, + "seed": seed, + }, + "results": results, + } + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + if markdown_output is not None: + markdown_output.parent.mkdir(parents=True, exist_ok=True) + markdown_output.write_text(render_markdown(payload), encoding="utf-8") + + +if __name__ == "__main__": + benchmark_command() diff --git a/setup.py b/setup.py index 78d5f12..0b4313f 100644 --- a/setup.py +++ b/setup.py @@ -121,6 +121,7 @@ class _CMakeBuildExt(_build_ext): "attn_prefill", "attn_paged_decode", "attn_paged_prefill", + "bf16_gemv", "rotary_emb", ) missing = [name for name in required if not any(lib_dir.glob(f"{name}.*.so"))] diff --git a/tests/extension/test_gemv.py b/tests/extension/test_gemv.py new file mode 100644 index 0000000..ff2881e --- /dev/null +++ b/tests/extension/test_gemv.py @@ -0,0 +1,155 @@ +import pytest +import torch +import torch.nn.functional as F + +from astrai.extension import bf16_gemv, is_available + +GEMV_AVAILABLE = ( + torch.cuda.is_available() + and is_available("bf16_gemv") + and torch.cuda.get_device_capability() >= (8, 0) +) +skip_no_gemv = pytest.mark.skipif( + not GEMV_AVAILABLE, + reason="BF16 GEMV requires a built kernel and compute capability 8.0+", +) + + +@skip_no_gemv +@pytest.mark.parametrize( + "n,k", + [(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912), (100000, 1536)], +) +def test_bf16_gemv_matches_linear_shape_families(n, k): + torch.manual_seed(17) + 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) + assert actual.shape == (n,) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + +@skip_no_gemv +@pytest.mark.parametrize("m", [2, 4, 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) + x = torch.randn(m, 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) + assert actual.shape == (m, n) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5) + + +@skip_no_gemv +def test_bf16_gemv_preserves_singleton_batch_and_fuses_bias(): + torch.manual_seed(23) + x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(1536, device="cuda", dtype=torch.bfloat16) + actual = bf16_gemv(x, weight, bias) + expected = F.linear(x, weight, bias) + assert actual.shape == (1, 1536) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + +@skip_no_gemv +def test_bf16_gemv_small_batch_fuses_bias(): + torch.manual_seed(25) + x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + bias = torch.randn(256, device="cuda", dtype=torch.bfloat16) + actual = bf16_gemv(x, weight, bias) + expected = F.linear(x, weight, bias) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + +@skip_no_gemv +def test_bf16_gemv_uses_current_stream(): + x = torch.randn(1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + actual = bf16_gemv(x, weight) + expected = F.linear(x, weight) + stream.synchronize() + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + +@skip_no_gemv +def test_bf16_gemv_cuda_graph_replay(): + torch.manual_seed(29) + x = torch.randn(1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + bf16_gemv(x, weight) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = bf16_gemv(x, weight) + + x.copy_(torch.randn_like(x)) + graph.replay() + expected = F.linear(x, weight) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + +@skip_no_gemv +def test_bf16_gemv_small_batch_cuda_graph_replay(): + torch.manual_seed(31) + x = torch.randn(8, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) + for _ in range(3): + bf16_gemv(x, weight) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = bf16_gemv(x, weight) + + x.copy_(torch.randn_like(x)) + graph.replay() + expected = F.linear(x, weight) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) + + +@skip_no_gemv +@pytest.mark.parametrize( + "make_args,error", + [ + ( + lambda: ( + torch.randn(3, 16, device="cuda", dtype=torch.bfloat16), + torch.randn(8, 16, device="cuda", dtype=torch.bfloat16), + ), + "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), + torch.randn(8, 16, device="cuda", dtype=torch.float16), + ), + "bf16", + ), + ( + lambda: ( + torch.randn( + 16, device="cuda", dtype=torch.bfloat16, requires_grad=True + ), + torch.randn(8, 16, device="cuda", dtype=torch.bfloat16), + ), + "autograd", + ), + ], +) +def test_bf16_gemv_rejects_unsupported_inputs(make_args, error): + with pytest.raises(RuntimeError, match=error): + bf16_gemv(*make_args()) diff --git a/tests/extension/test_linear_dispatch.py b/tests/extension/test_linear_dispatch.py new file mode 100644 index 0000000..2b6274b --- /dev/null +++ b/tests/extension/test_linear_dispatch.py @@ -0,0 +1,185 @@ +import logging + +import pytest +import torch +import torch.nn.functional as F + +from astrai.extension import explain, is_available, linear, op_backend +from astrai.extension.backend import linear as public_linear +from astrai.model.components.linear import Linear + +GEMV_AVAILABLE = ( + torch.cuda.is_available() + and is_available("bf16_gemv") + and torch.cuda.get_device_capability() >= (8, 0) +) +skip_no_gemv = pytest.mark.skipif( + not GEMV_AVAILABLE, + reason="BF16 GEMV requires a built kernel and compute capability 8.0+", +) + + +def test_linear_backend_is_public(): + assert linear is public_linear + + +def test_model_linear_routes_through_backend(monkeypatch): + sentinel = torch.randn(2, 4) + + def fake_linear(x, weight, bias): + assert x.shape == (2, 3) + assert weight.shape == (4, 3) + assert bias is None + return sentinel + + monkeypatch.setattr("astrai.model.components.linear.linear", fake_linear) + layer = Linear(3, 4) + assert layer(torch.randn(2, 3)) is sentinel + + +def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "1") + x = torch.randn(2, 8, requires_grad=True) + weight = torch.randn(4, 8, requires_grad=True) + actual = linear(x, weight) + expected = F.linear(x, weight) + torch.testing.assert_close(actual, expected) + actual.sum().backward() + assert x.grad is not None + assert weight.grad is not None + + +def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog): + monkeypatch.setenv("ASTRAI_GEMV", "invalid-test-mode") + with caplog.at_level(logging.WARNING): + trace = explain("linear", torch.randn(1, 8), torch.randn(4, 8)) + assert "using auto" in caplog.text + assert "=> torch" in trace + + +@skip_no_gemv +def test_mode_zero_disables_gemv(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "0") + x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + assert "=> torch" in explain("linear", x, weight) + torch.testing.assert_close(linear(x, weight), F.linear(x, weight)) + + +@skip_no_gemv +def test_mode_one_forces_capable_unmeasured_shape(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "1") + x = torch.randn(1, 64, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + assert "=> gemv" in explain("linear", x, weight) + torch.testing.assert_close( + linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25 + ) + + +@skip_no_gemv +@pytest.mark.parametrize("m", [2, 4, 8]) +def test_mode_one_dispatches_supported_small_batches(monkeypatch, m): + monkeypatch.setenv("ASTRAI_GEMV", "1") + x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + assert "=> gemv" in explain("linear", x, weight) + torch.testing.assert_close( + linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25 + ) + + +@skip_no_gemv +def test_auto_m1_falls_back_until_end_to_end_gate_passes(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "auto") + x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) + winning = torch.randn( + 1536, + 1536, + device="cuda", + dtype=torch.bfloat16, + requires_grad=True, + ) + with torch.no_grad(): + assert "=> torch" in explain("linear", x, winning) + torch.testing.assert_close(linear(x, winning), F.linear(x, winning)) + + +@skip_no_gemv +def test_auto_selects_measured_sm89_small_batch_winner(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "auto") + x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + trace = explain("linear", x, weight) + if torch.cuda.get_device_capability() == (8, 9): + assert "=> auto_gemv" in trace + else: + assert "=> torch" in trace + torch.testing.assert_close( + linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.5 + ) + + +@skip_no_gemv +@pytest.mark.parametrize( + "m,n,k", + [ + (2, 6912, 1536), # up/gate loses at every measured M + (2, 1536, 6912), # long-K accumulation changed checkpoint greedy output + (4, 100000, 1536), # LM head misses the 5% M=4 gate + (4, 1536, 6912), # long-K accumulation changed checkpoint greedy output + (8, 256, 1536), # remaining M=8 winners miss the 3% end-to-end gate + ], +) +def test_auto_rejects_measured_small_batch_losers(monkeypatch, m, n, k): + monkeypatch.setenv("ASTRAI_GEMV", "auto") + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + assert "=> torch" in explain("linear", x, weight) + + +@skip_no_gemv +def test_grad_enabled_and_unsupported_multirow_always_fall_back(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "1") + x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn( + 256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + assert "=> torch" in explain("linear", x, weight) + with torch.no_grad(): + multirow = x.expand(3, -1).contiguous() + assert "=> torch" in explain("linear", multirow, weight) + + +@skip_no_gemv +def test_explicit_gemv_selection_respects_capability(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "0") + x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(), op_backend(linear="gemv"): + assert "=> gemv" in explain("linear", x, weight) + torch.testing.assert_close( + linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25 + ) + + +@skip_no_gemv +def test_dispatched_linear_cuda_graph_replay(monkeypatch): + monkeypatch.setenv("ASTRAI_GEMV", "1") + x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) + weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) + with torch.no_grad(): + for _ in range(3): + linear(x, weight) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = linear(x, weight) + x.copy_(torch.randn_like(x)) + graph.replay() + expected = F.linear(x, weight) + torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)