ViperEkura
7540acb43e
- replace the per-shape auto tables in the linear backend with an M-banded rule (M in [2,4] on compute capability 8.0+) that measured at the HBM bandwidth floor across every family, and fold the capability check into the capable guard - drop the unreachable swiglu auto shape-table machinery so both backends share one env-mode ladder via the new dispatch.env_mode helper - add __all__ across extension modules, name the rotary registration records, and unify typing to the typing-module style - rewrite test_linear_dispatch.py around behavioral routing assertions and document the M-banded policy in the developer docs - Benchmark: L20 SM89, Python dispatch overhead 2.9us to 1.5us, auto now covers every projection shape at M in [2,4].
83 lines
2.9 KiB
Python
83 lines
2.9 KiB
Python
"""Inference-only dispatch for AstrAI linear layers.
|
|
|
|
The CUDA GEMV path is narrow by construction rather than by a measured
|
|
shape table: the kernel streams each weight exactly once, so automatic
|
|
selection is keyed on the decode batch size alone (M in [2, 4], where it
|
|
sits at the HBM bandwidth floor and beat the cuBLAS small-M path on every
|
|
measured family). Every training, prefill-sized, out-of-band, or
|
|
unsupported call falls back to PyTorch.
|
|
"""
|
|
|
|
from typing import Optional
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import Tensor
|
|
|
|
from astrai.extension.dispatch import env_mode
|
|
from astrai.extension.loader import is_available
|
|
from astrai.extension.ops.gemv import bf16_gemv
|
|
|
|
# M=1 keeps cuBLAS (its GEMV path is already at the bandwidth floor; only
|
|
# OPT 1.3B shapes ever passed the full gate). M >= 5 approaches the cuBLAS
|
|
# tensor-core crossover (M=8 regressed at wrapper level on every measured
|
|
# family, and cuBLAS clearly wins from M ~ 12).
|
|
_AUTO_GEMV_M = frozenset({2, 3, 4})
|
|
|
|
|
|
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,
|
|
)
|
|
|
|
|
|
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 not 1 <= x.shape[0] <= 8)
|
|
or x.shape[-1] != weight.shape[1]
|
|
or x.device != weight.device
|
|
or not x.is_contiguous()
|
|
or not weight.is_contiguous()
|
|
or torch.cuda.get_device_capability(x.get_device()) < (8, 0)
|
|
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 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 any M in ``{1, ..., 8}``, and ``auto`` (the
|
|
default) uses GEMV for decode batches with M in ``{2, 3, 4}``.
|
|
"""
|
|
mode = env_mode("ASTRAI_GEMV")
|
|
if mode != "0" and _gemv_capable(x, weight, bias):
|
|
m = 1 if x.ndim == 1 else x.shape[0]
|
|
if mode == "1" or m in _AUTO_GEMV_M:
|
|
return _inference_bf16_gemv(x, weight, bias)
|
|
return F.linear(x, weight, bias)
|
|
|
|
|
|
__all__ = ["linear"]
|