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
This commit is contained in:
0z5a
@@ -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",
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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"]
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,206 @@
|
||||
// Directly callable small-M BF16 GEMV primitive for decode-time linear layers.
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
|
||||
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 <int Rows>
|
||||
__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<const __nv_bfloat162*>(x);
|
||||
const auto* w2 =
|
||||
reinterpret_cast<const __nv_bfloat162*>(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 <int Rows>
|
||||
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<Rows><<<n, kThreads, 0, stream>>>(
|
||||
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<int>::max() &&
|
||||
n <= std::numeric_limits<int>::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::Tensor>();
|
||||
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<const __nv_bfloat16*>(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<const __nv_bfloat16*>(x.data_ptr());
|
||||
const auto* weight_ptr =
|
||||
reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr());
|
||||
auto* output_ptr = reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
|
||||
const int n_int = static_cast<int>(n);
|
||||
const int k_int = static_cast<int>(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"
|
||||
);
|
||||
}
|
||||
@@ -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/
|
||||
|
||||
@@ -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.
|
||||
@@ -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()
|
||||
@@ -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"))]
|
||||
|
||||
@@ -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())
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user