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