From 6709534d6499a7e849fb56678ec0be587461cb15 Mon Sep 17 00:00:00 2001 From: ViperEkura <3081035982@qq.com> Date: Sat, 5 Sep 2026 01:29:23 +0800 Subject: [PATCH] refactor: remove bf16 gemm and swiglu kernels and rebuild csrc benchmarks - 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 --- astrai/extension/__init__.py | 18 +- astrai/extension/backend/__init__.py | 4 - astrai/extension/backend/linear.py | 215 -------- astrai/extension/backend/swiglu.py | 67 --- astrai/extension/ops/__init__.py | 4 - astrai/extension/ops/gemm.py | 27 - astrai/extension/ops/swiglu.py | 22 - astrai/model/components/linear.py | 5 +- astrai/model/components/mlp.py | 3 +- csrc/CMakeLists.txt | 4 - csrc/bench/benchmark_attention.py | 649 +++++++++++++++++++++++ csrc/bench/benchmark_fp8.py | 394 ++++++++++++++ csrc/bench/benchmark_gemm.py | 333 ------------ csrc/bench/benchmark_gemm_common.py | 450 ---------------- csrc/bench/benchmark_rotary.py | 252 +++++++++ csrc/bench/benchmark_swiglu.py | 325 ------------ csrc/kernels/gemm.cu | 650 ------------------------ csrc/kernels/swiglu.cu | 297 ----------- docs/developer/cuda_kernels.md | 120 +---- setup.py | 2 - tests/extension/test_gemm.py | 396 --------------- tests/extension/test_linear_dispatch.py | 233 --------- tests/extension/test_swiglu.py | 112 ---- tests/extension/test_swiglu_dispatch.py | 118 ----- 24 files changed, 1306 insertions(+), 3394 deletions(-) delete mode 100644 astrai/extension/backend/linear.py delete mode 100644 astrai/extension/backend/swiglu.py delete mode 100644 astrai/extension/ops/gemm.py delete mode 100644 astrai/extension/ops/swiglu.py create mode 100644 csrc/bench/benchmark_attention.py create mode 100644 csrc/bench/benchmark_fp8.py delete mode 100644 csrc/bench/benchmark_gemm.py delete mode 100644 csrc/bench/benchmark_gemm_common.py create mode 100644 csrc/bench/benchmark_rotary.py delete mode 100644 csrc/bench/benchmark_swiglu.py delete mode 100644 csrc/kernels/gemm.cu delete mode 100644 csrc/kernels/swiglu.cu delete mode 100644 tests/extension/test_gemm.py delete mode 100644 tests/extension/test_linear_dispatch.py delete mode 100644 tests/extension/test_swiglu.py delete mode 100644 tests/extension/test_swiglu_dispatch.py diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index 96f6164..eaac3c5 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -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", diff --git a/astrai/extension/backend/__init__.py b/astrai/extension/backend/__init__.py index 8c4127a..189f8e8 100644 --- a/astrai/extension/backend/__init__.py +++ b/astrai/extension/backend/__init__.py @@ -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", ] diff --git a/astrai/extension/backend/linear.py b/astrai/extension/backend/linear.py deleted file mode 100644 index 9f9fb06..0000000 --- a/astrai/extension/backend/linear.py +++ /dev/null @@ -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"] diff --git a/astrai/extension/backend/swiglu.py b/astrai/extension/backend/swiglu.py deleted file mode 100644 index 90be251..0000000 --- a/astrai/extension/backend/swiglu.py +++ /dev/null @@ -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"] diff --git a/astrai/extension/ops/__init__.py b/astrai/extension/ops/__init__.py index 9b9f4cf..ac0009a 100644 --- a/astrai/extension/ops/__init__.py +++ b/astrai/extension/ops/__init__.py @@ -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", ] diff --git a/astrai/extension/ops/gemm.py b/astrai/extension/ops/gemm.py deleted file mode 100644 index 7293d64..0000000 --- a/astrai/extension/ops/gemm.py +++ /dev/null @@ -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"] diff --git a/astrai/extension/ops/swiglu.py b/astrai/extension/ops/swiglu.py deleted file mode 100644 index e1bfc24..0000000 --- a/astrai/extension/ops/swiglu.py +++ /dev/null @@ -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"] diff --git a/astrai/model/components/linear.py b/astrai/model/components/linear.py index 3e50e62..f3a9aeb 100644 --- a/astrai/model/components/linear.py +++ b/astrai/model/components/linear.py @@ -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) diff --git a/astrai/model/components/mlp.py b/astrai/model/components/mlp.py index defd8d4..083dad3 100644 --- a/astrai/model/components/mlp.py +++ b/astrai/model/components/mlp.py @@ -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} diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index 9b8cafd..b08b91f 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -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 ) diff --git a/csrc/bench/benchmark_attention.py b/csrc/bench/benchmark_attention.py new file mode 100644 index 0000000..02c8c35 --- /dev/null +++ b/csrc/bench/benchmark_attention.py @@ -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() diff --git a/csrc/bench/benchmark_fp8.py b/csrc/bench/benchmark_fp8.py new file mode 100644 index 0000000..208742c --- /dev/null +++ b/csrc/bench/benchmark_fp8.py @@ -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() diff --git a/csrc/bench/benchmark_gemm.py b/csrc/bench/benchmark_gemm.py deleted file mode 100644 index 4cc6c35..0000000 --- a/csrc/bench/benchmark_gemm.py +++ /dev/null @@ -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() diff --git a/csrc/bench/benchmark_gemm_common.py b/csrc/bench/benchmark_gemm_common.py deleted file mode 100644 index 3ae272a..0000000 --- a/csrc/bench/benchmark_gemm_common.py +++ /dev/null @@ -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() diff --git a/csrc/bench/benchmark_rotary.py b/csrc/bench/benchmark_rotary.py new file mode 100644 index 0000000..47e126a --- /dev/null +++ b/csrc/bench/benchmark_rotary.py @@ -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() diff --git a/csrc/bench/benchmark_swiglu.py b/csrc/bench/benchmark_swiglu.py deleted file mode 100644 index 2a5502f..0000000 --- a/csrc/bench/benchmark_swiglu.py +++ /dev/null @@ -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() diff --git a/csrc/kernels/gemm.cu b/csrc/kernels/gemm.cu deleted file mode 100644 index 7b95b48..0000000 --- a/csrc/kernels/gemm.cu +++ /dev/null @@ -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 -#include -#include -#include -#include - -#include -#include - -// 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 -__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(output_index) * k; - const unsigned whead_raw = - ((16u - (reinterpret_cast(wrow) & 15u)) & 15u) >> 1; - const int whead = static_cast(min(whead_raw, static_cast(k))); - const int wvecs = (k - whead) / 8; - const int wtail_start = whead + wvecs * 8; - const uint4* __restrict__ w4 = reinterpret_cast(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(x) + 2u * static_cast(whead)) & 15u) == 0u) { - const auto* x4 = reinterpret_cast(x + whead); - for (int v = threadIdx.x; v < wvecs; v += blockDim.x) { - const uint4 wv_raw = w4[v]; - const auto* wv = - reinterpret_cast(&wv_raw); -#pragma unroll - for (int row = 0; row < Rows; ++row) { - const uint4 xv_raw = - x4[(static_cast(row) * (k / 8)) + v]; - const auto* xv = - reinterpret_cast(&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(&wv_raw); -#pragma unroll - for (int row = 0; row < Rows; ++row) { - const __nv_bfloat16* xv = - x + static_cast(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(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(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 -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 - <<>>( - 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 -__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 -__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(r, c8 * 16); - if (ok) { - astrai::cp_async_16( - reinterpret_cast(dst), - src + (int64_t)(row0 + r) * k + kbase + c8 * 8 - ); - } else { - unsigned* w = reinterpret_cast(dst); -#pragma unroll - for (int i = 0; i < 4; ++i) - w[i] = 0u; - } - } -} - -template -__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( - a_slot(s), x, m, k, m0, tid, s - ); - stage_tile(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(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(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(); - __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( - 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( - a_slot(pf), x, m, k, m0, tid, pf - ); - stage_tile(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 -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, - cudaFuncAttributeMaxDynamicSharedMemorySize, smem - )); - return true; - }(); - (void)opted_in; - } - dim3 grid((m + BM - 1) / BM, (n + BN - 1) / BN); - tiled_gemm_kernel - <<>>( - 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::max() && - n <= std::numeric_limits::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_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(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(x.data_ptr()); - const auto* weight_ptr = - reinterpret_cast(weight.data_ptr()); - auto* output_ptr = reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); - const int m_int = static_cast(m); - const int n_int = static_cast(n); - const int k_int = static_cast(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(x.data_ptr()) & 15) == 0u && - (reinterpret_cast(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)" - ); -} diff --git a/csrc/kernels/swiglu.cu b/csrc/kernels/swiglu.cu deleted file mode 100644 index b054241..0000000 --- a/csrc/kernels/swiglu.cu +++ /dev/null @@ -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 -#include -#include -#include -#include - -#include -#include - -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 -__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(x); - const auto* up4 = reinterpret_cast( - up_weight + static_cast(output_index) * k - ); - const auto* gate4 = reinterpret_cast( - gate_weight + static_cast(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(&up_raw); - const auto* gate_values = - reinterpret_cast(&gate_raw); - -#pragma unroll - for (int row = 0; row < Rows; ++row) { - const uint4 x_raw = - x4[static_cast(row) * vector_count + vector_index]; - const auto* x_values = - reinterpret_cast(&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(row) * n + output_index] = - __float2bfloat16_rn(up * silu); - } - } - } -} - -template -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<<>>( - 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(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(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(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::max() && - n <= std::numeric_limits::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(x.data_ptr()); - const auto* up_ptr = - reinterpret_cast(up_weight.data_ptr()); - const auto* gate_ptr = - reinterpret_cast(gate_weight.data_ptr()); - auto* output_ptr = - reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); - const int n_int = static_cast(n); - const int k_int = static_cast(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" - ); -} diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index 766db09..50285f8 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -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 ``` diff --git a/setup.py b/setup.py index efc5907..78d5f12 100644 --- a/setup.py +++ b/setup.py @@ -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"))] diff --git a/tests/extension/test_gemm.py b/tests/extension/test_gemm.py deleted file mode 100644 index 9705102..0000000 --- a/tests/extension/test_gemm.py +++ /dev/null @@ -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) diff --git a/tests/extension/test_linear_dispatch.py b/tests/extension/test_linear_dispatch.py deleted file mode 100644 index 3cf1d3c..0000000 --- a/tests/extension/test_linear_dispatch.py +++ /dev/null @@ -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__() diff --git a/tests/extension/test_swiglu.py b/tests/extension/test_swiglu.py deleted file mode 100644 index f5b1c2d..0000000 --- a/tests/extension/test_swiglu.py +++ /dev/null @@ -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) diff --git a/tests/extension/test_swiglu_dispatch.py b/tests/extension/test_swiglu_dispatch.py deleted file mode 100644 index 8b040d4..0000000 --- a/tests/extension/test_swiglu_dispatch.py +++ /dev/null @@ -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)