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

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

Browse Source 
- delete csrc/kernels/gemm.cu and swiglu.cu and drop their CMake and setup.py registration
- remove the ops wrappers plus backend/linear.py and backend/swiglu.py so Linear and MLP call F.linear directly
- drop the four gemm and swiglu kernel test files and prune the stale cuda_kernels.md sections
- add csrc/bench benchmarks for the remaining kernels: attention decode prefill paged decode paged prefill versus single-launch SDPA references, rotary versus the torch fallback, fp8 quantize and mm_fp8 versus torch baselines
- attention, rotary_emb, and fp8_ops kernels are unchanged
This commit is contained in:
ViperEkura
2026-09-05 01:38:10 +08:00
parent a77e35dd51
commit 6709534d64
24 changed files with 1306 additions and 3394 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/__init__.py
+7 -11
View File
@@ -1,11 +1,10 @@
"""CUDA kernel wrappers, operator dispatch, and backend selection.
Public API:
- ``attention``, ``linear``, ``swiglu``, ``apply_rotary_emb`` — op
families with safe torch fallbacks (see ``astrai.extension.backend``)
- ``attention``, ``apply_rotary_emb`` — op families with safe torch
fallbacks (see ``astrai.extension.backend``)
- ``attn_decode`` / ``attn_prefill`` / ``attn_paged_decode`` /
``attn_paged_prefill`` — direct attention kernel wrappers
- ``bf16_gemm`` / ``bf16_swiglu`` — directly callable linear/MLP kernels
- ``AttentionBackend`` / ``TorchNativeBackend`` / ``CudaBackend`` /
``FlashAttnBackend`` — attention backend strategies
- ``resolve`` / ``explain`` / ``op_backend`` / ``env_mode`` — the shared
@@ -14,6 +13,9 @@ Public API:
Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
(blhd). Scale is always ``1/sqrt(head_dim)``. Wrapper functions call their
compiled CUDA kernels directly; fallback is the backend's responsibility.
Linear projections and dense-MLP SwiGLU run plain torch (``F.linear`` /
``Linear`` / ``MLP``); the former bf16_gemm / bf16_swiglu kernels and their
backends were removed.
"""
from astrai.extension.backend import (
@@ -27,8 +29,6 @@ from astrai.extension.backend import (
attention,
attn_backend,
get_backend,
linear,
swiglu,
)
from astrai.extension.dispatch import (
Axes,
@@ -52,9 +52,8 @@ from astrai.extension.ops import (
TensorLayout,
attn_decode,
attn_paged_decode,
attn_paged_prefill,
attn_prefill,
bf16_gemm,
bf16_swiglu,
)
__all__ = [
@@ -68,13 +67,10 @@ __all__ = [
"attention",
"attn_backend",
"get_backend",
"linear",
"swiglu",
"attn_decode",
"attn_paged_decode",
"attn_prefill",
"bf16_gemm",
"bf16_swiglu",
"attn_paged_prefill",
"is_available",
"KERNEL_NAMES",
"apply_rotary_emb",
astrai/extension/backend/__init__.py
-4
View File
@@ -11,9 +11,7 @@ from astrai.extension.backend.attention import (
attn_backend,
get_backend,
)
from astrai.extension.backend.linear import linear
from astrai.extension.backend.rotary import apply_rotary_emb
from astrai.extension.backend.swiglu import swiglu
__all__ = [
"ATTN_BACKEND",
@@ -26,6 +24,4 @@ __all__ = [
"attention",
"attn_backend",
"get_backend",
"linear",
"swiglu",
]
astrai/extension/backend/linear.py
-215
View File
@@ -1,215 +0,0 @@
"""Inference-only dispatch for AstrAI linear layers.
The CUDA GEMM path is sized by the decode batch M. M in [1, 8] uses the
register-resident GEMV kernel (any K); M in (8, 64] uses the tiled
kernel (K % 8 == 0, 16-byte-aligned tensors). Automatic mode
selects GEMM for M in [1, 64] where the primitive is capable; every
training, prefill-sized, or unsupported call falls back to PyTorch.
The family stays registered with the shared operator dispatcher, so
``op_backend(linear=...)``, ``ASTR_OPS=linear=...``, and ``resolve`` /
``explain`` keep working. The per-layer hot path only consults the
dispatcher when one of those selections is active.
"""
from typing import Any, Dict, List, Optional
import torch
import torch.nn.functional as F
from torch import Tensor
from astrai.extension.dispatch import (
ImplRecord,
Spec,
axis,
env_mode,
env_selection,
get_override,
register_family,
resolve,
tensor_axes,
)
from astrai.extension.loader import is_available
from astrai.extension.ops.gemm import bf16_gemm
def _torch_linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
return F.linear(x, weight, bias)
def _inference_bf16_gemm(
x: Tensor, weight: Tensor, bias: Optional[Tensor] = None
) -> Tensor:
return bf16_gemm(
x.detach(),
weight.detach(),
bias.detach() if bias is not None else None,
)
def _gemm_capable(x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> bool:
"""Check whether bf16_gemm can safely handle the call."""
if (
torch.is_grad_enabled()
or not x.is_cuda
or x.dtype != torch.bfloat16
or weight.dtype != torch.bfloat16
or weight.ndim != 2
or x.ndim not in (1, 2)
or x.shape[-1] != weight.shape[1]
or x.device != weight.device
or not x.is_contiguous()
or not weight.is_contiguous()
or torch.cuda.get_device_capability(x.get_device()) < (8, 0)
or not is_available("bf16_gemm")
):
return False
m = 1 if x.ndim == 1 else x.shape[0]
# M <= 32 is where the kernel wins: L2-rotation measurements on L20
# show every production shape at M=24-32 winning or tying, while
# M=48-64 loses the long-K down_proj by 8-10% (cuBLAS switches to a
# wider tile there). The kernel itself still accepts M <= 64 when
# called directly through astrai.extension.ops.gemm.
if not (1 <= m <= 32):
return False
# Vocabulary-sized lm_head weights (N in the tens of thousands+) stream
# better through cuBLAS: our skinny path ties it at M<=8 and the tiled
# path loses ~4% at M=9-16 (L2-rotation measurements on L20). Gate the
# whole shape family out instead of splitting hairs per M band.
if weight.shape[0] > 32768:
return False
# M > 8 tiled path requires K % 8 == 0 and 16-byte alignment.
k = x.shape[-1]
if m > 8 and (
k % 8 != 0 or (x.data_ptr() & 15) != 0 or (weight.data_ptr() & 15) != 0
):
return False
return bias is None or (
bias.device == x.device
and bias.dtype == torch.bfloat16
and bias.ndim == 1
and bias.shape[0] == weight.shape[0]
and bias.is_contiguous()
)
def _axes(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Dict[str, Any]:
weight_shape = tuple(weight.shape)
m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None)
supported_m = m is not None and 1 <= m <= 64
shape_matches = (
weight.ndim == 2
and x.ndim in (1, 2)
and bool(x.shape)
and x.shape[-1] == weight_shape[-1]
)
same_device = x.device == weight.device and (
bias is None or bias.device == x.device
)
bias_supported = bias is None or (
bias.ndim == 1
and weight.ndim == 2
and bias.shape[0] == weight_shape[0]
and bias.dtype == torch.bfloat16
and bias.is_contiguous()
)
capability = torch.cuda.get_device_capability(x.device) if x.is_cuda else None
return tensor_axes(
x,
mode=env_mode("ASTRAI_GEMM"),
m=m,
supported_m=supported_m,
auto_m=supported_m,
shape_matches=shape_matches,
same_device=same_device,
weight_dtype=weight.dtype,
x_contiguous=x.is_contiguous(),
weight_contiguous=weight.is_contiguous(),
bias_supported=bias_supported,
capability=capability,
)
_SPEC_CAPABLE = (
axis("device_cuda").truthy()
& axis("dtype").in_(torch.bfloat16)
& axis("weight_dtype").in_(torch.bfloat16)
& axis("grad_enabled").eq(False)
& axis("supported_m").truthy()
& axis("shape_matches").truthy()
& axis("same_device").truthy()
& axis("x_contiguous").truthy()
& axis("weight_contiguous").truthy()
& axis("bias_supported").truthy()
& Spec.of(
lambda ax: ax.get("capability") is not None and ax.get("capability") >= (8, 0),
"capability>=sm_80",
)
)
_SPEC_AUTO = _SPEC_CAPABLE & axis("auto_m").truthy()
def _linear_records() -> List[ImplRecord]:
mode = env_mode("ASTRAI_GEMM")
gemm_priority = 0 if mode == "1" else 100
auto_priority = 0 if mode == "auto" else 90
torch_priority = 0 if mode == "0" else 50
return [
ImplRecord(
family="linear",
name="gemm",
obj=_inference_bf16_gemm,
spec=_SPEC_CAPABLE,
available=lambda: is_available("bf16_gemm"),
priority=gemm_priority,
),
ImplRecord(
family="linear",
name="auto_gemm",
obj=_inference_bf16_gemm,
spec=_SPEC_AUTO,
available=lambda: is_available("bf16_gemm"),
priority=auto_priority,
),
ImplRecord(
family="linear",
name="torch",
obj=_torch_linear,
spec=Spec.always(),
priority=torch_priority,
),
]
def _fallback_record() -> ImplRecord:
return ImplRecord(
family="linear",
name="torch",
obj=_torch_linear,
spec=Spec.always(),
priority=999,
)
register_family("linear", _axes, _linear_records, _fallback_record)
def linear(x: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
"""Apply a linear projection with safe inference-only GEMM dispatch.
``ASTRAI_GEMM=0`` always uses PyTorch, ``1`` forces GEMM whenever the
primitive can safely handle the call (M in [1, 64], K % 8 == 0 and
16-byte-aligned for M > 8), and ``auto`` (the default) selects GEMM
for all capable decode batches.
"""
if get_override("linear") is not None or env_selection("linear") is not None:
return resolve("linear", x, weight, bias).record.obj(x, weight, bias)
mode = env_mode("ASTRAI_GEMM")
if mode != "0" and _gemm_capable(x, weight, bias):
return _inference_bf16_gemm(x, weight, bias)
return _torch_linear(x, weight, bias)
__all__ = ["linear"]
astrai/extension/backend/swiglu.py
-67
View File
@@ -1,67 +0,0 @@
"""Inference-only fused SwiGLU selection for dense MLP layers."""
import torch
import torch.nn.functional as F
from torch import Tensor
from astrai.extension.backend.linear import linear
from astrai.extension.dispatch import env_mode
from astrai.extension.loader import is_available
from astrai.extension.ops.swiglu import bf16_swiglu
def _unfused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
# Keep the existing linear backend in the fallback chain. This preserves
# any independently qualified GEMV batches instead of making the fusion
# decision suppress linear-level optimizations.
return linear(x, up_weight) * F.silu(linear(x, gate_weight))
def _fused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
return bf16_swiglu(x.detach(), up_weight.detach(), gate_weight.detach())
def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool:
return not (
torch.is_grad_enabled()
or not x.is_cuda
or x.dtype != torch.bfloat16
or up_weight.dtype != torch.bfloat16
or gate_weight.dtype != torch.bfloat16
or x.ndim not in (1, 2)
or up_weight.ndim != 2
or gate_weight.ndim != 2
or (x.ndim == 2 and not 1 <= x.shape[0] <= 8)
or up_weight.shape != gate_weight.shape
or x.shape[-1] != up_weight.shape[1]
or x.shape[-1] % 8 != 0
or x.device != up_weight.device
or x.device != gate_weight.device
or not x.is_contiguous()
or not up_weight.is_contiguous()
or not gate_weight.is_contiguous()
# The fused kernel reads all streams as uint4; contiguous-but-offset
# views are routed to the unfused chain instead of failing.
or (x.data_ptr() & 15) != 0
or (up_weight.data_ptr() & 15) != 0
or (gate_weight.data_ptr() & 15) != 0
or not is_available("bf16_swiglu")
)
def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor:
"""Apply the dense-MLP SwiGLU projection with a safe torch fallback.
``ASTRAI_SWIGLU=0`` keeps the unfused linear-backend chain; ``1`` forces
the fused primitive for supported inputs; ``auto`` (the default) uses
the fused primitive for decode batches with M in ``{1, ..., 8}``. The
fused kernel reads x once and covers both projections plus the SiLU
gate-multiply in a single launch, measured 9-15% faster than the
unfused chain per MLP call on L20 with L2-thrashing weight rotation.
"""
if env_mode("ASTRAI_SWIGLU") != "0" and _swiglu_capable(x, up_weight, gate_weight):
return _fused_swiglu(x, up_weight, gate_weight)
return _unfused_swiglu(x, up_weight, gate_weight)
__all__ = ["swiglu"]
astrai/extension/ops/__init__.py
-4
View File
@@ -7,9 +7,7 @@ from astrai.extension.ops.attention import (
attn_paged_prefill,
attn_prefill,
)
from astrai.extension.ops.gemm import bf16_gemm
from astrai.extension.ops.rotary import rotary_emb
from astrai.extension.ops.swiglu import bf16_swiglu
__all__ = [
"TensorLayout",
@@ -17,7 +15,5 @@ __all__ = [
"attn_paged_decode",
"attn_paged_prefill",
"attn_prefill",
"bf16_gemm",
"bf16_swiglu",
"rotary_emb",
]
astrai/extension/ops/gemm.py
-27
View File
@@ -1,27 +0,0 @@
"""Stateless wrapper for the directly callable BF16 GEMM primitive."""
from typing import Optional
import torch
from astrai.extension.loader import get_module
def bf16_gemm(
x: torch.Tensor,
weight: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
"""Compute ``F.linear(x, weight, bias)`` for up to 64 BF16 rows.
``x`` must have shape ``[K]`` or ``[M, K]`` with M in ``[1, 64]``, and
``weight`` must be a contiguous row-major ``[N, K]`` tensor. M in
``[1, 8]`` uses the register-resident skinny GEMM kernel (any K);
larger M uses the tiled kernel (K must be a multiple of 8 with
16-byte-aligned tensors). This primitive is inference-only and
intentionally performs no fallback or model-level dispatch.
"""
return get_module("bf16_gemm").bf16_gemm(x, weight, bias)
__all__ = ["bf16_gemm"]
astrai/extension/ops/swiglu.py
-22
View File
@@ -1,22 +0,0 @@
"""Stateless wrapper for the directly callable fused BF16 SwiGLU primitive."""
import torch
from astrai.extension.loader import get_module
def bf16_swiglu(
x: torch.Tensor,
up_weight: torch.Tensor,
gate_weight: torch.Tensor,
) -> torch.Tensor:
"""Compute ``linear(x, up) * silu(linear(x, gate))`` for M in [1, 8].
Inputs must be contiguous BF16 CUDA tensors. Both weights use row-major
``[N, K]`` storage with identical shapes, and K must be divisible by 8.
The primitive is inference-only and performs no fallback.
"""
return get_module("bf16_swiglu").bf16_swiglu(x, up_weight, gate_weight)
__all__ = ["bf16_swiglu"]
astrai/model/components/linear.py
+2 -3
View File
@@ -1,9 +1,8 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from astrai.extension.backend.linear import linear
class Linear(nn.Module):
def __init__(
@@ -22,4 +21,4 @@ class Linear(nn.Module):
nn.init.uniform_(self.bias, -bound, bound)
def forward(self, x: Tensor) -> Tensor:
return linear(x, self.weight, self.bias)
return F.linear(x, self.weight, self.bias)
astrai/model/components/mlp.py
+1 -2
View File
@@ -5,7 +5,6 @@ import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from astrai.extension.backend.swiglu import swiglu
from astrai.factory import BaseFactory
from astrai.model.components.linear import Linear
@@ -39,7 +38,7 @@ class MLP(nn.Module):
self.down = Linear(dim_ffn, dim, init_std=down_init_std)
def forward(self, x: Tensor) -> FFNOutput:
gated = swiglu(x, self.up.weight, self.gate.weight)
gated = self.up(x) * F.silu(self.gate(x))
out = self.down(gated)
return {"hidden_states": out, "aux_loss": None, "router_stats": None}
csrc/CMakeLists.txt
-4
View File
@@ -61,8 +61,6 @@ set(KERNEL_NAMES
attn_prefill
attn_paged_decode
attn_paged_prefill
bf16_gemm
bf16_swiglu
rotary_emb
)
set(KERNEL_SRCS
@@ -70,8 +68,6 @@ set(KERNEL_SRCS
attention/prefill.cu
attention/paged_decode.cu
attention/paged_prefill.cu
gemm.cu
swiglu.cu
rotary_emb.cu
)
csrc/bench/benchmark_attention.py
+649
View File
@@ -0,0 +1,649 @@
"""Benchmark the four attention kernels against single-launch torch SDPA.
Suites (--suite): decode, prefill, paged_decode, paged_prefill, all. The
torch side times one SDPA call per step over dense tensors: GQA expansion,
page-table gathers, padding, and masks are built once outside the timed
region, and masked calls prefer the cuDNN backend (the default masked path
is the slow math backend). The kernel's timed work still includes its fused
paged reads and current-token K/V append. Agreement is checked against the
same reference.
"""
from __future__ import annotations
import json
import math
import statistics
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Optional
import click
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
from astrai.extension import is_available
from astrai.extension.ops import (
attn_decode,
attn_paged_decode,
attn_paged_prefill,
attn_prefill,
)
from astrai.inference.workspace import MAX_SPLITS, Q_TILE_ROWS
@dataclass(frozen=True)
class GqaConfig:
"""One model family's attention geometry (llama-style GQA)."""
name: str
hq: int
hkv: int
head_dim: int
DEFAULT_CONFIGS = (
GqaConfig("llama2_7b", 32, 8, 128),
GqaConfig("llama3_70b", 64, 8, 128),
GqaConfig("qwen2_7b", 28, 4, 128),
GqaConfig("llama3_8b_d64", 32, 8, 64),
)
# (batch, per-request context length); context includes the token being
# decoded (kv_len = context, the last slot written in-kernel).
DECODE_CASES = ((1, 4096), (8, 4096), (32, 2048), (64, 1024))
# (batch, q_len); prefill from scratch so kv_len == q_len.
PREFILL_CASES = ((1, 2048), (1, 4096), (4, 1024), (8, 512))
def parse_config(value: str) -> GqaConfig:
parts = value.split(":")
if len(parts) != 4 or not parts[0]:
raise click.BadParameter("config must use NAME:HQ:HKV:HEAD_DIM")
try:
hq, hkv, head_dim = (int(item) for item in parts[1:])
except ValueError as exc:
raise click.BadParameter("HQ/HKV/HEAD_DIM must be integers") from exc
if hq <= 0 or hkv <= 0 or head_dim <= 0 or hq % hkv or head_dim % 32:
raise click.BadParameter(
"HQ/HKV positive with HQ % HKV == 0; HEAD_DIM % 32 == 0"
)
return GqaConfig(parts[0], hq, hkv, head_dim)
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iterations):
operation()
end.record()
end.synchronize()
return start.elapsed_time(end) / iterations
def summarize(values: list[float]) -> dict[str, float]:
ordered = sorted(values)
return {
"median_ms": statistics.median(ordered),
"p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
}
def measure_operations(
operations: dict[str, Callable[[], torch.Tensor]],
*,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, list[float]]:
for operation in operations.values():
for _ in range(warmup):
operation()
torch.cuda.synchronize()
samples: dict[str, list[float]] = {name: [] for name in operations}
order = tuple(operations)
# A-B-B-A order balances cache, clock, and temperature drift.
for _ in range(trials):
for name in (*order, *reversed(order)):
samples[name].append(time_operation(operations[name], iterations))
return samples
def repeat_kv_heads(x: torch.Tensor, n_rep: int) -> torch.Tensor:
"""Expand [*, n_kv_heads, head_dim] to [*, n_kv_heads * n_rep, head_dim]
with the backend's grouping (kv head = q head // n_rep)."""
if n_rep == 1:
return x
n_heads, head_dim = x.shape[-2:]
return (
x.unsqueeze(-2)
.expand(*x.shape[:-2], n_heads, n_rep, head_dim)
.reshape(*x.shape[:-2], n_heads * n_rep, head_dim)
)
def sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs) -> torch.Tensor:
"""SDPA over blhd tensors: [batch, seq, heads, head_dim] -> blhd."""
out = F.scaled_dot_product_attention(
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), **kwargs
)
return out.transpose(1, 2)
def prefer_cudnn_sdpa(
call: Callable[[], torch.Tensor],
) -> Callable[[], torch.Tensor]:
"""Return a closure running the masked SDPA ``call`` on cuDNN attention
when the backend accepts the bool mask, else torch's default (the
default masked path falls back to the much slower math backend)."""
def with_cudnn() -> torch.Tensor:
with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
return call()
try:
with_cudnn()
except Exception:
return call
return with_cudnn
def ragged_lens(batch: int, span: int) -> list[int]:
"""Deterministic mixed lengths spanning [span // 2, span]."""
if batch == 1:
return [span]
step = max(span // 2 // (batch - 1), 1)
return [span - (batch - 1 - i) * step for i in range(batch)]
def cumsum_indptr(lens: list[int]) -> torch.Tensor:
return torch.tensor(
[0, *torch.tensor(lens).cumsum(0).tolist()], dtype=torch.int32, device="cuda"
)
@dataclass(frozen=True)
class PagedInputs:
"""Standalone replicas of the PagePool / InferenceWorkspace tensors."""
q: torch.Tensor
k_cache: torch.Tensor
v_cache: torch.Tensor
req_to_token: torch.Tensor
req_pool_indices: torch.Tensor
kv_indptr: torch.Tensor
def build_paged_inputs(
config: GqaConfig, kv_lens: list[int], q_lens: Optional[list[int]]
) -> PagedInputs:
"""Flat pool + page table: request ``i`` owns the contiguous slot range
``[offset_i, offset_i + kv_len_i)``. Q is packed across requests when
``q_lens`` is given (ragged prefill), else [B, Hq, D] (decode)."""
batch = len(kv_lens)
pool = torch.randn(
sum(kv_lens),
config.hkv,
config.head_dim,
device="cuda",
dtype=torch.bfloat16,
)
req_to_token = torch.zeros(batch, max(kv_lens), dtype=torch.int32, device="cuda")
offset = 0
for i, length in enumerate(kv_lens):
req_to_token[i, :length] = torch.arange(
offset, offset + length, dtype=torch.int32, device="cuda"
)
offset += length
return PagedInputs(
q=torch.randn(
sum(q_lens) if q_lens is not None else batch,
config.hq,
config.head_dim,
device="cuda",
dtype=torch.bfloat16,
),
k_cache=pool,
v_cache=torch.randn_like(pool),
req_to_token=req_to_token,
req_pool_indices=torch.arange(batch, dtype=torch.int32, device="cuda"),
kv_indptr=cumsum_indptr(kv_lens),
)
def report_result(
suite: str,
config: GqaConfig,
case: dict[str, int],
operations: dict[str, Callable[[], torch.Tensor]],
samples: dict[str, list[float]],
io_bytes: int,
reference: torch.Tensor,
actual: torch.Tensor,
) -> dict[str, object]:
difference = actual.float() - reference.float()
result: dict[str, object] = {
"suite": suite,
"config": config.name,
"agreement": {
"max_abs_error": float(difference.abs().max()),
"cosine_similarity": float(
F.cosine_similarity(
actual.float().flatten(), reference.float().flatten(), dim=0
)
),
},
"estimated_io_bytes": io_bytes,
**case,
}
for name, operation in operations.items():
latency = summarize(samples[name])
result[name] = {
"effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
**latency,
}
speedup = (result["torch"]["median_ms"] / result["cuda"]["median_ms"] - 1.0) * 100.0
label = f"B={case.get('batch')}" + (
f" ctx={case['context']}" if "context" in case else f" q={case['q_len']}"
)
print(
f"{suite},{config.name},{label},{result['torch']['median_ms']:.4f},"
f"{result['cuda']['median_ms']:.4f},{speedup:+.1f}%,"
f"{result['agreement']['max_abs_error']:.4f}"
)
return result
# ---------------------------------------------------------------------------
# Suites
# ---------------------------------------------------------------------------
def benchmark_decode(
config: GqaConfig,
batch: int,
context: int,
*,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, object]:
q = torch.randn(
batch, 1, config.hq, config.head_dim, device="cuda", dtype=torch.bfloat16
)
k = torch.randn(
batch,
context,
config.hkv,
config.head_dim,
device="cuda",
dtype=torch.bfloat16,
)
v = torch.randn_like(k)
# GQA expansion is data preparation, not attention compute — build it
# once so the timed torch side is a single SDPA launch.
k_expanded = repeat_kv_heads(k, config.hq // config.hkv)
v_expanded = repeat_kv_heads(v, config.hq // config.hkv)
def torch_op() -> torch.Tensor:
return sdpa(q, k_expanded, v_expanded)
def cuda_op() -> torch.Tensor:
return attn_decode(q, k, v, is_causal=True)
operations = {"torch": torch_op, "cuda": cuda_op}
samples = measure_operations(
operations, warmup=warmup, iterations=iterations, trials=trials
)
io_bytes = (2 * q.numel() + 2 * k.numel()) * q.element_size()
return report_result(
"decode",
config,
{"batch": batch, "context": context},
operations,
samples,
io_bytes,
torch_op(),
cuda_op(),
)
def benchmark_prefill(
config: GqaConfig,
batch: int,
q_len: int,
*,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, object]:
q = torch.randn(
batch, q_len, config.hq, config.head_dim, device="cuda", dtype=torch.bfloat16
)
k = torch.randn(
batch, q_len, config.hkv, config.head_dim, device="cuda", dtype=torch.bfloat16
)
v = torch.randn_like(k)
k_expanded = repeat_kv_heads(k, config.hq // config.hkv)
v_expanded = repeat_kv_heads(v, config.hq // config.hkv)
def torch_op() -> torch.Tensor:
return sdpa(q, k_expanded, v_expanded, is_causal=True)
def cuda_op() -> torch.Tensor:
return attn_prefill(q, k, v, is_causal=True)
operations = {"torch": torch_op, "cuda": cuda_op}
samples = measure_operations(
operations, warmup=warmup, iterations=iterations, trials=trials
)
io_bytes = (2 * q.numel() + 2 * k.numel()) * q.element_size()
return report_result(
"prefill",
config,
{"batch": batch, "q_len": q_len},
operations,
samples,
io_bytes,
torch_op(),
cuda_op(),
)
def benchmark_paged_decode(
config: GqaConfig,
batch: int,
context: int,
*,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, object]:
n_rep = config.hq // config.hkv
kv_lens = [length + 1 for length in ragged_lens(batch, context)]
inputs = build_paged_inputs(config, kv_lens, None)
new_k = torch.randn(
batch, config.hkv, config.head_dim, device="cuda", dtype=torch.bfloat16
)
new_v = torch.randn_like(new_k)
o_part = torch.empty(
batch,
config.hq,
MAX_SPLITS,
config.head_dim,
dtype=torch.float32,
device="cuda",
)
ml_part = torch.empty(
batch, config.hq, MAX_SPLITS, 2, dtype=torch.float32, device="cuda"
)
out_buf = torch.empty(
batch, config.hq, config.head_dim, dtype=torch.bfloat16, device="cuda"
)
def cuda_op() -> torch.Tensor:
return attn_paged_decode(
inputs.q,
inputs.k_cache,
inputs.v_cache,
inputs.req_to_token,
inputs.req_pool_indices,
inputs.kv_indptr,
new_k=new_k,
new_v=new_v,
is_causal=True,
o_part_buf=o_part,
ml_part_buf=ml_part,
out_buf=out_buf,
)
# Reference-side data preparation happens once, outside the timed
# closure: append the current-token K/V into the pool (the kernel does
# this fused inside its launch), gather padded K/V, expand GQA heads.
max_len = max(kv_lens)
slots = inputs.req_to_token[:, :max_len].long()
last_slots = inputs.req_to_token[
torch.arange(batch, device="cuda"), torch.tensor(kv_lens) - 1
].long()
inputs.k_cache[last_slots] = new_k
inputs.v_cache[last_slots] = new_v
k_expanded = repeat_kv_heads(inputs.k_cache[slots], n_rep)
v_expanded = repeat_kv_heads(inputs.v_cache[slots], n_rep)
position = torch.arange(max_len, device="cuda")
lengths = torch.tensor(kv_lens, device="cuda", dtype=torch.long)
keep_mask = (position[None, :] < lengths[:, None])[:, None, None, :]
q_batched = inputs.q.unsqueeze(1)
sdpa_call = prefer_cudnn_sdpa(
lambda: sdpa(q_batched, k_expanded, v_expanded, attn_mask=keep_mask)
)
def torch_op() -> torch.Tensor:
return sdpa_call().squeeze(1) # [B, Hq, D]
operations = {"torch": torch_op, "cuda": cuda_op}
samples = measure_operations(
operations, warmup=warmup, iterations=iterations, trials=trials
)
io_bytes = (
2 * inputs.q.numel() # q read + out write
+ 2 * sum(kv_lens) * config.hkv * config.head_dim # k/v reads
+ 2 * new_k.numel() # new k/v writes
) * inputs.q.element_size()
return report_result(
"paged_decode",
config,
{"batch": batch, "context": context},
operations,
samples,
io_bytes,
torch_op(),
cuda_op(),
)
def benchmark_paged_prefill(
config: GqaConfig,
batch: int,
q_len: int,
*,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, object]:
n_rep = config.hq // config.hkv
q_lens = ragged_lens(batch, q_len)
inputs = build_paged_inputs(config, q_lens, q_lens)
qo_indptr = cumsum_indptr(q_lens)
tile_batches, tile_indices = [], []
for request, length in enumerate(q_lens):
n_tiles = (length + Q_TILE_ROWS - 1) // Q_TILE_ROWS
tile_batches.extend([request] * n_tiles)
tile_indices.extend(range(n_tiles))
q_tile_to_batch = torch.tensor(tile_batches, dtype=torch.int32, device="cuda")
q_tile_to_index = torch.tensor(tile_indices, dtype=torch.int32, device="cuda")
def cuda_op() -> torch.Tensor:
return attn_paged_prefill(
inputs.q,
inputs.k_cache,
inputs.v_cache,
inputs.req_to_token,
inputs.req_pool_indices,
inputs.kv_indptr,
qo_indptr,
q_tile_to_batch,
q_tile_to_index,
is_causal=True,
)
# Same once-only preparation: gather padded K/V, expand GQA heads, pad Q,
# build the causal + validity mask. The timed reference is one SDPA call
# plus the packed-row unpack.
max_len = max(q_lens)
slots = inputs.req_to_token[:, :max_len].long()
k_expanded = repeat_kv_heads(inputs.k_cache[slots], n_rep)
v_expanded = repeat_kv_heads(inputs.v_cache[slots], n_rep)
position = torch.arange(max_len, device="cuda")
lengths = torch.tensor(q_lens, device="cuda", dtype=torch.long)
causal = position[None, :, None] >= position[None, None, :]
keep = position[None, None, :] < lengths[:, None, None]
attn_mask = (causal & keep).unsqueeze(1)
q_padded = torch.zeros(
batch,
max_len,
config.hq,
config.head_dim,
device="cuda",
dtype=inputs.q.dtype,
)
for i, length in enumerate(q_lens):
q_padded[i, :length] = inputs.q[int(qo_indptr[i]) : int(qo_indptr[i + 1])]
sdpa_call = prefer_cudnn_sdpa(
lambda: sdpa(q_padded, k_expanded, v_expanded, attn_mask=attn_mask)
)
def torch_op() -> torch.Tensor:
out = sdpa_call()
return torch.cat([out[i, :length] for i, length in enumerate(q_lens)])
operations = {"torch": torch_op, "cuda": cuda_op}
samples = measure_operations(
operations, warmup=warmup, iterations=iterations, trials=trials
)
io_bytes = (
2 * inputs.q.numel() + 2 * sum(q_lens) * config.hkv * config.head_dim
) * inputs.q.element_size()
return report_result(
"paged_prefill",
config,
{"batch": batch, "q_len": q_len},
operations,
samples,
io_bytes,
torch_op(),
cuda_op(),
)
@click.command(help=__doc__)
@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
@click.option(
"--suite",
"suites",
type=click.Choice(("decode", "prefill", "paged_decode", "paged_prefill", "all")),
multiple=True,
default=("all",),
show_default=True,
)
@click.option(
"--config",
"config_values",
multiple=True,
help="Filter defaults by bare name, or add/override with NAME:HQ:HKV:HEAD_DIM.",
)
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
@click.option("--iterations", type=click.IntRange(min=1), default=50, show_default=True)
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
@click.option("--seed", type=int, default=0, show_default=True)
def benchmark_command(
output: Path | None,
suites: tuple[str, ...],
config_values: tuple[str, ...],
warmup: int,
iterations: int,
trials: int,
seed: int,
) -> None:
if not torch.cuda.is_available():
raise click.ClickException("CUDA is required")
kernel_for_suite = {
"decode": "attn_decode",
"prefill": "attn_prefill",
"paged_decode": "attn_paged_decode",
"paged_prefill": "attn_paged_prefill",
}
selected = (
tuple(kernel_for_suite) if "all" in suites else tuple(dict.fromkeys(suites))
)
missing = [
kernel_for_suite[suite]
for suite in selected
if not is_available(kernel_for_suite[suite])
]
if missing:
raise click.ClickException(f"built kernels required: {', '.join(missing)}")
# A bare name filters the matching default; a full spec overrides or appends.
chosen: dict[str, GqaConfig] = {}
for value in config_values:
config = (
next((c for c in DEFAULT_CONFIGS if c.name == value), None)
if ":" not in value
else parse_config(value)
)
if config is None:
raise click.BadParameter(f"unknown default config {value!r}")
chosen[config.name] = config
configs = tuple(chosen.values()) or DEFAULT_CONFIGS
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
runners = {
"decode": (benchmark_decode, DECODE_CASES),
"prefill": (benchmark_prefill, PREFILL_CASES),
"paged_decode": (benchmark_paged_decode, DECODE_CASES),
"paged_prefill": (benchmark_paged_prefill, PREFILL_CASES),
}
print("suite,config,case,torch_ms,cuda_ms,speedup,max_abs")
results = []
with torch.inference_mode():
for suite in selected:
runner, cases = runners[suite]
for config in configs:
for case in cases:
results.append(
runner(
config,
*case,
warmup=warmup,
iterations=iterations,
trials=trials,
)
)
torch.cuda.empty_cache()
if output is not None:
props = torch.cuda.get_device_properties(0)
payload = {
"metadata": {
"gpu_name": props.name,
"compute_capability": f"{props.major}.{props.minor}",
"torch_version": torch.__version__,
"cuda_version": torch.version.cuda,
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
},
"settings": {
"warmup": warmup,
"iterations": iterations,
"trials": trials,
"seed": seed,
"order": "A-B-B-A",
"suites": list(selected),
"configs": [asdict(config) for config in configs],
},
"results": results,
}
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
if __name__ == "__main__":
benchmark_command()
csrc/bench/benchmark_fp8.py
+394
View File
@@ -0,0 +1,394 @@
"""Benchmark the FP8 quantize and GEMM kernels against torch baselines.
Suites (--suite): quantize (plain / delayed-scaling ring / dual-orientation
entries vs the aten float8 cast) and gemm (pre-quantized ``mm_fp8`` in the
NT orientation the fp8 linear path uses, vs bf16 ``F.linear``). GEMM
agreement reports both kernel error (vs the fp32 dequantized fp8 product)
and format error (that product vs the bf16 matmul). FP8 MMA requires
compute capability 89+.
"""
from __future__ import annotations
import json
import math
import statistics
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable
import click
import torch
import torch.nn.functional as F
from astrai.extension import is_available
from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
@dataclass(frozen=True)
class MatrixShape:
name: str
rows: int
cols: int
QUANTIZE_SHAPES = (
MatrixShape("astrai_1b_act", 2048, 1536),
MatrixShape("llama2_7b_act", 2048, 4096),
MatrixShape("llama2_7b_down_w", 4096, 11008),
MatrixShape("llama3_70b_act", 2048, 8192),
)
# GEMM shapes as (N, K) weight mats; M comes from --m-values.
GEMM_SHAPES = (
MatrixShape("llama2_7b_qkv", 4096, 4096),
MatrixShape("llama2_7b_up_gate", 11008, 4096),
MatrixShape("llama2_7b_down", 4096, 11008),
MatrixShape("llama3_70b_up_gate", 28672, 8192),
)
def parse_positive_ints(value: str) -> tuple[int, ...]:
try:
values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
except ValueError as exc:
raise click.BadParameter("expected comma-separated integers") from exc
if not values or any(item <= 0 for item in values):
raise click.BadParameter("values must be positive integers")
return values
def parse_shape(value: str) -> MatrixShape:
parts = value.split(":")
if len(parts) != 3 or not parts[0]:
raise click.BadParameter("shape must use NAME:ROWS:COLS")
try:
rows, cols = (int(item) for item in parts[1:])
except ValueError as exc:
raise click.BadParameter("ROWS and COLS must be integers") from exc
if rows <= 0 or cols <= 0:
raise click.BadParameter("ROWS and COLS must be positive")
return MatrixShape(parts[0], rows, cols)
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iterations):
operation()
end.record()
end.synchronize()
return start.elapsed_time(end) / iterations
def summarize(values: list[float]) -> dict[str, float]:
ordered = sorted(values)
return {
"median_ms": statistics.median(ordered),
"p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
}
def measure_operations(
operations: dict[str, Callable[[], torch.Tensor]],
*,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, list[float]]:
for operation in operations.values():
for _ in range(warmup):
operation()
torch.cuda.synchronize()
samples: dict[str, list[float]] = {name: [] for name in operations}
order = tuple(operations)
# A-B-C-C-B-A order balances cache, clock, and temperature drift.
for _ in range(trials):
for name in (*order, *reversed(order)):
samples[name].append(time_operation(operations[name], iterations))
return samples
def quant_step(x: torch.Tensor, fmt: str) -> torch.Tensor:
"""Quantization step (dequant scale) from the current amax; ``quantize``
takes its reciprocal as the multiplier."""
amax = x.abs().amax().to(torch.float32).clamp_min(1e-12)
return amax / FP8_MAX[fmt]
def benchmark_quantize(
shape: MatrixShape,
fmt: str,
*,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, object]:
x = torch.randn(shape.rows, shape.cols, device="cuda", dtype=torch.bfloat16) * 0.1
multiplier = quant_step(x, fmt).reciprocal()
fp8_dtype = torch.float8_e4m3fn if fmt == "e4m3" else torch.float8_e5m2
ring_state = torch.zeros(16 + 4, dtype=torch.float32, device="cuda")
operations: dict[str, Callable[[], torch.Tensor]] = {
"torch_cast": lambda: x.to(fp8_dtype),
"plain": lambda: quantize(x, multiplier, fmt)[0],
"ring": lambda: quantize(
x,
multiplier,
fmt,
ring_state=ring_state,
hist_idx=0,
fp8_max=FP8_MAX[fmt],
)[0],
"dual": lambda: quantize_dual(x, multiplier, fmt)[0],
}
samples = measure_operations(
operations, warmup=warmup, iterations=iterations, trials=trials
)
with torch.no_grad():
step = quant_step(x, fmt)
x8, _ = quantize(x, multiplier, fmt)
dequant_error = float((x8.to(torch.float32) * step - x.float()).abs().max())
dual_t = quantize_dual(x, multiplier, fmt)[1]
dual_matches = bool(torch.equal(dual_t.t().contiguous(), x8.contiguous()))
io_bytes = 3 * x.numel() + 4 # bf16 read + fp8 write + f32 amax
result: dict[str, object] = {
"suite": "quantize",
"shape": shape.name,
"rows": shape.rows,
"cols": shape.cols,
"fmt": fmt,
"estimated_io_bytes": io_bytes,
"dequant_max_abs_error": dequant_error,
"dual_transpose_matches": dual_matches,
}
for name, samples_ms in samples.items():
latency = summarize(samples_ms)
bytes_per_call = io_bytes + x.numel() if name == "dual" else io_bytes
result[name] = {
"effective_bandwidth_gbps": bytes_per_call
/ (latency["median_ms"] / 1000)
/ 1e9,
**latency,
}
speedup = (
result["torch_cast"]["median_ms"] / result["plain"]["median_ms"] - 1.0
) * 100.0
print(
f"quantize,{shape.name},{shape.rows}x{shape.cols},"
f"{result['torch_cast']['median_ms']:.4f},{result['plain']['median_ms']:.4f},"
f"{result['ring']['median_ms']:.4f},{result['dual']['median_ms']:.4f},"
f"{speedup:+.1f}%,{dequant_error:.5f},{dual_matches}"
)
return result
def benchmark_gemm(
shape: MatrixShape,
m: int,
fmt: str,
*,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, object]:
x = torch.randn(m, shape.cols, device="cuda", dtype=torch.bfloat16) * 0.1
w = (
torch.randn(shape.rows, shape.cols, device="cuda", dtype=torch.bfloat16)
* shape.cols**-0.5
)
sx, sw = quant_step(x, fmt), quant_step(w, fmt)
with torch.no_grad():
x8, _ = quantize(x, sx.reciprocal(), fmt)
w8, _ = quantize(w, sw.reciprocal(), fmt)
dequant_scale = sx * sw
def torch_op() -> torch.Tensor:
return F.linear(x, w)
def fp8_op() -> torch.Tensor:
return mm_fp8(x8, w8, dequant_scale, trans_b=True)
operations = {"torch_bf16": torch_op, "fp8": fp8_op}
samples = measure_operations(
operations, warmup=warmup, iterations=iterations, trials=trials
)
with torch.no_grad():
actual = fp8_op().float()
dequant_reference = (
x8.to(torch.float32) @ w8.to(torch.float32).t() * dequant_scale
)
bf16_reference = torch_op().float()
kernel_difference = actual - dequant_reference
format_difference = dequant_reference - bf16_reference
io_bytes = m * shape.cols + shape.rows * shape.cols + 2 * m * shape.rows
result: dict[str, object] = {
"suite": "gemm",
"shape": shape.name,
"m": m,
"n": shape.rows,
"k": shape.cols,
"fmt": fmt,
"estimated_io_bytes": io_bytes,
"kernel_max_abs": float(kernel_difference.abs().max()),
"kernel_rel_l2": float(
kernel_difference.norm() / dequant_reference.norm().clamp_min(1e-12)
),
"format_rel_l2": float(
format_difference.norm() / bf16_reference.norm().clamp_min(1e-12)
),
}
for name, samples_ms in samples.items():
latency = summarize(samples_ms)
result[name] = {
"effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
**latency,
}
speedup = (
result["torch_bf16"]["median_ms"] / result["fp8"]["median_ms"] - 1.0
) * 100.0
print(
f"gemm,{shape.name},{m}x{shape.rows}x{shape.cols},"
f"{result['torch_bf16']['median_ms']:.4f},{result['fp8']['median_ms']:.4f},"
f"{speedup:+.1f}%,{result['kernel_max_abs']:.4f},"
f"{result['kernel_rel_l2']:.6f},{result['format_rel_l2']:.6f}"
)
return result
@click.command(help=__doc__)
@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
@click.option(
"--suite",
"suites",
type=click.Choice(("quantize", "gemm", "all")),
multiple=True,
default=("all",),
show_default=True,
)
@click.option("--fmt", type=click.Choice(("e4m3", "e5m2")), default="e4m3")
@click.option("--m-values", default="512,2048,4096", show_default=True)
@click.option(
"--shape",
"shape_values",
multiple=True,
help="Filter defaults by bare name (either suite), or add/override with "
"NAME:ROWS:COLS.",
)
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
@click.option(
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
)
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
@click.option("--seed", type=int, default=0, show_default=True)
def benchmark_command(
output: Path | None,
suites: tuple[str, ...],
fmt: str,
m_values: str,
shape_values: tuple[str, ...],
warmup: int,
iterations: int,
trials: int,
seed: int,
) -> None:
if not torch.cuda.is_available():
raise click.ClickException("CUDA is required")
if not is_available("fp8_ops"):
raise click.ClickException(
"the built fp8_ops kernel is required (compute capability 89+)"
)
selected = ("quantize", "gemm") if "all" in suites else tuple(dict.fromkeys(suites))
m_values_parsed = parse_positive_ints(m_values)
# Any --shape selection replaces the defaults for both suites: a bare
# name keeps that suite's matching default, a NAME:ROWS:COLS spec
# overrides the same-name default or adds a new one.
bare_names = {value for value in shape_values if ":" not in value}
known = {shape.name for shape in QUANTIZE_SHAPES + GEMM_SHAPES}
unknown = sorted(bare_names - known)
if unknown:
raise click.BadParameter(f"unknown default shape names: {', '.join(unknown)}")
specs = [parse_shape(value) for value in shape_values if ":" in value]
def resolve_shapes(defaults: tuple[MatrixShape, ...]) -> list[MatrixShape]:
if not shape_values:
return list(defaults)
by_name = {shape.name: shape for shape in defaults if shape.name in bare_names}
for spec in specs:
by_name[spec.name] = spec
return list(by_name.values())
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
results = []
with torch.inference_mode():
if "quantize" in selected:
print(
"suite,shape,size,cast_ms,plain_ms,ring_ms,dual_ms,vs_cast,"
"dequant_max,dual_ok"
)
for shape in resolve_shapes(QUANTIZE_SHAPES):
results.append(
benchmark_quantize(
shape, fmt, warmup=warmup, iterations=iterations, trials=trials
)
)
torch.cuda.empty_cache()
if "gemm" in selected:
print(
"suite,shape,mxn_xk,bf16_ms,fp8_ms,speedup,kernel_max,"
"kernel_rel_l2,format_rel_l2"
)
for shape in resolve_shapes(GEMM_SHAPES):
for m in m_values_parsed:
results.append(
benchmark_gemm(
shape,
m,
fmt,
warmup=warmup,
iterations=iterations,
trials=trials,
)
)
torch.cuda.empty_cache()
if output is not None:
props = torch.cuda.get_device_properties(0)
payload = {
"metadata": {
"gpu_name": props.name,
"compute_capability": f"{props.major}.{props.minor}",
"torch_version": torch.__version__,
"cuda_version": torch.version.cuda,
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
"fmt": fmt,
},
"settings": {
"warmup": warmup,
"iterations": iterations,
"trials": trials,
"seed": seed,
"order": "A-B-C-C-B-A",
"suites": list(selected),
"m_values": list(m_values_parsed),
},
"results": results,
}
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
if __name__ == "__main__":
benchmark_command()
csrc/bench/benchmark_gemm.py
-333
View File
@@ -1,333 +0,0 @@
"""Benchmark decode-time linear shapes before enabling custom GEMM dispatch.
The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It
establishes the per-architecture cuBLAS baseline that later GEMM primitives and
dispatch decisions must beat.
"""
from __future__ import annotations
import json
import math
import statistics
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Iterable
import click
import torch
import torch.nn.functional as F
@dataclass(frozen=True)
class LinearShape:
name: str
n: int
k: int
DEFAULT_SHAPES = (
LinearShape("q_proj", 1536, 1536),
LinearShape("k_proj", 256, 1536),
LinearShape("v_proj", 256, 1536),
LinearShape("attn_out", 1536, 1536),
LinearShape("mlp_up", 6912, 1536),
LinearShape("mlp_gate", 6912, 1536),
LinearShape("mlp_down", 1536, 6912),
LinearShape("lm_head", 100000, 1536),
)
DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16}
def parse_positive_ints(value: str) -> tuple[int, ...]:
"""Parse a comma-separated, duplicate-free list of positive integers."""
try:
values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
except ValueError as exc:
raise click.BadParameter("expected comma-separated integers") from exc
if not values or any(item <= 0 for item in values):
raise click.BadParameter("values must be positive integers")
return values
def parse_shape(value: str) -> LinearShape:
"""Parse NAME:N:K into a benchmark shape."""
parts = value.split(":")
if len(parts) != 3 or not parts[0]:
raise click.BadParameter("shape must use NAME:N:K")
try:
n, k = (int(item) for item in parts[1:])
except ValueError as exc:
raise click.BadParameter("N and K must be integers") from exc
if n <= 0 or k <= 0:
raise click.BadParameter("N and K must be positive")
return LinearShape(parts[0], n, k)
def estimate_io_bytes(
m: int, n: int, k: int, element_size: int, *, has_bias: bool
) -> int:
"""Estimate bytes touched once by Y[M,N] = X[M,K] @ W[N,K].T."""
elements = m * k + n * k + m * n
if has_bias:
elements += n
return elements * element_size
def percentile(values: Iterable[float], quantile: float) -> float:
ordered = sorted(values)
if not ordered:
raise ValueError("percentile requires at least one sample")
rank = (len(ordered) - 1) * quantile
lower = math.floor(rank)
upper = math.ceil(rank)
if lower == upper:
return ordered[lower]
fraction = rank - lower
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
def summarize_latency(samples_ms: list[float]) -> dict[str, float]:
return {
"median_ms": statistics.median(samples_ms),
"p90_ms": percentile(samples_ms, 0.90),
"p99_ms": percentile(samples_ms, 0.99),
"min_ms": min(samples_ms),
"max_ms": max(samples_ms),
}
def measure_cuda_ms(
operation: Callable[[], torch.Tensor], *, warmup: int, iterations: int, trials: int
) -> list[float]:
for _ in range(warmup):
operation()
torch.cuda.synchronize()
samples = []
for _ in range(trials):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iterations):
operation()
end.record()
end.synchronize()
samples.append(start.elapsed_time(end) / iterations)
return samples
def count_cuda_kernels(
operation: Callable[[], torch.Tensor], repeats: int = 5
) -> float:
"""Profile a few calls and return the average device events per call."""
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA,
],
acc_events=True,
) as profile:
for _ in range(repeats):
operation()
torch.cuda.synchronize()
device_type = torch.autograd.DeviceType.CUDA
events = [event for event in profile.events() if event.device_type == device_type]
return len(events) / repeats
def capture_linear(
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None
) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
for _ in range(3):
F.linear(x, weight, bias)
torch.cuda.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
output = F.linear(x, weight, bias)
return graph, output
def benchmark_case(
shape: LinearShape,
m: int,
*,
dtype: torch.dtype,
mode: str,
bias_enabled: bool,
warmup: int,
iterations: int,
trials: int,
) -> dict[str, object]:
x = torch.randn((m, shape.k), device="cuda", dtype=dtype)
weight = torch.randn((shape.n, shape.k), device="cuda", dtype=dtype)
bias = torch.randn(shape.n, device="cuda", dtype=dtype) if bias_enabled else None
graph = None
graph_output = None
if mode == "graph":
graph, graph_output = capture_linear(x, weight, bias)
def operation() -> torch.Tensor:
graph.replay()
return graph_output
else:
def operation() -> torch.Tensor:
return F.linear(x, weight, bias)
samples_ms = measure_cuda_ms(
operation, warmup=warmup, iterations=iterations, trials=trials
)
latency = summarize_latency(samples_ms)
io_bytes = estimate_io_bytes(
m, shape.n, shape.k, x.element_size(), has_bias=bias is not None
)
median_seconds = latency["median_ms"] / 1000
result: dict[str, object] = {
"name": shape.name,
"m": m,
"n": shape.n,
"k": shape.k,
"mode": mode,
"bias": bias is not None,
"estimated_io_bytes": io_bytes,
"effective_bandwidth_gbps": io_bytes / median_seconds / 1e9,
"cuda_kernel_launches_per_call": count_cuda_kernels(operation),
**latency,
"samples_ms": samples_ms,
}
return result
def render_markdown(payload: dict[str, object]) -> str:
metadata = payload["metadata"]
assert isinstance(metadata, dict)
results = payload["results"]
assert isinstance(results, list)
lines = [
"# Decode linear baseline",
"",
f"- GPU: {metadata['gpu_name']}",
f"- Compute capability: {metadata['compute_capability']}",
f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
f"- Dtype: {metadata['dtype']}",
"",
"| Layer | M | N | K | Mode | Median (ms) | p99 (ms) | GB/s | CUDA kernels/call |",
"|---|---:|---:|---:|---|---:|---:|---:|---:|",
]
for item in results:
assert isinstance(item, dict)
lines.append(
"| {name} | {m} | {n} | {k} | {mode} | {median_ms:.4f} | "
"{p99_ms:.4f} | {effective_bandwidth_gbps:.1f} | "
"{cuda_kernel_launches_per_call:.2f} |".format(**item)
)
lines.append("")
return "\n".join(lines)
def device_metadata(dtype_name: str) -> dict[str, object]:
props = torch.cuda.get_device_properties(0)
return {
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
"gpu_name": props.name,
"compute_capability": f"{props.major}.{props.minor}",
"total_memory_bytes": props.total_memory,
"torch_version": torch.__version__,
"cuda_version": torch.version.cuda,
"dtype": dtype_name,
}
@click.command(help=__doc__)
@click.option("--output", type=click.Path(path_type=Path), required=True)
@click.option("--markdown-output", type=click.Path(path_type=Path))
@click.option("--m-values", default="1,2,4,8,16,32", show_default=True)
@click.option(
"--shape",
"shape_values",
multiple=True,
help="Override defaults with repeatable NAME:N:K shapes.",
)
@click.option("--dtype", type=click.Choice(tuple(DTYPES)), default="bfloat16")
@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both")
@click.option("--bias/--no-bias", default=False)
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
@click.option(
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
)
@click.option("--trials", type=click.IntRange(min=1), default=20, show_default=True)
@click.option("--seed", type=int, default=0, show_default=True)
def benchmark_command(
output: Path,
markdown_output: Path | None,
m_values: str,
shape_values: tuple[str, ...],
dtype: str,
mode: str,
bias: bool,
warmup: int,
iterations: int,
trials: int,
seed: int,
) -> None:
if not torch.cuda.is_available():
raise click.ClickException("CUDA is required")
parsed_m = parse_positive_ints(m_values)
shapes = tuple(parse_shape(item) for item in shape_values) or DEFAULT_SHAPES
modes = ("eager", "graph") if mode == "both" else (mode,)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
results = []
for shape in shapes:
for m in parsed_m:
for current_mode in modes:
click.echo(
f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}"
)
results.append(
benchmark_case(
shape,
m,
dtype=DTYPES[dtype],
mode=current_mode,
bias_enabled=bias,
warmup=warmup,
iterations=iterations,
trials=trials,
)
)
payload: dict[str, object] = {
"schema_version": 1,
"metadata": device_metadata(dtype),
"parameters": {
"m_values": list(parsed_m),
"shapes": [asdict(shape) for shape in shapes],
"modes": list(modes),
"bias": bias,
"warmup": warmup,
"iterations": iterations,
"trials": trials,
"seed": seed,
},
"results": results,
}
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
if markdown_output is not None:
markdown_output.parent.mkdir(parents=True, exist_ok=True)
markdown_output.write_text(render_markdown(payload), encoding="utf-8")
if __name__ == "__main__":
benchmark_command()
csrc/bench/benchmark_gemm_common.py
-450
View File
@@ -1,450 +0,0 @@
"""Benchmark the BF16 GEMM primitive and guarded linear dispatcher.
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;
it measures dispatcher overhead and dependent MLP work, but is deliberately
not presented as a whole-model throughput benchmark.
"""
import argparse
import gc
import json
import math
import os
import statistics
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
import torch
import torch.nn.functional as F
from astrai.extension import bf16_gemm, is_available, linear
@dataclass(frozen=True)
class Shape:
label: str
n: int
k: int
@dataclass(frozen=True)
class Chain:
label: str
hidden: int
kv: int
intermediate: int
fused_qkv: bool = False
gated_mlp: bool = True
@dataclass(frozen=True)
class Timing:
median_ms: float
p90_ms: float
ASTRAI_SHAPES = (
Shape("astrai_qkv", 256, 1536),
Shape("astrai_square", 1536, 1536),
Shape("astrai_up_gate", 6912, 1536),
Shape("astrai_down", 1536, 6912),
Shape("astrai_lm_head", 100000, 1536),
)
TRADITIONAL_SHAPES = (
Shape("llama2_7b_qo", 4096, 4096),
Shape("llama2_7b_up_gate", 11008, 4096),
Shape("llama2_7b_down", 4096, 11008),
Shape("llama3_8b_kv", 1024, 4096),
Shape("llama3_8b_up_gate", 14336, 4096),
Shape("llama3_8b_down", 4096, 14336),
Shape("llama2_13b_qo", 5120, 5120),
Shape("llama2_13b_up_gate", 13824, 5120),
Shape("llama2_13b_down", 5120, 13824),
Shape("gpt_neox_up", 16384, 4096),
Shape("gpt_neox_down", 4096, 16384),
Shape("qwen2_7b_kv", 512, 3584),
Shape("qwen2_7b_qo", 3584, 3584),
Shape("qwen2_7b_up_gate", 18944, 3584),
Shape("qwen2_7b_down", 3584, 18944),
Shape("llama3_70b_kv", 1024, 8192),
Shape("llama3_70b_qo", 8192, 8192),
Shape("llama3_70b_up_gate", 28672, 8192),
Shape("llama3_70b_down", 8192, 28672),
Shape("opt_1_3b_qkvo", 2048, 2048),
Shape("opt_1_3b_up", 8192, 2048),
Shape("opt_1_3b_down", 2048, 8192),
)
CHAINS = (
Chain("llama2_7b", 4096, 4096, 11008),
Chain("llama3_8b", 4096, 1024, 14336),
Chain("llama2_13b", 5120, 5120, 13824),
Chain("gpt_neox_20b", 4096, 4096, 16384, fused_qkv=True),
Chain("qwen2_7b", 3584, 512, 18944),
Chain("llama3_70b", 8192, 1024, 28672),
Chain("opt_1_3b", 2048, 2048, 8192, gated_mlp=False),
)
def _elapsed_ms(fn: Callable[[], torch.Tensor], inner: int) -> float:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(inner):
fn()
end.record()
end.synchronize()
return start.elapsed_time(end) / inner
def _timing(values: list[float]) -> Timing:
ordered = sorted(values)
p90_index = max(0, math.ceil(0.9 * len(ordered)) - 1)
return Timing(statistics.median(ordered), ordered[p90_index])
def _measure_pair(
baseline: Callable[[], torch.Tensor],
candidate: Callable[[], torch.Tensor],
*,
warmup: int,
samples: int,
inner: int,
prepare_baseline: Callable[[], None] = lambda: None,
prepare_candidate: Callable[[], None] = lambda: None,
) -> tuple[Timing, Timing]:
cases = (
("baseline", prepare_baseline, baseline),
("candidate", prepare_candidate, candidate),
)
for iteration in range(warmup):
_, prepare, fn = cases[iteration % 2]
prepare()
fn()
torch.cuda.synchronize()
values: dict[str, list[float]] = {"baseline": [], "candidate": []}
for sample in range(samples):
order = cases if sample % 2 == 0 else tuple(reversed(cases))
for label, prepare, fn in order:
prepare()
values[label].append(_elapsed_ms(fn, inner))
return _timing(values["baseline"]), _timing(values["candidate"])
def _print_header() -> None:
print(
"suite,label,m,n,k,torch_median_ms,torch_p90_ms,"
"candidate_median_ms,candidate_p90_ms,speedup_pct,"
"max_abs,relative_l2,argmax_equal"
)
def _print_result(
suite: str,
label: str,
m: int,
n: int,
k: int,
baseline: Timing,
candidate: Timing,
reference: torch.Tensor,
actual: torch.Tensor,
) -> dict[str, object]:
difference = actual.float() - reference.float()
max_abs = difference.abs().max().item()
relative_l2 = difference.norm().item() / max(reference.float().norm().item(), 1e-12)
argmax_equal = torch.equal(actual.argmax(dim=-1), reference.argmax(dim=-1))
speedup = (baseline.median_ms / candidate.median_ms - 1.0) * 100.0
result: dict[str, object] = {
"suite": suite,
"label": label,
"m": m,
"n": n,
"k": k,
"torch_median_ms": baseline.median_ms,
"torch_p90_ms": baseline.p90_ms,
"candidate_median_ms": candidate.median_ms,
"candidate_p90_ms": candidate.p90_ms,
"speedup_pct": speedup,
"max_abs": max_abs,
"relative_l2": relative_l2,
"argmax_equal": argmax_equal,
}
print(
f"{suite},{label},{m},{n},{k},"
f"{baseline.median_ms:.6f},{baseline.p90_ms:.6f},"
f"{candidate.median_ms:.6f},{candidate.p90_ms:.6f},"
f"{speedup:+.2f},{max_abs:.6f},{relative_l2:.8f},"
f"{str(argmax_equal).lower()}",
flush=True,
)
return result
def _weight(n: int, k: int, device: torch.device, std: float) -> torch.Tensor:
weight = torch.empty((n, k), device=device, dtype=torch.bfloat16)
weight.normal_(mean=0.0, std=std)
return weight.requires_grad_(True)
def _kernel_functions(
x: torch.Tensor, weight: torch.Tensor
) -> tuple[Callable[[], torch.Tensor], Callable[[], torch.Tensor]]:
def baseline() -> torch.Tensor:
return F.linear(x, weight)
def candidate() -> torch.Tensor:
return bf16_gemm(x, weight.detach())
return baseline, candidate
def benchmark_kernels(
args: argparse.Namespace, device: torch.device
) -> list[dict[str, object]]:
if args.family == "astrai":
shapes = ASTRAI_SHAPES
elif args.family == "traditional":
shapes = TRADITIONAL_SHAPES
else:
shapes = ASTRAI_SHAPES + TRADITIONAL_SHAPES
if args.shape_label:
requested = set(args.shape_label)
shapes = tuple(shape for shape in shapes if shape.label in requested)
missing = requested - {shape.label for shape in shapes}
if missing:
raise ValueError(f"unknown shape labels: {', '.join(sorted(missing))}")
results: list[dict[str, object]] = []
for shape in shapes:
weight = _weight(shape.n, shape.k, device, args.weight_std)
for m in args.m:
x = torch.randn((m, shape.k), device=device, dtype=torch.bfloat16)
baseline_fn, candidate_fn = _kernel_functions(x, weight)
with torch.inference_mode():
reference = baseline_fn()
actual = candidate_fn()
baseline, candidate = _measure_pair(
baseline_fn,
candidate_fn,
warmup=args.warmup,
samples=args.samples,
inner=args.inner,
)
results.append(
_print_result(
"kernel",
shape.label,
m,
shape.n,
shape.k,
baseline,
candidate,
reference,
actual,
)
)
del baseline_fn, candidate_fn, x, reference, actual
del weight
gc.collect()
torch.cuda.empty_cache()
return results
def _set_mode(mode: str) -> None:
os.environ["ASTRAI_GEMM"] = mode
def _chain_weights(
spec: Chain, device: torch.device, std: float
) -> dict[str, torch.Tensor]:
weights = {
"o": _weight(spec.hidden, spec.hidden, device, std),
"up": _weight(spec.intermediate, spec.hidden, device, std),
"down": _weight(spec.hidden, spec.intermediate, device, std),
}
if spec.fused_qkv:
weights["qkv"] = _weight(3 * spec.hidden, spec.hidden, device, std)
else:
weights.update(
{
"q": _weight(spec.hidden, spec.hidden, device, std),
"k": _weight(spec.kv, spec.hidden, device, std),
"v": _weight(spec.kv, spec.hidden, device, std),
}
)
if spec.gated_mlp:
weights["gate"] = _weight(spec.intermediate, spec.hidden, device, std)
return weights
def _chain_fn(
x: torch.Tensor, weights: dict[str, torch.Tensor], spec: Chain
) -> Callable[[], torch.Tensor]:
def run() -> torch.Tensor:
output_projection = linear(x, weights["o"])
up = linear(x, weights["up"])
if spec.fused_qkv:
attention_projection = linear(x, weights["qkv"])[..., : x.shape[-1]]
hidden = F.gelu(up)
else:
attention_projection = linear(x, weights["q"])
linear(x, weights["k"])
linear(x, weights["v"])
if spec.gated_mlp:
gate = linear(x, weights["gate"])
hidden = F.silu(gate) * up
else:
hidden = F.gelu(up)
down = linear(hidden, weights["down"])
return attention_projection + output_projection + down
return run
def benchmark_chains(
args: argparse.Namespace, device: torch.device
) -> list[dict[str, object]]:
results: list[dict[str, object]] = []
chains = CHAINS
if args.chain_label:
requested = set(args.chain_label)
chains = tuple(chain for chain in chains if chain.label in requested)
missing = requested - {chain.label for chain in chains}
if missing:
raise ValueError(f"unknown chain labels: {', '.join(sorted(missing))}")
for spec in chains:
weights = _chain_weights(spec, device, args.weight_std)
for m in args.m:
x = torch.randn((m, spec.hidden), device=device, dtype=torch.bfloat16)
run = _chain_fn(x, weights, spec)
with torch.inference_mode():
_set_mode("0")
reference = run()
_set_mode(args.candidate_mode)
actual = run()
baseline, candidate = _measure_pair(
run,
run,
warmup=args.warmup,
samples=args.samples,
inner=args.chain_inner,
prepare_baseline=lambda: _set_mode("0"),
prepare_candidate=lambda: _set_mode(args.candidate_mode),
)
results.append(
_print_result(
"synthetic_chain",
spec.label,
m,
spec.hidden,
spec.intermediate,
baseline,
candidate,
reference,
actual,
)
)
del x, reference, actual
del weights
gc.collect()
torch.cuda.empty_cache()
return results
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--suite", choices=("kernel", "chain", "all"), default="all")
parser.add_argument(
"--family", choices=("astrai", "traditional", "all"), default="all"
)
parser.add_argument(
"--m", type=int, nargs="+", choices=(1, 2, 4, 8), default=(1, 2, 4, 8)
)
parser.add_argument(
"--shape-label",
action="append",
help="limit the kernel suite to one or more named shape labels",
)
parser.add_argument(
"--chain-label",
action="append",
help="limit the chain suite to one or more named model families",
)
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--warmup", type=int, default=20)
parser.add_argument("--samples", type=int, default=9)
parser.add_argument("--inner", type=int, default=100)
parser.add_argument("--chain-inner", type=int, default=20)
parser.add_argument(
"--candidate-mode",
choices=("auto", "1"),
default="auto",
help="dispatcher mode for the candidate side of the chain suite",
)
parser.add_argument("--weight-std", type=float, default=0.02)
parser.add_argument("--seed", type=int, default=20260902)
parser.add_argument(
"--output",
type=Path,
help="optional JSON output; stdout always retains the compact CSV table",
)
return parser.parse_args()
def main() -> None:
args = parse_args()
if not torch.cuda.is_available() or not is_available("bf16_gemm"):
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:
raise ValueError("warmup must be non-negative and sample/inner counts positive")
torch.cuda.set_device(args.device)
device = torch.device("cuda", args.device)
torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)
properties = torch.cuda.get_device_properties(device)
print(
f"# device={properties.name}, capability={properties.major}.{properties.minor}, "
f"seed={args.seed}, weight_std={args.weight_std}"
)
_print_header()
results: list[dict[str, object]] = []
if args.suite in ("kernel", "all"):
results.extend(benchmark_kernels(args, device))
if args.suite in ("chain", "all"):
results.extend(benchmark_chains(args, device))
if args.output is not None:
payload = {
"environment": {
"device": properties.name,
"capability": f"{properties.major}.{properties.minor}",
"torch": torch.__version__,
"cuda": torch.version.cuda,
},
"parameters": {
"suite": args.suite,
"family": args.family,
"m": args.m,
"shape_labels": args.shape_label,
"chain_labels": args.chain_label,
"candidate_mode": args.candidate_mode,
"seed": args.seed,
"weight_std": args.weight_std,
"warmup": args.warmup,
"samples": args.samples,
"inner": args.inner,
"chain_inner": args.chain_inner,
},
"results": results,
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(payload, indent=2) + "\n")
if __name__ == "__main__":
main()
csrc/bench/benchmark_rotary.py
+252
View File
@@ -0,0 +1,252 @@
"""Benchmark the fused rotary-embedding kernel against the torch fallback.
The baseline is the complex-multiply fallback from
``astrai.extension.backend.rotary``. Layouts mirror the production call
shapes: packed 3D [tokens, n_heads, head_dim] and dense 4D
[batch, seq_len, n_heads, head_dim]; positions are random integers so every
row exercises a distinct cos/sin gather.
"""
from __future__ import annotations
import json
import math
import statistics
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable
import click
import torch
import torch.nn.functional as F
from astrai.extension import is_available
from astrai.extension.ops import rotary_emb
@dataclass(frozen=True)
class RotaryCase:
name: str
layout: str
batch: int
seq_len: int
heads: int
head_dim: int
DEFAULT_CASES = (
RotaryCase("decode_bs1", "packed", 1, 1, 32, 128),
RotaryCase("decode_bs32", "packed", 32, 1, 32, 128),
RotaryCase("prefill_4k_llama7b", "packed", 1, 4096, 32, 128),
RotaryCase("prefill_4k_llama70b", "packed", 1, 4096, 64, 128),
RotaryCase("train_8x2k_llama7b", "dense", 8, 2048, 32, 128),
RotaryCase("train_4x1k_d64", "dense", 4, 1024, 32, 64),
)
def parse_case(value: str) -> RotaryCase:
parts = value.split(":")
if len(parts) != 6 or not parts[0]:
raise click.BadParameter("case must use NAME:LAYOUT:BATCH:SEQ:HEADS:HEAD_DIM")
name, layout, batch, seq_len, heads, head_dim = parts
try:
fields = (int(batch), int(seq_len), int(heads), int(head_dim))
except ValueError as exc:
raise click.BadParameter("fields must be integers") from exc
if layout not in ("packed", "dense"):
raise click.BadParameter("layout must be 'packed' or 'dense'")
if any(field <= 0 for field in fields) or head_dim % 2:
raise click.BadParameter("fields must be positive; HEAD_DIM even")
return RotaryCase(name, layout, fields[0], fields[1], fields[2], fields[3])
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iterations):
operation()
end.record()
end.synchronize()
return start.elapsed_time(end) / iterations
def summarize(values: list[float]) -> dict[str, float]:
ordered = sorted(values)
return {
"median_ms": statistics.median(ordered),
"p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
}
def torch_apply(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
"""The complex-multiply fallback (mirrors backend.rotary._torch_apply)."""
cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
dtype = x.dtype
x_ = x.float().reshape(*x.shape[:-1], -1, 2)
x_complex = torch.view_as_complex(x_)
freqs_cis_complex = torch.complex(cos, sin).unsqueeze(-2)
x_rotated = x_complex * freqs_cis_complex
return torch.view_as_real(x_rotated).flatten(-2).to(dtype)
def build_freqs(head_dim: int, positions: torch.Tensor) -> torch.Tensor:
"""[cos, sin] pairs for the given positions, laid out [..., head_dim/2, 2]."""
theta = 10000.0 ** (
-torch.arange(0, head_dim, 2, dtype=torch.float64, device=positions.device)
/ head_dim
)
freqs = positions.double().unsqueeze(-1) * theta
return torch.stack([freqs.cos(), freqs.sin()], dim=-1).float()
def benchmark_case(
case: RotaryCase, *, warmup: int, iterations: int, trials: int
) -> dict[str, object]:
if case.layout == "packed":
tokens = case.batch * case.seq_len
x = torch.randn(
tokens, case.heads, case.head_dim, device="cuda", dtype=torch.bfloat16
)
positions = torch.randint(0, 65536, (tokens,), device="cuda")
else:
x = torch.randn(
case.batch,
case.seq_len,
case.heads,
case.head_dim,
device="cuda",
dtype=torch.bfloat16,
)
positions = torch.randint(0, 65536, (case.batch, case.seq_len), device="cuda")
freqs_cis = build_freqs(case.head_dim, positions)
operations: dict[str, Callable[[], torch.Tensor]] = {
"torch": lambda: torch_apply(x, freqs_cis),
"cuda": lambda: rotary_emb(x, freqs_cis),
}
for operation in operations.values():
for _ in range(warmup):
operation()
torch.cuda.synchronize()
samples: dict[str, list[float]] = {name: [] for name in operations}
order = tuple(operations)
# A-B-B-A order balances cache, clock, and temperature drift.
for _ in range(trials):
for name in (*order, *reversed(order)):
samples[name].append(time_operation(operations[name], iterations))
with torch.no_grad():
expected = operations["torch"]().float()
actual = operations["cuda"]().float()
difference = actual - expected
io_bytes = (
2 * x.numel() * x.element_size() + freqs_cis.numel() * freqs_cis.element_size()
)
result: dict[str, object] = {
"case": case.name,
"layout": case.layout,
"heads": case.heads,
"head_dim": case.head_dim,
"estimated_io_bytes": io_bytes,
"max_abs_error": float(difference.abs().max()),
"cosine_similarity": float(
F.cosine_similarity(actual.flatten(), expected.flatten(), dim=0)
),
}
for name, samples_ms in samples.items():
latency = summarize(samples_ms)
result[name] = {
"effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
**latency,
}
speedup = (result["torch"]["median_ms"] / result["cuda"]["median_ms"] - 1.0) * 100.0
print(
f"{case.name},{case.layout},{result['torch']['median_ms']:.5f},"
f"{result['cuda']['median_ms']:.5f},{speedup:+.1f}%,"
f"{result['max_abs_error']:.5f}"
)
return result
@click.command(help=__doc__)
@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
@click.option(
"--case",
"case_values",
multiple=True,
help="Filter defaults by bare name, or add/override with "
"NAME:LAYOUT:BATCH:SEQ:HEADS:HEAD_DIM.",
)
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
@click.option(
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
)
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
@click.option("--seed", type=int, default=0, show_default=True)
def benchmark_command(
output: Path | None,
case_values: tuple[str, ...],
warmup: int,
iterations: int,
trials: int,
seed: int,
) -> None:
if not torch.cuda.is_available():
raise click.ClickException("CUDA is required")
if not is_available("rotary_emb"):
raise click.ClickException("the built rotary_emb kernel is required")
# A bare name filters the matching default; a full spec overrides or appends.
chosen: dict[str, RotaryCase] = {}
for value in case_values:
case = (
next((c for c in DEFAULT_CASES if c.name == value), None)
if ":" not in value
else parse_case(value)
)
if case is None:
raise click.BadParameter(f"unknown default case {value!r}")
chosen[case.name] = case
cases = tuple(chosen.values()) or DEFAULT_CASES
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
print("case,layout,torch_ms,cuda_ms,speedup,max_abs")
results = []
with torch.inference_mode():
for case in cases:
results.append(
benchmark_case(
case, warmup=warmup, iterations=iterations, trials=trials
)
)
torch.cuda.empty_cache()
if output is not None:
props = torch.cuda.get_device_properties(0)
payload = {
"metadata": {
"gpu_name": props.name,
"compute_capability": f"{props.major}.{props.minor}",
"torch_version": torch.__version__,
"cuda_version": torch.version.cuda,
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
},
"settings": {
"warmup": warmup,
"iterations": iterations,
"trials": trials,
"seed": seed,
"order": "A-B-B-A",
},
"results": results,
}
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
if __name__ == "__main__":
benchmark_command()
csrc/bench/benchmark_swiglu.py
-325
View File
@@ -1,325 +0,0 @@
"""Benchmark fused BF16 SwiGLU against torch and unfused GEMM chains."""
from __future__ import annotations
import json
import math
import statistics
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Callable, Iterable
import click
import torch
import torch.nn.functional as F
from astrai.extension import bf16_gemm, bf16_swiglu, is_available
@dataclass(frozen=True)
class SwiGLUShape:
name: str
n: int
k: int
DEFAULT_SHAPES = (
SwiGLUShape("astrai_1b", 6912, 1536),
SwiGLUShape("llama2_7b", 11008, 4096),
SwiGLUShape("llama3_8b", 14336, 4096),
SwiGLUShape("llama2_13b", 13824, 5120),
SwiGLUShape("gpt_neox_20b", 16384, 6144),
)
def parse_positive_ints(value: str) -> tuple[int, ...]:
try:
values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
except ValueError as exc:
raise click.BadParameter("expected comma-separated integers") from exc
if not values or any(item <= 0 for item in values):
raise click.BadParameter("values must be positive integers")
return values
def parse_shape(value: str) -> SwiGLUShape:
parts = value.split(":")
if len(parts) != 3 or not parts[0]:
raise click.BadParameter("shape must use NAME:N:K")
try:
n, k = (int(item) for item in parts[1:])
except ValueError as exc:
raise click.BadParameter("N and K must be integers") from exc
if n <= 0 or k <= 0 or k % 8:
raise click.BadParameter("N must be positive and K positive/divisible by 8")
return SwiGLUShape(parts[0], n, k)
def percentile(values: Iterable[float], quantile: float) -> float:
ordered = sorted(values)
rank = (len(ordered) - 1) * quantile
lower = math.floor(rank)
upper = math.ceil(rank)
if lower == upper:
return ordered[lower]
fraction = rank - lower
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
def summarize(values: list[float]) -> dict[str, float]:
return {
"median_ms": statistics.median(values),
"p90_ms": percentile(values, 0.90),
"p99_ms": percentile(values, 0.99),
"min_ms": min(values),
"max_ms": max(values),
}
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
for _ in range(iterations):
operation()
end.record()
end.synchronize()
return start.elapsed_time(end) / iterations
def count_cuda_kernels(
operation: Callable[[], torch.Tensor], repeats: int = 5
) -> float:
with torch.profiler.profile(
activities=[
torch.profiler.ProfilerActivity.CPU,
torch.profiler.ProfilerActivity.CUDA,
],
acc_events=True,
) as profile:
for _ in range(repeats):
operation()
torch.cuda.synchronize()
device_type = torch.autograd.DeviceType.CUDA
events = [event for event in profile.events() if event.device_type == device_type]
return len(events) / repeats
def capture(operation: Callable[[], torch.Tensor]):
for _ in range(3):
operation()
torch.cuda.synchronize()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
output = operation()
def replay() -> torch.Tensor:
graph.replay()
return output
return replay
def make_operations(x, up_weight, gate_weight, mode: str):
operations: dict[str, Callable[[], torch.Tensor]] = {
"torch": lambda: F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)),
"gemm_chain": lambda: (
bf16_gemm(x, up_weight) * F.silu(bf16_gemm(x, gate_weight))
),
"fused": lambda: bf16_swiglu(x, up_weight, gate_weight),
}
if mode == "graph":
operations = {name: capture(op) for name, op in operations.items()}
return operations
def benchmark_case(
shape: SwiGLUShape,
m: int,
mode: str,
*,
warmup: int,
iterations: int,
trials: int,
) -> list[dict[str, object]]:
x = torch.randn((m, shape.k), device="cuda", dtype=torch.bfloat16) * 0.1
scale = shape.k**-0.5
up_weight = (
torch.randn((shape.n, shape.k), device="cuda", dtype=torch.bfloat16) * scale
)
gate_weight = (
torch.randn((shape.n, shape.k), device="cuda", dtype=torch.bfloat16) * scale
)
operations = make_operations(x, up_weight, gate_weight, mode)
for operation in operations.values():
for _ in range(warmup):
operation()
torch.cuda.synchronize()
samples = {name: [] for name in operations}
forward_order = tuple(operations)
# A-B-C-C-B-A order balances cache, clock, and temperature drift.
for _ in range(trials):
for name in (*forward_order, *reversed(forward_order)):
samples[name].append(time_operation(operations[name], iterations))
with torch.no_grad():
expected = operations["torch"]().clone()
actual = operations["fused"]().clone()
difference = (actual.float() - expected.float()).abs()
max_abs_error = float(difference.max())
mean_abs_error = float(difference.mean())
cosine_similarity = float(
F.cosine_similarity(actual.float().flatten(), expected.float().flatten(), dim=0)
)
results = []
for name, operation in operations.items():
result: dict[str, object] = {
"shape": shape.name,
"m": m,
"n": shape.n,
"k": shape.k,
"mode": mode,
"implementation": name,
"cuda_kernel_launches_per_call": count_cuda_kernels(operation),
**summarize(samples[name]),
}
if name == "fused":
result.update(
max_abs_error=max_abs_error,
mean_abs_error=mean_abs_error,
cosine_similarity=cosine_similarity,
)
results.append(result)
return results
def device_metadata() -> dict[str, object]:
props = torch.cuda.get_device_properties(0)
return {
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
"gpu_name": props.name,
"compute_capability": f"{props.major}.{props.minor}",
"total_memory_bytes": props.total_memory,
"torch_version": torch.__version__,
"cuda_version": torch.version.cuda,
"dtype": "bfloat16",
}
def render_markdown(payload: dict[str, object]) -> str:
metadata = payload["metadata"]
results = payload["results"]
assert isinstance(metadata, dict)
assert isinstance(results, list)
by_case = {
(item["shape"], item["m"], item["mode"], item["implementation"]): item
for item in results
}
cases = sorted({(item["shape"], item["m"], item["mode"]) for item in results})
lines = [
"# Fused SwiGLU benchmark",
"",
f"- GPU: {metadata['gpu_name']}",
f"- Compute capability: {metadata['compute_capability']}",
f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
"",
"| Shape | M | Mode | torch ms | GEMM chain ms | fused ms | "
"vs best unfused | fused kernels | max abs | cosine |",
"|---|---:|---|---:|---:|---:|---:|---:|---:|---:|",
]
for shape, m, mode in cases:
torch_item = by_case[(shape, m, mode, "torch")]
gemm_item = by_case[(shape, m, mode, "gemm_chain")]
fused_item = by_case[(shape, m, mode, "fused")]
best = min(torch_item["median_ms"], gemm_item["median_ms"])
improvement = (best / fused_item["median_ms"] - 1) * 100
lines.append(
f"| {shape} | {m} | {mode} | {torch_item['median_ms']:.5f} | "
f"{gemm_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | "
f"{improvement:+.2f}% | "
f"{fused_item['cuda_kernel_launches_per_call']:.1f} | "
f"{fused_item['max_abs_error']:.5f} | "
f"{fused_item['cosine_similarity']:.8f} |"
)
lines.append("")
return "\n".join(lines)
@click.command(help=__doc__)
@click.option("--output", type=click.Path(path_type=Path), required=True)
@click.option("--markdown-output", type=click.Path(path_type=Path))
@click.option("--m-values", default="1,2,4,8", show_default=True)
@click.option("--shape", "shape_values", multiple=True, help="Repeat NAME:N:K.")
@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both")
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
@click.option(
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
)
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
@click.option("--seed", type=int, default=0, show_default=True)
def benchmark_command(
output: Path,
markdown_output: Path | None,
m_values: str,
shape_values: tuple[str, ...],
mode: str,
warmup: int,
iterations: int,
trials: int,
seed: int,
) -> None:
if not torch.cuda.is_available():
raise click.ClickException("CUDA is required")
if not is_available("bf16_gemm") or not is_available("bf16_swiglu"):
raise click.ClickException("built bf16_gemm and bf16_swiglu are required")
shapes = tuple(parse_shape(value) for value in shape_values) or DEFAULT_SHAPES
m_values_parsed = parse_positive_ints(m_values)
if any(m > 8 for m in m_values_parsed):
raise click.BadParameter("fused primitive supports M up to 8")
modes = ("eager", "graph") if mode == "both" else (mode,)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
results = []
with torch.inference_mode():
for shape in shapes:
for m in m_values_parsed:
for current_mode in modes:
click.echo(
f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}"
)
results.extend(
benchmark_case(
shape,
m,
current_mode,
warmup=warmup,
iterations=iterations,
trials=trials,
)
)
torch.cuda.empty_cache()
payload: dict[str, object] = {
"metadata": device_metadata(),
"settings": {
"warmup": warmup,
"iterations": iterations,
"trials": trials,
"seed": seed,
"order": "A-B-C-C-B-A",
},
"results": results,
}
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(payload, indent=2) + "\n")
if markdown_output is not None:
markdown_output.parent.mkdir(parents=True, exist_ok=True)
markdown_output.write_text(render_markdown(payload))
if __name__ == "__main__":
benchmark_command()
csrc/kernels/gemm.cu
-650
View File
@@ -1,650 +0,0 @@
// 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/swiglu.cu
-297
View File
@@ -1,297 +0,0 @@
// 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
// all decode rows. Bandwidth-bound in the cold-HBM decode regime, so variant
// selection beyond the M=8 block-size rule is noise.
#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 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;
}
__device__ __forceinline__ float round_bf16(float value) {
return __bfloat162float(__float2bfloat16_rn(value));
}
template <int Threads, int Rows>
__global__ void bf16_swiglu_kernel(
const __nv_bfloat16* __restrict__ x,
const __nv_bfloat16* __restrict__ up_weight,
const __nv_bfloat16* __restrict__ gate_weight,
__nv_bfloat16* __restrict__ output,
int n,
int k
) {
constexpr int kWarps = Threads / kWarpSize;
const int output_index = blockIdx.x;
const int lane = threadIdx.x & (kWarpSize - 1);
const int warp = threadIdx.x / kWarpSize;
const int vector_count = k / 8;
float up_sums[Rows] = {};
float gate_sums[Rows] = {};
__shared__ float up_warp_sums[Rows][kWarps];
__shared__ float gate_warp_sums[Rows][kWarps];
const auto* x4 = reinterpret_cast<const uint4*>(x);
const auto* up4 = reinterpret_cast<const uint4*>(
up_weight + static_cast<int64_t>(output_index) * k
);
const auto* gate4 = reinterpret_cast<const uint4*>(
gate_weight + static_cast<int64_t>(output_index) * k
);
// Read each pair of up/gate weight chunks once per CTA, then reuse it for
// every active decode row. The fused epilogue removes two [M, N]
// intermediates and the standalone SiLU and multiply launches.
for (int vector_index = threadIdx.x;
vector_index < vector_count;
vector_index += blockDim.x) {
const uint4 up_raw = up4[vector_index];
const uint4 gate_raw = gate4[vector_index];
const auto* up_values =
reinterpret_cast<const __nv_bfloat162*>(&up_raw);
const auto* gate_values =
reinterpret_cast<const __nv_bfloat162*>(&gate_raw);
#pragma unroll
for (int row = 0; row < Rows; ++row) {
const uint4 x_raw =
x4[static_cast<int64_t>(row) * vector_count + vector_index];
const auto* x_values =
reinterpret_cast<const __nv_bfloat162*>(&x_raw);
#pragma unroll
for (int pair = 0; pair < 4; ++pair) {
const float2 xv = __bfloat1622float2(x_values[pair]);
const float2 uv = __bfloat1622float2(up_values[pair]);
const float2 gv = __bfloat1622float2(gate_values[pair]);
up_sums[row] = fmaf(xv.x, uv.x, up_sums[row]);
up_sums[row] = fmaf(xv.y, uv.y, up_sums[row]);
gate_sums[row] = fmaf(xv.x, gv.x, gate_sums[row]);
gate_sums[row] = fmaf(xv.y, gv.y, gate_sums[row]);
}
}
}
#pragma unroll
for (int row = 0; row < Rows; ++row) {
up_sums[row] = warp_sum(up_sums[row]);
gate_sums[row] = warp_sum(gate_sums[row]);
}
if (lane == 0) {
#pragma unroll
for (int row = 0; row < Rows; ++row) {
up_warp_sums[row][warp] = up_sums[row];
gate_warp_sums[row][warp] = gate_sums[row];
}
}
__syncthreads();
if (warp == 0) {
#pragma unroll
for (int row = 0; row < Rows; ++row) {
float up = lane < kWarps ? up_warp_sums[row][lane] : 0.0f;
float gate = lane < kWarps ? gate_warp_sums[row][lane] : 0.0f;
up = warp_sum(up);
gate = warp_sum(gate);
if (lane == 0) {
// Match the public composition's BF16 rounding boundaries:
// BF16 linear outputs, BF16 SiLU output, then BF16 multiply.
up = round_bf16(up);
gate = round_bf16(gate);
const float silu = round_bf16(gate / (1.0f + expf(-gate)));
output[static_cast<int64_t>(row) * n + output_index] =
__float2bfloat16_rn(up * silu);
}
}
}
}
template <int Threads, int Rows>
void launch_bf16_swiglu(
const __nv_bfloat16* x,
const __nv_bfloat16* up_weight,
const __nv_bfloat16* gate_weight,
__nv_bfloat16* output,
int n,
int k,
cudaStream_t stream
) {
bf16_swiglu_kernel<Threads, Rows><<<n, Threads, 0, stream>>>(
x, up_weight, gate_weight, output, n, k
);
}
torch::Tensor bf16_swiglu(
torch::Tensor x,
torch::Tensor up_weight,
torch::Tensor gate_weight
) {
TORCH_CHECK(
x.is_cuda() && up_weight.is_cuda() && gate_weight.is_cuda(),
"x, up_weight, and gate_weight must be CUDA tensors"
);
TORCH_CHECK(
x.device() == up_weight.device() && x.device() == gate_weight.device(),
"x and weights must share a device"
);
TORCH_CHECK(
x.scalar_type() == torch::kBFloat16 &&
up_weight.scalar_type() == torch::kBFloat16 &&
gate_weight.scalar_type() == torch::kBFloat16,
"x and weights must be bf16"
);
TORCH_CHECK(
x.dim() == 1 || x.dim() == 2,
"x must have shape [K] or [M, K]"
);
TORCH_CHECK(
up_weight.dim() == 2 && gate_weight.dim() == 2,
"weights must have shape [N, K]"
);
TORCH_CHECK(
x.is_contiguous() && up_weight.is_contiguous() &&
gate_weight.is_contiguous(),
"x and weights must be contiguous"
);
// The kernel loads all three streams as uint4; contiguous-but-offset
// views would fault with an opaque "misaligned address" CUDA error, so
// reject them here with an actionable message.
TORCH_CHECK(
(reinterpret_cast<uintptr_t>(x.data_ptr()) & 15u) == 0u,
"bf16_swiglu requires 16-byte aligned x (storage_offset must keep "
"data_ptr divisible by 16); clone the tensor or use the torch path"
);
TORCH_CHECK(
(reinterpret_cast<uintptr_t>(up_weight.data_ptr()) & 15u) == 0u,
"bf16_swiglu requires 16-byte aligned up_weight (storage_offset "
"must keep data_ptr divisible by 16); clone the tensor or use the "
"torch path"
);
TORCH_CHECK(
(reinterpret_cast<uintptr_t>(gate_weight.data_ptr()) & 15u) == 0u,
"bf16_swiglu requires 16-byte aligned gate_weight (storage_offset "
"must keep data_ptr divisible by 16); clone the tensor or use the "
"torch path"
);
TORCH_CHECK(
!x.requires_grad() && !up_weight.requires_grad() &&
!gate_weight.requires_grad(),
"bf16_swiglu 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 = up_weight.size(0);
TORCH_CHECK(m >= 1 && m <= 8, "M must be in [1, 8]");
TORCH_CHECK(
gate_weight.sizes() == up_weight.sizes(),
"up_weight and gate_weight must have identical shapes"
);
TORCH_CHECK(up_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 % 8 == 0, "K must be divisible by 8");
TORCH_CHECK(
k <= std::numeric_limits<int>::max() &&
n <= std::numeric_limits<int>::max(),
"N or K exceeds the CUDA launcher limit"
);
const at::cuda::OptionalCUDAGuard guard(x.device());
const auto* properties = at::cuda::getDeviceProperties(x.device().index());
TORCH_CHECK(
properties->major >= 8,
"bf16_swiglu 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* up_ptr =
reinterpret_cast<const __nv_bfloat16*>(up_weight.data_ptr());
const auto* gate_ptr =
reinterpret_cast<const __nv_bfloat16*>(gate_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);
// Block size 256 keeps the weight streams at the HBM bandwidth floor for
// M in [1, 7]; M=8 halves the CTA so each thread owns more of the row
// and the shared-memory reduction tree shrinks (measured on L20 with
// rotated cold weights; larger CTAs only add idle warps).
switch (m) {
case 1:
launch_bf16_swiglu<256, 1>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 2:
launch_bf16_swiglu<256, 2>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 3:
launch_bf16_swiglu<256, 3>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 4:
launch_bf16_swiglu<256, 4>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 5:
launch_bf16_swiglu<256, 5>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 6:
launch_bf16_swiglu<256, 6>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 7:
launch_bf16_swiglu<256, 7>(
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
);
break;
case 8:
launch_bf16_swiglu<128, 8>(
x_ptr, up_ptr, gate_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_swiglu",
&bf16_swiglu,
py::arg("x"),
py::arg("up_weight"),
py::arg("gate_weight"),
"M in [1, 8] fused BF16 up/gate projection and SwiGLU"
);
}
docs/developer/cuda_kernels.md
+1 -119
View File
@@ -1,9 +1,7 @@
# CUDA Kernels
AstrAI includes optional custom CUDA kernels for attention, rotary embedding,
BF16 GEMM/SwiGLU, and FP8 GEMM. These are built when `nvcc` is available and
CUDA is detected. BF16 GEMM and SwiGLU are directly callable and can be
selected by guarded model dispatchers described below.
and FP8 GEMM. These are built when `nvcc` is available and CUDA is detected.
## Overview
@@ -14,121 +12,8 @@ selected by guarded model dispatchers described below.
| `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention |
| `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) |
| `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
| `bf16_gemm` | `gemm.cu` | M=1..64 BF16 linear, skinny GEMM path for M<=8 and tensor-core tiled path above (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+) |
### BF16 GEMM primitive
`astrai.extension.bf16_gemm(x, weight, bias=None)` accepts a contiguous BF16
input shaped `[K]` or `[M, K]`, with `M` in `[1, 64]`, and row-major weights
`[N, K]`. One entry point selects the internal path by M:
- `M` in `[1, 8]` — register-resident GEMV kernel, any positive `K`: one
flat 128-thread CTA per weight row computes the output column for all M
tokens together. Under HBM-streaming conditions (weights rotated through
L2, as in real decode) the kernel is bandwidth-bound; on L20 the flat
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
`ASTRAI_GEMM=0` for an unconditional `F.linear` fallback, `1` to force the
kernel for any capable M (1-64, with K%8==0 and alignment for M>8), or
`auto` (the default). On compute capability 8.0+, `auto` routes all capable
decode batches M in [1, 64] through the kernel. The GEMV path (M≤8, any K)
measured universal wins at the HBM bandwidth floor (AstrAI 1B end-to-end
+8.8% at M=1 rising to +17% at M=8). The tiled path (M 9-64) measured on
L20 across five shape families wins or ties four (q/o, kv, gate/up, lm_head)
and regresses one (down-projection M>16: +25% latency; narrow-N long-K where
cuBLAS is strong). Out-of-band, training, prefill-sized, or unsupported
calls fall back to PyTorch.
The measurements below date from the original `[2, 4]` automatic band; the
band has since widened to `{1, ..., 8}` with the flat 128-thread kernel:
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
main-versus-warp-tiling run used identical interleaved settings; for the four
M=4 selected shapes, candidate latency changed from 0.016292 to 0.016108 ms
for `(4096,4096)`, 0.037939 to 0.028539 ms for `(11008,4096)`, 0.043407 to
0.039803 ms for `(4096,11008)`, and 0.006697 to 0.006390 ms for
`(1024,4096)`.
The dependent projection-chain gate, which includes Python dispatch and
rotates through distinct weights instead of repeatedly warming one matrix,
measured:
| Synthetic chain | M=2 | M=4 | Row argmax parity |
|---|---:|---:|---|
| LLaMA 2 7B | +8.49% | fallback (M=4 bands excluded) | exact |
| LLaMA 3 8B | +8.50% | +6.44% | exact |
| LLaMA 2 13B | +5.66% | +5.67% | exact |
| GPT-NeoX 20B | +6.95% | +5.93% | exact |
| Qwen2 7B | +7.48% | +7.48% | exact |
| LLaMA 3 70B | +7.77% | +7.69% | exact |
| OPT 1.3B | +25.20% | fallback (M=4 up projection regresses) | exact |
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.
These are synthetic projection-chain measurements, not whole-model throughput
claims. Reproduce them with `csrc/bench/benchmark_gemm_common.py`.
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
greedy-stable winners missed the 3% end-to-end gate. Long-K MLP-down bands are
also excluded because their valid BF16 error changed a checkpoint greedy
argmax; the enabled M=2/4 bands matched the baseline greedy output exactly.
Training, prefill, unmeasured architectures, and losing shape bands always
remain on PyTorch. Use mode `1` only for explicit A/B runs outside this table.
The primitive remains directly callable and deliberately has no internal
`F.linear` fallback. The model-level backend owns fallback and dispatch policy.
### BF16 SwiGLU primitive
`astrai.extension.bf16_swiglu(x, up_weight, gate_weight)` fuses the two dense
MLP projections with `up * silu(gate)` into one CUDA launch for contiguous
BF16 inputs with `M` in `[1, 8]` and K divisible by 8. It preserves the BF16
rounding boundaries of the two projection outputs, SiLU output, and final
product while accumulating dot products in FP32.
The kernel is a single CTA-reuse tiling: one CTA per output column reads each
up/gate weight chunk once and applies it to all M rows. Block size is 256
threads for M in [1, 7] and 128 for M=8, where the shorter shared-memory
reduction tree wins under cold-HBM decode traffic. An earlier per-shape
`(6912,1536)` warp-per-row variant and its dispatch table were removed: HBM
measurements with rotated weights showed the table was tuned against L2-cache
regime timing and was up to 6% slower than CTA reuse at M=2/4; the kernel is
bandwidth-bound, so finer variant selection is noise.
Dense `MLP` modules route through the SwiGLU backend. `ASTRAI_SWIGLU=0` keeps
the unfused linear backend, and `1` explicitly forces the fused primitive.
`auto` is the default but currently has no enabled bands: although direct
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
M=1/2/4. Automatic dispatch therefore remains numerically identical to the
existing path.
Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
| Variant | File | Optimization |
@@ -318,13 +203,10 @@ astrai/extension/
├── ops/
│ ├── attention.py # Stateless attention kernel wrappers
│ ├── rotary.py # Stateless rotary kernel wrapper
│ ├── gemm.py # Stateless BF16 GEMM primitive
│ ├── swiglu.py # Stateless fused BF16 SwiGLU primitive
│ └── fp8.py # Stateless FP8 primitives (custom_op)
├── fp8.py # FP8 strategy layer (fp8_autocast, recipes)
└── backend/
├── attention.py # Backend selection, KV cache I/O, and fallback
├── swiglu.py # Inference-only fused/unfused SwiGLU policy
└── rotary.py # Per-call CUDA/torch rotary dispatch
```
setup.py
-2
View File
@@ -121,8 +121,6 @@ class _CMakeBuildExt(_build_ext):
"attn_prefill",
"attn_paged_decode",
"attn_paged_prefill",
"bf16_gemm",
"bf16_swiglu",
"rotary_emb",
)
missing = [name for name in required if not any(lib_dir.glob(f"{name}.*.so"))]
tests/extension/test_gemm.py
-396
View File
@@ -1,396 +0,0 @@
import pytest
import torch
import torch.nn.functional as F
from astrai.extension import bf16_gemm, is_available
GEMM_AVAILABLE = (
torch.cuda.is_available()
and is_available("bf16_gemm")
and torch.cuda.get_device_capability() >= (8, 0)
)
skip_no_gemm = pytest.mark.skipif(
not GEMM_AVAILABLE,
reason="BF16 GEMM requires a built kernel and compute capability 8.0+",
)
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(
"n,k",
[(256, 1536), (1536, 1536), (6912, 1536), (1536, 6912), (100000, 1536)],
)
def test_bf16_gemm_matches_linear_shape_families(n, k):
torch.manual_seed(17)
x = torch.randn(k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemm(x, weight)
expected = F.linear(x, weight)
assert actual.shape == (n,)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemm
@pytest.mark.parametrize("m", [2, 3, 4, 5, 6, 7, 8])
@pytest.mark.parametrize("n,k", [(256, 1536), (1536, 1536), (1536, 6912)])
def test_bf16_gemm_matches_small_decode_batches(m, n, k):
torch.manual_seed(19 + m)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemm(x, weight)
expected = F.linear(x, weight)
assert actual.shape == (m, n)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.5)
@skip_no_gemm
@pytest.mark.parametrize("m", [2, 4])
@pytest.mark.parametrize(
"n,k",
[
(1024, 4096),
(4096, 4096),
(11008, 4096),
(4096, 11008),
(14336, 4096),
(4096, 14336),
(5120, 5120),
(13824, 5120),
(5120, 13824),
(16384, 4096),
(4096, 16384),
(512, 3584),
(3584, 3584),
(18944, 3584),
(3584, 18944),
(1024, 8192),
(8192, 8192),
(28672, 8192),
(8192, 28672),
(2048, 2048),
(8192, 2048),
(2048, 8192),
],
)
def test_bf16_gemm_matches_common_transformer_shapes(m, n, k):
torch.manual_seed(2026 + m + n + k)
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)
expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemm
@pytest.mark.parametrize(
"m,n,k",
[
(1, 8192, 2048),
(8, 4096, 11008),
(8, 512, 3584),
(8, 1024, 8192),
(8, 2048, 8192),
],
)
def test_bf16_gemm_matches_m8_edge_bands(m, n, k):
torch.manual_seed(2026 + m + n + k)
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)
expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemm
def test_bf16_gemm_preserves_singleton_batch_and_fuses_bias():
torch.manual_seed(23)
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
bias = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemm(x, weight, bias)
expected = F.linear(x, weight, bias)
assert actual.shape == (1, 1536)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemm
def test_bf16_gemm_small_batch_fuses_bias():
torch.manual_seed(25)
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
bias = torch.randn(256, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemm(x, weight, bias)
expected = F.linear(x, weight, bias)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemm
def test_bf16_gemm_uses_current_stream():
x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
torch.cuda.synchronize()
stream = torch.cuda.Stream()
with torch.cuda.stream(stream):
actual = bf16_gemm(x, weight)
expected = F.linear(x, weight)
stream.synchronize()
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemm
def test_bf16_gemm_cuda_graph_replay():
torch.manual_seed(29)
x = torch.randn(1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
for _ in range(3):
bf16_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()
expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemm
@pytest.mark.parametrize("n,k", [(64, 7), (64, 12), (33, 100), (256, 1534)])
def test_bf16_gemm_handles_unaligned_k(n, k):
torch.manual_seed(29)
x = torch.randn(k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16)
actual = bf16_gemm(x, weight)
expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
x3 = torch.randn(3, k, device="cuda", dtype=torch.bfloat16)
actual3 = bf16_gemm(x3, weight)
torch.testing.assert_close(actual3, F.linear(x3, weight), rtol=0.02, atol=0.5)
@skip_no_gemm
@pytest.mark.parametrize("m", [1, 2, 3, 4])
def test_bf16_gemm_accepts_complementary_misalignment(m):
"""Misaligned weight rows plus an x base chosen so the vectorized 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
``x`` itself, faulting with a misaligned-address CUDA error)."""
torch.manual_seed(37)
n, k = 256, 1536
# offset 5 elements = +10 bytes: weight rows land at 10 % 16 (whead=3)
# and x at 10 % 16, so (x + 2*whead) % 16 == 0 selects the fast path.
big_w = torch.randn(n * k + 8, device="cuda", dtype=torch.bfloat16)
weight = big_w[5 : 5 + n * k].view(n, k)
big_x = torch.randn(m * k + 8, device="cuda", dtype=torch.bfloat16)
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
actual = bf16_gemm(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)
@skip_no_gemm
def test_bf16_gemm_scalar_path_handles_misaligned_weight_only():
"""Weight rows misaligned while x stays 16B-aligned take the scalar-x
middle and must stay exact."""
torch.manual_seed(41)
n, k = 256, 1536
big_w = torch.randn(n * k + 8, device="cuda", dtype=torch.bfloat16)
weight = big_w[5 : 5 + n * k].view(n, k)
x = torch.randn(2, k, device="cuda", dtype=torch.bfloat16)
assert (weight.data_ptr() & 15) == 10 and (x.data_ptr() & 15) == 0
actual = bf16_gemm(x, weight)
torch.testing.assert_close(actual, F.linear(x, weight), rtol=0.02, atol=0.5)
@skip_no_gemm
def test_bf16_gemm_small_batch_cuda_graph_replay():
torch.manual_seed(31)
x = torch.randn(8, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
for _ in range(3):
bf16_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()
expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
@skip_no_gemm
@pytest.mark.parametrize(
"make_args,error",
[
(
lambda: (
torch.randn(65, 16, device="cuda", dtype=torch.bfloat16),
torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
),
"M must",
),
(
lambda: (
torch.randn(16, device="cuda", dtype=torch.float16),
torch.randn(8, 16, device="cuda", dtype=torch.float16),
),
"bf16",
),
(
lambda: (
torch.randn(
16, device="cuda", dtype=torch.bfloat16, requires_grad=True
),
torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
),
"autograd",
),
],
)
def test_bf16_gemm_rejects_unsupported_inputs(make_args, error):
with pytest.raises(RuntimeError, match=error):
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
-233
View File
@@ -1,233 +0,0 @@
import importlib
import logging
import pytest
import torch
import torch.nn.functional as F
from astrai.extension import is_available, linear
from astrai.extension.dispatch import explain, op_backend, resolve
# The package attribute ``linear`` is the dispatched function; reach the
# module object explicitly for monkeypatching its private helpers.
linear_module = importlib.import_module("astrai.extension.backend.linear")
GEMM_AVAILABLE = (
torch.cuda.is_available()
and is_available("bf16_gemm")
and torch.cuda.get_device_capability() >= (8, 0)
)
skip_no_gemm = pytest.mark.skipif(
not GEMM_AVAILABLE,
reason="BF16 GEMM requires a built kernel and compute capability 8.0+",
)
def _routes_to_gemm(monkeypatch, x, weight, bias=None) -> bool:
"""Patch the GEMM entry point to a sentinel and report whether
``linear`` selected it (torch fallback would compute a real tensor)."""
sentinel = object()
def fake_gemm(x, weight, bias):
return sentinel
monkeypatch.setattr(linear_module, "_inference_bf16_gemm", fake_gemm)
return linear(x, weight, bias) is sentinel
def test_model_linear_routes_through_backend(monkeypatch):
sentinel = torch.randn(2, 4)
def fake_linear(x, weight, bias):
assert x.shape == (2, 3)
assert weight.shape == (4, 3)
assert bias is None
return sentinel
monkeypatch.setattr("astrai.model.components.linear.linear", fake_linear)
from astrai.model.components.linear import Linear
layer = Linear(3, 4)
assert layer(torch.randn(2, 3)) is sentinel
def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
monkeypatch.setenv("ASTRAI_GEMM", "invalid-test-mode")
x = torch.randn(2, 8)
weight = torch.randn(4, 8)
with caplog.at_level(logging.WARNING):
actual = linear(x, weight)
assert "using auto" in caplog.text
torch.testing.assert_close(actual, F.linear(x, weight))
def test_cpu_and_training_calls_fall_back_to_torch(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(2, 8, requires_grad=True)
weight = torch.randn(4, 8, requires_grad=True)
actual = linear(x, weight)
expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected)
actual.sum().backward()
assert x.grad is not None
assert weight.grad is not None
@skip_no_gemm
@pytest.mark.parametrize("m", [1, 2, 3, 4, 5, 6, 7, 8])
def test_auto_selects_small_decode_batches(monkeypatch, m):
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", [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):
"""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)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
)
@skip_no_gemm
def test_mode_zero_disables_gemm(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMM", "0")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(linear(x, weight), F.linear(x, weight))
@skip_no_gemm
@pytest.mark.parametrize("m", [1, 2, 8, 16, 32])
def test_mode_one_forces_every_capable_batch(monkeypatch, m):
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(m, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemm
def test_mode_one_rejects_oversized_batch_and_grad(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMM", "1")
weight = torch.randn(
256, 1536, device="cuda", dtype=torch.bfloat16, requires_grad=True
)
with torch.no_grad():
oversized = torch.randn(65, 1536, device="cuda", dtype=torch.bfloat16)
assert not _routes_to_gemm(monkeypatch, oversized, weight)
assert not _routes_to_gemm(
monkeypatch, torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16), weight
)
@skip_no_gemm
def test_mode_one_supports_bias_and_vector_input(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(1, 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)
with torch.no_grad():
assert _routes_to_gemm(monkeypatch, x, weight, bias)
monkeypatch.undo()
torch.testing.assert_close(
linear(x, weight, bias),
F.linear(x, weight, bias),
rtol=0.02,
atol=0.25,
)
@skip_no_gemm
def test_dispatched_linear_cuda_graph_replay(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
for _ in range(3):
linear(x, weight)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
actual = linear(x, weight)
x.copy_(torch.randn_like(x))
graph.replay()
expected = F.linear(x, weight)
torch.testing.assert_close(actual, expected, rtol=0.02, atol=0.25)
def test_linear_family_is_registered_with_shared_dispatcher():
from astrai.extension.dispatch import _FAMILIES
assert "linear" in _FAMILIES
x = torch.randn(2, 8)
weight = torch.randn(4, 8)
resolution = resolve("linear", x, weight)
assert resolution.record.family == "linear"
assert resolution.origin in ("chain", "fallback")
assert "linear" in explain("linear", x, weight)
@skip_no_gemm
def test_ops_env_override_forces_torch_for_capable_call(monkeypatch):
"""ASTR_OPS=linear=torch must keep working after the M-band rewrite
(regression: the family was silently dropped from the dispatcher, so
the override warned, fell through, and the gemm kernel still ran)."""
monkeypatch.setenv("ASTR_OPS", "linear=torch")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(1536, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
)
@skip_no_gemm
def test_ops_env_override_forces_gemm(monkeypatch):
monkeypatch.setenv("ASTR_OPS", "linear=gemm")
x = torch.randn(1, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad():
# M=1 is outside the auto band but inside the forced gemm record.
assert _routes_to_gemm(monkeypatch, x, weight)
@skip_no_gemm
def test_op_backend_context_selects_torch(monkeypatch):
monkeypatch.setenv("ASTRAI_GEMM", "1")
x = torch.randn(2, 1536, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16)
with torch.no_grad(), op_backend(linear="torch"):
assert not _routes_to_gemm(monkeypatch, x, weight)
torch.testing.assert_close(
linear(x, weight), F.linear(x, weight), rtol=0.02, atol=0.25
)
# The override is scoped: the forced mode applies again afterwards.
with torch.no_grad():
assert _routes_to_gemm(monkeypatch, x, weight)
def test_op_backend_rejects_unknown_linear_handle():
with pytest.raises(ValueError, match="Unknown linear implementation"):
op_backend(linear="nonexistent").__enter__()
tests/extension/test_swiglu.py
-112
View File
@@ -1,112 +0,0 @@
import pytest
import torch
import torch.nn.functional as F
from astrai.extension import bf16_swiglu, is_available
SWIGLU_AVAILABLE = (
torch.cuda.is_available()
and is_available("bf16_swiglu")
and torch.cuda.get_device_capability() >= (8, 0)
)
skip_no_swiglu = pytest.mark.skipif(
not SWIGLU_AVAILABLE,
reason="BF16 SwiGLU requires a built kernel and compute capability 8.0+",
)
def reference_swiglu(x, up_weight, gate_weight):
return F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight))
@skip_no_swiglu
@pytest.mark.parametrize("m", [1, 2, 4, 8])
@pytest.mark.parametrize("n,k", [(6912, 1536), (4096, 4096), (11008, 4096)])
def test_bf16_swiglu_matches_common_dense_mlp_shapes(m, n, k):
torch.manual_seed(37 + m)
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.1
up_weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
gate_weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5)
actual = bf16_swiglu(x, up_weight, gate_weight)
expected = reference_swiglu(x, up_weight, gate_weight)
assert actual.shape == (m, n)
torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)
@skip_no_swiglu
def test_bf16_swiglu_preserves_vector_shape():
x = torch.randn(1536, device="cuda", dtype=torch.bfloat16) * 0.1
up_weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) * 0.02
gate_weight = torch.randn_like(up_weight) * 0.02
actual = bf16_swiglu(x, up_weight, gate_weight)
expected = reference_swiglu(x, up_weight, gate_weight)
assert actual.shape == (256,)
torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)
@skip_no_swiglu
def test_bf16_swiglu_uses_current_stream_and_cuda_graph():
torch.manual_seed(43)
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1
up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02
gate_weight = torch.randn_like(up_weight) * 0.02
stream = torch.cuda.Stream()
with torch.cuda.stream(stream):
for _ in range(3):
bf16_swiglu(x, up_weight, gate_weight)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
actual = bf16_swiglu(x, up_weight, gate_weight)
x.copy_(torch.randn_like(x) * 0.1)
graph.replay()
stream.synchronize()
expected = reference_swiglu(x, up_weight, gate_weight)
torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)
@skip_no_swiglu
@pytest.mark.parametrize(
"make_args,error",
[
(
lambda: (
torch.randn(9, 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",
),
(
lambda: (
torch.randn(2, 15, device="cuda", dtype=torch.bfloat16),
torch.randn(8, 15, device="cuda", dtype=torch.bfloat16),
torch.randn(8, 15, device="cuda", dtype=torch.bfloat16),
),
"divisible by 8",
),
(
lambda: (
torch.randn(2, 16, device="cuda", dtype=torch.bfloat16),
torch.randn(8, 16, device="cuda", dtype=torch.bfloat16),
torch.randn(7, 16, device="cuda", dtype=torch.bfloat16),
),
"identical shapes",
),
],
)
def test_bf16_swiglu_rejects_unsupported_inputs(make_args, error):
with pytest.raises(RuntimeError, match=error):
bf16_swiglu(*make_args())
@skip_no_swiglu
def test_bf16_swiglu_rejects_misaligned_storage_with_clear_error():
"""Contiguous-but-offset views must fail the wrapper's TORCH_CHECK with
an actionable message instead of a sticky CUDA misaligned-address error
(regression: the kernel casts x directly to uint4 without checking)."""
k = 1536
base = torch.randn(k + 1, device="cuda", dtype=torch.bfloat16)
x = base[1:] # +2 bytes: contiguous but not 16B-aligned
weights = torch.randn(8, k, device="cuda", dtype=torch.bfloat16)
with pytest.raises(RuntimeError, match="16-byte aligned"):
bf16_swiglu(x, weights, weights)
tests/extension/test_swiglu_dispatch.py
-118
View File
@@ -1,118 +0,0 @@
import logging
import pytest
import torch
import torch.nn.functional as F
from astrai.extension import is_available, swiglu
from astrai.model.components.mlp import MLP
SWIGLU_AVAILABLE = (
torch.cuda.is_available()
and is_available("bf16_swiglu")
and torch.cuda.get_device_capability() >= (8, 0)
)
skip_no_swiglu = pytest.mark.skipif(
not SWIGLU_AVAILABLE,
reason="BF16 SwiGLU requires a built kernel and compute capability 8.0+",
)
def reference_swiglu(x, up_weight, gate_weight):
return F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight))
def test_cpu_and_training_calls_fall_back_with_gradients(monkeypatch):
monkeypatch.setenv("ASTRAI_SWIGLU", "1")
x = torch.randn(2, 8, requires_grad=True)
up_weight = torch.randn(4, 8, requires_grad=True)
gate_weight = torch.randn(4, 8, requires_grad=True)
actual = swiglu(x, up_weight, gate_weight)
expected = reference_swiglu(x, up_weight, gate_weight)
torch.testing.assert_close(actual, expected)
actual.sum().backward()
assert x.grad is not None
assert up_weight.grad is not None
assert gate_weight.grad is not None
def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog):
monkeypatch.setenv("ASTRAI_SWIGLU", "invalid-test-mode")
x = torch.randn(2, 8)
up_weight = torch.randn(4, 8)
gate_weight = torch.randn(4, 8)
with caplog.at_level(logging.WARNING):
actual = swiglu(x, up_weight, gate_weight)
assert "using auto" in caplog.text
torch.testing.assert_close(actual, reference_swiglu(x, up_weight, gate_weight))
def test_mlp_routes_through_swiglu_backend(monkeypatch):
sentinel = torch.randn(2, 4)
def fake_swiglu(x, up_weight, gate_weight):
assert x.shape == (2, 3)
assert up_weight.shape == gate_weight.shape == (4, 3)
return sentinel
monkeypatch.setattr("astrai.model.components.mlp.swiglu", fake_swiglu)
layer = MLP(3, 4)
output = layer(torch.randn(2, 3))
assert output["hidden_states"].shape == (2, 3)
@skip_no_swiglu
def test_mode_zero_disables_fused_kernel(monkeypatch):
monkeypatch.setenv("ASTRAI_SWIGLU", "0")
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1
up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02
gate_weight = torch.randn_like(up_weight) * 0.02
with torch.no_grad():
actual = swiglu(x, up_weight, gate_weight)
expected = reference_swiglu(x, up_weight, gate_weight)
torch.testing.assert_close(actual, expected)
@skip_no_swiglu
def test_mode_one_forces_supported_shape(monkeypatch):
monkeypatch.setenv("ASTRAI_SWIGLU", "1")
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1
up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02
gate_weight = torch.randn_like(up_weight) * 0.02
with torch.no_grad():
actual = swiglu(x, up_weight, gate_weight)
expected = reference_swiglu(x, up_weight, gate_weight)
torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)
@skip_no_swiglu
def test_auto_uses_fused_chain_for_decode_batches(monkeypatch):
# Auto adopts the fused primitive for the decode band; numerics match
# the unfused linear-backend chain within BF16 accumulation-order noise.
monkeypatch.setenv("ASTRAI_SWIGLU", "auto")
x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16)
up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16)
gate_weight = torch.randn_like(up_weight)
with torch.no_grad():
actual = swiglu(x, up_weight, gate_weight)
expected = reference_swiglu(x, up_weight, gate_weight)
torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.1)
@skip_no_swiglu
def test_mode_one_falls_back_for_misaligned_storage(monkeypatch):
"""Contiguous-but-offset views must route to the unfused torch chain
even with ASTRAI_SWIGLU=1 instead of reaching the uint4-only kernel
(regression: the fused primitive faulted with a misaligned-address
CUDA error for such inputs)."""
monkeypatch.setenv("ASTRAI_SWIGLU", "1")
k = 1536
x_base = torch.randn(2 * k + 8, device="cuda", dtype=torch.bfloat16) * 0.1
x = x_base[1 : 1 + 2 * k].view(2, k)
assert x.is_contiguous() and (x.data_ptr() & 15) != 0
up_weight = torch.randn(64, k, device="cuda", dtype=torch.bfloat16) * 0.02
gate_weight = torch.randn_like(up_weight)
with torch.no_grad():
actual = swiglu(x, up_weight, gate_weight)
expected = reference_swiglu(x, up_weight, gate_weight)
torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01)