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

refactor: remove bf16 gemm and swiglu kernels and rebuild csrc benchmarks

Browse Source 
- delete csrc/kernels/gemm.cu and swiglu.cu and drop their CMake and setup.py registration
- remove the ops wrappers plus backend/linear.py and backend/swiglu.py so Linear and MLP call F.linear directly
- drop the four gemm and swiglu kernel test files and prune the stale cuda_kernels.md sections
- add csrc/bench benchmarks for the remaining kernels: attention decode prefill paged decode paged prefill versus single-launch SDPA references, rotary versus the torch fallback, fp8 quantize and mm_fp8 versus torch baselines
- attention, rotary_emb, and fp8_ops kernels are unchanged
This commit is contained in:
ViperEkura
2026-09-05 01:38:10 +08:00
parent a77e35dd51
commit 6709534d64
24 changed files with 1306 additions and 3394 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/__init__.py
+7 -11
View File
@@ -1,11 +1,10 @@
"""CUDA kernel wrappers, operator dispatch, and backend selection.
Public API:
- ``attention``, ``linear``, ``swiglu``, ``apply_rotary_emb`` — op
families with safe torch fallbacks (see ``astrai.extension.backend``)
- ``attention``, ``apply_rotary_emb`` — op families with safe torch
fallbacks (see ``astrai.extension.backend``)
- ``attn_decode`` / ``attn_prefill`` / ``attn_paged_decode`` /
``attn_paged_prefill`` — direct attention kernel wrappers
- ``bf16_gemm`` / ``bf16_swiglu`` — directly callable linear/MLP kernels
- ``AttentionBackend`` / ``TorchNativeBackend`` / ``CudaBackend`` /
``FlashAttnBackend`` — attention backend strategies
- ``resolve`` / ``explain`` / ``op_backend`` / ``env_mode`` — the shared
@@ -14,6 +13,9 @@ Public API:
Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
(blhd). Scale is always ``1/sqrt(head_dim)``. Wrapper functions call their
compiled CUDA kernels directly; fallback is the backend's responsibility.
Linear projections and dense-MLP SwiGLU run plain torch (``F.linear`` /
``Linear`` / ``MLP``); the former bf16_gemm / bf16_swiglu kernels and their
backends were removed.
"""
from astrai.extension.backend import (
@@ -27,8 +29,6 @@ from astrai.extension.backend import (
attention,
attn_backend,
get_backend,
linear,
swiglu,
)
from astrai.extension.dispatch import (
Axes,
@@ -52,9 +52,8 @@ from astrai.extension.ops import (
TensorLayout,
attn_decode,
attn_paged_decode,
attn_paged_prefill,
attn_prefill,
bf16_gemm,
bf16_swiglu,
)
__all__ = [
@@ -68,13 +67,10 @@ __all__ = [
"attention",
"attn_backend",
"get_backend",
"linear",
"swiglu",
"attn_decode",
"attn_paged_decode",
"attn_prefill",
"bf16_gemm",
"bf16_swiglu",
"attn_paged_prefill",
"is_available",
"KERNEL_NAMES",
"apply_rotary_emb",
astrai/extension/backend/__init__.py
-4
View File
@@ -11,9 +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
from astrai.extension.backend.swiglu import swiglu
__all__ = [
"ATTN_BACKEND",
@@ -26,6 +24,4 @@ __all__ = [
"attention",
"attn_backend",
"get_backend",
"linear",
"swiglu",
]
astrai/extension/backend/linear.py
-215
View File
@@ -1,215 +0,0 @@
"""Inference-only dispatch for AstrAI linear layers.
The CUDA GEMM path is sized by the decode batch M. M in [1, 8] uses the
register-resident GEMV kernel (any K); M in (8, 64] uses the tiled
kernel (K % 8 == 0, 16-byte-aligned tensors). Automatic mode
selects GEMM for M in [1, 64] where the primitive is capable; every
training, prefill-sized, or unsupported call falls back to PyTorch.
The family stays registered with the shared operator dispatcher, so
``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` /
``explain`` keep working. The per-layer hot path only consults the
dispatcher when one of those selections is active.
"""
from typing import Any, Dict, List, Optional
import torch
import torch.nn.functional as F
from torch import Tensor
from astrai.extension.dispatch import (
ImplRecord,
Spec,
axis,
env_mode,
env_selection,
get_override,
register_family,
resolve,
tensor_axes,
)
from astrai.extension.loader import is_available
from astrai.extension.ops.gemm import bf16_gemm
def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
return F.linear(x, weight, bias)
def _inference_bf16_gemm(
x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
) -> Tensor:
return bf16_gemm(
x.detach(),
weight.detach(),
bias.detach() if bias is not None else None,
)
def _gemm_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
"""Check whether bf16_gemm can safely handle the call."""
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.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_gemm")
):
return False
m = 1 if x.ndim == 1 else x.shape[0]
# M <= 32 is where the kernel wins: L2-rotation measurements on L20
# show every production shape at M=24-32 winning or tying, while
# M=48-64 loses the long-K down_proj by 8-10% (cuBLAS switches to a
# wider tile there). The kernel itself still accepts M <= 64 when
# called directly through astrai.extension.ops.gemm.
if not (1 <= m <= 32):
return False
# Vocabulary-sized lm_head weights (N in the tens of thousands+) stream
# better through cuBLAS: our skinny path ties it at M<=8 and the tiled
# path loses ~4% at M=9-16 (L2-rotation measurements on L20). Gate the
# whole shape family out instead of splitting hairs per M band.
if weight.shape[0] > 32768:
return False
# M > 8 tiled path requires K % 8 == 0 and 16-byte alignment.
k = x.shape[-1]
if m > 8 and (
k % 8 != 0 or (x.data_ptr() & 15) != 0 or (weight.data_ptr() & 15) != 0
):
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 _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]:
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 <= 64
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
return tensor_axes(
x,
mode=env_mode("ASTRAI_GEMM"),
m=m,
supported_m=supported_m,
auto_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,
capability=capability,
)
_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.of(
lambda ax: ax.get("capability") is not None and ax.get("capability") >= (8, 0),
"capability>=sm_80",
)
)
_SPEC_AUTO = _SPEC_CAPABLE & axis("auto_m").truthy()
def _linear_records() -> List[ImplRecord]:
mode = env_mode("ASTRAI_GEMM")
gemm_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="gemm",
obj=_inference_bf16_gemm,
spec=_SPEC_CAPABLE,
available=lambda: is_available("bf16_gemm"),
priority=gemm_priority,
),
ImplRecord(
family="linear",
name="auto_gemm",
obj=_inference_bf16_gemm,
spec=_SPEC_AUTO,
available=lambda: is_available("bf16_gemm"),
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 GEMM dispatch.
``ASTRAI_GEMM=0`` always uses PyTorch, ``1`` forces GEMM whenever the
primitive can safely handle the call (M in [1, 64], K % 8 == 0 and
16-byte-aligned for M > 8), and ``auto`` (the default) selects GEMM
for all capable decode batches.
"""
if get_override("linear") is not None or env_selection("linear") is not None:
return resolve("linear", x, weight, bias).record.obj(x, weight, bias)
mode = env_mode("ASTRAI_GEMM")
if mode != "0" and _gemm_capable(x, weight, bias):
return _inference_bf16_gemm(x, weight, bias)
return _torch_linear(x, weight, bias)
__all__ = ["linear"]
astrai/extension/backend/swiglu.py
-67
View File
@@ -1,67 +0,0 @@
"""Inference-only fused SwiGLU selection for dense MLP layers."""
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
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 batches instead of making the fusion
# decision suppress linear-level optimizations.
return linear(x, up_weight) * F.silu(linear(x, gate_weight))
def _fused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
return bf16_swiglu(x.detach(), up_weight.detach(), gate_weight.detach())
def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
return not (
torch.is_grad_enabled()
or not x.is_cuda
or x.dtype != torch.bfloat16
or up_weight.dtype != torch.bfloat16
or gate_weight.dtype != torch.bfloat16
or x.ndim not in (1, 2)
or up_weight.ndim != 2
or gate_weight.ndim != 2
or (x.ndim == 2 and not 1 <= x.shape[0] <= 8)
or up_weight.shape != gate_weight.shape
or x.shape[-1] != up_weight.shape[1]
or x.shape[-1] % 8 != 0
or x.device != up_weight.device
or x.device != gate_weight.device
or not x.is_contiguous()
or not up_weight.is_contiguous()
or not gate_weight.is_contiguous()
# The fused kernel reads all streams as uint4; contiguous-but-offset
# views are routed to the unfused chain instead of failing.
or (x.data_ptr() & 15) != 0
or (up_weight.data_ptr() & 15) != 0
or (gate_weight.data_ptr() & 15) != 0
or not is_available("bf16_swiglu")
)
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; ``auto`` (the default) uses
the fused primitive for decode batches with M in ``{1, ..., 8}``. The
fused kernel reads x once and covers both projections plus the SiLU
gate-multiply in a single launch, measured 9-15% faster than the
unfused chain per MLP call on L20 with L2-thrashing weight rotation.
"""
if env_mode("ASTRAI_SWIGLU") != "0" and _swiglu_capable(x, up_weight, gate_weight):
return _fused_swiglu(x, up_weight, gate_weight)
return _unfused_swiglu(x, up_weight, gate_weight)
__all__ = ["swiglu"]
astrai/extension/ops/__init__.py
-4
View File
@@ -7,9 +7,7 @@ from astrai.extension.ops.attention import (
attn_paged_prefill,
attn_prefill,
)
from astrai.extension.ops.gemm import bf16_gemm
from astrai.extension.ops.rotary import rotary_emb
from astrai.extension.ops.swiglu import bf16_swiglu
__all__ = [
"TensorLayout",
@@ -17,7 +15,5 @@ __all__ = [
"attn_paged_decode",
"attn_paged_prefill",
"attn_prefill",
"bf16_gemm",
"bf16_swiglu",
"rotary_emb",
]
astrai/extension/ops/gemm.py
-27
View File
@@ -1,27 +0,0 @@
"""Stateless wrapper for the directly callable BF16 GEMM primitive."""
from typing import Optional
import torch
from astrai.extension.loader import get_module
def bf16_gemm(
x: torch.Tensor,
weight: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Compute ``F.linear(x, weight, bias)`` for up to 64 BF16 rows.
``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``[1, 64]``, and
``weight`` must be a contiguous row-major ``[N, K]`` tensor. M in
``[1, 8]`` uses the register-resident skinny GEMM kernel (any K);
larger M uses the tiled kernel (K must be a multiple of 8 with
16-byte-aligned tensors). This primitive is inference-only and
intentionally performs no fallback or model-level dispatch.
"""
return get_module("bf16_gemm").bf16_gemm(x, weight, bias)
__all__ = ["bf16_gemm"]
astrai/extension/ops/swiglu.py
-22
View File
@@ -1,22 +0,0 @@
"""Stateless wrapper for the directly callable fused BF16 SwiGLU primitive."""
import torch
from astrai.extension.loader import get_module
def bf16_swiglu(
x: torch.Tensor,
up_weight: torch.Tensor,
gate_weight: torch.Tensor,
) -> torch.Tensor:
"""Compute ``linear(x, up) * silu(linear(x, gate))`` for M in [1, 8].
Inputs must be contiguous BF16 CUDA tensors. Both weights use row-major
``[N, K]`` storage with identical shapes, and K must be divisible by 8.
The primitive is inference-only and performs no fallback.
"""
return get_module("bf16_swiglu").bf16_swiglu(x, up_weight, gate_weight)
__all__ = ["bf16_swiglu"]