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

perf: rebuild decode gemm dispatch around shape-driven tile configs

Browse Source 
- split-K removed entirely: tiled kernel walks K in one pass, no partials/semas workspace, no memset, single launch per call
- skinny GEMM (M<=8) dispatch table replaces the hand-written switch
- shape-driven four-family table replaces plan_gemm: wide-N (n>=4096) default {16,64,64,3,128} with BM=32 at M>16; narrow-N deep-K rings {16,32,256,2,64} while the grid fits one wave, {16,32,128,2,64} past it
- narrow-N is K-serial: widening the grid measurably does nothing (BN 64->32 ties, doubled m_tiles tie, kv at 4 blocks ties q/o at 24); deeper K chunks win until 72KB smem forces one CTA per SM and past one wave the 2-wave quantization loses to BK=128
- launch-check macros in common/launch.cuh; smem opt-in for the 72KB/60KB rings
- rename kernels/bf16_*.cu to gemm.cu/swiglu.cu; module names unchanged
- Python gate: lm_head (N>32768) falls back to cuBLAS, band narrows to M<=32
- drop the stale per-op benchmark narratives; fold the live numbers into cuda_kernels.md

Benchmark: NVIDIA L20 (sm_89, 92 SMs), CUDA 12.8, bf16, L2-thrash weight rotation, per-call medians at M=16: q/o 9.5us, kv 8.6us, gate/up 33.3us, down 33.7us (down -29% vs prior default). End-to-end 1B decode (gen 128, 3 trials, tokens/s vs cuBLAS): B=1 260 vs 252, B=8 1660 vs 1446, B=16 2464 vs 2437, B=32 3620 vs 3690. Prior split-K dispatch measured B=16 2243 / B=32 3393.
This commit is contained in:
ViperEkura
2026-09-04 22:41:39 +08:00
parent 8e39d9d8c9
commit 1798474316
21 changed files with 1098 additions and 697 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/__init__.py
+3 -3
View File
@@ -5,7 +5,7 @@ Public API:
families with safe torch fallbacks (see ``astrai.extension.backend``) families with safe torch fallbacks (see ``astrai.extension.backend``)
- ``attn_decode`` / ``attn_prefill`` / ``attn_paged_decode`` / - ``attn_decode`` / ``attn_prefill`` / ``attn_paged_decode`` /
``attn_paged_prefill`` — direct attention kernel wrappers ``attn_paged_prefill`` — direct attention kernel wrappers
- ``bf16_gemv`` / ``bf16_swiglu`` — directly callable linear/MLP kernels - ``bf16_gemm`` / ``bf16_swiglu`` — directly callable linear/MLP kernels
- ``AttentionBackend`` / ``TorchNativeBackend`` / ``CudaBackend`` / - ``AttentionBackend`` / ``TorchNativeBackend`` / ``CudaBackend`` /
``FlashAttnBackend`` — attention backend strategies ``FlashAttnBackend`` — attention backend strategies
- ``resolve`` / ``explain`` / ``op_backend`` / ``env_mode`` — the shared - ``resolve`` / ``explain`` / ``op_backend`` / ``env_mode`` — the shared
@@ -53,7 +53,7 @@ from astrai.extension.ops import (
attn_decode, attn_decode,
attn_paged_decode, attn_paged_decode,
attn_prefill, attn_prefill,
bf16_gemv, bf16_gemm,
bf16_swiglu, bf16_swiglu,
) )
@@ -73,7 +73,7 @@ __all__ = [
"attn_decode", "attn_decode",
"attn_paged_decode", "attn_paged_decode",
"attn_prefill", "attn_prefill",
"bf16_gemv", "bf16_gemm",
"bf16_swiglu", "bf16_swiglu",
"is_available", "is_available",
"KERNEL_NAMES", "KERNEL_NAMES",
astrai/extension/backend/linear.py
+53 -48
View File
@@ -1,17 +1,15 @@
"""Inference-only dispatch for AstrAI linear layers. """Inference-only dispatch for AstrAI linear layers.
The CUDA GEMV path is narrow by construction rather than by a measured The CUDA GEMM path is sized by the decode batch M. M in [1, 8] uses the
shape table: the kernel streams each weight exactly once, so automatic register-resident GEMV kernel (any K); M in (8, 64] uses the tiled
selection is keyed on the decode batch size alone (M in [2, 4], where it kernel (K % 8 == 0, 16-byte-aligned tensors). Automatic mode
sits at the HBM bandwidth floor and beat the cuBLAS small-M path on every selects GEMM for M in [1, 64] where the primitive is capable; every
measured family). Every training, prefill-sized, out-of-band, or training, prefill-sized, or unsupported call falls back to PyTorch.
unsupported call falls back to PyTorch.
The family stays registered with the shared operator dispatcher, so The family stays registered with the shared operator dispatcher, so
``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` / ``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` /
``explain`` keep working like for attention and rotary. The per-layer ``explain`` keep working. The per-layer hot path only consults the
hot path only consults the dispatcher when one of those selections is dispatcher when one of those selections is active.
active, keeping it free of axes dictionaries and record sorting.
""" """
from typing import Any, Dict, List, Optional from typing import Any, Dict, List, Optional
@@ -32,33 +30,25 @@ from astrai.extension.dispatch import (
tensor_axes, tensor_axes,
) )
from astrai.extension.loader import is_available from astrai.extension.loader import is_available
from astrai.extension.ops.gemv import bf16_gemv from astrai.extension.ops.gemm import bf16_gemm
# 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 _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor: def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
return F.linear(x, weight, bias) return F.linear(x, weight, bias)
def _inference_bf16_gemv( def _inference_bf16_gemm(
x: Tensor, weight: Tensor, bias: Optional[Tensor] = None x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
) -> Tensor: ) -> Tensor:
# Model parameters retain requires_grad=True after eval(). Dispatch is return bf16_gemm(
# 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(), x.detach(),
weight.detach(), weight.detach(),
bias.detach() if bias is not None else None, bias.detach() if bias is not None else None,
) )
def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool: def _gemm_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
"""Check whether bf16_gemm can safely handle the call."""
if ( if (
torch.is_grad_enabled() torch.is_grad_enabled()
or not x.is_cuda or not x.is_cuda
@@ -66,13 +56,32 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
or weight.dtype != torch.bfloat16 or weight.dtype != torch.bfloat16
or weight.ndim != 2 or weight.ndim != 2
or x.ndim not in (1, 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.shape[-1] != weight.shape[1]
or x.device != weight.device or x.device != weight.device
or not x.is_contiguous() or not x.is_contiguous()
or not weight.is_contiguous() or not weight.is_contiguous()
or torch.cuda.get_device_capability(x.get_device()) < (8, 0) or torch.cuda.get_device_capability(x.get_device()) < (8, 0)
or not is_available("bf16_gemv") 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 False
return bias is None or ( return bias is None or (
@@ -87,7 +96,7 @@ def _gemv_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]: def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]:
weight_shape = tuple(weight.shape) weight_shape = tuple(weight.shape)
m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None) 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 supported_m = m is not None and 1 <= m <= 64
shape_matches = ( shape_matches = (
weight.ndim == 2 weight.ndim == 2
and x.ndim in (1, 2) and x.ndim in (1, 2)
@@ -107,10 +116,10 @@ def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str,
capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None
return tensor_axes( return tensor_axes(
x, x,
mode=env_mode("ASTRAI_GEMV"), mode=env_mode("ASTRAI_GEMM"),
m=m, m=m,
supported_m=supported_m, supported_m=supported_m,
auto_m=m in _AUTO_GEMV_M, auto_m=supported_m,
shape_matches=shape_matches, shape_matches=shape_matches,
same_device=same_device, same_device=same_device,
weight_dtype=weight.dtype, weight_dtype=weight.dtype,
@@ -142,25 +151,25 @@ _SPEC_AUTO = _SPEC_CAPABLE & axis("auto_m").truthy()
def _linear_records() -> List[ImplRecord]: def _linear_records() -> List[ImplRecord]:
mode = env_mode("ASTRAI_GEMV") mode = env_mode("ASTRAI_GEMM")
gemv_priority = 0 if mode == "1" else 100 gemm_priority = 0 if mode == "1" else 100
auto_priority = 0 if mode == "auto" else 90 auto_priority = 0 if mode == "auto" else 90
torch_priority = 0 if mode == "0" else 50 torch_priority = 0 if mode == "0" else 50
return [ return [
ImplRecord( ImplRecord(
family="linear", family="linear",
name="gemv", name="gemm",
obj=_inference_bf16_gemv, obj=_inference_bf16_gemm,
spec=_SPEC_CAPABLE, spec=_SPEC_CAPABLE,
available=lambda: is_available("bf16_gemv"), available=lambda: is_available("bf16_gemm"),
priority=gemv_priority, priority=gemm_priority,
), ),
ImplRecord( ImplRecord(
family="linear", family="linear",
name="auto_gemv", name="auto_gemm",
obj=_inference_bf16_gemv, obj=_inference_bf16_gemm,
spec=_SPEC_AUTO, spec=_SPEC_AUTO,
available=lambda: is_available("bf16_gemv"), available=lambda: is_available("bf16_gemm"),
priority=auto_priority, priority=auto_priority,
), ),
ImplRecord( ImplRecord(
@@ -187,23 +196,19 @@ register_family("linear", _axes, _linear_records, _fallback_record)
def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor: def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
"""Apply a linear projection with safe inference-only GEMV dispatch. """Apply a linear projection with safe inference-only GEMM dispatch.
``ASTRAI_GEMV=0`` always uses PyTorch, ``1`` forces GEMV whenever the ``ASTRAI_GEMM=0`` always uses PyTorch, ``1`` forces GEMM whenever the
primitive can safely handle any M in ``{1, ..., 8}``, and ``auto`` (the primitive can safely handle the call (M in [1, 64], K % 8 == 0 and
default) uses GEMV for decode batches with M in ``{2, 3, 4}``. 16-byte-aligned for M > 8), and ``auto`` (the default) selects GEMM
for all capable decode batches.
""" """
# Route through the shared dispatcher whenever a selection is active so
# explicit/context/env overrides stay honored; otherwise keep the hot
# path free of axes dictionaries and record sorting.
if get_override("linear") is not None or env_selection("linear") is not None: 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) return resolve("linear", x, weight, bias).record.obj(x, weight, bias)
mode = env_mode("ASTRAI_GEMV") mode = env_mode("ASTRAI_GEMM")
if mode != "0" and _gemv_capable(x, weight, bias): if mode != "0" and _gemm_capable(x, weight, bias):
m = 1 if x.ndim == 1 else x.shape[0] return _inference_bf16_gemm(x, weight, bias)
if mode == "1" or m in _AUTO_GEMV_M:
return _inference_bf16_gemv(x, weight, bias)
return _torch_linear(x, weight, bias) return _torch_linear(x, weight, bias)
astrai/extension/backend/swiglu.py
+8 -8
View File
@@ -52,16 +52,16 @@ def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor: def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
"""Apply the dense-MLP SwiGLU projection with a safe torch fallback. """Apply the dense-MLP SwiGLU projection with a safe torch fallback.
``ASTRAI_SWIGLU=0`` and ``auto`` keep the unfused linear-backend chain; ``ASTRAI_SWIGLU=0`` keeps the unfused linear-backend chain; ``1`` forces
``1`` forces the fused primitive for supported inputs. Auto will adopt the fused primitive for supported inputs; ``auto`` (the default) uses
an M-banded rule mirroring the linear backend once end-to-end evidence the fused primitive for decode batches with M in ``{1, ..., 8}``. The
qualifies one. 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") != "1" or not _swiglu_capable( if env_mode("ASTRAI_SWIGLU") != "0" and _swiglu_capable(x, up_weight, gate_weight):
x, up_weight, gate_weight
):
return _unfused_swiglu(x, up_weight, gate_weight)
return _fused_swiglu(x, up_weight, gate_weight) return _fused_swiglu(x, up_weight, gate_weight)
return _unfused_swiglu(x, up_weight, gate_weight)
__all__ = ["swiglu"] __all__ = ["swiglu"]
astrai/extension/ops/__init__.py
+2 -2
View File
@@ -7,7 +7,7 @@ from astrai.extension.ops.attention import (
attn_paged_prefill, attn_paged_prefill,
attn_prefill, attn_prefill,
) )
from astrai.extension.ops.gemv import bf16_gemv from astrai.extension.ops.gemm import bf16_gemm
from astrai.extension.ops.rotary import rotary_emb from astrai.extension.ops.rotary import rotary_emb
from astrai.extension.ops.swiglu import bf16_swiglu from astrai.extension.ops.swiglu import bf16_swiglu
@@ -17,7 +17,7 @@ __all__ = [
"attn_paged_decode", "attn_paged_decode",
"attn_paged_prefill", "attn_paged_prefill",
"attn_prefill", "attn_prefill",
"bf16_gemv", "bf16_gemm",
"bf16_swiglu", "bf16_swiglu",
"rotary_emb", "rotary_emb",
] ]
astrai/extension/ops/gemm.py
+27
View File
@@ -0,0 +1,27 @@
"""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/gemv.py
-26
View File
@@ -1,26 +0,0 @@
"""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, 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)
__all__ = ["bf16_gemv"]
csrc/CMakeLists.txt
+3 -3
View File
@@ -61,7 +61,7 @@ set(KERNEL_NAMES
attn_prefill attn_prefill
attn_paged_decode attn_paged_decode
attn_paged_prefill attn_paged_prefill
bf16_gemv bf16_gemm
bf16_swiglu bf16_swiglu
rotary_emb rotary_emb
) )
@@ -70,8 +70,8 @@ set(KERNEL_SRCS
attention/prefill.cu attention/prefill.cu
attention/paged_decode.cu attention/paged_decode.cu
attention/paged_prefill.cu attention/paged_prefill.cu
bf16_gemv.cu gemm.cu
bf16_swiglu.cu swiglu.cu
rotary_emb.cu rotary_emb.cu
) )
csrc/bench/benchmark_gemv.py → csrc/bench/benchmark_gemm.py
+2 -2
View File
@@ -1,7 +1,7 @@
"""Benchmark decode-time linear shapes before enabling custom GEMV dispatch. """Benchmark decode-time linear shapes before enabling custom GEMM dispatch.
The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It
establishes the per-architecture cuBLAS baseline that later GEMV primitives and establishes the per-architecture cuBLAS baseline that later GEMM primitives and
dispatch decisions must beat. dispatch decisions must beat.
""" """
csrc/bench/benchmark_gemv_common.py → csrc/bench/benchmark_gemm_common.py
+6 -6
View File
@@ -1,4 +1,4 @@
"""Benchmark the BF16 GEMV primitive and guarded linear dispatcher. """Benchmark the BF16 GEMM primitive and guarded linear dispatcher.
The kernel suite covers AstrAI's native projections plus common LLaMA and The kernel suite covers AstrAI's native projections plus common LLaMA and
GPT-NeoX matrix shapes. The chain suite is a synthetic projection/MLP chain; GPT-NeoX matrix shapes. The chain suite is a synthetic projection/MLP chain;
@@ -19,7 +19,7 @@ from pathlib import Path
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
from astrai.extension import bf16_gemv, is_available, linear from astrai.extension import bf16_gemm, is_available, linear
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -198,7 +198,7 @@ def _kernel_functions(
return F.linear(x, weight) return F.linear(x, weight)
def candidate() -> torch.Tensor: def candidate() -> torch.Tensor:
return bf16_gemv(x, weight.detach()) return bf16_gemm(x, weight.detach())
return baseline, candidate return baseline, candidate
@@ -256,7 +256,7 @@ def benchmark_kernels(
def _set_mode(mode: str) -> None: def _set_mode(mode: str) -> None:
os.environ["ASTRAI_GEMV"] = mode os.environ["ASTRAI_GEMM"] = mode
def _chain_weights( def _chain_weights(
@@ -398,8 +398,8 @@ def parse_args() -> argparse.Namespace:
def main() -> None: def main() -> None:
args = parse_args() args = parse_args()
if not torch.cuda.is_available() or not is_available("bf16_gemv"): if not torch.cuda.is_available() or not is_available("bf16_gemm"):
raise RuntimeError("benchmark requires CUDA and the built bf16_gemv extension") raise RuntimeError("benchmark requires CUDA and the built bf16_gemm extension")
if args.warmup < 0 or args.samples < 1 or args.inner < 1 or args.chain_inner < 1: if args.warmup < 0 or args.samples < 1 or args.inner < 1 or args.chain_inner < 1:
raise ValueError("warmup must be non-negative and sample/inner counts positive") raise ValueError("warmup must be non-negative and sample/inner counts positive")
csrc/bench/benchmark_swiglu.py
+10 -10
View File
@@ -1,4 +1,4 @@
"""Benchmark fused BF16 SwiGLU against torch and unfused GEMV chains.""" """Benchmark fused BF16 SwiGLU against torch and unfused GEMM chains."""
from __future__ import annotations from __future__ import annotations
@@ -14,7 +14,7 @@ import click
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
from astrai.extension import bf16_gemv, bf16_swiglu, is_available from astrai.extension import bf16_gemm, bf16_swiglu, is_available
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -124,8 +124,8 @@ def capture(operation: Callable[[], torch.Tensor]):
def make_operations(x, up_weight, gate_weight, mode: str): def make_operations(x, up_weight, gate_weight, mode: str):
operations: dict[str, Callable[[], torch.Tensor]] = { operations: dict[str, Callable[[], torch.Tensor]] = {
"torch": lambda: F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)), "torch": lambda: F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)),
"gemv_chain": lambda: ( "gemm_chain": lambda: (
bf16_gemv(x, up_weight) * F.silu(bf16_gemv(x, gate_weight)) bf16_gemm(x, up_weight) * F.silu(bf16_gemm(x, gate_weight))
), ),
"fused": lambda: bf16_swiglu(x, up_weight, gate_weight), "fused": lambda: bf16_swiglu(x, up_weight, gate_weight),
} }
@@ -226,19 +226,19 @@ def render_markdown(payload: dict[str, object]) -> str:
f"- Compute capability: {metadata['compute_capability']}", f"- Compute capability: {metadata['compute_capability']}",
f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}", f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
"", "",
"| Shape | M | Mode | torch ms | GEMV chain ms | fused ms | " "| Shape | M | Mode | torch ms | GEMM chain ms | fused ms | "
"vs best unfused | fused kernels | max abs | cosine |", "vs best unfused | fused kernels | max abs | cosine |",
"|---|---:|---|---:|---:|---:|---:|---:|---:|---:|", "|---|---:|---|---:|---:|---:|---:|---:|---:|---:|",
] ]
for shape, m, mode in cases: for shape, m, mode in cases:
torch_item = by_case[(shape, m, mode, "torch")] torch_item = by_case[(shape, m, mode, "torch")]
gemv_item = by_case[(shape, m, mode, "gemv_chain")] gemm_item = by_case[(shape, m, mode, "gemm_chain")]
fused_item = by_case[(shape, m, mode, "fused")] fused_item = by_case[(shape, m, mode, "fused")]
best = min(torch_item["median_ms"], gemv_item["median_ms"]) best = min(torch_item["median_ms"], gemm_item["median_ms"])
improvement = (best / fused_item["median_ms"] - 1) * 100 improvement = (best / fused_item["median_ms"] - 1) * 100
lines.append( lines.append(
f"| {shape} | {m} | {mode} | {torch_item['median_ms']:.5f} | " f"| {shape} | {m} | {mode} | {torch_item['median_ms']:.5f} | "
f"{gemv_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | " f"{gemm_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | "
f"{improvement:+.2f}% | " f"{improvement:+.2f}% | "
f"{fused_item['cuda_kernel_launches_per_call']:.1f} | " f"{fused_item['cuda_kernel_launches_per_call']:.1f} | "
f"{fused_item['max_abs_error']:.5f} | " f"{fused_item['max_abs_error']:.5f} | "
@@ -273,8 +273,8 @@ def benchmark_command(
) -> None: ) -> None:
if not torch.cuda.is_available(): if not torch.cuda.is_available():
raise click.ClickException("CUDA is required") raise click.ClickException("CUDA is required")
if not is_available("bf16_gemv") or not is_available("bf16_swiglu"): if not is_available("bf16_gemm") or not is_available("bf16_swiglu"):
raise click.ClickException("built bf16_gemv and bf16_swiglu are required") raise click.ClickException("built bf16_gemm and bf16_swiglu are required")
shapes = tuple(parse_shape(value) for value in shape_values) or DEFAULT_SHAPES shapes = tuple(parse_shape(value) for value in shape_values) or DEFAULT_SHAPES
m_values_parsed = parse_positive_ints(m_values) m_values_parsed = parse_positive_ints(m_values)
if any(m > 8 for m in m_values_parsed): if any(m > 8 for m in m_values_parsed):
csrc/kernels/bf16_gemv.cu
-316
View File
@@ -1,316 +0,0 @@
// 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 kHalfCtaThreads = 128;
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, int Threads>
__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;
float sums[Rows] = {};
__shared__ float warp_sums[Rows][Threads / kWarpSize];
// Weight row: scalar head/tail around a 16-byte-aligned uint4 middle so
// any K is accepted while keeping 128-bit weight loads, which dominate
// bandwidth on decode shapes. x pairs with scalar loads: it is a tiny
// L1/L2-resident matrix, consecutive threads still touch contiguous
// addresses, and no per-row alignment case analysis is needed.
const __nv_bfloat16* __restrict__ wrow =
weight + static_cast<int64_t>(output_index) * k;
const unsigned whead_raw =
((16u - (reinterpret_cast<uintptr_t>(wrow) & 15u)) & 15u) >> 1;
const int whead = static_cast<int>(min(whead_raw, static_cast<unsigned>(k)));
const int wvecs = (k - whead) / 8;
const int wtail_start = whead + wvecs * 8;
const uint4* __restrict__ w4 = reinterpret_cast<const uint4*>(wrow + whead);
// x chunks pair element-for-element with the aligned weight middle:
// the uint4 view is rooted at ``x + whead`` (16-byte aligned by the
// branch guard), and each row strides by ``k / 8`` vectors because its
// first middle element sits ``whead`` scalars past ``row * k``. When
// K % 8 == 0 and the weight row is already aligned (whead == 0, the
// production case) this reduces to one pure uint4 loop with an empty
// head/tail. Otherwise per-row uint4 loads are not 16-byte addressable,
// and scalar x pairing keeps the kernel correct for any K while the
// weight stream stays vectorized.
if (k % 8 == 0 &&
((reinterpret_cast<uintptr_t>(x) + 2u * static_cast<unsigned>(whead)) & 15u) == 0u) {
const auto* x4 = reinterpret_cast<const uint4*>(x + whead);
for (int v = threadIdx.x; v < wvecs; v += blockDim.x) {
const uint4 wv_raw = w4[v];
const auto* wv =
reinterpret_cast<const __nv_bfloat162*>(&wv_raw);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
const uint4 xv_raw =
x4[(static_cast<int64_t>(row) * (k / 8)) + v];
const auto* xv =
reinterpret_cast<const __nv_bfloat162*>(&xv_raw);
#pragma unroll
for (int p = 0; p < 4; ++p) {
sums[row] = fmaf(
__bfloat162float(__low2bfloat16(xv[p])),
__bfloat162float(__low2bfloat16(wv[p])),
sums[row]
);
sums[row] = fmaf(
__bfloat162float(__high2bfloat16(xv[p])),
__bfloat162float(__high2bfloat16(wv[p])),
sums[row]
);
}
}
}
} else {
for (int v = threadIdx.x; v < wvecs; v += blockDim.x) {
const uint4 wv_raw = w4[v];
const __nv_bfloat16* wv_s =
reinterpret_cast<const __nv_bfloat16*>(&wv_raw);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
const __nv_bfloat16* xv =
x + static_cast<int64_t>(row) * k + whead + 8 * v;
#pragma unroll
for (int s = 0; s < 8; ++s) {
sums[row] = fmaf(
__bfloat162float(xv[s]),
__bfloat162float(wv_s[s]),
sums[row]
);
}
}
}
}
// Head and tail remainders: plain scalar pairing, at most 14 elements.
for (int i = threadIdx.x; i < whead; i += blockDim.x) {
const float wv = __bfloat162float(wrow[i]);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
sums[row] = fmaf(
__bfloat162float(x[static_cast<int64_t>(row) * k + i]),
wv,
sums[row]
);
}
}
for (int i = wtail_start + threadIdx.x; i < k; i += blockDim.x) {
const float wv = __bfloat162float(wrow[i]);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
sums[row] = fmaf(
__bfloat162float(x[static_cast<int64_t>(row) * k + i]),
wv,
sums[row]
);
}
}
#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 < (Threads / 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
) {
// Decode is HBM weight-streaming bound: with weights rotated through L2,
// 128/256-thread CTAs measure within noise on L20 except for small
// weight matrices at the largest decode batch, where the smaller CTA
// wins 5-9% (see docs/developer/decode_linear_benchmark.md).
constexpr int64_t kSmallWeightLimit = int64_t{12} << 20;
if constexpr (Rows == 8) {
if (k % 8 == 0 &&
(reinterpret_cast<uintptr_t>(x) & 15u) == 0u &&
(reinterpret_cast<uintptr_t>(weight) & 15u) == 0u &&
static_cast<int64_t>(n) * k <= kSmallWeightLimit) {
bf16_gemv_kernel<Rows, kHalfCtaThreads>
<<<n, kHalfCtaThreads, 0, stream>>>(
x, weight, bias, output, n, k
);
return;
}
}
bf16_gemv_kernel<Rows, kThreads><<<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 <= 8,
"M must be in [1, 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 <= 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 3:
launch_bf16_gemv<3>(
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 5:
launch_bf16_gemv<5>(
x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 6:
launch_bf16_gemv<6>(
x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 7:
launch_bf16_gemv<7>(
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, 8] BF16 GEMV with FP32 accumulation and optional fused bias"
);
}
csrc/kernels/common/launch.cuh
+36
View File
@@ -0,0 +1,36 @@
// Launch-and-check macros — pure CUDA, no torch, so out-of-tree harnesses
// (tile sweeps, csrc/tests) share the exact production launch discipline.
//
// Include order matters for overrides: define ASTRAI_LAUNCH_FAIL before
// including this header (directly or via another kernel header) to swap
// print+abort for a throwing check, as the torch entry units do with
// C10_CUDA_CHECK.
#pragma once
#include <cstdio>
#include <cstdlib>
#include <cuda_runtime.h>
#define ASTRAI_LAUNCH_FAIL(err, what) \
do { \
std::fprintf( \
stderr, "ASTRAI: %s failed: %s (%s:%d)\n", what, \
cudaGetErrorString(err), __FILE__, __LINE__ \
); \
std::abort(); \
} while (0)
#define ASTRAI_CUDA_CHECK(expr) \
do { \
cudaError_t astrai_err_ = (expr); \
if (astrai_err_ != cudaSuccess) { \
ASTRAI_LAUNCH_FAIL(astrai_err_, #expr); \
} \
} while (0)
// Check a kernel launch. Wrap the raw <<<>>> with this on the next line;
// a rejected configuration must fail loudly instead of silently measuring
// as a constant ~3us no-op (the tile-sweep lesson).
#define ASTRAI_LAUNCH_CHECK() ASTRAI_CUDA_CHECK(cudaGetLastError())
csrc/kernels/gemm.cu
+650
View File
@@ -0,0 +1,650 @@
// BF16 decode-time GEMM primitive for linear layers: one entry point whose
// internal path is sized by the decode batch M.
//
// M in [1, 8] — one CTA per weight row, the M rows held in registers
// (any K).
// M in (8, 32] — BM=16 tensor-core tiles; shape-driven configs (see
// the dispatch table at the entry point) cover every
// production shape (K % 8 == 0 and 16-byte-aligned
// tensors).
//
// Both paths take row-major [N, K] weights, accumulate in FP32, fuse an
// optional bias, and are inference-only.
#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>
// Route kernel-launch failures through the torch error check instead of
// common/launch.cuh's print+abort default. Must precede the kernel code.
#define ASTRAI_LAUNCH_FAIL(err, what) C10_CUDA_CHECK(err)
#include "common/cp_async.cuh"
#include "common/launch.cuh"
#include "common/mma.cuh"
namespace {
constexpr int kHalfCtaThreads = 128;
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, int Threads>
__global__ void skinny_gemm_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;
float sums[Rows] = {};
__shared__ float warp_sums[Rows][Threads / kWarpSize];
// Weight row: scalar head/tail around a 16-byte-aligned uint4 middle so
// any K is accepted while keeping 128-bit weight loads, which dominate
// bandwidth on decode shapes. x pairs with scalar loads: it is a tiny
// L1/L2-resident matrix, consecutive threads still touch contiguous
// addresses, and no per-row alignment case analysis is needed.
const __nv_bfloat16* __restrict__ wrow =
weight + static_cast<int64_t>(output_index) * k;
const unsigned whead_raw =
((16u - (reinterpret_cast<uintptr_t>(wrow) & 15u)) & 15u) >> 1;
const int whead = static_cast<int>(min(whead_raw, static_cast<unsigned>(k)));
const int wvecs = (k - whead) / 8;
const int wtail_start = whead + wvecs * 8;
const uint4* __restrict__ w4 = reinterpret_cast<const uint4*>(wrow + whead);
// x chunks pair element-for-element with the aligned weight middle:
// the uint4 view is rooted at ``x + whead`` (16-byte aligned by the
// branch guard), and each row strides by ``k / 8`` vectors because its
// first middle element sits ``whead`` scalars past ``row * k``. When
// K % 8 == 0 and the weight row is already aligned (whead == 0, the
// production case) this reduces to one pure uint4 loop with an empty
// head/tail. Otherwise per-row uint4 loads are not 16-byte addressable,
// and scalar x pairing keeps the kernel correct for any K while the
// weight stream stays vectorized.
if (k % 8 == 0 &&
((reinterpret_cast<uintptr_t>(x) + 2u * static_cast<unsigned>(whead)) & 15u) == 0u) {
const auto* x4 = reinterpret_cast<const uint4*>(x + whead);
for (int v = threadIdx.x; v < wvecs; v += blockDim.x) {
const uint4 wv_raw = w4[v];
const auto* wv =
reinterpret_cast<const __nv_bfloat162*>(&wv_raw);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
const uint4 xv_raw =
x4[(static_cast<int64_t>(row) * (k / 8)) + v];
const auto* xv =
reinterpret_cast<const __nv_bfloat162*>(&xv_raw);
#pragma unroll
for (int p = 0; p < 4; ++p) {
sums[row] = fmaf(
__bfloat162float(__low2bfloat16(xv[p])),
__bfloat162float(__low2bfloat16(wv[p])),
sums[row]
);
sums[row] = fmaf(
__bfloat162float(__high2bfloat16(xv[p])),
__bfloat162float(__high2bfloat16(wv[p])),
sums[row]
);
}
}
}
} else {
for (int v = threadIdx.x; v < wvecs; v += blockDim.x) {
const uint4 wv_raw = w4[v];
const __nv_bfloat16* wv_s =
reinterpret_cast<const __nv_bfloat16*>(&wv_raw);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
const __nv_bfloat16* xv =
x + static_cast<int64_t>(row) * k + whead + 8 * v;
#pragma unroll
for (int s = 0; s < 8; ++s) {
sums[row] = fmaf(
__bfloat162float(xv[s]),
__bfloat162float(wv_s[s]),
sums[row]
);
}
}
}
}
// Head and tail remainders: plain scalar pairing, at most 14 elements.
for (int i = threadIdx.x; i < whead; i += blockDim.x) {
const float wv = __bfloat162float(wrow[i]);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
sums[row] = fmaf(
__bfloat162float(x[static_cast<int64_t>(row) * k + i]),
wv,
sums[row]
);
}
}
for (int i = wtail_start + threadIdx.x; i < k; i += blockDim.x) {
const float wv = __bfloat162float(wrow[i]);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
sums[row] = fmaf(
__bfloat162float(x[static_cast<int64_t>(row) * k + i]),
wv,
sums[row]
);
}
}
#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 < (Threads / 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, int Threads>
void launch_skinny_gemm(
const __nv_bfloat16* x,
const __nv_bfloat16* weight,
const __nv_bfloat16* bias,
__nv_bfloat16* output,
int n,
int k,
cudaStream_t stream
) {
// Split-K was tried and rejected for the small-N shapes (GQA kv
// projections): their ~48K uint4 loads already saturate thread-level
// parallelism one load deep, so they sit on the launch+HBM latency
// floor, and the fence+atomic+last-CTA partial round trip adds ~0.6us
// of fixed sync cost (measured -27% at M=1, -113% at M=8 on L20).
// The real fix for those shapes is fusing the QKV projections so the
// tiny kv rows stop launching as standalone kernels at all.
//
// Decode is HBM weight-streaming bound. Every shape launches with a
// single 128-thread CTA size: measured on L20 (sm_89) with L2-thrashing
// weight rotation, flat 128t is within ~1% of a per-shape tuned
// 128/256/512 mix at M=1 and M=8 and gives up at most ~2% at M=2-4
// (512t down_proj, 256t gate/up/lm_head cells), while dodging the
// Rows>=7 register cliff on the 307MB lm_head (+21% DRAM throughput
// vs 256t at M=8). Re-measured
// after the table refactor: 256t on wide-N gate/up wins only ~2.8% at
// M=2-4 and ties at M=1/6, inside run-to-run drift. Simplicity keeps
// winning over the last ~2%.
skinny_gemm_kernel<Rows, Threads>
<<<n, Threads, 0, stream>>>(
x, weight, bias, output, n, k
);
ASTRAI_LAUNCH_CHECK();
}
// Compile-time dispatch tables: replace a hand-written switch over M with
// function-pointer tables indexed by M-1. Adding a Rows variant means adding
// one table entry, not a new case block at the call site.
using SkinnyGemmFn = void (*)(
const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*,
__nv_bfloat16*, int, int, cudaStream_t
);
constexpr SkinnyGemmFn kSkinnyGemm[8] = {
&launch_skinny_gemm<1, kHalfCtaThreads>,
&launch_skinny_gemm<2, kHalfCtaThreads>,
&launch_skinny_gemm<3, kHalfCtaThreads>,
&launch_skinny_gemm<4, kHalfCtaThreads>,
&launch_skinny_gemm<5, kHalfCtaThreads>,
&launch_skinny_gemm<6, kHalfCtaThreads>,
&launch_skinny_gemm<7, kHalfCtaThreads>,
&launch_skinny_gemm<8, kHalfCtaThreads>,
};
// ---------------------------------------------------------------------------
// M > 8 path: small-M (M <= 64) tiled GEMM.
//
// F.linear for decode batches in (8, 64]: cuBLAS tiles the small M as a
// single tile row, which starves the grid (measured on L20, sm_89: 24-54
// CTAs on 92 SMs, tensor pipe 24-42%, DRAM <= 72%). This kernel keeps the
// M rows in one CTA tile and fills the SMs along N and K instead.
//
// CUTLASS-style configuration: the kernel is parameterized by template
// parameters — CTA tile (BM x BN x BK), pipeline depth, thread count —
// with the warp layout derived inside (kMt M fragments, kNt n16 tiles per
// warp). Four families are instantiated (see the dispatch table at the
// entry point): the default (BN=64, BK=64, 3 stages) and its BM=32 wide-N
// variant, plus narrow-N deep-K rings (BK=256/128, BN=32, 64 threads); a
// new shape is an instantiation, not a rewrite.
//
// Operand staging reuses the FP8 GEMM's scheme (see fp8/gemm/*.cuh and
// docs/developer/cuda_kernels.md): 16B chunks XOR-swizzled with row & 7,
// one barrier per k-tile, and a kStages+1 ring whose prefetch for tile
// i+kStages lands in the slot tile i-1 released — no post-compute barrier.
// ---------------------------------------------------------------------------
using bf16 = __nv_bfloat16;
// Logical (row, byte column) -> byte offset in one flat [rows * BK*2B]
// staging tile. The 16B chunk index is XORed with row & (chunks - 1) so an
// ldmatrix fragment load (8 consecutive rows x 16B) hits all 32 banks once.
template <int BK>
__device__ __forceinline__ int tile_off(int row, int byte_col) {
constexpr int kRowBytes = BK * 2;
constexpr int kChunks = kRowBytes / 16;
static_assert(
(kChunks & (kChunks - 1)) == 0, "swizzle needs a power-of-two chunk count"
);
return row * kRowBytes +
(((byte_col >> 4) ^ (row & (kChunks - 1))) << 4) + (byte_col & 15);
}
// Predicated staging of one [RowsTile x BK] operand slice from a
// row-major [total_rows, K] tensor into its swizzled ring slot. Chunks past
// the row count or past K zero-fill (the wrapper guarantees K % 8 == 0 and
// 16B-aligned rows, so a misaligned *valid* chunk cannot occur).
template <int RowsTile, int BK, int kThreads>
__device__ __forceinline__ void stage_tile(
char* slot, const bf16* __restrict__ src, int total_rows, int64_t k,
int row0, int tid, int kt
) {
constexpr int kRowBytes = BK * 2;
constexpr int kChunks = kRowBytes / 16;
#pragma unroll
for (int c = tid; c < RowsTile * kChunks; c += kThreads) {
const int r = c / kChunks;
const int c8 = c % kChunks;
const int64_t kbase = (int64_t)kt * BK;
const bool ok = row0 + r < total_rows && kbase + c8 * 8 + 8 <= k;
char* dst = slot + tile_off<BK>(r, c8 * 16);
if (ok) {
astrai::cp_async_16(
reinterpret_cast<bf16*>(dst),
src + (int64_t)(row0 + r) * k + kbase + c8 * 8
);
} else {
unsigned* w = reinterpret_cast<unsigned*>(dst);
#pragma unroll
for (int i = 0; i < 4; ++i)
w[i] = 0u;
}
}
}
template <int BM, int BN, int BK, int kStages, int kThreads>
__global__ __launch_bounds__(kThreads, 1) void tiled_gemm_kernel(
const bf16* __restrict__ x,
const bf16* __restrict__ w,
const bf16* __restrict__ bias,
bf16* __restrict__ out,
int m,
int n,
int k
) {
constexpr int kMt = BM / 16; // m16 fragments
constexpr int kWarps = kThreads / 32;
constexpr int kNt = BN / (kWarps * 16); // n16 tiles per warp
constexpr int kRowBytes = BK * 2;
constexpr int kRing = kStages + 1; // ring slots per operand
constexpr int kSegs = BK / 16; // m16k16 segments per tile
constexpr int kSegXor = 32; // bytes: +2 chunks per segment
static_assert(BM % 16 == 0, "BM must be a multiple of 16");
static_assert(BN % (kWarps * 16) == 0, "warps must tile BN in n16 units");
extern __shared__ __align__(16) char smem[];
char* const a_ring = smem;
char* const b_ring = smem + kRing * BM * kRowBytes;
const int tid = threadIdx.x;
const int warp = tid >> 5;
const int lane = tid & 31;
// Standard 2D grid mapping: grid.x covers M tiles, grid.y covers N
// tiles. The grid fills the SMs along N; every block walks the whole
// K range in a single pass.
const int m0 = blockIdx.x * BM;
const int n0 = blockIdx.y * BN;
const int total_tiles = (k + BK - 1) / BK;
auto a_slot = [&](int t) {
return a_ring + (t % kRing) * BM * kRowBytes;
};
auto b_slot = [&](int t) {
return b_ring + (t % kRing) * BN * kRowBytes;
};
#pragma unroll
for (int s = 0; s < kStages; ++s) {
if (s < total_tiles) {
stage_tile<BM, BK, kThreads>(
a_slot(s), x, m, k, m0, tid, s
);
stage_tile<BN, BK, kThreads>(b_slot(s), w, n, k, n0, tid, s);
}
astrai::cp_async_commit_group();
}
// Per-lane ldmatrix fragment addresses (relative to each ring slot):
// A x4: lanes 0-7 mat0 (rows 0-7, chunk k-lo), 8-15 mat1 (rows 8-15,
// k-lo), 16-23 mat2 (rows 0-7, k-hi), 24-31 mat3 (rows 8-15, k-hi).
// B x4 pair: lanes 0-7/8-15 the first n8 tile's k-lo/k-hi chunks,
// 16-23/24-31 the second n8 tile's. Warp w owns the n16 tiles at
// (w + j * kWarps) * 16 for j in [0, kNt).
const int a_row = ((lane >> 3) & 1) * 8 + (lane & 7);
const unsigned a_off = tile_off<BK>(a_row, (lane >> 4) * 16);
unsigned b_off[kNt];
#pragma unroll
for (int j = 0; j < kNt; ++j) {
const int b_row =
(warp + j * kWarps) * 16 + (lane & 7) + (lane >> 4) * 8;
b_off[j] = tile_off<BK>(b_row, ((lane >> 3) & 1) * 16);
}
const unsigned a_base0 = __cvta_generic_to_shared(a_ring) + a_off;
unsigned b_base0[kNt];
#pragma unroll
for (int j = 0; j < kNt; ++j)
b_base0[j] = __cvta_generic_to_shared(b_ring) + b_off[j];
float acc[kMt][kNt][2][4] = {};
for (int i = 0; i < total_tiles; ++i) {
astrai::cp_async_wait_group<kStages - 1>();
__syncthreads();
const unsigned a_base =
a_base0 + (unsigned)((i % kRing) * BM * kRowBytes);
unsigned b_base[kNt];
#pragma unroll
for (int j = 0; j < kNt; ++j)
b_base[j] = b_base0[j] + (unsigned)((i % kRing) * BN * kRowBytes);
#pragma unroll
for (int seg = 0; seg < kSegs; ++seg) {
const unsigned a_seg = a_base ^ (unsigned)(seg * kSegXor);
unsigned a4[kMt][4], b4[kNt][4];
#pragma unroll
for (int mt = 0; mt < kMt; ++mt)
astrai::ldmatrix_x4_lane(
a4[mt], a_seg + (unsigned)(mt * 16 * kRowBytes)
);
#pragma unroll
for (int j = 0; j < kNt; ++j)
astrai::ldmatrix_x4_lane(
b4[j], (b_base[j] ^ (unsigned)(seg * kSegXor))
);
#pragma unroll
for (int mt = 0; mt < kMt; ++mt)
#pragma unroll
for (int j = 0; j < kNt; ++j)
#pragma unroll
for (int nt = 0; nt < 2; ++nt)
astrai::mma_sync<bf16>(
acc[mt][j][nt], a4[mt], b4[j] + nt * 2,
acc[mt][j][nt]
);
}
// Prefetch tile i+kStages into the slot tile i-1 released. The
// barrier at the top of the next iteration separates every
// thread's reads of that slot (iteration i-1) from these writes.
const int pf = i + kStages;
if (pf < total_tiles) {
stage_tile<BM, BK, kThreads>(
a_slot(pf), x, m, k, m0, tid, pf
);
stage_tile<BN, BK, kThreads>(b_slot(pf), w, n, k, n0, tid, pf);
}
astrai::cp_async_commit_group();
}
const int row0 = lane >> 2;
const int col0 = (lane & 3) * 2;
#pragma unroll
for (int mt = 0; mt < kMt; ++mt) {
if (m0 + mt * 16 + row0 >= m)
continue;
const int64_t orow = (int64_t)(m0 + mt * 16 + row0) * n;
#pragma unroll
for (int j = 0; j < kNt; ++j) {
#pragma unroll
for (int nt = 0; nt < 2; ++nt) {
const int col =
n0 + (warp + j * kWarps) * 16 + col0 + nt * 8;
if (col >= n)
continue;
float2 v, v8;
v.x = acc[mt][j][nt][0];
v.y = acc[mt][j][nt][1];
v8.x = acc[mt][j][nt][2];
v8.y = acc[mt][j][nt][3];
if (bias != nullptr) {
v.x += __bfloat162float(bias[col]);
v.y += __bfloat162float(bias[col + 1]);
v8.x += __bfloat162float(bias[col]);
v8.y += __bfloat162float(bias[col + 1]);
}
if (col + 1 < n) {
*reinterpret_cast<__nv_bfloat162*>(out + orow + col) =
__floats2bfloat162_rn(v.x, v.y);
if (m0 + mt * 16 + row0 + 8 < m)
*reinterpret_cast<__nv_bfloat162*>(
out + orow + 8 * n + col) =
__floats2bfloat162_rn(v8.x, v8.y);
} else {
out[orow + col] = __float2bfloat16_rn(v.x);
if (m0 + mt * 16 + row0 + 8 < m)
out[orow + 8 * n + col] =
__float2bfloat16_rn(v8.x);
}
}
}
}
}
template <int BM, int BN, int BK, int kStages, int kThreads>
void launch_tiled_gemm(
const bf16* x,
const bf16* w,
const bf16* bias,
bf16* out,
int m,
int n,
int k,
cudaStream_t stream
) {
constexpr int kRing = kStages + 1;
constexpr int smem = kRing * (BM + BN) * BK * 2;
// 99KB is the sm_86/89 per-block opt-in ceiling; configs above 48KB
// (long-K BK=128) run one CTA per SM and pay a one-time attribute opt-in.
static_assert(smem <= 99 * 1024, "family must fit the sm_86/89 opt-in ceiling");
if constexpr (smem > 48 * 1024) {
static const bool opted_in = [] {
ASTRAI_CUDA_CHECK(cudaFuncSetAttribute(
tiled_gemm_kernel<BM, BN, BK, kStages, kThreads>,
cudaFuncAttributeMaxDynamicSharedMemorySize, smem
));
return true;
}();
(void)opted_in;
}
dim3 grid((m + BM - 1) / BM, (n + BN - 1) / BN);
tiled_gemm_kernel<BM, BN, BK, kStages, kThreads>
<<<grid, kThreads, smem, stream>>>(
x, w, bias, out, m, n, k
);
ASTRAI_LAUNCH_CHECK();
}
using TiledGemmFn = void (*)(
const bf16*, const bf16*, const bf16*, bf16*, int, int, int,
cudaStream_t
);
// CTA tile configuration pairing the shape parameters with their matched
// instantiation. Field order is canonical everywhere it appears — template
// arguments, this struct, the dispatch table — as BM, BN, BK, then
// pipeline depth (stages) and CTA size (threads). A new family is one
// row here plus one select branch.
struct TileConfig {
int bm;
int bn;
int bk;
int stages;
int threads;
TiledGemmFn launch;
};
// Shape -> tile config, measured on L20 with L2-thrashing weight rotation.
// The selector trades grid fill against K-loop serial latency:
// - Wide N (n >= 4096): ceil(n/64) tiles already cover the SMs, the GEMM
// is HBM-bound and the default config wins; BM widens to 32 at M>16 to
// halve the re-staged weight stream.
// - Narrow N: few N tiles leave the grid K-serial — widening the grid
// does not help (measured: BN 64->32 ties, doubled m_tiles tie, kv at
// 4 blocks ties q/o at 24); fewer, deeper K chunks do. BK=256 with a
// 72KB two-stage ring is the winner while the grid fits one wave
// (72KB smem means one CTA per SM); past one wave its 2-wave
// quantization loses to BK=128's 36.9KB two-CTA ring.
inline TileConfig select_tile_config(int m, int n, int k) {
if (n >= 4096) {
if (m > 16) {
return {32, 64, 64, 3, 128,
&launch_tiled_gemm<32, 64, 64, 3, 128>};
}
return {16, 64, 64, 3, 128, &launch_tiled_gemm<16, 64, 64, 3, 128>};
}
const int n_tiles = (n + 31) / 32;
const int m_tiles = (m + 15) / 16;
if (n_tiles * m_tiles <= 92) { // one wave on the 92-SM L20
return {16, 32, 256, 2, 64, &launch_tiled_gemm<16, 32, 256, 2, 64>};
}
return {16, 32, 128, 2, 64, &launch_tiled_gemm<16, 32, 128, 2, 64>};
}
} // namespace
// Single entry point: M in [1, 8] routes to the register-resident skinny
// GEMM kernel (any K), M in (8, 64] to the tiled kernel (K % 8 == 0 and
// 16-byte-aligned tensors, checked at the branch).
torch::Tensor bf16_gemm(
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_gemm 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(weight.size(1) == k, "weight K must match x K");
TORCH_CHECK(m >= 1 && m <= 64, "M must be in [1, 64]");
TORCH_CHECK(k > 0 && n > 0, "N and K must be positive");
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_gemm 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_gemm 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 m_int = static_cast<int>(m);
const int n_int = static_cast<int>(n);
const int k_int = static_cast<int>(k);
if (m <= 8) {
kSkinnyGemm[m_int - 1](
x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int,
stream.stream()
);
} else {
TORCH_CHECK(
k % 8 == 0 &&
(reinterpret_cast<uintptr_t>(x.data_ptr()) & 15) == 0u &&
(reinterpret_cast<uintptr_t>(weight.data_ptr()) & 15) == 0u,
"M > 8 requires K to be a multiple of 8 and x/weight 16-byte aligned"
);
const TileConfig cfg = select_tile_config(m_int, n_int, k_int);
cfg.launch(
x_ptr, weight_ptr, bias_ptr, output_ptr, m_int, n_int, k_int,
stream.stream()
);
}
C10_CUDA_CHECK(cudaGetLastError());
return output;
}
PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
module.def(
"bf16_gemm",
&bf16_gemm,
py::arg("x"),
py::arg("weight"),
py::arg("bias") = py::none(),
"M in [1, 64] BF16 GEMM with optional fused bias "
"(register-resident skinny GEMM path for M <= 8, tensor-core tiles above)"
);
}
csrc/kernels/bf16_swiglu.cu → csrc/kernels/swiglu.cu
+1 -2
View File
@@ -1,8 +1,7 @@
// Fused small-M BF16 SwiGLU primitive for decode-time dense MLP layers. // Fused small-M BF16 SwiGLU primitive for decode-time dense MLP layers.
// One CTA per output column; each weight pair is read once and reused across // One CTA per output column; each weight pair is read once and reused across
// all decode rows. Bandwidth-bound in the cold-HBM decode regime, so variant // all decode rows. Bandwidth-bound in the cold-HBM decode regime, so variant
// selection beyond the M=8 block-size rule is noise (see // selection beyond the M=8 block-size rule is noise.
// docs/developer/swiglu_benchmark.md).
#include <ATen/cuda/CUDAContext.h> #include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h> #include <c10/cuda/CUDAGuard.h>
docs/developer/cuda_kernels.md
+47 -42
View File
@@ -1,8 +1,8 @@
# CUDA Kernels # CUDA Kernels
AstrAI includes optional custom CUDA kernels for attention, rotary embedding, AstrAI includes optional custom CUDA kernels for attention, rotary embedding,
BF16 GEMV/SwiGLU, and FP8 GEMM. These are built when `nvcc` is available and BF16 GEMM/SwiGLU, and FP8 GEMM. These are built when `nvcc` is available and
CUDA is detected. BF16 GEMV and SwiGLU are directly callable and can be CUDA is detected. BF16 GEMM and SwiGLU are directly callable and can be
selected by guarded model dispatchers described below. selected by guarded model dispatchers described below.
## Overview ## Overview
@@ -14,49 +14,55 @@ selected by guarded model dispatchers described below.
| `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention | | `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) | | `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) | | `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
| `bf16_gemv` | `bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (sm_80+) | | `bf16_gemm` | `gemm.cu` | M=1..64 BF16 linear, skinny GEMM path for M<=8 and tensor-core tiled path above (sm_80+) |
| `bf16_swiglu` | `bf16_swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) | | `bf16_swiglu` | `swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) |
| `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) | | `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) |
### BF16 GEMV primitive ### BF16 GEMM primitive
`astrai.extension.bf16_gemv(x, weight, bias=None)` accepts a contiguous BF16 `astrai.extension.bf16_gemm(x, weight, bias=None)` accepts a contiguous BF16
input shaped `[K]` or `[M, K]`, with `M` in `[1, 8]` and any positive `K`, and input shaped `[K]` or `[M, K]`, with `M` in `[1, 64]`, and row-major weights
row-major weights `[N, K]`. One CTA computes an output row for all M tokens `[N, K]`. One entry point selects the internal path by M:
together, reusing the weight row across tokens. CTA size is 256 threads,
except for small weight matrices (`N*K <= 12 MiB`) at `M=8`, where a
128-thread CTA measured 5-9% faster on L20. Variant selection is otherwise
intentionally shape-free: under HBM-streaming conditions (weights rotated
through L2, as in real decode) the kernel is bandwidth-bound and block-size
choice measures within noise, so earlier per-shape variant tables were
removed along with the warp-tiled kernel.
The weight stream uses 128-bit vectorized loads anchored at each row's first - `M` in `[1, 8]` — register-resident GEMV kernel, any positive `K`: one
16-byte-aligned address with scalar head/tail sweeps for unaligned remainders, flat 128-thread CTA per weight row computes the output column for all M
so arbitrary `K` and storage offsets stay correct. Accumulation is FP32; the tokens together. Under HBM-streaming conditions (weights rotated through
optional BF16 bias is fused before the BF16 store. The launcher uses the L2, as in real decode) the kernel is bandwidth-bound; on L20 the flat
current CUDA stream, is CUDA Graph capture-safe, and requires sm_80 or newer. 128-thread CTA measured within ~1% of a per-shape tuned 128/256/512 mix at
M=1 and M=8 and at most ~2% below it at M=2-4, while dodging the register
cliff the larger CTAs hit on the 307MB lm_head at M>=7 (+21% DRAM
throughput vs 256t). Simplicity was chosen over the last ~2%.
- `M` in `(8, 64]` — CUTLASS-style tiled kernel (`tiled_gemm_kernel`,
fully parameterized by template parameters BM/BN/BK/stages/threads):
multistage cp.async staging with XOR-swizzled 16B chunks, no split-K —
every block walks the whole K range in one pass and the grid fills the
SMs along N. Dispatch is shape-driven: wide N (n >= 4096, grid already
full) uses the bandwidth-optimal default (BN=64, BK=64, 3 stages; BM
widens to 32 at M>16); narrow N is K-serial, where widening the grid
measurably does nothing and deeper K chunks win (BN=32, BK=256 while
the grid fits one wave, BK=128 past it). Requires `K % 8 == 0` and
16-byte-aligned tensors.
Both paths accumulate in FP32, fuse the optional BF16 bias before the BF16
store, use the current CUDA stream, are CUDA Graph capture-safe, and require
sm_80 or newer. The GEMV path anchors 128-bit weight loads at each row's
first 16-byte-aligned address with scalar head/tail sweeps for unaligned
remainders, so arbitrary `K` and storage offsets stay correct.
Model `Linear` calls route through the lightweight linear backend. Set Model `Linear` calls route through the lightweight linear backend. Set
`ASTRAI_GEMV=0` for an unconditional `F.linear` fallback, `1` to force the `ASTRAI_GEMM=0` for an unconditional `F.linear` fallback, `1` to force the
kernel for any supported M in [1, 8], or `auto` (the default). Automatic kernel for any capable M (1-64, with K%8==0 and alignment for M>8), or
dispatch is keyed on the decode batch size alone: the kernel streams each `auto` (the default). On compute capability 8.0+, `auto` routes all capable
weight exactly once, so once a batch size is profitable it is profitable decode batches M in [1, 64] through the kernel. The GEMV path (M≤8, any K)
across projection shapes. On compute capability 8.0+, `auto` selects the measured universal wins at the HBM bandwidth floor (AstrAI 1B end-to-end
kernel for `M` in [2, 4], where every measured model family beat the cuBLAS +8.8% at M=1 rising to +17% at M=8). The tiled path (M 9-64) measured on
small-M path at the HBM bandwidth floor (AstrAI 1B chain +11.8% to +14.0%, L20 across five shape families wins or ties four (q/o, kv, gate/up, lm_head)
common LLaMA/Qwen/OPT chains +5.66% to +25.20%). `M=1` keeps cuBLAS, whose and regresses one (down-projection M>16: +25% latency; narrow-N long-K where
GEMV path is already at the floor, and `M >= 5` approaches the cuBLAS cuBLAS is strong). Out-of-band, training, prefill-sized, or unsupported
tensor-core crossover (M=8 regressed at wrapper level in every measured calls fall back to PyTorch.
family). Out-of-band, training, prefill-sized, or unsupported calls fall
back to PyTorch.
Inside the primitive, a templated cooperative kernel uses 256 threads, The measurements below date from the original `[2, 4]` automatic band; the
except at the largest decode batch where a 128-thread CTA wins 5-9% on band has since widened to `{1, ..., 8}` with the flat 128-thread kernel:
small weight matrices: `M=8` with 16-byte-aligned inputs, `K % 8 == 0`,
and `N*K <= 12 MiB` selects the smaller CTA. This internal selector is
separate from model automatic dispatch, whose Python/wrapper overhead is
included in the gates above.
On NVIDIA L20 (SM89), the common-shape microbenchmark reports +5.37% to On NVIDIA L20 (SM89), the common-shape microbenchmark reports +5.37% to
+114.39% for M=2 and +5.38% to +115.26% for M=4 versus `F.linear`. The paired +114.39% for M=2 and +5.38% to +115.26% for M=4 versus `F.linear`. The paired
@@ -84,7 +90,7 @@ OPT 1.3B M=1 is +4.54%. Qwen2 and LLaMA 3 70B M=1, and all three new
families at M=8, remain exact PyTorch fallbacks. families at M=8, remain exact PyTorch fallbacks.
These are synthetic projection-chain measurements, not whole-model throughput These are synthetic projection-chain measurements, not whole-model throughput
claims. Reproduce them with `csrc/bench/benchmark_gemv_common.py`. claims. Reproduce them with `csrc/bench/benchmark_gemm_common.py`.
The AstrAI 1B A→B→B→A results use the real `InferenceEngine`, including scheduler, The AstrAI 1B A→B→B→A results use the real `InferenceEngine`, including scheduler,
sampling, and CUDA Graph. M=8 stays on PyTorch because its remaining sampling, and CUDA Graph. M=8 stays on PyTorch because its remaining
@@ -121,8 +127,7 @@ the unfused linear backend, and `1` explicitly forces the fused primitive.
errors are small (maximum absolute error at most 2.4e-4 in the L20 matrix), errors are small (maximum absolute error at most 2.4e-4 in the L20 matrix),
the different FP32 reduction order changed greedy checkpoint output for the different FP32 reduction order changed greedy checkpoint output for
M=1/2/4. Automatic dispatch therefore remains numerically identical to the M=1/2/4. Automatic dispatch therefore remains numerically identical to the
existing path. See [the benchmark protocol](./swiglu_benchmark.md) for raw existing path.
operator, engine, and checkpoint evidence.
Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist: Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
@@ -313,7 +318,7 @@ astrai/extension/
├── ops/ ├── ops/
│ ├── attention.py # Stateless attention kernel wrappers │ ├── attention.py # Stateless attention kernel wrappers
│ ├── rotary.py # Stateless rotary kernel wrapper │ ├── rotary.py # Stateless rotary kernel wrapper
│ ├── gemv.py # Stateless BF16 GEMV primitive │ ├── gemm.py # Stateless BF16 GEMM primitive
│ ├── swiglu.py # Stateless fused BF16 SwiGLU primitive │ ├── swiglu.py # Stateless fused BF16 SwiGLU primitive
│ └── fp8.py # Stateless FP8 primitives (custom_op) │ └── fp8.py # Stateless FP8 primitives (custom_op)
├── fp8.py # FP8 strategy layer (fp8_autocast, recipes) ├── fp8.py # FP8 strategy layer (fp8_autocast, recipes)
docs/developer/decode_linear_benchmark.md
-42
View File
@@ -1,42 +0,0 @@
# Decode linear shape benchmark
`csrc/bench/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 csrc/bench/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.
For direct A/B coverage of the custom kernel and guarded dispatcher across
traditional LLaMA and GPT-NeoX decode shapes, use:
```bash
CUDA_VISIBLE_DEVICES=0 PYTHONPATH=. python csrc/bench/benchmark_gemv_common.py \
--suite all --family traditional --m 2 4 \
--output results/gemv_common.json
```
The kernel suite compares the directly callable primitive with `F.linear`.
Use repeatable `--shape-label` and `--chain-label` filters for a focused run.
The synthetic-chain suite alternates `ASTRAI_GEMV=0` and `auto`, includes
dependent MLP work and Python dispatch, and rotates through distinct weights.
Automatic dispatch is keyed on the decode batch size alone (`M` in `[2, 4]` on
compute capability 8.0+); use `--candidate-mode 1` to characterize a family
before widening that band. The checked-in final evidence always uses `auto`.
It is deliberately not labeled a whole-model throughput benchmark. Both
suites report median/p90 CUDA-event latency plus maximum absolute error,
relative L2 error, and row-wise argmax parity.
docs/developer/swiglu_benchmark.md
-83
View File
@@ -1,83 +0,0 @@
# Fused SwiGLU benchmark
`csrc/bench/benchmark_swiglu.py` compares the directly callable fused BF16
SwiGLU primitive with both `F.linear` and the existing two-GEMV chain. It covers
the native AstrAI 1B MLP plus LLaMA 2 7B/13B, LLaMA 3 8B, and GPT-NeoX 20B
up/gate shapes at M=1/2/4/8 in eager and CUDA Graph modes.
```bash
CUDA_VISIBLE_DEVICES=0 python csrc/bench/benchmark_swiglu.py \
--output results/swiglu.json \
--markdown-output results/swiglu.md \
--m-values 1,2,4,8 --mode both \
--warmup 20 --iterations 100 --trials 10
```
Each trial uses A-B-C-C-B-A ordering to balance clock, cache, and temperature
drift. The generated JSON records every timing sample, p50/p90/p99, CUDA launch
count, maximum/mean absolute error, and cosine similarity.
## L20 findings
Hardware was one NVIDIA L20 (sm_89), PyTorch 2.11.0+cu128, CUDA 12.8. The
existing GPU5 inference service remained resident (15.4 GiB) but idle at the
sampling boundaries; no process or container was stopped.
For AstrAI 1B `(N,K)=(6912,1536)`, CUDA Graph medians were:
| M | torch (ms) | GEMV chain (ms) | fused (ms) | vs best unfused |
|---:|---:|---:|---:|---:|
| 1 | 0.02564 | 0.02298 | 0.01375 | +67.13% |
| 2 | 0.02484 | 0.02628 | 0.01416 | +75.40% |
| 4 | 0.02507 | 0.03839 | 0.01806 | +38.82% |
| 8 | 0.02563 | 0.07007 | 0.03339 | -23.24% |
The wide traditional shapes are weight-bandwidth dominated. CTA reuse keeps
the fused primitive within roughly -1.2% to +0.9% of the best unfused chain,
so none is eligible for automatic selection. This negative crossover is kept
in the raw evidence rather than hidden by a favorable subset.
The real 24-layer AstrAI checkpoint was then run through `InferenceEngine`,
including scheduler, sampling, and CUDA Graph. A-B-B-A medians were:
| Batch | unfused (ms/step) | forced fused (ms/step) | throughput gain |
|---:|---:|---:|---:|
| 1 | 4.125 | 3.925 | +5.10% |
| 2 | 4.245 | 4.055 | +4.69% |
| 4 | 4.475 | 4.305 | +3.95% |
## Dispatch decision
Direct correctness stayed close (`max_abs <= 2.4e-4`, cosine approximately
1.0), but deterministic greedy generations changed at M=1, M=2, and M=4.
For that reason no SM89 shape is enabled in `auto`. The default path stays on
the existing unfused linear backend, including any independently qualified
GEMV dispatch. `ASTRAI_SWIGLU=1` remains an explicit benchmark/experimentation
switch for callers that accept normal BF16 reduction-order variation. A future
automatic band must repeat both the performance and checkpoint-output gates.
## HBM re-measurement and kernel simplification
The operator numbers above are L2-resident: the AstrAI pair is 40.5 MB,
smaller than the 96 MB L2, so a tight timing loop re-reads warm weights
(13.75 us implies ~3.1 TB/s, far above the 864 GB/s spec). Real decode rotates
~1 GB of per-layer weights through L2 every step, so every call is cold.
Re-measuring with rotated weight copies (>= 240 MB working set) on the same
L20 showed:
- The fused CTA-reuse kernel sits at the dual-stream cold-read floor
(702 vs 699 GB/s at (6912,1536); 369 vs 370 GB/s at (11008,4096)). Wide
LLaMA matrices cap at ~370-400 GB/s regardless of kernel, even for a
pure-read loop, so the old per-variant gaps there were noise.
- The `(6912,1536)` warp-per-row variant (formerly M=2/4/8) is 2-6% slower
than CTA reuse at M=2/4 under cold weights and no longer wins at M=8 once
the CTA drops to 128 threads. It and its dispatch table were deleted.
- New rule: 256 threads for M in [1, 7], 128 threads for M=8. End-to-end
through the built module at (6912,1536): 738-752 GB/s for M in [1, 4] and
702 GB/s at M=8 (+6% over the removed warp path).
The M=8 CUDA-Graph regression reported above (`-23.24%`) does not survive the
cold-weight regime: cuBLAS reaches L2 bandwidth in the warm loop while both
fused paths converge to the same HBM floor.
setup.py
+1 -1
View File
@@ -121,7 +121,7 @@ class _CMakeBuildExt(_build_ext):
"attn_prefill", "attn_prefill",
"attn_paged_decode", "attn_paged_decode",
"attn_paged_prefill", "attn_paged_prefill",
"bf16_gemv", "bf16_gemm",
"bf16_swiglu", "bf16_swiglu",
"rotary_emb", "rotary_emb",
) )
tests/extension/test_gemv.py → tests/extension/test_gemm.py
+184 -49
View File
@@ -2,48 +2,65 @@ import pytest
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
from astrai.extension import bf16_gemv, is_available from astrai.extension import bf16_gemm, is_available
GEMV_AVAILABLE = ( GEMM_AVAILABLE = (
torch.cuda.is_available() torch.cuda.is_available()
and is_available("bf16_gemv") and is_available("bf16_gemm")
and torch.cuda.get_device_capability() >= (8, 0) and torch.cuda.get_device_capability() >= (8, 0)
) )
skip_no_gemv = pytest.mark.skipif( skip_no_gemm = pytest.mark.skipif(
not GEMV_AVAILABLE, not GEMM_AVAILABLE,
reason="BF16 GEMV requires a built kernel and compute capability 8.0+", reason="BF16 GEMM requires a built kernel and compute capability 8.0+",
) )
@skip_no_gemv def _assert_close_fp64(actual, x, weight, bias=None):
"""Compare bf16 kernel output vs fp64-exact with ulp-scaled tolerance.
Avoids false failures from cuBLAS default bf16 split-K partial reduction
(which can introduce ~2 ulp diffs on near-tie rounding at long K). Used
for tiled-path tests (M > 8, K >= 4096) where the bf16 accumulation tie
pattern may differ from cuBLAS's."""
exact = x.double() @ weight.double().T
if bias is not None:
exact = exact + bias.double()
ulp = (exact.abs() * 2**-9).clamp(min=2**-9)
max_ulp = ((actual.double() - exact).abs() / ulp).max().item()
assert max_ulp < 8, (
f"max_ulp={max_ulp:.1f} exceeds 8 (exact fp32 accumulation should stay within ~2 ulps)"
)
@skip_no_gemm
@pytest.mark.parametrize( @pytest.mark.parametrize(
"n,k", "n,k",
[(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912), (100000, 1536)], [(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912), (100000, 1536)],
) )
def test_bf16_gemv_matches_linear_shape_families(n, k): def test_bf16_gemm_matches_linear_shape_families(n, k):
torch.manual_seed(17) torch.manual_seed(17)
x = torch.randn(k, device="cuda", dtype=torch.bfloat16) x = torch.randn(k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
expected = F.linear(x, weight) expected = F.linear(x, weight)
assert actual.shape == (n,) assert actual.shape == (n,)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize("m", [2, 3, 4, 5, 6, 7, 8]) @pytest.mark.parametrize("m", [2, 3, 4, 5, 6, 7, 8])
@pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (1536, 6912)]) @pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (1536, 6912)])
def test_bf16_gemv_matches_small_decode_batches(m, n, k): def test_bf16_gemm_matches_small_decode_batches(m, n, k):
torch.manual_seed(19 + m) torch.manual_seed(19 + m)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
expected = F.linear(x, weight) expected = F.linear(x, weight)
assert actual.shape == (m, n) assert actual.shape == (m, n)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5)
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize("m", [2, 4]) @pytest.mark.parametrize("m", [2, 4])
@pytest.mark.parametrize( @pytest.mark.parametrize(
"n,k", "n,k",
@@ -72,17 +89,17 @@ def test_bf16_gemv_matches_small_decode_batches(m, n, k):
(2048, 8192), (2048, 8192),
], ],
) )
def test_bf16_gemv_matches_common_transformer_shapes(m, n, k): def test_bf16_gemm_matches_common_transformer_shapes(m, n, k):
torch.manual_seed(2026 + m + n + k) torch.manual_seed(2026 + m + n + k)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16) weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)
weight.normal_(mean=0.0, std=0.02) weight.normal_(mean=0.0, std=0.02)
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
expected = F.linear(x, weight) expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize( @pytest.mark.parametrize(
"m,n,k", "m,n,k",
[ [
@@ -93,63 +110,63 @@ def test_bf16_gemv_matches_common_transformer_shapes(m, n, k):
(8, 2048, 8192), (8, 2048, 8192),
], ],
) )
def test_bf16_gemv_matches_m8_edge_bands(m, n, k): def test_bf16_gemm_matches_m8_edge_bands(m, n, k):
torch.manual_seed(2026 + m + n + k) torch.manual_seed(2026 + m + n + k)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16) weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)
weight.normal_(mean=0.0, std=0.02) weight.normal_(mean=0.0, std=0.02)
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
expected = F.linear(x, weight) expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv @skip_no_gemm
def test_bf16_gemv_preserves_singleton_batch_and_fuses_bias(): def test_bf16_gemm_preserves_singleton_batch_and_fuses_bias():
torch.manual_seed(23) torch.manual_seed(23)
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 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) bias = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemv(x, weight, bias) actual = bf16_gemm(x, weight, bias)
expected = F.linear(x, weight, bias) expected = F.linear(x, weight, bias)
assert actual.shape == (1, 1536) assert actual.shape == (1, 1536)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv @skip_no_gemm
def test_bf16_gemv_small_batch_fuses_bias(): def test_bf16_gemm_small_batch_fuses_bias():
torch.manual_seed(25) torch.manual_seed(25)
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 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) bias = torch.randn(256, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemv(x, weight, bias) actual = bf16_gemm(x, weight, bias)
expected = F.linear(x, weight, bias) expected = F.linear(x, weight, bias)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv @skip_no_gemm
def test_bf16_gemv_uses_current_stream(): def test_bf16_gemm_uses_current_stream():
x = torch.randn(1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
torch.cuda.synchronize() torch.cuda.synchronize()
stream = torch.cuda.Stream() stream = torch.cuda.Stream()
with torch.cuda.stream(stream): with torch.cuda.stream(stream):
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
expected = F.linear(x, weight) expected = F.linear(x, weight)
stream.synchronize() stream.synchronize()
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv @skip_no_gemm
def test_bf16_gemv_cuda_graph_replay(): def test_bf16_gemm_cuda_graph_replay():
torch.manual_seed(29) torch.manual_seed(29)
x = torch.randn(1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
for _ in range(3): for _ in range(3):
bf16_gemv(x, weight) bf16_gemm(x, weight)
graph = torch.cuda.CUDAGraph() graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph): with torch.cuda.graph(graph):
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
x.copy_(torch.randn_like(x)) x.copy_(torch.randn_like(x))
graph.replay() graph.replay()
@@ -157,24 +174,24 @@ def test_bf16_gemv_cuda_graph_replay():
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize("n,k", [(64, 7), (64, 12), (33, 100), (256, 1534)]) @pytest.mark.parametrize("n,k", [(64, 7), (64, 12), (33, 100), (256, 1534)])
def test_bf16_gemv_handles_unaligned_k(n, k): def test_bf16_gemm_handles_unaligned_k(n, k):
torch.manual_seed(29) torch.manual_seed(29)
x = torch.randn(k, device="cuda", dtype=torch.bfloat16) x = torch.randn(k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
expected = F.linear(x, weight) expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
x3 = torch.randn(3, k, device="cuda", dtype=torch.bfloat16) x3 = torch.randn(3, k, device="cuda", dtype=torch.bfloat16)
actual3 = bf16_gemv(x3, weight) actual3 = bf16_gemm(x3, weight)
torch.testing.assert_close(actual3, F.linear(x3, weight), rtol=0.02, atol=0.5) torch.testing.assert_close(actual3, F.linear(x3, weight), rtol=0.02, atol=0.5)
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize("m", [1, 2, 3, 4]) @pytest.mark.parametrize("m", [1, 2, 3, 4])
def test_bf16_gemv_accepts_complementary_misalignment(m): def test_bf16_gemm_accepts_complementary_misalignment(m):
"""Misaligned weight rows plus an x base chosen so the vectorized branch """Misaligned weight rows plus an x base chosen so the vectorized branch
is entered with a non-16B-aligned ``x`` pointer (regression: the branch is entered with a non-16B-aligned ``x`` pointer (regression: the branch
guard checked ``x + whead`` alignment but the uint4 view was rooted at guard checked ``x + whead`` alignment but the uint4 view was rooted at
@@ -189,13 +206,13 @@ def test_bf16_gemv_accepts_complementary_misalignment(m):
x = big_x[5 : 5 + m * k].view(m, k) if m > 1 else big_x[5 : 5 + k] x = big_x[5 : 5 + m * k].view(m, k) if m > 1 else big_x[5 : 5 + k]
assert (x.data_ptr() & 15) == 10 and (weight.data_ptr() & 15) == 10 assert (x.data_ptr() & 15) == 10 and (weight.data_ptr() & 15) == 10
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
expected = F.linear(x, weight) expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5 if m > 1 else 0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5 if m > 1 else 0.25)
@skip_no_gemv @skip_no_gemm
def test_bf16_gemv_scalar_path_handles_misaligned_weight_only(): def test_bf16_gemm_scalar_path_handles_misaligned_weight_only():
"""Weight rows misaligned while x stays 16B-aligned take the scalar-x """Weight rows misaligned while x stays 16B-aligned take the scalar-x
middle and must stay exact.""" middle and must stay exact."""
torch.manual_seed(41) torch.manual_seed(41)
@@ -205,21 +222,21 @@ def test_bf16_gemv_scalar_path_handles_misaligned_weight_only():
x = torch.randn(2, k, device="cuda", dtype=torch.bfloat16) x = torch.randn(2, k, device="cuda", dtype=torch.bfloat16)
assert (weight.data_ptr() & 15) == 10 and (x.data_ptr() & 15) == 0 assert (weight.data_ptr() & 15) == 10 and (x.data_ptr() & 15) == 0
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
torch.testing.assert_close(actual, F.linear(x, weight), rtol=0.02, atol=0.5) torch.testing.assert_close(actual, F.linear(x, weight), rtol=0.02, atol=0.5)
@skip_no_gemv @skip_no_gemm
def test_bf16_gemv_small_batch_cuda_graph_replay(): def test_bf16_gemm_small_batch_cuda_graph_replay():
torch.manual_seed(31) torch.manual_seed(31)
x = torch.randn(8, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(8, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
for _ in range(3): for _ in range(3):
bf16_gemv(x, weight) bf16_gemm(x, weight)
graph = torch.cuda.CUDAGraph() graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph): with torch.cuda.graph(graph):
actual = bf16_gemv(x, weight) actual = bf16_gemm(x, weight)
x.copy_(torch.randn_like(x)) x.copy_(torch.randn_like(x))
graph.replay() graph.replay()
@@ -227,13 +244,13 @@ def test_bf16_gemv_small_batch_cuda_graph_replay():
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25) torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize( @pytest.mark.parametrize(
"make_args,error", "make_args,error",
[ [
( (
lambda: ( lambda: (
torch.randn(9, 16, device="cuda", dtype=torch.bfloat16), torch.randn(65, 16, device="cuda", dtype=torch.bfloat16),
torch.randn(8, 16, device="cuda", dtype=torch.bfloat16), torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
), ),
"M must", "M must",
@@ -256,6 +273,124 @@ def test_bf16_gemv_small_batch_cuda_graph_replay():
), ),
], ],
) )
def test_bf16_gemv_rejects_unsupported_inputs(make_args, error): def test_bf16_gemm_rejects_unsupported_inputs(make_args, error):
with pytest.raises(RuntimeError, match=error): with pytest.raises(RuntimeError, match=error):
bf16_gemv(*make_args()) bf16_gemm(*make_args())
# ---------------------------------------------------------------------------
# Tiled path: M in (8, 64]
# ---------------------------------------------------------------------------
@skip_no_gemm
@pytest.mark.parametrize("m", [9, 12, 16, 17, 24, 32, 33, 48, 64])
@pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912)])
def test_bf16_gemm_tiled_matches_decode_batches(m, n, k):
torch.manual_seed(31 + m)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.empty(n, k, device="cuda", dtype=torch.bfloat16)
weight.normal_(mean=0.0, std=0.02)
actual = bf16_gemm(x, weight)
assert actual.shape == (m, n)
_assert_close_fp64(actual, x, weight)
@skip_no_gemm
@pytest.mark.parametrize("m", [12, 64])
def test_bf16_gemm_tiled_matches_lm_head(m):
# N=100000 fills the SMs with N tiles alone: the splits=1 epilogue.
torch.manual_seed(37 + m)
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.empty(100000, 1536, device="cuda", dtype=torch.bfloat16)
weight.normal_(mean=0.0, std=0.02)
actual = bf16_gemm(x, weight)
_assert_close_fp64(actual, x, weight)
@skip_no_gemm
@pytest.mark.parametrize(
"m,n,k",
[
(12, 100, 72),
(12, 96, 1536),
(12, 1632, 1536),
(64, 100, 72),
(17, 200, 8),
(33, 160, 152),
],
)
def test_bf16_gemm_tiled_handles_remainder_tiles(m, n, k):
# N not a multiple of 64 (predicated epilogue columns) and K not a
# multiple of 64 (zero-filled staging chunks).
torch.manual_seed(41 + m + n + k)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemm(x, weight)
_assert_close_fp64(actual, x, weight)
@skip_no_gemm
@pytest.mark.parametrize("m,n,k", [(12, 1536, 1536), (64, 6912, 1536)])
def test_bf16_gemm_tiled_fuses_bias(m, n, k):
# (12, 1536) exercises the narrow-N deep-K config; (64, 6912) the
# wide-N default.
torch.manual_seed(43 + m)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
bias = torch.randn(n, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemm(x, weight, bias)
_assert_close_fp64(actual, x, weight, bias)
@skip_no_gemm
def test_bf16_gemm_tiled_deterministic_across_runs():
# Single-pass K accumulation with no atomics: reruns are bitwise
# identical — CUDA Graph replay relies on this.
torch.manual_seed(47)
x = torch.randn(16, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
first = bf16_gemm(x, weight)
second = bf16_gemm(x, weight)
assert torch.equal(first, second)
@skip_no_gemm
def test_bf16_gemm_tiled_cuda_graph_replay():
torch.manual_seed(53)
x = torch.randn(16, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
for _ in range(3):
bf16_gemm(x, weight)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
actual = bf16_gemm(x, weight)
x.copy_(torch.randn_like(x))
graph.replay()
_assert_close_fp64(actual, x, weight)
@skip_no_gemm
def test_bf16_gemm_tiled_rejects_k_not_multiple_of_8():
x = torch.randn(12, 12, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(64, 12, device="cuda", dtype=torch.bfloat16)
with pytest.raises(RuntimeError, match="multiple of 8"):
bf16_gemm(x, weight)
@skip_no_gemm
def test_bf16_gemm_tiled_rejects_misaligned_x():
# A 2-byte storage offset breaks the 16B alignment the tiled path
# stages chunks on; the M <= 8 GEMV path still accepts it.
k = 1536
storage = torch.randn(12 * k + 1, device="cuda", dtype=torch.bfloat16)
x8 = storage[1 : 1 + 8 * k].view(8, k)
x12 = storage[1 : 1 + 12 * k].view(12, k)
weight = torch.randn(1536, k, device="cuda", dtype=torch.bfloat16)
torch.testing.assert_close(
bf16_gemm(x8, weight), F.linear(x8, weight), rtol=0.02, atol=0.5
)
with pytest.raises(RuntimeError, match="16-byte"):
bf16_gemm(x12, weight)
tests/extension/test_linear_dispatch.py
+61 -49
View File
@@ -12,26 +12,26 @@ from astrai.extension.dispatch import explain, op_backend, resolve
# module object explicitly for monkeypatching its private helpers. # module object explicitly for monkeypatching its private helpers.
linear_module = importlib.import_module("astrai.extension.backend.linear") linear_module = importlib.import_module("astrai.extension.backend.linear")
GEMV_AVAILABLE = ( GEMM_AVAILABLE = (
torch.cuda.is_available() torch.cuda.is_available()
and is_available("bf16_gemv") and is_available("bf16_gemm")
and torch.cuda.get_device_capability() >= (8, 0) and torch.cuda.get_device_capability() >= (8, 0)
) )
skip_no_gemv = pytest.mark.skipif( skip_no_gemm = pytest.mark.skipif(
not GEMV_AVAILABLE, not GEMM_AVAILABLE,
reason="BF16 GEMV requires a built kernel and compute capability 8.0+", reason="BF16 GEMM requires a built kernel and compute capability 8.0+",
) )
def _routes_to_gemv(monkeypatch, x, weight, bias=None) -> bool: def _routes_to_gemm(monkeypatch, x, weight, bias=None) -> bool:
"""Patch the GEMV entry point to a sentinel and report whether """Patch the GEMM entry point to a sentinel and report whether
``linear`` selected it (torch fallback would compute a real tensor).""" ``linear`` selected it (torch fallback would compute a real tensor)."""
sentinel = object() sentinel = object()
def fake_gemv(x, weight, bias): def fake_gemm(x, weight, bias):
return sentinel return sentinel
monkeypatch.setattr(linear_module, "_inference_bf16_gemv", fake_gemv) monkeypatch.setattr(linear_module, "_inference_bf16_gemm", fake_gemm)
return linear(x, weight, bias) is sentinel return linear(x, weight, bias) is sentinel
@@ -52,7 +52,7 @@ def test_model_linear_routes_through_backend(monkeypatch):
def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog): def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
monkeypatch.setenv("ASTRAI_GEMV", "invalid-test-mode") monkeypatch.setenv("ASTRAI_GEMM", "invalid-test-mode")
x = torch.randn(2, 8) x = torch.randn(2, 8)
weight = torch.randn(4, 8) weight = torch.randn(4, 8)
with caplog.at_level(logging.WARNING): with caplog.at_level(logging.WARNING):
@@ -62,7 +62,7 @@ def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch): def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1") monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(2, 8, requires_grad=True) x = torch.randn(2, 8, requires_grad=True)
weight = torch.randn(4, 8, requires_grad=True) weight = torch.randn(4, 8, requires_grad=True)
actual = linear(x, weight) actual = linear(x, weight)
@@ -73,71 +73,83 @@ def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
assert weight.grad is not None assert weight.grad is not None
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize("m", [2, 3, 4]) @pytest.mark.parametrize("m", [1, 2, 3, 4, 5, 6, 7, 8])
def test_auto_selects_small_decode_batches(monkeypatch, m): def test_auto_selects_small_decode_batches(monkeypatch, m):
monkeypatch.setenv("ASTRAI_GEMV", "auto") monkeypatch.setenv("ASTRAI_GEMM", "auto")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(): with torch.no_grad():
assert _routes_to_gemv(monkeypatch, x, weight) assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize("m", [1, 5, 8, 9]) @pytest.mark.parametrize("m", [12, 16, 24, 32])
def test_auto_selects_larger_decode_batches(monkeypatch, m):
"""Auto covers M up to 32; 48+ loses to cuBLAS on long-K shapes."""
monkeypatch.setenv("ASTRAI_GEMM", "auto")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemm
@pytest.mark.parametrize("m", [48, 64, 65])
def test_auto_falls_back_outside_band(monkeypatch, m): def test_auto_falls_back_outside_band(monkeypatch, m):
monkeypatch.setenv("ASTRAI_GEMV", "auto") """M beyond 32 falls back to cuBLAS (measured regression at M=48+)."""
monkeypatch.setenv("ASTRAI_GEMM", "auto")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(): with torch.no_grad():
assert not _routes_to_gemv(monkeypatch, x, weight) assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close( torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25 linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
) )
@skip_no_gemv @skip_no_gemm
def test_mode_zero_disables_gemv(monkeypatch): def test_mode_zero_disables_gemm(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "0") monkeypatch.setenv("ASTRAI_GEMM", "0")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(): with torch.no_grad():
assert not _routes_to_gemv(monkeypatch, x, weight) assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(linear(x, weight), F.linear(x, weight)) torch.testing.assert_close(linear(x, weight), F.linear(x, weight))
@skip_no_gemv @skip_no_gemm
@pytest.mark.parametrize("m", [1, 2, 8]) @pytest.mark.parametrize("m", [1, 2, 8, 16, 32])
def test_mode_one_forces_every_capable_batch(monkeypatch, m): def test_mode_one_forces_every_capable_batch(monkeypatch, m):
monkeypatch.setenv("ASTRAI_GEMV", "1") monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(): with torch.no_grad():
assert _routes_to_gemv(monkeypatch, x, weight) assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemv @skip_no_gemm
def test_mode_one_rejects_oversized_batch_and_grad(monkeypatch): def test_mode_one_rejects_oversized_batch_and_grad(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1") monkeypatch.setenv("ASTRAI_GEMM", "1")
weight = torch.randn( weight = torch.randn(
256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True 256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True
) )
with torch.no_grad(): with torch.no_grad():
oversized = torch.randn(9, 1536, device="cuda", dtype=torch.bfloat16) oversized = torch.randn(65, 1536, device="cuda", dtype=torch.bfloat16)
assert not _routes_to_gemv(monkeypatch, oversized, weight) assert not _routes_to_gemm(monkeypatch, oversized, weight)
assert not _routes_to_gemv( assert not _routes_to_gemm(
monkeypatch, torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16), weight monkeypatch, torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16), weight
) )
@skip_no_gemv @skip_no_gemm
def test_mode_one_supports_bias_and_vector_input(monkeypatch): def test_mode_one_supports_bias_and_vector_input(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1") monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 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) bias = torch.randn(256, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(): with torch.no_grad():
assert _routes_to_gemv(monkeypatch, x, weight, bias) assert _routes_to_gemm(monkeypatch, x, weight, bias)
monkeypatch.undo() monkeypatch.undo()
torch.testing.assert_close( torch.testing.assert_close(
linear(x, weight, bias), linear(x, weight, bias),
@@ -147,9 +159,9 @@ def test_mode_one_supports_bias_and_vector_input(monkeypatch):
) )
@skip_no_gemv @skip_no_gemm
def test_dispatched_linear_cuda_graph_replay(monkeypatch): def test_dispatched_linear_cuda_graph_replay(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1") monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(): with torch.no_grad():
@@ -176,44 +188,44 @@ def test_linear_family_is_registered_with_shared_dispatcher():
assert "linear" in explain("linear", x, weight) assert "linear" in explain("linear", x, weight)
@skip_no_gemv @skip_no_gemm
def test_ops_env_override_forces_torch_for_capable_call(monkeypatch): def test_ops_env_override_forces_torch_for_capable_call(monkeypatch):
"""ASTR_OPS=linear=torch must keep working after the M-band rewrite """ASTR_OPS=linear=torch must keep working after the M-band rewrite
(regression: the family was silently dropped from the dispatcher, so (regression: the family was silently dropped from the dispatcher, so
the override warned, fell through, and the gemv kernel still ran).""" the override warned, fell through, and the gemm kernel still ran)."""
monkeypatch.setenv("ASTR_OPS", "linear=torch") monkeypatch.setenv("ASTR_OPS", "linear=torch")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(): with torch.no_grad():
assert not _routes_to_gemv(monkeypatch, x, weight) assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close( torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25 linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
) )
@skip_no_gemv @skip_no_gemm
def test_ops_env_override_forces_gemv(monkeypatch): def test_ops_env_override_forces_gemm(monkeypatch):
monkeypatch.setenv("ASTR_OPS", "linear=gemv") monkeypatch.setenv("ASTR_OPS", "linear=gemm")
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(): with torch.no_grad():
# M=1 is outside the auto band but inside the forced gemv record. # M=1 is outside the auto band but inside the forced gemm record.
assert _routes_to_gemv(monkeypatch, x, weight) assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemv @skip_no_gemm
def test_op_backend_context_selects_torch(monkeypatch): def test_op_backend_context_selects_torch(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMV", "1") monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(), op_backend(linear="torch"): with torch.no_grad(), op_backend(linear="torch"):
assert not _routes_to_gemv(monkeypatch, x, weight) assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close( torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25 linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
) )
# The override is scoped: the forced mode applies again afterwards. # The override is scoped: the forced mode applies again afterwards.
with torch.no_grad(): with torch.no_grad():
assert _routes_to_gemv(monkeypatch, x, weight) assert _routes_to_gemm(monkeypatch, x, weight)
def test_op_backend_rejects_unknown_linear_handle(): def test_op_backend_rejects_unknown_linear_handle():
tests/extension/test_swiglu_dispatch.py
+3 -4
View File
@@ -86,10 +86,9 @@ def test_mode_one_forces_supported_shape(monkeypatch):
@skip_no_swiglu @skip_no_swiglu
def test_auto_uses_unfused_chain_until_shape_is_qualified(monkeypatch): def test_auto_uses_fused_chain_for_decode_batches(monkeypatch):
# The fusion table is empty, so auto keeps the unfused linear-backend # Auto adopts the fused primitive for the decode band; numerics match
# chain. The linear backend may still dispatch its own GEMV for M=4, # the unfused linear-backend chain within BF16 accumulation-order noise.
# hence the relaxed tolerance versus the pure-torch reference.
monkeypatch.setenv("ASTRAI_SWIGLU", "auto") monkeypatch.setenv("ASTRAI_SWIGLU", "auto")
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16)