ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: dispatch linear gemv by decode batch size and unify extension style

Browse Source 
- 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].
This commit is contained in:
ViperEkura
2026-09-03 07:23:13 +08:00
parent 27abb7c5e7
commit 7540acb43e
14 changed files with 225 additions and 536 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/backend/linear.py
+18 -249
View File
@@ -1,191 +1,28 @@
"""Inference-only dispatch for AstrAI linear layers.
The CUDA GEMV path is deliberately narrow: automatic selection is enabled
only for small decode batches and BF16 shapes measured to beat ``F.linear``
on a supported architecture. Every training, prefill, unsupported-layout,
and unmeasured call falls back to PyTorch.
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.
"""
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.dispatch import env_mode
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 projection-chain/engine
# >=3% gates pass and output argmax remains stable. M=1 is limited to OPT 1.3B;
# M=8 remains empty because at least one projection in each measured family
# misses the per-shape gate even when its aggregate chain result is positive.
_COMMON_TRANSFORMER_SM89_SHAPES = frozenset(
{
(1024, 4096), # LLaMA 3 8B K/V
(4096, 4096), # LLaMA 2/3 7B/8B Q/O
(11008, 4096), # LLaMA 2 7B gate/up
(4096, 11008), # LLaMA 2 7B down
(14336, 4096), # LLaMA 3 8B gate/up
(4096, 14336), # LLaMA 3 8B down
(5120, 5120), # LLaMA 2 13B Q/K/V/O
(13824, 5120), # LLaMA 2 13B gate/up
(5120, 13824), # LLaMA 2 13B down
(16384, 4096), # GPT-NeoX MLP up
(4096, 16384), # GPT-NeoX MLP down
}
)
_COMMON_TRANSFORMER_SM89_M4_SHAPES = _COMMON_TRANSFORMER_SM89_SHAPES - {
(4096, 4096),
(11008, 4096),
(4096, 11008),
}
_QWEN2_7B_SM89_SHAPES = frozenset(
{
(512, 3584), # K/V
(3584, 3584), # Q/O
(18944, 3584), # gate/up
(3584, 18944), # down
}
)
_LLAMA3_70B_SM89_SHAPES = frozenset(
{
(1024, 8192), # K/V
(8192, 8192), # Q/O
(28672, 8192), # gate/up
(8192, 28672), # down
}
)
_OPT_1_3B_SM89_SHAPES = frozenset(
{
(2048, 2048), # Q/K/V/O
(8192, 2048), # MLP up
(2048, 8192), # MLP down
}
)
_AUTO_GEMV_SHAPES: dict[tuple[int, int], dict[int, frozenset[tuple[int, int]]]] = {
(8, 9): {
1: _OPT_1_3B_SM89_SHAPES,
2: _COMMON_TRANSFORMER_SM89_SHAPES
| _QWEN2_7B_SM89_SHAPES
| _LLAMA3_70B_SM89_SHAPES
| _OPT_1_3B_SM89_SHAPES
| frozenset(
{
(256, 1536),
(1536, 1536),
(100000, 1536),
}
),
4: _COMMON_TRANSFORMER_SM89_M4_SHAPES
| _QWEN2_7B_SM89_SHAPES
| _LLAMA3_70B_SM89_SHAPES
| 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 is not None and 1 <= m <= 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,
)
_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()
)
_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)
# 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(
@@ -201,11 +38,6 @@ def _inference_bf16_gemv(
)
@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()
@@ -219,6 +51,7 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
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
@@ -231,83 +64,19 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
)
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 any M in ``{1, ..., 8}``, and ``auto`` (the
default) uses only architecture/shape bands backed by benchmark evidence.
default) uses GEMV for decode batches with M in ``{2, 3, 4}``.
"""
# 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)
mode = env_mode("ASTRAI_GEMV")
if mode != "0" and _gemv_capable(x, weight, bias):
if mode == "1" or _auto_gemv_shape(x, weight):
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 _torch_linear(x, weight, bias)
return F.linear(x, weight, bias)
__all__ = ["linear"]
astrai/extension/backend/rotary.py
+15 -7
View File
@@ -9,6 +9,8 @@ Layout: x is [batch, seq_len, n_heads, head_dim] (bf16).
freqs_cis is [batch, seq_len, dim/2, 2] (f32) — [cos, sin] pairs.
"""
from typing import Any, Dict, List
import torch
from torch import Tensor
@@ -41,7 +43,7 @@ def _torch_apply(x: Tensor, freqs_cis: Tensor) -> Tensor:
return x_out.to(dtype)
def _rotary_records() -> list:
def _rotary_records() -> List[ImplRecord]:
return [
ImplRecord(
family="rotary",
@@ -61,12 +63,15 @@ def _rotary_records() -> list:
]
register_family(
"rotary",
lambda x, freqs_cis: tensor_axes(x),
_rotary_records,
lambda: _rotary_records()[-1],
)
def _axes(x: Tensor, freqs_cis: Tensor) -> Dict[str, Any]:
return tensor_axes(x)
def _fallback_record() -> ImplRecord:
return _rotary_records()[-1]
register_family("rotary", _axes, _rotary_records, _fallback_record)
def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
@@ -80,3 +85,6 @@ def apply_rotary_emb(x: Tensor, freqs_cis: Tensor) -> Tensor:
[batch, seq_len, n_heads, head_dim] (bf16)
"""
return resolve("rotary", x, freqs_cis).record.obj(x, freqs_cis)
__all__ = ["apply_rotary_emb"]
astrai/extension/backend/swiglu.py
+10 -57
View File
@@ -1,46 +1,18 @@
"""Inference-only fused SwiGLU selection for dense MLP layers."""
import logging
import os
from functools import cache
import torch
import torch.nn.functional as F
from torch import Tensor
from astrai.extension.backend.linear import linear
from astrai.extension.dispatch import env_mode
from astrai.extension.loader import is_available
from astrai.extension.ops.swiglu import bf16_swiglu
logger = logging.getLogger(__name__)
# Shape keys are (N, K) for the paired up/gate projections. Automatic entries
# are populated only after the primitive, MLP chain, and greedy checkpoint
# gates pass on that architecture.
_AUTO_SWIGLU_SHAPES: dict[tuple[int, int], dict[int, frozenset[tuple[int, int]]]] = {}
_AUTO_SWIGLU_M = frozenset(
m for architecture in _AUTO_SWIGLU_SHAPES.values() for m in architecture
)
_VALID_MODES = {"0", "1", "auto"}
_WARNED_MODES: set[str] = set()
def _swiglu_mode() -> str:
mode = os.environ.get("ASTRAI_SWIGLU", "auto").strip().lower()
if mode in _VALID_MODES:
return mode
if mode not in _WARNED_MODES:
_WARNED_MODES.add(mode)
logger.warning(
"ASTRAI_SWIGLU=%r is invalid; expected 0, 1, or auto; using auto",
mode,
)
return "auto"
def _unfused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
# Keep the existing linear backend in the fallback chain. This preserves
# any independently qualified GEMV shapes instead of making the fusion
# any independently qualified GEMV batches instead of making the fusion
# decision suppress linear-level optimizations.
return linear(x, up_weight) * F.silu(linear(x, gate_weight))
@@ -49,11 +21,6 @@ def _fused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
return bf16_swiglu(x.detach(), up_weight.detach(), gate_weight.detach())
@cache
def _device_capability(device_index: int) -> tuple[int, int]:
return torch.cuda.get_device_capability(device_index)
def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
return not (
torch.is_grad_enabled()
@@ -77,33 +44,19 @@ def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
)
def _auto_swiglu_shape(x: Tensor, up_weight: Tensor) -> bool:
capability = _device_capability(x.get_device())
m = 1 if x.ndim == 1 else x.shape[0]
return (up_weight.shape[0], up_weight.shape[1]) in _AUTO_SWIGLU_SHAPES.get(
capability, {}
).get(m, ())
def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
"""Apply the dense-MLP SwiGLU projection with a safe torch fallback.
``ASTRAI_SWIGLU=0`` keeps the unfused linear-backend chain, ``1`` forces
the fused primitive for supported inputs, and ``auto`` uses only
architecture/shape bands backed by benchmark and checkpoint evidence.
``ASTRAI_SWIGLU=0`` and ``auto`` keep the unfused linear-backend chain;
``1`` forces the fused primitive for supported inputs. Auto will adopt
an M-banded rule mirroring the linear backend once end-to-end evidence
qualifies one.
"""
mode = _swiglu_mode()
if mode == "0" or (mode == "auto" and not _AUTO_SWIGLU_SHAPES):
return _unfused_swiglu(x, up_weight, gate_weight)
if mode == "auto":
m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
if m not in _AUTO_SWIGLU_M:
return _unfused_swiglu(x, up_weight, gate_weight)
if _swiglu_capable(x, up_weight, gate_weight) and (
mode == "1" or _auto_swiglu_shape(x, up_weight)
if env_mode("ASTRAI_SWIGLU") != "1" or not _swiglu_capable(
x, up_weight, gate_weight
):
return _fused_swiglu(x, up_weight, gate_weight)
return _unfused_swiglu(x, up_weight, gate_weight)
return _unfused_swiglu(x, up_weight, gate_weight)
return _fused_swiglu(x, up_weight, gate_weight)
__all__ = ["swiglu"]