refactor: remove bf16 gemm and swiglu kernels and rebuild csrc benchmarks
- delete csrc/kernels/gemm.cu and swiglu.cu and drop their CMake and setup.py registration - remove the ops wrappers plus backend/linear.py and backend/swiglu.py so Linear and MLP call F.linear directly - drop the four gemm and swiglu kernel test files and prune the stale cuda_kernels.md sections - add csrc/bench benchmarks for the remaining kernels: attention decode prefill paged decode paged prefill versus single-launch SDPA references, rotary versus the torch fallback, fp8 quantize and mm_fp8 versus torch baselines - attention, rotary_emb, and fp8_ops kernels are unchanged
This commit is contained in:
@@ -61,8 +61,6 @@ set(KERNEL_NAMES
|
||||
attn_prefill
|
||||
attn_paged_decode
|
||||
attn_paged_prefill
|
||||
bf16_gemm
|
||||
bf16_swiglu
|
||||
rotary_emb
|
||||
)
|
||||
set(KERNEL_SRCS
|
||||
@@ -70,8 +68,6 @@ set(KERNEL_SRCS
|
||||
attention/prefill.cu
|
||||
attention/paged_decode.cu
|
||||
attention/paged_prefill.cu
|
||||
gemm.cu
|
||||
swiglu.cu
|
||||
rotary_emb.cu
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,649 @@
|
||||
"""Benchmark the four attention kernels against single-launch torch SDPA.
|
||||
|
||||
Suites (--suite): decode, prefill, paged_decode, paged_prefill, all. The
|
||||
torch side times one SDPA call per step over dense tensors: GQA expansion,
|
||||
page-table gathers, padding, and masks are built once outside the timed
|
||||
region, and masked calls prefer the cuDNN backend (the default masked path
|
||||
is the slow math backend). The kernel's timed work still includes its fused
|
||||
paged reads and current-token K/V append. Agreement is checked against the
|
||||
same reference.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
import click
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.attention import SDPBackend, sdpa_kernel
|
||||
|
||||
from astrai.extension import is_available
|
||||
from astrai.extension.ops import (
|
||||
attn_decode,
|
||||
attn_paged_decode,
|
||||
attn_paged_prefill,
|
||||
attn_prefill,
|
||||
)
|
||||
from astrai.inference.workspace import MAX_SPLITS, Q_TILE_ROWS
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GqaConfig:
|
||||
"""One model family's attention geometry (llama-style GQA)."""
|
||||
|
||||
name: str
|
||||
hq: int
|
||||
hkv: int
|
||||
head_dim: int
|
||||
|
||||
|
||||
DEFAULT_CONFIGS = (
|
||||
GqaConfig("llama2_7b", 32, 8, 128),
|
||||
GqaConfig("llama3_70b", 64, 8, 128),
|
||||
GqaConfig("qwen2_7b", 28, 4, 128),
|
||||
GqaConfig("llama3_8b_d64", 32, 8, 64),
|
||||
)
|
||||
|
||||
# (batch, per-request context length); context includes the token being
|
||||
# decoded (kv_len = context, the last slot written in-kernel).
|
||||
DECODE_CASES = ((1, 4096), (8, 4096), (32, 2048), (64, 1024))
|
||||
# (batch, q_len); prefill from scratch so kv_len == q_len.
|
||||
PREFILL_CASES = ((1, 2048), (1, 4096), (4, 1024), (8, 512))
|
||||
|
||||
|
||||
def parse_config(value: str) -> GqaConfig:
|
||||
parts = value.split(":")
|
||||
if len(parts) != 4 or not parts[0]:
|
||||
raise click.BadParameter("config must use NAME:HQ:HKV:HEAD_DIM")
|
||||
try:
|
||||
hq, hkv, head_dim = (int(item) for item in parts[1:])
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("HQ/HKV/HEAD_DIM must be integers") from exc
|
||||
if hq <= 0 or hkv <= 0 or head_dim <= 0 or hq % hkv or head_dim % 32:
|
||||
raise click.BadParameter(
|
||||
"HQ/HKV positive with HQ % HKV == 0; HEAD_DIM % 32 == 0"
|
||||
)
|
||||
return GqaConfig(parts[0], hq, hkv, head_dim)
|
||||
|
||||
|
||||
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
operation()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / iterations
|
||||
|
||||
|
||||
def summarize(values: list[float]) -> dict[str, float]:
|
||||
ordered = sorted(values)
|
||||
return {
|
||||
"median_ms": statistics.median(ordered),
|
||||
"p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
|
||||
}
|
||||
|
||||
|
||||
def measure_operations(
|
||||
operations: dict[str, Callable[[], torch.Tensor]],
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, list[float]]:
|
||||
for operation in operations.values():
|
||||
for _ in range(warmup):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
samples: dict[str, list[float]] = {name: [] for name in operations}
|
||||
order = tuple(operations)
|
||||
# A-B-B-A order balances cache, clock, and temperature drift.
|
||||
for _ in range(trials):
|
||||
for name in (*order, *reversed(order)):
|
||||
samples[name].append(time_operation(operations[name], iterations))
|
||||
return samples
|
||||
|
||||
|
||||
def repeat_kv_heads(x: torch.Tensor, n_rep: int) -> torch.Tensor:
|
||||
"""Expand [*, n_kv_heads, head_dim] to [*, n_kv_heads * n_rep, head_dim]
|
||||
with the backend's grouping (kv head = q head // n_rep)."""
|
||||
if n_rep == 1:
|
||||
return x
|
||||
n_heads, head_dim = x.shape[-2:]
|
||||
return (
|
||||
x.unsqueeze(-2)
|
||||
.expand(*x.shape[:-2], n_heads, n_rep, head_dim)
|
||||
.reshape(*x.shape[:-2], n_heads * n_rep, head_dim)
|
||||
)
|
||||
|
||||
|
||||
def sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs) -> torch.Tensor:
|
||||
"""SDPA over blhd tensors: [batch, seq, heads, head_dim] -> blhd."""
|
||||
out = F.scaled_dot_product_attention(
|
||||
q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), **kwargs
|
||||
)
|
||||
return out.transpose(1, 2)
|
||||
|
||||
|
||||
def prefer_cudnn_sdpa(
|
||||
call: Callable[[], torch.Tensor],
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
"""Return a closure running the masked SDPA ``call`` on cuDNN attention
|
||||
when the backend accepts the bool mask, else torch's default (the
|
||||
default masked path falls back to the much slower math backend)."""
|
||||
|
||||
def with_cudnn() -> torch.Tensor:
|
||||
with sdpa_kernel([SDPBackend.CUDNN_ATTENTION]):
|
||||
return call()
|
||||
|
||||
try:
|
||||
with_cudnn()
|
||||
except Exception:
|
||||
return call
|
||||
return with_cudnn
|
||||
|
||||
|
||||
def ragged_lens(batch: int, span: int) -> list[int]:
|
||||
"""Deterministic mixed lengths spanning [span // 2, span]."""
|
||||
if batch == 1:
|
||||
return [span]
|
||||
step = max(span // 2 // (batch - 1), 1)
|
||||
return [span - (batch - 1 - i) * step for i in range(batch)]
|
||||
|
||||
|
||||
def cumsum_indptr(lens: list[int]) -> torch.Tensor:
|
||||
return torch.tensor(
|
||||
[0, *torch.tensor(lens).cumsum(0).tolist()], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PagedInputs:
|
||||
"""Standalone replicas of the PagePool / InferenceWorkspace tensors."""
|
||||
|
||||
q: torch.Tensor
|
||||
k_cache: torch.Tensor
|
||||
v_cache: torch.Tensor
|
||||
req_to_token: torch.Tensor
|
||||
req_pool_indices: torch.Tensor
|
||||
kv_indptr: torch.Tensor
|
||||
|
||||
|
||||
def build_paged_inputs(
|
||||
config: GqaConfig, kv_lens: list[int], q_lens: Optional[list[int]]
|
||||
) -> PagedInputs:
|
||||
"""Flat pool + page table: request ``i`` owns the contiguous slot range
|
||||
``[offset_i, offset_i + kv_len_i)``. Q is packed across requests when
|
||||
``q_lens`` is given (ragged prefill), else [B, Hq, D] (decode)."""
|
||||
batch = len(kv_lens)
|
||||
pool = torch.randn(
|
||||
sum(kv_lens),
|
||||
config.hkv,
|
||||
config.head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
req_to_token = torch.zeros(batch, max(kv_lens), dtype=torch.int32, device="cuda")
|
||||
offset = 0
|
||||
for i, length in enumerate(kv_lens):
|
||||
req_to_token[i, :length] = torch.arange(
|
||||
offset, offset + length, dtype=torch.int32, device="cuda"
|
||||
)
|
||||
offset += length
|
||||
return PagedInputs(
|
||||
q=torch.randn(
|
||||
sum(q_lens) if q_lens is not None else batch,
|
||||
config.hq,
|
||||
config.head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
),
|
||||
k_cache=pool,
|
||||
v_cache=torch.randn_like(pool),
|
||||
req_to_token=req_to_token,
|
||||
req_pool_indices=torch.arange(batch, dtype=torch.int32, device="cuda"),
|
||||
kv_indptr=cumsum_indptr(kv_lens),
|
||||
)
|
||||
|
||||
|
||||
def report_result(
|
||||
suite: str,
|
||||
config: GqaConfig,
|
||||
case: dict[str, int],
|
||||
operations: dict[str, Callable[[], torch.Tensor]],
|
||||
samples: dict[str, list[float]],
|
||||
io_bytes: int,
|
||||
reference: torch.Tensor,
|
||||
actual: torch.Tensor,
|
||||
) -> dict[str, object]:
|
||||
difference = actual.float() - reference.float()
|
||||
result: dict[str, object] = {
|
||||
"suite": suite,
|
||||
"config": config.name,
|
||||
"agreement": {
|
||||
"max_abs_error": float(difference.abs().max()),
|
||||
"cosine_similarity": float(
|
||||
F.cosine_similarity(
|
||||
actual.float().flatten(), reference.float().flatten(), dim=0
|
||||
)
|
||||
),
|
||||
},
|
||||
"estimated_io_bytes": io_bytes,
|
||||
**case,
|
||||
}
|
||||
for name, operation in operations.items():
|
||||
latency = summarize(samples[name])
|
||||
result[name] = {
|
||||
"effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
|
||||
**latency,
|
||||
}
|
||||
speedup = (result["torch"]["median_ms"] / result["cuda"]["median_ms"] - 1.0) * 100.0
|
||||
label = f"B={case.get('batch')}" + (
|
||||
f" ctx={case['context']}" if "context" in case else f" q={case['q_len']}"
|
||||
)
|
||||
print(
|
||||
f"{suite},{config.name},{label},{result['torch']['median_ms']:.4f},"
|
||||
f"{result['cuda']['median_ms']:.4f},{speedup:+.1f}%,"
|
||||
f"{result['agreement']['max_abs_error']:.4f}"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Suites
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def benchmark_decode(
|
||||
config: GqaConfig,
|
||||
batch: int,
|
||||
context: int,
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, object]:
|
||||
q = torch.randn(
|
||||
batch, 1, config.hq, config.head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
k = torch.randn(
|
||||
batch,
|
||||
context,
|
||||
config.hkv,
|
||||
config.head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
v = torch.randn_like(k)
|
||||
# GQA expansion is data preparation, not attention compute — build it
|
||||
# once so the timed torch side is a single SDPA launch.
|
||||
k_expanded = repeat_kv_heads(k, config.hq // config.hkv)
|
||||
v_expanded = repeat_kv_heads(v, config.hq // config.hkv)
|
||||
|
||||
def torch_op() -> torch.Tensor:
|
||||
return sdpa(q, k_expanded, v_expanded)
|
||||
|
||||
def cuda_op() -> torch.Tensor:
|
||||
return attn_decode(q, k, v, is_causal=True)
|
||||
|
||||
operations = {"torch": torch_op, "cuda": cuda_op}
|
||||
samples = measure_operations(
|
||||
operations, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
io_bytes = (2 * q.numel() + 2 * k.numel()) * q.element_size()
|
||||
return report_result(
|
||||
"decode",
|
||||
config,
|
||||
{"batch": batch, "context": context},
|
||||
operations,
|
||||
samples,
|
||||
io_bytes,
|
||||
torch_op(),
|
||||
cuda_op(),
|
||||
)
|
||||
|
||||
|
||||
def benchmark_prefill(
|
||||
config: GqaConfig,
|
||||
batch: int,
|
||||
q_len: int,
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, object]:
|
||||
q = torch.randn(
|
||||
batch, q_len, config.hq, config.head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
k = torch.randn(
|
||||
batch, q_len, config.hkv, config.head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
v = torch.randn_like(k)
|
||||
k_expanded = repeat_kv_heads(k, config.hq // config.hkv)
|
||||
v_expanded = repeat_kv_heads(v, config.hq // config.hkv)
|
||||
|
||||
def torch_op() -> torch.Tensor:
|
||||
return sdpa(q, k_expanded, v_expanded, is_causal=True)
|
||||
|
||||
def cuda_op() -> torch.Tensor:
|
||||
return attn_prefill(q, k, v, is_causal=True)
|
||||
|
||||
operations = {"torch": torch_op, "cuda": cuda_op}
|
||||
samples = measure_operations(
|
||||
operations, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
io_bytes = (2 * q.numel() + 2 * k.numel()) * q.element_size()
|
||||
return report_result(
|
||||
"prefill",
|
||||
config,
|
||||
{"batch": batch, "q_len": q_len},
|
||||
operations,
|
||||
samples,
|
||||
io_bytes,
|
||||
torch_op(),
|
||||
cuda_op(),
|
||||
)
|
||||
|
||||
|
||||
def benchmark_paged_decode(
|
||||
config: GqaConfig,
|
||||
batch: int,
|
||||
context: int,
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, object]:
|
||||
n_rep = config.hq // config.hkv
|
||||
kv_lens = [length + 1 for length in ragged_lens(batch, context)]
|
||||
inputs = build_paged_inputs(config, kv_lens, None)
|
||||
new_k = torch.randn(
|
||||
batch, config.hkv, config.head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
new_v = torch.randn_like(new_k)
|
||||
o_part = torch.empty(
|
||||
batch,
|
||||
config.hq,
|
||||
MAX_SPLITS,
|
||||
config.head_dim,
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
)
|
||||
ml_part = torch.empty(
|
||||
batch, config.hq, MAX_SPLITS, 2, dtype=torch.float32, device="cuda"
|
||||
)
|
||||
out_buf = torch.empty(
|
||||
batch, config.hq, config.head_dim, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
|
||||
def cuda_op() -> torch.Tensor:
|
||||
return attn_paged_decode(
|
||||
inputs.q,
|
||||
inputs.k_cache,
|
||||
inputs.v_cache,
|
||||
inputs.req_to_token,
|
||||
inputs.req_pool_indices,
|
||||
inputs.kv_indptr,
|
||||
new_k=new_k,
|
||||
new_v=new_v,
|
||||
is_causal=True,
|
||||
o_part_buf=o_part,
|
||||
ml_part_buf=ml_part,
|
||||
out_buf=out_buf,
|
||||
)
|
||||
|
||||
# Reference-side data preparation happens once, outside the timed
|
||||
# closure: append the current-token K/V into the pool (the kernel does
|
||||
# this fused inside its launch), gather padded K/V, expand GQA heads.
|
||||
max_len = max(kv_lens)
|
||||
slots = inputs.req_to_token[:, :max_len].long()
|
||||
last_slots = inputs.req_to_token[
|
||||
torch.arange(batch, device="cuda"), torch.tensor(kv_lens) - 1
|
||||
].long()
|
||||
inputs.k_cache[last_slots] = new_k
|
||||
inputs.v_cache[last_slots] = new_v
|
||||
k_expanded = repeat_kv_heads(inputs.k_cache[slots], n_rep)
|
||||
v_expanded = repeat_kv_heads(inputs.v_cache[slots], n_rep)
|
||||
position = torch.arange(max_len, device="cuda")
|
||||
lengths = torch.tensor(kv_lens, device="cuda", dtype=torch.long)
|
||||
keep_mask = (position[None, :] < lengths[:, None])[:, None, None, :]
|
||||
q_batched = inputs.q.unsqueeze(1)
|
||||
|
||||
sdpa_call = prefer_cudnn_sdpa(
|
||||
lambda: sdpa(q_batched, k_expanded, v_expanded, attn_mask=keep_mask)
|
||||
)
|
||||
|
||||
def torch_op() -> torch.Tensor:
|
||||
return sdpa_call().squeeze(1) # [B, Hq, D]
|
||||
|
||||
operations = {"torch": torch_op, "cuda": cuda_op}
|
||||
samples = measure_operations(
|
||||
operations, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
io_bytes = (
|
||||
2 * inputs.q.numel() # q read + out write
|
||||
+ 2 * sum(kv_lens) * config.hkv * config.head_dim # k/v reads
|
||||
+ 2 * new_k.numel() # new k/v writes
|
||||
) * inputs.q.element_size()
|
||||
return report_result(
|
||||
"paged_decode",
|
||||
config,
|
||||
{"batch": batch, "context": context},
|
||||
operations,
|
||||
samples,
|
||||
io_bytes,
|
||||
torch_op(),
|
||||
cuda_op(),
|
||||
)
|
||||
|
||||
|
||||
def benchmark_paged_prefill(
|
||||
config: GqaConfig,
|
||||
batch: int,
|
||||
q_len: int,
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, object]:
|
||||
n_rep = config.hq // config.hkv
|
||||
q_lens = ragged_lens(batch, q_len)
|
||||
inputs = build_paged_inputs(config, q_lens, q_lens)
|
||||
qo_indptr = cumsum_indptr(q_lens)
|
||||
|
||||
tile_batches, tile_indices = [], []
|
||||
for request, length in enumerate(q_lens):
|
||||
n_tiles = (length + Q_TILE_ROWS - 1) // Q_TILE_ROWS
|
||||
tile_batches.extend([request] * n_tiles)
|
||||
tile_indices.extend(range(n_tiles))
|
||||
q_tile_to_batch = torch.tensor(tile_batches, dtype=torch.int32, device="cuda")
|
||||
q_tile_to_index = torch.tensor(tile_indices, dtype=torch.int32, device="cuda")
|
||||
|
||||
def cuda_op() -> torch.Tensor:
|
||||
return attn_paged_prefill(
|
||||
inputs.q,
|
||||
inputs.k_cache,
|
||||
inputs.v_cache,
|
||||
inputs.req_to_token,
|
||||
inputs.req_pool_indices,
|
||||
inputs.kv_indptr,
|
||||
qo_indptr,
|
||||
q_tile_to_batch,
|
||||
q_tile_to_index,
|
||||
is_causal=True,
|
||||
)
|
||||
|
||||
# Same once-only preparation: gather padded K/V, expand GQA heads, pad Q,
|
||||
# build the causal + validity mask. The timed reference is one SDPA call
|
||||
# plus the packed-row unpack.
|
||||
max_len = max(q_lens)
|
||||
slots = inputs.req_to_token[:, :max_len].long()
|
||||
k_expanded = repeat_kv_heads(inputs.k_cache[slots], n_rep)
|
||||
v_expanded = repeat_kv_heads(inputs.v_cache[slots], n_rep)
|
||||
position = torch.arange(max_len, device="cuda")
|
||||
lengths = torch.tensor(q_lens, device="cuda", dtype=torch.long)
|
||||
causal = position[None, :, None] >= position[None, None, :]
|
||||
keep = position[None, None, :] < lengths[:, None, None]
|
||||
attn_mask = (causal & keep).unsqueeze(1)
|
||||
q_padded = torch.zeros(
|
||||
batch,
|
||||
max_len,
|
||||
config.hq,
|
||||
config.head_dim,
|
||||
device="cuda",
|
||||
dtype=inputs.q.dtype,
|
||||
)
|
||||
for i, length in enumerate(q_lens):
|
||||
q_padded[i, :length] = inputs.q[int(qo_indptr[i]) : int(qo_indptr[i + 1])]
|
||||
|
||||
sdpa_call = prefer_cudnn_sdpa(
|
||||
lambda: sdpa(q_padded, k_expanded, v_expanded, attn_mask=attn_mask)
|
||||
)
|
||||
|
||||
def torch_op() -> torch.Tensor:
|
||||
out = sdpa_call()
|
||||
return torch.cat([out[i, :length] for i, length in enumerate(q_lens)])
|
||||
|
||||
operations = {"torch": torch_op, "cuda": cuda_op}
|
||||
samples = measure_operations(
|
||||
operations, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
io_bytes = (
|
||||
2 * inputs.q.numel() + 2 * sum(q_lens) * config.hkv * config.head_dim
|
||||
) * inputs.q.element_size()
|
||||
return report_result(
|
||||
"paged_prefill",
|
||||
config,
|
||||
{"batch": batch, "q_len": q_len},
|
||||
operations,
|
||||
samples,
|
||||
io_bytes,
|
||||
torch_op(),
|
||||
cuda_op(),
|
||||
)
|
||||
|
||||
|
||||
@click.command(help=__doc__)
|
||||
@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
|
||||
@click.option(
|
||||
"--suite",
|
||||
"suites",
|
||||
type=click.Choice(("decode", "prefill", "paged_decode", "paged_prefill", "all")),
|
||||
multiple=True,
|
||||
default=("all",),
|
||||
show_default=True,
|
||||
)
|
||||
@click.option(
|
||||
"--config",
|
||||
"config_values",
|
||||
multiple=True,
|
||||
help="Filter defaults by bare name, or add/override with NAME:HQ:HKV:HEAD_DIM.",
|
||||
)
|
||||
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option("--iterations", type=click.IntRange(min=1), default=50, show_default=True)
|
||||
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option("--seed", type=int, default=0, show_default=True)
|
||||
def benchmark_command(
|
||||
output: Path | None,
|
||||
suites: tuple[str, ...],
|
||||
config_values: tuple[str, ...],
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise click.ClickException("CUDA is required")
|
||||
kernel_for_suite = {
|
||||
"decode": "attn_decode",
|
||||
"prefill": "attn_prefill",
|
||||
"paged_decode": "attn_paged_decode",
|
||||
"paged_prefill": "attn_paged_prefill",
|
||||
}
|
||||
selected = (
|
||||
tuple(kernel_for_suite) if "all" in suites else tuple(dict.fromkeys(suites))
|
||||
)
|
||||
missing = [
|
||||
kernel_for_suite[suite]
|
||||
for suite in selected
|
||||
if not is_available(kernel_for_suite[suite])
|
||||
]
|
||||
if missing:
|
||||
raise click.ClickException(f"built kernels required: {', '.join(missing)}")
|
||||
|
||||
# A bare name filters the matching default; a full spec overrides or appends.
|
||||
chosen: dict[str, GqaConfig] = {}
|
||||
for value in config_values:
|
||||
config = (
|
||||
next((c for c in DEFAULT_CONFIGS if c.name == value), None)
|
||||
if ":" not in value
|
||||
else parse_config(value)
|
||||
)
|
||||
if config is None:
|
||||
raise click.BadParameter(f"unknown default config {value!r}")
|
||||
chosen[config.name] = config
|
||||
configs = tuple(chosen.values()) or DEFAULT_CONFIGS
|
||||
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
runners = {
|
||||
"decode": (benchmark_decode, DECODE_CASES),
|
||||
"prefill": (benchmark_prefill, PREFILL_CASES),
|
||||
"paged_decode": (benchmark_paged_decode, DECODE_CASES),
|
||||
"paged_prefill": (benchmark_paged_prefill, PREFILL_CASES),
|
||||
}
|
||||
print("suite,config,case,torch_ms,cuda_ms,speedup,max_abs")
|
||||
results = []
|
||||
with torch.inference_mode():
|
||||
for suite in selected:
|
||||
runner, cases = runners[suite]
|
||||
for config in configs:
|
||||
for case in cases:
|
||||
results.append(
|
||||
runner(
|
||||
config,
|
||||
*case,
|
||||
warmup=warmup,
|
||||
iterations=iterations,
|
||||
trials=trials,
|
||||
)
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if output is not None:
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
payload = {
|
||||
"metadata": {
|
||||
"gpu_name": props.name,
|
||||
"compute_capability": f"{props.major}.{props.minor}",
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
||||
},
|
||||
"settings": {
|
||||
"warmup": warmup,
|
||||
"iterations": iterations,
|
||||
"trials": trials,
|
||||
"seed": seed,
|
||||
"order": "A-B-B-A",
|
||||
"suites": list(selected),
|
||||
"configs": [asdict(config) for config in configs],
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_command()
|
||||
@@ -0,0 +1,394 @@
|
||||
"""Benchmark the FP8 quantize and GEMM kernels against torch baselines.
|
||||
|
||||
Suites (--suite): quantize (plain / delayed-scaling ring / dual-orientation
|
||||
entries vs the aten float8 cast) and gemm (pre-quantized ``mm_fp8`` in the
|
||||
NT orientation the fp8 linear path uses, vs bf16 ``F.linear``). GEMM
|
||||
agreement reports both kernel error (vs the fp32 dequantized fp8 product)
|
||||
and format error (that product vs the bf16 matmul). FP8 MMA requires
|
||||
compute capability 89+.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import click
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension import is_available
|
||||
from astrai.extension.ops.fp8 import mm_fp8, quantize, quantize_dual
|
||||
|
||||
FP8_MAX = {"e4m3": 448.0, "e5m2": 57344.0}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MatrixShape:
|
||||
name: str
|
||||
rows: int
|
||||
cols: int
|
||||
|
||||
|
||||
QUANTIZE_SHAPES = (
|
||||
MatrixShape("astrai_1b_act", 2048, 1536),
|
||||
MatrixShape("llama2_7b_act", 2048, 4096),
|
||||
MatrixShape("llama2_7b_down_w", 4096, 11008),
|
||||
MatrixShape("llama3_70b_act", 2048, 8192),
|
||||
)
|
||||
|
||||
# GEMM shapes as (N, K) weight mats; M comes from --m-values.
|
||||
GEMM_SHAPES = (
|
||||
MatrixShape("llama2_7b_qkv", 4096, 4096),
|
||||
MatrixShape("llama2_7b_up_gate", 11008, 4096),
|
||||
MatrixShape("llama2_7b_down", 4096, 11008),
|
||||
MatrixShape("llama3_70b_up_gate", 28672, 8192),
|
||||
)
|
||||
|
||||
|
||||
def parse_positive_ints(value: str) -> tuple[int, ...]:
|
||||
try:
|
||||
values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("expected comma-separated integers") from exc
|
||||
if not values or any(item <= 0 for item in values):
|
||||
raise click.BadParameter("values must be positive integers")
|
||||
return values
|
||||
|
||||
|
||||
def parse_shape(value: str) -> MatrixShape:
|
||||
parts = value.split(":")
|
||||
if len(parts) != 3 or not parts[0]:
|
||||
raise click.BadParameter("shape must use NAME:ROWS:COLS")
|
||||
try:
|
||||
rows, cols = (int(item) for item in parts[1:])
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("ROWS and COLS must be integers") from exc
|
||||
if rows <= 0 or cols <= 0:
|
||||
raise click.BadParameter("ROWS and COLS must be positive")
|
||||
return MatrixShape(parts[0], rows, cols)
|
||||
|
||||
|
||||
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
operation()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / iterations
|
||||
|
||||
|
||||
def summarize(values: list[float]) -> dict[str, float]:
|
||||
ordered = sorted(values)
|
||||
return {
|
||||
"median_ms": statistics.median(ordered),
|
||||
"p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
|
||||
}
|
||||
|
||||
|
||||
def measure_operations(
|
||||
operations: dict[str, Callable[[], torch.Tensor]],
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, list[float]]:
|
||||
for operation in operations.values():
|
||||
for _ in range(warmup):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
samples: dict[str, list[float]] = {name: [] for name in operations}
|
||||
order = tuple(operations)
|
||||
# A-B-C-C-B-A order balances cache, clock, and temperature drift.
|
||||
for _ in range(trials):
|
||||
for name in (*order, *reversed(order)):
|
||||
samples[name].append(time_operation(operations[name], iterations))
|
||||
return samples
|
||||
|
||||
|
||||
def quant_step(x: torch.Tensor, fmt: str) -> torch.Tensor:
|
||||
"""Quantization step (dequant scale) from the current amax; ``quantize``
|
||||
takes its reciprocal as the multiplier."""
|
||||
amax = x.abs().amax().to(torch.float32).clamp_min(1e-12)
|
||||
return amax / FP8_MAX[fmt]
|
||||
|
||||
|
||||
def benchmark_quantize(
|
||||
shape: MatrixShape,
|
||||
fmt: str,
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, object]:
|
||||
x = torch.randn(shape.rows, shape.cols, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
multiplier = quant_step(x, fmt).reciprocal()
|
||||
fp8_dtype = torch.float8_e4m3fn if fmt == "e4m3" else torch.float8_e5m2
|
||||
ring_state = torch.zeros(16 + 4, dtype=torch.float32, device="cuda")
|
||||
|
||||
operations: dict[str, Callable[[], torch.Tensor]] = {
|
||||
"torch_cast": lambda: x.to(fp8_dtype),
|
||||
"plain": lambda: quantize(x, multiplier, fmt)[0],
|
||||
"ring": lambda: quantize(
|
||||
x,
|
||||
multiplier,
|
||||
fmt,
|
||||
ring_state=ring_state,
|
||||
hist_idx=0,
|
||||
fp8_max=FP8_MAX[fmt],
|
||||
)[0],
|
||||
"dual": lambda: quantize_dual(x, multiplier, fmt)[0],
|
||||
}
|
||||
samples = measure_operations(
|
||||
operations, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
step = quant_step(x, fmt)
|
||||
x8, _ = quantize(x, multiplier, fmt)
|
||||
dequant_error = float((x8.to(torch.float32) * step - x.float()).abs().max())
|
||||
dual_t = quantize_dual(x, multiplier, fmt)[1]
|
||||
dual_matches = bool(torch.equal(dual_t.t().contiguous(), x8.contiguous()))
|
||||
|
||||
io_bytes = 3 * x.numel() + 4 # bf16 read + fp8 write + f32 amax
|
||||
result: dict[str, object] = {
|
||||
"suite": "quantize",
|
||||
"shape": shape.name,
|
||||
"rows": shape.rows,
|
||||
"cols": shape.cols,
|
||||
"fmt": fmt,
|
||||
"estimated_io_bytes": io_bytes,
|
||||
"dequant_max_abs_error": dequant_error,
|
||||
"dual_transpose_matches": dual_matches,
|
||||
}
|
||||
for name, samples_ms in samples.items():
|
||||
latency = summarize(samples_ms)
|
||||
bytes_per_call = io_bytes + x.numel() if name == "dual" else io_bytes
|
||||
result[name] = {
|
||||
"effective_bandwidth_gbps": bytes_per_call
|
||||
/ (latency["median_ms"] / 1000)
|
||||
/ 1e9,
|
||||
**latency,
|
||||
}
|
||||
speedup = (
|
||||
result["torch_cast"]["median_ms"] / result["plain"]["median_ms"] - 1.0
|
||||
) * 100.0
|
||||
print(
|
||||
f"quantize,{shape.name},{shape.rows}x{shape.cols},"
|
||||
f"{result['torch_cast']['median_ms']:.4f},{result['plain']['median_ms']:.4f},"
|
||||
f"{result['ring']['median_ms']:.4f},{result['dual']['median_ms']:.4f},"
|
||||
f"{speedup:+.1f}%,{dequant_error:.5f},{dual_matches}"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def benchmark_gemm(
|
||||
shape: MatrixShape,
|
||||
m: int,
|
||||
fmt: str,
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, object]:
|
||||
x = torch.randn(m, shape.cols, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
w = (
|
||||
torch.randn(shape.rows, shape.cols, device="cuda", dtype=torch.bfloat16)
|
||||
* shape.cols**-0.5
|
||||
)
|
||||
sx, sw = quant_step(x, fmt), quant_step(w, fmt)
|
||||
with torch.no_grad():
|
||||
x8, _ = quantize(x, sx.reciprocal(), fmt)
|
||||
w8, _ = quantize(w, sw.reciprocal(), fmt)
|
||||
dequant_scale = sx * sw
|
||||
|
||||
def torch_op() -> torch.Tensor:
|
||||
return F.linear(x, w)
|
||||
|
||||
def fp8_op() -> torch.Tensor:
|
||||
return mm_fp8(x8, w8, dequant_scale, trans_b=True)
|
||||
|
||||
operations = {"torch_bf16": torch_op, "fp8": fp8_op}
|
||||
samples = measure_operations(
|
||||
operations, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
actual = fp8_op().float()
|
||||
dequant_reference = (
|
||||
x8.to(torch.float32) @ w8.to(torch.float32).t() * dequant_scale
|
||||
)
|
||||
bf16_reference = torch_op().float()
|
||||
kernel_difference = actual - dequant_reference
|
||||
format_difference = dequant_reference - bf16_reference
|
||||
io_bytes = m * shape.cols + shape.rows * shape.cols + 2 * m * shape.rows
|
||||
|
||||
result: dict[str, object] = {
|
||||
"suite": "gemm",
|
||||
"shape": shape.name,
|
||||
"m": m,
|
||||
"n": shape.rows,
|
||||
"k": shape.cols,
|
||||
"fmt": fmt,
|
||||
"estimated_io_bytes": io_bytes,
|
||||
"kernel_max_abs": float(kernel_difference.abs().max()),
|
||||
"kernel_rel_l2": float(
|
||||
kernel_difference.norm() / dequant_reference.norm().clamp_min(1e-12)
|
||||
),
|
||||
"format_rel_l2": float(
|
||||
format_difference.norm() / bf16_reference.norm().clamp_min(1e-12)
|
||||
),
|
||||
}
|
||||
for name, samples_ms in samples.items():
|
||||
latency = summarize(samples_ms)
|
||||
result[name] = {
|
||||
"effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
|
||||
**latency,
|
||||
}
|
||||
speedup = (
|
||||
result["torch_bf16"]["median_ms"] / result["fp8"]["median_ms"] - 1.0
|
||||
) * 100.0
|
||||
print(
|
||||
f"gemm,{shape.name},{m}x{shape.rows}x{shape.cols},"
|
||||
f"{result['torch_bf16']['median_ms']:.4f},{result['fp8']['median_ms']:.4f},"
|
||||
f"{speedup:+.1f}%,{result['kernel_max_abs']:.4f},"
|
||||
f"{result['kernel_rel_l2']:.6f},{result['format_rel_l2']:.6f}"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@click.command(help=__doc__)
|
||||
@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
|
||||
@click.option(
|
||||
"--suite",
|
||||
"suites",
|
||||
type=click.Choice(("quantize", "gemm", "all")),
|
||||
multiple=True,
|
||||
default=("all",),
|
||||
show_default=True,
|
||||
)
|
||||
@click.option("--fmt", type=click.Choice(("e4m3", "e5m2")), default="e4m3")
|
||||
@click.option("--m-values", default="512,2048,4096", show_default=True)
|
||||
@click.option(
|
||||
"--shape",
|
||||
"shape_values",
|
||||
multiple=True,
|
||||
help="Filter defaults by bare name (either suite), or add/override with "
|
||||
"NAME:ROWS:COLS.",
|
||||
)
|
||||
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option(
|
||||
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
|
||||
)
|
||||
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option("--seed", type=int, default=0, show_default=True)
|
||||
def benchmark_command(
|
||||
output: Path | None,
|
||||
suites: tuple[str, ...],
|
||||
fmt: str,
|
||||
m_values: str,
|
||||
shape_values: tuple[str, ...],
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise click.ClickException("CUDA is required")
|
||||
if not is_available("fp8_ops"):
|
||||
raise click.ClickException(
|
||||
"the built fp8_ops kernel is required (compute capability 89+)"
|
||||
)
|
||||
selected = ("quantize", "gemm") if "all" in suites else tuple(dict.fromkeys(suites))
|
||||
m_values_parsed = parse_positive_ints(m_values)
|
||||
|
||||
# Any --shape selection replaces the defaults for both suites: a bare
|
||||
# name keeps that suite's matching default, a NAME:ROWS:COLS spec
|
||||
# overrides the same-name default or adds a new one.
|
||||
bare_names = {value for value in shape_values if ":" not in value}
|
||||
known = {shape.name for shape in QUANTIZE_SHAPES + GEMM_SHAPES}
|
||||
unknown = sorted(bare_names - known)
|
||||
if unknown:
|
||||
raise click.BadParameter(f"unknown default shape names: {', '.join(unknown)}")
|
||||
specs = [parse_shape(value) for value in shape_values if ":" in value]
|
||||
|
||||
def resolve_shapes(defaults: tuple[MatrixShape, ...]) -> list[MatrixShape]:
|
||||
if not shape_values:
|
||||
return list(defaults)
|
||||
by_name = {shape.name: shape for shape in defaults if shape.name in bare_names}
|
||||
for spec in specs:
|
||||
by_name[spec.name] = spec
|
||||
return list(by_name.values())
|
||||
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
results = []
|
||||
with torch.inference_mode():
|
||||
if "quantize" in selected:
|
||||
print(
|
||||
"suite,shape,size,cast_ms,plain_ms,ring_ms,dual_ms,vs_cast,"
|
||||
"dequant_max,dual_ok"
|
||||
)
|
||||
for shape in resolve_shapes(QUANTIZE_SHAPES):
|
||||
results.append(
|
||||
benchmark_quantize(
|
||||
shape, fmt, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
if "gemm" in selected:
|
||||
print(
|
||||
"suite,shape,mxn_xk,bf16_ms,fp8_ms,speedup,kernel_max,"
|
||||
"kernel_rel_l2,format_rel_l2"
|
||||
)
|
||||
for shape in resolve_shapes(GEMM_SHAPES):
|
||||
for m in m_values_parsed:
|
||||
results.append(
|
||||
benchmark_gemm(
|
||||
shape,
|
||||
m,
|
||||
fmt,
|
||||
warmup=warmup,
|
||||
iterations=iterations,
|
||||
trials=trials,
|
||||
)
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if output is not None:
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
payload = {
|
||||
"metadata": {
|
||||
"gpu_name": props.name,
|
||||
"compute_capability": f"{props.major}.{props.minor}",
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"fmt": fmt,
|
||||
},
|
||||
"settings": {
|
||||
"warmup": warmup,
|
||||
"iterations": iterations,
|
||||
"trials": trials,
|
||||
"seed": seed,
|
||||
"order": "A-B-C-C-B-A",
|
||||
"suites": list(selected),
|
||||
"m_values": list(m_values_parsed),
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_command()
|
||||
@@ -1,333 +0,0 @@
|
||||
"""Benchmark decode-time linear shapes before enabling custom GEMM dispatch.
|
||||
|
||||
The benchmark deliberately calls ``torch.nn.functional.linear`` directly. It
|
||||
establishes the per-architecture cuBLAS baseline that later GEMM primitives and
|
||||
dispatch decisions must beat.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from dataclasses import asdict, dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Callable, Iterable
|
||||
|
||||
import click
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LinearShape:
|
||||
name: str
|
||||
n: int
|
||||
k: int
|
||||
|
||||
|
||||
DEFAULT_SHAPES = (
|
||||
LinearShape("q_proj", 1536, 1536),
|
||||
LinearShape("k_proj", 256, 1536),
|
||||
LinearShape("v_proj", 256, 1536),
|
||||
LinearShape("attn_out", 1536, 1536),
|
||||
LinearShape("mlp_up", 6912, 1536),
|
||||
LinearShape("mlp_gate", 6912, 1536),
|
||||
LinearShape("mlp_down", 1536, 6912),
|
||||
LinearShape("lm_head", 100000, 1536),
|
||||
)
|
||||
DTYPES = {"bfloat16": torch.bfloat16, "float16": torch.float16}
|
||||
|
||||
|
||||
def parse_positive_ints(value: str) -> tuple[int, ...]:
|
||||
"""Parse a comma-separated, duplicate-free list of positive integers."""
|
||||
try:
|
||||
values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("expected comma-separated integers") from exc
|
||||
if not values or any(item <= 0 for item in values):
|
||||
raise click.BadParameter("values must be positive integers")
|
||||
return values
|
||||
|
||||
|
||||
def parse_shape(value: str) -> LinearShape:
|
||||
"""Parse NAME:N:K into a benchmark shape."""
|
||||
parts = value.split(":")
|
||||
if len(parts) != 3 or not parts[0]:
|
||||
raise click.BadParameter("shape must use NAME:N:K")
|
||||
try:
|
||||
n, k = (int(item) for item in parts[1:])
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("N and K must be integers") from exc
|
||||
if n <= 0 or k <= 0:
|
||||
raise click.BadParameter("N and K must be positive")
|
||||
return LinearShape(parts[0], n, k)
|
||||
|
||||
|
||||
def estimate_io_bytes(
|
||||
m: int, n: int, k: int, element_size: int, *, has_bias: bool
|
||||
) -> int:
|
||||
"""Estimate bytes touched once by Y[M,N] = X[M,K] @ W[N,K].T."""
|
||||
elements = m * k + n * k + m * n
|
||||
if has_bias:
|
||||
elements += n
|
||||
return elements * element_size
|
||||
|
||||
|
||||
def percentile(values: Iterable[float], quantile: float) -> float:
|
||||
ordered = sorted(values)
|
||||
if not ordered:
|
||||
raise ValueError("percentile requires at least one sample")
|
||||
rank = (len(ordered) - 1) * quantile
|
||||
lower = math.floor(rank)
|
||||
upper = math.ceil(rank)
|
||||
if lower == upper:
|
||||
return ordered[lower]
|
||||
fraction = rank - lower
|
||||
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
|
||||
|
||||
|
||||
def summarize_latency(samples_ms: list[float]) -> dict[str, float]:
|
||||
return {
|
||||
"median_ms": statistics.median(samples_ms),
|
||||
"p90_ms": percentile(samples_ms, 0.90),
|
||||
"p99_ms": percentile(samples_ms, 0.99),
|
||||
"min_ms": min(samples_ms),
|
||||
"max_ms": max(samples_ms),
|
||||
}
|
||||
|
||||
|
||||
def measure_cuda_ms(
|
||||
operation: Callable[[], torch.Tensor], *, warmup: int, iterations: int, trials: int
|
||||
) -> list[float]:
|
||||
for _ in range(warmup):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
samples = []
|
||||
for _ in range(trials):
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
operation()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
samples.append(start.elapsed_time(end) / iterations)
|
||||
return samples
|
||||
|
||||
|
||||
def count_cuda_kernels(
|
||||
operation: Callable[[], torch.Tensor], repeats: int = 5
|
||||
) -> float:
|
||||
"""Profile a few calls and return the average device events per call."""
|
||||
with torch.profiler.profile(
|
||||
activities=[
|
||||
torch.profiler.ProfilerActivity.CPU,
|
||||
torch.profiler.ProfilerActivity.CUDA,
|
||||
],
|
||||
acc_events=True,
|
||||
) as profile:
|
||||
for _ in range(repeats):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
device_type = torch.autograd.DeviceType.CUDA
|
||||
events = [event for event in profile.events() if event.device_type == device_type]
|
||||
return len(events) / repeats
|
||||
|
||||
|
||||
def capture_linear(
|
||||
x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None
|
||||
) -> tuple[torch.cuda.CUDAGraph, torch.Tensor]:
|
||||
for _ in range(3):
|
||||
F.linear(x, weight, bias)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
output = F.linear(x, weight, bias)
|
||||
return graph, output
|
||||
|
||||
|
||||
def benchmark_case(
|
||||
shape: LinearShape,
|
||||
m: int,
|
||||
*,
|
||||
dtype: torch.dtype,
|
||||
mode: str,
|
||||
bias_enabled: bool,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> dict[str, object]:
|
||||
x = torch.randn((m, shape.k), device="cuda", dtype=dtype)
|
||||
weight = torch.randn((shape.n, shape.k), device="cuda", dtype=dtype)
|
||||
bias = torch.randn(shape.n, device="cuda", dtype=dtype) if bias_enabled else None
|
||||
|
||||
graph = None
|
||||
graph_output = None
|
||||
if mode == "graph":
|
||||
graph, graph_output = capture_linear(x, weight, bias)
|
||||
|
||||
def operation() -> torch.Tensor:
|
||||
graph.replay()
|
||||
return graph_output
|
||||
|
||||
else:
|
||||
|
||||
def operation() -> torch.Tensor:
|
||||
return F.linear(x, weight, bias)
|
||||
|
||||
samples_ms = measure_cuda_ms(
|
||||
operation, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
latency = summarize_latency(samples_ms)
|
||||
io_bytes = estimate_io_bytes(
|
||||
m, shape.n, shape.k, x.element_size(), has_bias=bias is not None
|
||||
)
|
||||
median_seconds = latency["median_ms"] / 1000
|
||||
|
||||
result: dict[str, object] = {
|
||||
"name": shape.name,
|
||||
"m": m,
|
||||
"n": shape.n,
|
||||
"k": shape.k,
|
||||
"mode": mode,
|
||||
"bias": bias is not None,
|
||||
"estimated_io_bytes": io_bytes,
|
||||
"effective_bandwidth_gbps": io_bytes / median_seconds / 1e9,
|
||||
"cuda_kernel_launches_per_call": count_cuda_kernels(operation),
|
||||
**latency,
|
||||
"samples_ms": samples_ms,
|
||||
}
|
||||
return result
|
||||
|
||||
|
||||
def render_markdown(payload: dict[str, object]) -> str:
|
||||
metadata = payload["metadata"]
|
||||
assert isinstance(metadata, dict)
|
||||
results = payload["results"]
|
||||
assert isinstance(results, list)
|
||||
|
||||
lines = [
|
||||
"# Decode linear baseline",
|
||||
"",
|
||||
f"- GPU: {metadata['gpu_name']}",
|
||||
f"- Compute capability: {metadata['compute_capability']}",
|
||||
f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
|
||||
f"- Dtype: {metadata['dtype']}",
|
||||
"",
|
||||
"| Layer | M | N | K | Mode | Median (ms) | p99 (ms) | GB/s | CUDA kernels/call |",
|
||||
"|---|---:|---:|---:|---|---:|---:|---:|---:|",
|
||||
]
|
||||
for item in results:
|
||||
assert isinstance(item, dict)
|
||||
lines.append(
|
||||
"| {name} | {m} | {n} | {k} | {mode} | {median_ms:.4f} | "
|
||||
"{p99_ms:.4f} | {effective_bandwidth_gbps:.1f} | "
|
||||
"{cuda_kernel_launches_per_call:.2f} |".format(**item)
|
||||
)
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def device_metadata(dtype_name: str) -> dict[str, object]:
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
return {
|
||||
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"gpu_name": props.name,
|
||||
"compute_capability": f"{props.major}.{props.minor}",
|
||||
"total_memory_bytes": props.total_memory,
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"dtype": dtype_name,
|
||||
}
|
||||
|
||||
|
||||
@click.command(help=__doc__)
|
||||
@click.option("--output", type=click.Path(path_type=Path), required=True)
|
||||
@click.option("--markdown-output", type=click.Path(path_type=Path))
|
||||
@click.option("--m-values", default="1,2,4,8,16,32", show_default=True)
|
||||
@click.option(
|
||||
"--shape",
|
||||
"shape_values",
|
||||
multiple=True,
|
||||
help="Override defaults with repeatable NAME:N:K shapes.",
|
||||
)
|
||||
@click.option("--dtype", type=click.Choice(tuple(DTYPES)), default="bfloat16")
|
||||
@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both")
|
||||
@click.option("--bias/--no-bias", default=False)
|
||||
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option(
|
||||
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
|
||||
)
|
||||
@click.option("--trials", type=click.IntRange(min=1), default=20, show_default=True)
|
||||
@click.option("--seed", type=int, default=0, show_default=True)
|
||||
def benchmark_command(
|
||||
output: Path,
|
||||
markdown_output: Path | None,
|
||||
m_values: str,
|
||||
shape_values: tuple[str, ...],
|
||||
dtype: str,
|
||||
mode: str,
|
||||
bias: bool,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise click.ClickException("CUDA is required")
|
||||
|
||||
parsed_m = parse_positive_ints(m_values)
|
||||
shapes = tuple(parse_shape(item) for item in shape_values) or DEFAULT_SHAPES
|
||||
modes = ("eager", "graph") if mode == "both" else (mode,)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
results = []
|
||||
for shape in shapes:
|
||||
for m in parsed_m:
|
||||
for current_mode in modes:
|
||||
click.echo(
|
||||
f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}"
|
||||
)
|
||||
results.append(
|
||||
benchmark_case(
|
||||
shape,
|
||||
m,
|
||||
dtype=DTYPES[dtype],
|
||||
mode=current_mode,
|
||||
bias_enabled=bias,
|
||||
warmup=warmup,
|
||||
iterations=iterations,
|
||||
trials=trials,
|
||||
)
|
||||
)
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"schema_version": 1,
|
||||
"metadata": device_metadata(dtype),
|
||||
"parameters": {
|
||||
"m_values": list(parsed_m),
|
||||
"shapes": [asdict(shape) for shape in shapes],
|
||||
"modes": list(modes),
|
||||
"bias": bias,
|
||||
"warmup": warmup,
|
||||
"iterations": iterations,
|
||||
"trials": trials,
|
||||
"seed": seed,
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
||||
if markdown_output is not None:
|
||||
markdown_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
markdown_output.write_text(render_markdown(payload), encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_command()
|
||||
@@ -1,450 +0,0 @@
|
||||
"""Benchmark the BF16 GEMM primitive and guarded linear dispatcher.
|
||||
|
||||
The kernel suite covers AstrAI's native projections plus common LLaMA and
|
||||
GPT-NeoX matrix shapes. The chain suite is a synthetic projection/MLP chain;
|
||||
it measures dispatcher overhead and dependent MLP work, but is deliberately
|
||||
not presented as a whole-model throughput benchmark.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import statistics
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension import bf16_gemm, is_available, linear
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Shape:
|
||||
label: str
|
||||
n: int
|
||||
k: int
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Chain:
|
||||
label: str
|
||||
hidden: int
|
||||
kv: int
|
||||
intermediate: int
|
||||
fused_qkv: bool = False
|
||||
gated_mlp: bool = True
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Timing:
|
||||
median_ms: float
|
||||
p90_ms: float
|
||||
|
||||
|
||||
ASTRAI_SHAPES = (
|
||||
Shape("astrai_qkv", 256, 1536),
|
||||
Shape("astrai_square", 1536, 1536),
|
||||
Shape("astrai_up_gate", 6912, 1536),
|
||||
Shape("astrai_down", 1536, 6912),
|
||||
Shape("astrai_lm_head", 100000, 1536),
|
||||
)
|
||||
|
||||
TRADITIONAL_SHAPES = (
|
||||
Shape("llama2_7b_qo", 4096, 4096),
|
||||
Shape("llama2_7b_up_gate", 11008, 4096),
|
||||
Shape("llama2_7b_down", 4096, 11008),
|
||||
Shape("llama3_8b_kv", 1024, 4096),
|
||||
Shape("llama3_8b_up_gate", 14336, 4096),
|
||||
Shape("llama3_8b_down", 4096, 14336),
|
||||
Shape("llama2_13b_qo", 5120, 5120),
|
||||
Shape("llama2_13b_up_gate", 13824, 5120),
|
||||
Shape("llama2_13b_down", 5120, 13824),
|
||||
Shape("gpt_neox_up", 16384, 4096),
|
||||
Shape("gpt_neox_down", 4096, 16384),
|
||||
Shape("qwen2_7b_kv", 512, 3584),
|
||||
Shape("qwen2_7b_qo", 3584, 3584),
|
||||
Shape("qwen2_7b_up_gate", 18944, 3584),
|
||||
Shape("qwen2_7b_down", 3584, 18944),
|
||||
Shape("llama3_70b_kv", 1024, 8192),
|
||||
Shape("llama3_70b_qo", 8192, 8192),
|
||||
Shape("llama3_70b_up_gate", 28672, 8192),
|
||||
Shape("llama3_70b_down", 8192, 28672),
|
||||
Shape("opt_1_3b_qkvo", 2048, 2048),
|
||||
Shape("opt_1_3b_up", 8192, 2048),
|
||||
Shape("opt_1_3b_down", 2048, 8192),
|
||||
)
|
||||
|
||||
CHAINS = (
|
||||
Chain("llama2_7b", 4096, 4096, 11008),
|
||||
Chain("llama3_8b", 4096, 1024, 14336),
|
||||
Chain("llama2_13b", 5120, 5120, 13824),
|
||||
Chain("gpt_neox_20b", 4096, 4096, 16384, fused_qkv=True),
|
||||
Chain("qwen2_7b", 3584, 512, 18944),
|
||||
Chain("llama3_70b", 8192, 1024, 28672),
|
||||
Chain("opt_1_3b", 2048, 2048, 8192, gated_mlp=False),
|
||||
)
|
||||
|
||||
|
||||
def _elapsed_ms(fn: Callable[[], torch.Tensor], inner: int) -> float:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(inner):
|
||||
fn()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / inner
|
||||
|
||||
|
||||
def _timing(values: list[float]) -> Timing:
|
||||
ordered = sorted(values)
|
||||
p90_index = max(0, math.ceil(0.9 * len(ordered)) - 1)
|
||||
return Timing(statistics.median(ordered), ordered[p90_index])
|
||||
|
||||
|
||||
def _measure_pair(
|
||||
baseline: Callable[[], torch.Tensor],
|
||||
candidate: Callable[[], torch.Tensor],
|
||||
*,
|
||||
warmup: int,
|
||||
samples: int,
|
||||
inner: int,
|
||||
prepare_baseline: Callable[[], None] = lambda: None,
|
||||
prepare_candidate: Callable[[], None] = lambda: None,
|
||||
) -> tuple[Timing, Timing]:
|
||||
cases = (
|
||||
("baseline", prepare_baseline, baseline),
|
||||
("candidate", prepare_candidate, candidate),
|
||||
)
|
||||
for iteration in range(warmup):
|
||||
_, prepare, fn = cases[iteration % 2]
|
||||
prepare()
|
||||
fn()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
values: dict[str, list[float]] = {"baseline": [], "candidate": []}
|
||||
for sample in range(samples):
|
||||
order = cases if sample % 2 == 0 else tuple(reversed(cases))
|
||||
for label, prepare, fn in order:
|
||||
prepare()
|
||||
values[label].append(_elapsed_ms(fn, inner))
|
||||
return _timing(values["baseline"]), _timing(values["candidate"])
|
||||
|
||||
|
||||
def _print_header() -> None:
|
||||
print(
|
||||
"suite,label,m,n,k,torch_median_ms,torch_p90_ms,"
|
||||
"candidate_median_ms,candidate_p90_ms,speedup_pct,"
|
||||
"max_abs,relative_l2,argmax_equal"
|
||||
)
|
||||
|
||||
|
||||
def _print_result(
|
||||
suite: str,
|
||||
label: str,
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
baseline: Timing,
|
||||
candidate: Timing,
|
||||
reference: torch.Tensor,
|
||||
actual: torch.Tensor,
|
||||
) -> dict[str, object]:
|
||||
difference = actual.float() - reference.float()
|
||||
max_abs = difference.abs().max().item()
|
||||
relative_l2 = difference.norm().item() / max(reference.float().norm().item(), 1e-12)
|
||||
argmax_equal = torch.equal(actual.argmax(dim=-1), reference.argmax(dim=-1))
|
||||
speedup = (baseline.median_ms / candidate.median_ms - 1.0) * 100.0
|
||||
result: dict[str, object] = {
|
||||
"suite": suite,
|
||||
"label": label,
|
||||
"m": m,
|
||||
"n": n,
|
||||
"k": k,
|
||||
"torch_median_ms": baseline.median_ms,
|
||||
"torch_p90_ms": baseline.p90_ms,
|
||||
"candidate_median_ms": candidate.median_ms,
|
||||
"candidate_p90_ms": candidate.p90_ms,
|
||||
"speedup_pct": speedup,
|
||||
"max_abs": max_abs,
|
||||
"relative_l2": relative_l2,
|
||||
"argmax_equal": argmax_equal,
|
||||
}
|
||||
print(
|
||||
f"{suite},{label},{m},{n},{k},"
|
||||
f"{baseline.median_ms:.6f},{baseline.p90_ms:.6f},"
|
||||
f"{candidate.median_ms:.6f},{candidate.p90_ms:.6f},"
|
||||
f"{speedup:+.2f},{max_abs:.6f},{relative_l2:.8f},"
|
||||
f"{str(argmax_equal).lower()}",
|
||||
flush=True,
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def _weight(n: int, k: int, device: torch.device, std: float) -> torch.Tensor:
|
||||
weight = torch.empty((n, k), device=device, dtype=torch.bfloat16)
|
||||
weight.normal_(mean=0.0, std=std)
|
||||
return weight.requires_grad_(True)
|
||||
|
||||
|
||||
def _kernel_functions(
|
||||
x: torch.Tensor, weight: torch.Tensor
|
||||
) -> tuple[Callable[[], torch.Tensor], Callable[[], torch.Tensor]]:
|
||||
def baseline() -> torch.Tensor:
|
||||
return F.linear(x, weight)
|
||||
|
||||
def candidate() -> torch.Tensor:
|
||||
return bf16_gemm(x, weight.detach())
|
||||
|
||||
return baseline, candidate
|
||||
|
||||
|
||||
def benchmark_kernels(
|
||||
args: argparse.Namespace, device: torch.device
|
||||
) -> list[dict[str, object]]:
|
||||
if args.family == "astrai":
|
||||
shapes = ASTRAI_SHAPES
|
||||
elif args.family == "traditional":
|
||||
shapes = TRADITIONAL_SHAPES
|
||||
else:
|
||||
shapes = ASTRAI_SHAPES + TRADITIONAL_SHAPES
|
||||
if args.shape_label:
|
||||
requested = set(args.shape_label)
|
||||
shapes = tuple(shape for shape in shapes if shape.label in requested)
|
||||
missing = requested - {shape.label for shape in shapes}
|
||||
if missing:
|
||||
raise ValueError(f"unknown shape labels: {', '.join(sorted(missing))}")
|
||||
|
||||
results: list[dict[str, object]] = []
|
||||
for shape in shapes:
|
||||
weight = _weight(shape.n, shape.k, device, args.weight_std)
|
||||
for m in args.m:
|
||||
x = torch.randn((m, shape.k), device=device, dtype=torch.bfloat16)
|
||||
baseline_fn, candidate_fn = _kernel_functions(x, weight)
|
||||
with torch.inference_mode():
|
||||
reference = baseline_fn()
|
||||
actual = candidate_fn()
|
||||
baseline, candidate = _measure_pair(
|
||||
baseline_fn,
|
||||
candidate_fn,
|
||||
warmup=args.warmup,
|
||||
samples=args.samples,
|
||||
inner=args.inner,
|
||||
)
|
||||
results.append(
|
||||
_print_result(
|
||||
"kernel",
|
||||
shape.label,
|
||||
m,
|
||||
shape.n,
|
||||
shape.k,
|
||||
baseline,
|
||||
candidate,
|
||||
reference,
|
||||
actual,
|
||||
)
|
||||
)
|
||||
del baseline_fn, candidate_fn, x, reference, actual
|
||||
del weight
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return results
|
||||
|
||||
|
||||
def _set_mode(mode: str) -> None:
|
||||
os.environ["ASTRAI_GEMM"] = mode
|
||||
|
||||
|
||||
def _chain_weights(
|
||||
spec: Chain, device: torch.device, std: float
|
||||
) -> dict[str, torch.Tensor]:
|
||||
weights = {
|
||||
"o": _weight(spec.hidden, spec.hidden, device, std),
|
||||
"up": _weight(spec.intermediate, spec.hidden, device, std),
|
||||
"down": _weight(spec.hidden, spec.intermediate, device, std),
|
||||
}
|
||||
if spec.fused_qkv:
|
||||
weights["qkv"] = _weight(3 * spec.hidden, spec.hidden, device, std)
|
||||
else:
|
||||
weights.update(
|
||||
{
|
||||
"q": _weight(spec.hidden, spec.hidden, device, std),
|
||||
"k": _weight(spec.kv, spec.hidden, device, std),
|
||||
"v": _weight(spec.kv, spec.hidden, device, std),
|
||||
}
|
||||
)
|
||||
if spec.gated_mlp:
|
||||
weights["gate"] = _weight(spec.intermediate, spec.hidden, device, std)
|
||||
return weights
|
||||
|
||||
|
||||
def _chain_fn(
|
||||
x: torch.Tensor, weights: dict[str, torch.Tensor], spec: Chain
|
||||
) -> Callable[[], torch.Tensor]:
|
||||
def run() -> torch.Tensor:
|
||||
output_projection = linear(x, weights["o"])
|
||||
up = linear(x, weights["up"])
|
||||
if spec.fused_qkv:
|
||||
attention_projection = linear(x, weights["qkv"])[..., : x.shape[-1]]
|
||||
hidden = F.gelu(up)
|
||||
else:
|
||||
attention_projection = linear(x, weights["q"])
|
||||
linear(x, weights["k"])
|
||||
linear(x, weights["v"])
|
||||
if spec.gated_mlp:
|
||||
gate = linear(x, weights["gate"])
|
||||
hidden = F.silu(gate) * up
|
||||
else:
|
||||
hidden = F.gelu(up)
|
||||
down = linear(hidden, weights["down"])
|
||||
return attention_projection + output_projection + down
|
||||
|
||||
return run
|
||||
|
||||
|
||||
def benchmark_chains(
|
||||
args: argparse.Namespace, device: torch.device
|
||||
) -> list[dict[str, object]]:
|
||||
results: list[dict[str, object]] = []
|
||||
chains = CHAINS
|
||||
if args.chain_label:
|
||||
requested = set(args.chain_label)
|
||||
chains = tuple(chain for chain in chains if chain.label in requested)
|
||||
missing = requested - {chain.label for chain in chains}
|
||||
if missing:
|
||||
raise ValueError(f"unknown chain labels: {', '.join(sorted(missing))}")
|
||||
for spec in chains:
|
||||
weights = _chain_weights(spec, device, args.weight_std)
|
||||
for m in args.m:
|
||||
x = torch.randn((m, spec.hidden), device=device, dtype=torch.bfloat16)
|
||||
run = _chain_fn(x, weights, spec)
|
||||
with torch.inference_mode():
|
||||
_set_mode("0")
|
||||
reference = run()
|
||||
_set_mode(args.candidate_mode)
|
||||
actual = run()
|
||||
baseline, candidate = _measure_pair(
|
||||
run,
|
||||
run,
|
||||
warmup=args.warmup,
|
||||
samples=args.samples,
|
||||
inner=args.chain_inner,
|
||||
prepare_baseline=lambda: _set_mode("0"),
|
||||
prepare_candidate=lambda: _set_mode(args.candidate_mode),
|
||||
)
|
||||
results.append(
|
||||
_print_result(
|
||||
"synthetic_chain",
|
||||
spec.label,
|
||||
m,
|
||||
spec.hidden,
|
||||
spec.intermediate,
|
||||
baseline,
|
||||
candidate,
|
||||
reference,
|
||||
actual,
|
||||
)
|
||||
)
|
||||
del x, reference, actual
|
||||
del weights
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
return results
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--suite", choices=("kernel", "chain", "all"), default="all")
|
||||
parser.add_argument(
|
||||
"--family", choices=("astrai", "traditional", "all"), default="all"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--m", type=int, nargs="+", choices=(1, 2, 4, 8), default=(1, 2, 4, 8)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shape-label",
|
||||
action="append",
|
||||
help="limit the kernel suite to one or more named shape labels",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chain-label",
|
||||
action="append",
|
||||
help="limit the chain suite to one or more named model families",
|
||||
)
|
||||
parser.add_argument("--device", type=int, default=0)
|
||||
parser.add_argument("--warmup", type=int, default=20)
|
||||
parser.add_argument("--samples", type=int, default=9)
|
||||
parser.add_argument("--inner", type=int, default=100)
|
||||
parser.add_argument("--chain-inner", type=int, default=20)
|
||||
parser.add_argument(
|
||||
"--candidate-mode",
|
||||
choices=("auto", "1"),
|
||||
default="auto",
|
||||
help="dispatcher mode for the candidate side of the chain suite",
|
||||
)
|
||||
parser.add_argument("--weight-std", type=float, default=0.02)
|
||||
parser.add_argument("--seed", type=int, default=20260902)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=Path,
|
||||
help="optional JSON output; stdout always retains the compact CSV table",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if not torch.cuda.is_available() or not is_available("bf16_gemm"):
|
||||
raise RuntimeError("benchmark requires CUDA and the built bf16_gemm extension")
|
||||
if args.warmup < 0 or args.samples < 1 or args.inner < 1 or args.chain_inner < 1:
|
||||
raise ValueError("warmup must be non-negative and sample/inner counts positive")
|
||||
|
||||
torch.cuda.set_device(args.device)
|
||||
device = torch.device("cuda", args.device)
|
||||
torch.manual_seed(args.seed)
|
||||
torch.cuda.manual_seed_all(args.seed)
|
||||
properties = torch.cuda.get_device_properties(device)
|
||||
print(
|
||||
f"# device={properties.name}, capability={properties.major}.{properties.minor}, "
|
||||
f"seed={args.seed}, weight_std={args.weight_std}"
|
||||
)
|
||||
_print_header()
|
||||
results: list[dict[str, object]] = []
|
||||
if args.suite in ("kernel", "all"):
|
||||
results.extend(benchmark_kernels(args, device))
|
||||
if args.suite in ("chain", "all"):
|
||||
results.extend(benchmark_chains(args, device))
|
||||
if args.output is not None:
|
||||
payload = {
|
||||
"environment": {
|
||||
"device": properties.name,
|
||||
"capability": f"{properties.major}.{properties.minor}",
|
||||
"torch": torch.__version__,
|
||||
"cuda": torch.version.cuda,
|
||||
},
|
||||
"parameters": {
|
||||
"suite": args.suite,
|
||||
"family": args.family,
|
||||
"m": args.m,
|
||||
"shape_labels": args.shape_label,
|
||||
"chain_labels": args.chain_label,
|
||||
"candidate_mode": args.candidate_mode,
|
||||
"seed": args.seed,
|
||||
"weight_std": args.weight_std,
|
||||
"warmup": args.warmup,
|
||||
"samples": args.samples,
|
||||
"inner": args.inner,
|
||||
"chain_inner": args.chain_inner,
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
args.output.write_text(json.dumps(payload, indent=2) + "\n")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,252 @@
|
||||
"""Benchmark the fused rotary-embedding kernel against the torch fallback.
|
||||
|
||||
The baseline is the complex-multiply fallback from
|
||||
``astrai.extension.backend.rotary``. Layouts mirror the production call
|
||||
shapes: packed 3D [tokens, n_heads, head_dim] and dense 4D
|
||||
[batch, seq_len, n_heads, head_dim]; positions are random integers so every
|
||||
row exercises a distinct cos/sin gather.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Callable
|
||||
|
||||
import click
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension import is_available
|
||||
from astrai.extension.ops import rotary_emb
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class RotaryCase:
|
||||
name: str
|
||||
layout: str
|
||||
batch: int
|
||||
seq_len: int
|
||||
heads: int
|
||||
head_dim: int
|
||||
|
||||
|
||||
DEFAULT_CASES = (
|
||||
RotaryCase("decode_bs1", "packed", 1, 1, 32, 128),
|
||||
RotaryCase("decode_bs32", "packed", 32, 1, 32, 128),
|
||||
RotaryCase("prefill_4k_llama7b", "packed", 1, 4096, 32, 128),
|
||||
RotaryCase("prefill_4k_llama70b", "packed", 1, 4096, 64, 128),
|
||||
RotaryCase("train_8x2k_llama7b", "dense", 8, 2048, 32, 128),
|
||||
RotaryCase("train_4x1k_d64", "dense", 4, 1024, 32, 64),
|
||||
)
|
||||
|
||||
|
||||
def parse_case(value: str) -> RotaryCase:
|
||||
parts = value.split(":")
|
||||
if len(parts) != 6 or not parts[0]:
|
||||
raise click.BadParameter("case must use NAME:LAYOUT:BATCH:SEQ:HEADS:HEAD_DIM")
|
||||
name, layout, batch, seq_len, heads, head_dim = parts
|
||||
try:
|
||||
fields = (int(batch), int(seq_len), int(heads), int(head_dim))
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("fields must be integers") from exc
|
||||
if layout not in ("packed", "dense"):
|
||||
raise click.BadParameter("layout must be 'packed' or 'dense'")
|
||||
if any(field <= 0 for field in fields) or head_dim % 2:
|
||||
raise click.BadParameter("fields must be positive; HEAD_DIM even")
|
||||
return RotaryCase(name, layout, fields[0], fields[1], fields[2], fields[3])
|
||||
|
||||
|
||||
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
operation()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / iterations
|
||||
|
||||
|
||||
def summarize(values: list[float]) -> dict[str, float]:
|
||||
ordered = sorted(values)
|
||||
return {
|
||||
"median_ms": statistics.median(ordered),
|
||||
"p90_ms": ordered[max(0, math.ceil(0.9 * len(ordered)) - 1)],
|
||||
}
|
||||
|
||||
|
||||
def torch_apply(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
|
||||
"""The complex-multiply fallback (mirrors backend.rotary._torch_apply)."""
|
||||
cos, sin = freqs_cis[..., 0], freqs_cis[..., 1]
|
||||
dtype = x.dtype
|
||||
x_ = x.float().reshape(*x.shape[:-1], -1, 2)
|
||||
x_complex = torch.view_as_complex(x_)
|
||||
freqs_cis_complex = torch.complex(cos, sin).unsqueeze(-2)
|
||||
x_rotated = x_complex * freqs_cis_complex
|
||||
return torch.view_as_real(x_rotated).flatten(-2).to(dtype)
|
||||
|
||||
|
||||
def build_freqs(head_dim: int, positions: torch.Tensor) -> torch.Tensor:
|
||||
"""[cos, sin] pairs for the given positions, laid out [..., head_dim/2, 2]."""
|
||||
theta = 10000.0 ** (
|
||||
-torch.arange(0, head_dim, 2, dtype=torch.float64, device=positions.device)
|
||||
/ head_dim
|
||||
)
|
||||
freqs = positions.double().unsqueeze(-1) * theta
|
||||
return torch.stack([freqs.cos(), freqs.sin()], dim=-1).float()
|
||||
|
||||
|
||||
def benchmark_case(
|
||||
case: RotaryCase, *, warmup: int, iterations: int, trials: int
|
||||
) -> dict[str, object]:
|
||||
if case.layout == "packed":
|
||||
tokens = case.batch * case.seq_len
|
||||
x = torch.randn(
|
||||
tokens, case.heads, case.head_dim, device="cuda", dtype=torch.bfloat16
|
||||
)
|
||||
positions = torch.randint(0, 65536, (tokens,), device="cuda")
|
||||
else:
|
||||
x = torch.randn(
|
||||
case.batch,
|
||||
case.seq_len,
|
||||
case.heads,
|
||||
case.head_dim,
|
||||
device="cuda",
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
positions = torch.randint(0, 65536, (case.batch, case.seq_len), device="cuda")
|
||||
freqs_cis = build_freqs(case.head_dim, positions)
|
||||
operations: dict[str, Callable[[], torch.Tensor]] = {
|
||||
"torch": lambda: torch_apply(x, freqs_cis),
|
||||
"cuda": lambda: rotary_emb(x, freqs_cis),
|
||||
}
|
||||
for operation in operations.values():
|
||||
for _ in range(warmup):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
samples: dict[str, list[float]] = {name: [] for name in operations}
|
||||
order = tuple(operations)
|
||||
# A-B-B-A order balances cache, clock, and temperature drift.
|
||||
for _ in range(trials):
|
||||
for name in (*order, *reversed(order)):
|
||||
samples[name].append(time_operation(operations[name], iterations))
|
||||
|
||||
with torch.no_grad():
|
||||
expected = operations["torch"]().float()
|
||||
actual = operations["cuda"]().float()
|
||||
difference = actual - expected
|
||||
io_bytes = (
|
||||
2 * x.numel() * x.element_size() + freqs_cis.numel() * freqs_cis.element_size()
|
||||
)
|
||||
|
||||
result: dict[str, object] = {
|
||||
"case": case.name,
|
||||
"layout": case.layout,
|
||||
"heads": case.heads,
|
||||
"head_dim": case.head_dim,
|
||||
"estimated_io_bytes": io_bytes,
|
||||
"max_abs_error": float(difference.abs().max()),
|
||||
"cosine_similarity": float(
|
||||
F.cosine_similarity(actual.flatten(), expected.flatten(), dim=0)
|
||||
),
|
||||
}
|
||||
for name, samples_ms in samples.items():
|
||||
latency = summarize(samples_ms)
|
||||
result[name] = {
|
||||
"effective_bandwidth_gbps": io_bytes / (latency["median_ms"] / 1000) / 1e9,
|
||||
**latency,
|
||||
}
|
||||
speedup = (result["torch"]["median_ms"] / result["cuda"]["median_ms"] - 1.0) * 100.0
|
||||
print(
|
||||
f"{case.name},{case.layout},{result['torch']['median_ms']:.5f},"
|
||||
f"{result['cuda']['median_ms']:.5f},{speedup:+.1f}%,"
|
||||
f"{result['max_abs_error']:.5f}"
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
@click.command(help=__doc__)
|
||||
@click.option("--output", type=click.Path(path_type=Path), help="Optional JSON output.")
|
||||
@click.option(
|
||||
"--case",
|
||||
"case_values",
|
||||
multiple=True,
|
||||
help="Filter defaults by bare name, or add/override with "
|
||||
"NAME:LAYOUT:BATCH:SEQ:HEADS:HEAD_DIM.",
|
||||
)
|
||||
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option(
|
||||
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
|
||||
)
|
||||
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option("--seed", type=int, default=0, show_default=True)
|
||||
def benchmark_command(
|
||||
output: Path | None,
|
||||
case_values: tuple[str, ...],
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise click.ClickException("CUDA is required")
|
||||
if not is_available("rotary_emb"):
|
||||
raise click.ClickException("the built rotary_emb kernel is required")
|
||||
|
||||
# A bare name filters the matching default; a full spec overrides or appends.
|
||||
chosen: dict[str, RotaryCase] = {}
|
||||
for value in case_values:
|
||||
case = (
|
||||
next((c for c in DEFAULT_CASES if c.name == value), None)
|
||||
if ":" not in value
|
||||
else parse_case(value)
|
||||
)
|
||||
if case is None:
|
||||
raise click.BadParameter(f"unknown default case {value!r}")
|
||||
chosen[case.name] = case
|
||||
cases = tuple(chosen.values()) or DEFAULT_CASES
|
||||
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
print("case,layout,torch_ms,cuda_ms,speedup,max_abs")
|
||||
results = []
|
||||
with torch.inference_mode():
|
||||
for case in cases:
|
||||
results.append(
|
||||
benchmark_case(
|
||||
case, warmup=warmup, iterations=iterations, trials=trials
|
||||
)
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if output is not None:
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
payload = {
|
||||
"metadata": {
|
||||
"gpu_name": props.name,
|
||||
"compute_capability": f"{props.major}.{props.minor}",
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
||||
},
|
||||
"settings": {
|
||||
"warmup": warmup,
|
||||
"iterations": iterations,
|
||||
"trials": trials,
|
||||
"seed": seed,
|
||||
"order": "A-B-B-A",
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_command()
|
||||
@@ -1,325 +0,0 @@
|
||||
"""Benchmark fused BF16 SwiGLU against torch and unfused GEMM chains."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import math
|
||||
import statistics
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Callable, Iterable
|
||||
|
||||
import click
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from astrai.extension import bf16_gemm, bf16_swiglu, is_available
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SwiGLUShape:
|
||||
name: str
|
||||
n: int
|
||||
k: int
|
||||
|
||||
|
||||
DEFAULT_SHAPES = (
|
||||
SwiGLUShape("astrai_1b", 6912, 1536),
|
||||
SwiGLUShape("llama2_7b", 11008, 4096),
|
||||
SwiGLUShape("llama3_8b", 14336, 4096),
|
||||
SwiGLUShape("llama2_13b", 13824, 5120),
|
||||
SwiGLUShape("gpt_neox_20b", 16384, 6144),
|
||||
)
|
||||
|
||||
|
||||
def parse_positive_ints(value: str) -> tuple[int, ...]:
|
||||
try:
|
||||
values = tuple(dict.fromkeys(int(item.strip()) for item in value.split(",")))
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("expected comma-separated integers") from exc
|
||||
if not values or any(item <= 0 for item in values):
|
||||
raise click.BadParameter("values must be positive integers")
|
||||
return values
|
||||
|
||||
|
||||
def parse_shape(value: str) -> SwiGLUShape:
|
||||
parts = value.split(":")
|
||||
if len(parts) != 3 or not parts[0]:
|
||||
raise click.BadParameter("shape must use NAME:N:K")
|
||||
try:
|
||||
n, k = (int(item) for item in parts[1:])
|
||||
except ValueError as exc:
|
||||
raise click.BadParameter("N and K must be integers") from exc
|
||||
if n <= 0 or k <= 0 or k % 8:
|
||||
raise click.BadParameter("N must be positive and K positive/divisible by 8")
|
||||
return SwiGLUShape(parts[0], n, k)
|
||||
|
||||
|
||||
def percentile(values: Iterable[float], quantile: float) -> float:
|
||||
ordered = sorted(values)
|
||||
rank = (len(ordered) - 1) * quantile
|
||||
lower = math.floor(rank)
|
||||
upper = math.ceil(rank)
|
||||
if lower == upper:
|
||||
return ordered[lower]
|
||||
fraction = rank - lower
|
||||
return ordered[lower] * (1 - fraction) + ordered[upper] * fraction
|
||||
|
||||
|
||||
def summarize(values: list[float]) -> dict[str, float]:
|
||||
return {
|
||||
"median_ms": statistics.median(values),
|
||||
"p90_ms": percentile(values, 0.90),
|
||||
"p99_ms": percentile(values, 0.99),
|
||||
"min_ms": min(values),
|
||||
"max_ms": max(values),
|
||||
}
|
||||
|
||||
|
||||
def time_operation(operation: Callable[[], torch.Tensor], iterations: int) -> float:
|
||||
start = torch.cuda.Event(enable_timing=True)
|
||||
end = torch.cuda.Event(enable_timing=True)
|
||||
start.record()
|
||||
for _ in range(iterations):
|
||||
operation()
|
||||
end.record()
|
||||
end.synchronize()
|
||||
return start.elapsed_time(end) / iterations
|
||||
|
||||
|
||||
def count_cuda_kernels(
|
||||
operation: Callable[[], torch.Tensor], repeats: int = 5
|
||||
) -> float:
|
||||
with torch.profiler.profile(
|
||||
activities=[
|
||||
torch.profiler.ProfilerActivity.CPU,
|
||||
torch.profiler.ProfilerActivity.CUDA,
|
||||
],
|
||||
acc_events=True,
|
||||
) as profile:
|
||||
for _ in range(repeats):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
device_type = torch.autograd.DeviceType.CUDA
|
||||
events = [event for event in profile.events() if event.device_type == device_type]
|
||||
return len(events) / repeats
|
||||
|
||||
|
||||
def capture(operation: Callable[[], torch.Tensor]):
|
||||
for _ in range(3):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
output = operation()
|
||||
|
||||
def replay() -> torch.Tensor:
|
||||
graph.replay()
|
||||
return output
|
||||
|
||||
return replay
|
||||
|
||||
|
||||
def make_operations(x, up_weight, gate_weight, mode: str):
|
||||
operations: dict[str, Callable[[], torch.Tensor]] = {
|
||||
"torch": lambda: F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)),
|
||||
"gemm_chain": lambda: (
|
||||
bf16_gemm(x, up_weight) * F.silu(bf16_gemm(x, gate_weight))
|
||||
),
|
||||
"fused": lambda: bf16_swiglu(x, up_weight, gate_weight),
|
||||
}
|
||||
if mode == "graph":
|
||||
operations = {name: capture(op) for name, op in operations.items()}
|
||||
return operations
|
||||
|
||||
|
||||
def benchmark_case(
|
||||
shape: SwiGLUShape,
|
||||
m: int,
|
||||
mode: str,
|
||||
*,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
) -> list[dict[str, object]]:
|
||||
x = torch.randn((m, shape.k), device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
scale = shape.k**-0.5
|
||||
up_weight = (
|
||||
torch.randn((shape.n, shape.k), device="cuda", dtype=torch.bfloat16) * scale
|
||||
)
|
||||
gate_weight = (
|
||||
torch.randn((shape.n, shape.k), device="cuda", dtype=torch.bfloat16) * scale
|
||||
)
|
||||
operations = make_operations(x, up_weight, gate_weight, mode)
|
||||
for operation in operations.values():
|
||||
for _ in range(warmup):
|
||||
operation()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
samples = {name: [] for name in operations}
|
||||
forward_order = tuple(operations)
|
||||
# A-B-C-C-B-A order balances cache, clock, and temperature drift.
|
||||
for _ in range(trials):
|
||||
for name in (*forward_order, *reversed(forward_order)):
|
||||
samples[name].append(time_operation(operations[name], iterations))
|
||||
|
||||
with torch.no_grad():
|
||||
expected = operations["torch"]().clone()
|
||||
actual = operations["fused"]().clone()
|
||||
difference = (actual.float() - expected.float()).abs()
|
||||
max_abs_error = float(difference.max())
|
||||
mean_abs_error = float(difference.mean())
|
||||
cosine_similarity = float(
|
||||
F.cosine_similarity(actual.float().flatten(), expected.float().flatten(), dim=0)
|
||||
)
|
||||
|
||||
results = []
|
||||
for name, operation in operations.items():
|
||||
result: dict[str, object] = {
|
||||
"shape": shape.name,
|
||||
"m": m,
|
||||
"n": shape.n,
|
||||
"k": shape.k,
|
||||
"mode": mode,
|
||||
"implementation": name,
|
||||
"cuda_kernel_launches_per_call": count_cuda_kernels(operation),
|
||||
**summarize(samples[name]),
|
||||
}
|
||||
if name == "fused":
|
||||
result.update(
|
||||
max_abs_error=max_abs_error,
|
||||
mean_abs_error=mean_abs_error,
|
||||
cosine_similarity=cosine_similarity,
|
||||
)
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
|
||||
def device_metadata() -> dict[str, object]:
|
||||
props = torch.cuda.get_device_properties(0)
|
||||
return {
|
||||
"timestamp_utc": datetime.now(timezone.utc).isoformat(),
|
||||
"gpu_name": props.name,
|
||||
"compute_capability": f"{props.major}.{props.minor}",
|
||||
"total_memory_bytes": props.total_memory,
|
||||
"torch_version": torch.__version__,
|
||||
"cuda_version": torch.version.cuda,
|
||||
"dtype": "bfloat16",
|
||||
}
|
||||
|
||||
|
||||
def render_markdown(payload: dict[str, object]) -> str:
|
||||
metadata = payload["metadata"]
|
||||
results = payload["results"]
|
||||
assert isinstance(metadata, dict)
|
||||
assert isinstance(results, list)
|
||||
by_case = {
|
||||
(item["shape"], item["m"], item["mode"], item["implementation"]): item
|
||||
for item in results
|
||||
}
|
||||
cases = sorted({(item["shape"], item["m"], item["mode"]) for item in results})
|
||||
lines = [
|
||||
"# Fused SwiGLU benchmark",
|
||||
"",
|
||||
f"- GPU: {metadata['gpu_name']}",
|
||||
f"- Compute capability: {metadata['compute_capability']}",
|
||||
f"- PyTorch / CUDA: {metadata['torch_version']} / {metadata['cuda_version']}",
|
||||
"",
|
||||
"| Shape | M | Mode | torch ms | GEMM chain ms | fused ms | "
|
||||
"vs best unfused | fused kernels | max abs | cosine |",
|
||||
"|---|---:|---|---:|---:|---:|---:|---:|---:|---:|",
|
||||
]
|
||||
for shape, m, mode in cases:
|
||||
torch_item = by_case[(shape, m, mode, "torch")]
|
||||
gemm_item = by_case[(shape, m, mode, "gemm_chain")]
|
||||
fused_item = by_case[(shape, m, mode, "fused")]
|
||||
best = min(torch_item["median_ms"], gemm_item["median_ms"])
|
||||
improvement = (best / fused_item["median_ms"] - 1) * 100
|
||||
lines.append(
|
||||
f"| {shape} | {m} | {mode} | {torch_item['median_ms']:.5f} | "
|
||||
f"{gemm_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | "
|
||||
f"{improvement:+.2f}% | "
|
||||
f"{fused_item['cuda_kernel_launches_per_call']:.1f} | "
|
||||
f"{fused_item['max_abs_error']:.5f} | "
|
||||
f"{fused_item['cosine_similarity']:.8f} |"
|
||||
)
|
||||
lines.append("")
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
@click.command(help=__doc__)
|
||||
@click.option("--output", type=click.Path(path_type=Path), required=True)
|
||||
@click.option("--markdown-output", type=click.Path(path_type=Path))
|
||||
@click.option("--m-values", default="1,2,4,8", show_default=True)
|
||||
@click.option("--shape", "shape_values", multiple=True, help="Repeat NAME:N:K.")
|
||||
@click.option("--mode", type=click.Choice(("eager", "graph", "both")), default="both")
|
||||
@click.option("--warmup", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option(
|
||||
"--iterations", type=click.IntRange(min=1), default=100, show_default=True
|
||||
)
|
||||
@click.option("--trials", type=click.IntRange(min=1), default=10, show_default=True)
|
||||
@click.option("--seed", type=int, default=0, show_default=True)
|
||||
def benchmark_command(
|
||||
output: Path,
|
||||
markdown_output: Path | None,
|
||||
m_values: str,
|
||||
shape_values: tuple[str, ...],
|
||||
mode: str,
|
||||
warmup: int,
|
||||
iterations: int,
|
||||
trials: int,
|
||||
seed: int,
|
||||
) -> None:
|
||||
if not torch.cuda.is_available():
|
||||
raise click.ClickException("CUDA is required")
|
||||
if not is_available("bf16_gemm") or not is_available("bf16_swiglu"):
|
||||
raise click.ClickException("built bf16_gemm and bf16_swiglu are required")
|
||||
shapes = tuple(parse_shape(value) for value in shape_values) or DEFAULT_SHAPES
|
||||
m_values_parsed = parse_positive_ints(m_values)
|
||||
if any(m > 8 for m in m_values_parsed):
|
||||
raise click.BadParameter("fused primitive supports M up to 8")
|
||||
modes = ("eager", "graph") if mode == "both" else (mode,)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
results = []
|
||||
with torch.inference_mode():
|
||||
for shape in shapes:
|
||||
for m in m_values_parsed:
|
||||
for current_mode in modes:
|
||||
click.echo(
|
||||
f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}"
|
||||
)
|
||||
results.extend(
|
||||
benchmark_case(
|
||||
shape,
|
||||
m,
|
||||
current_mode,
|
||||
warmup=warmup,
|
||||
iterations=iterations,
|
||||
trials=trials,
|
||||
)
|
||||
)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
payload: dict[str, object] = {
|
||||
"metadata": device_metadata(),
|
||||
"settings": {
|
||||
"warmup": warmup,
|
||||
"iterations": iterations,
|
||||
"trials": trials,
|
||||
"seed": seed,
|
||||
"order": "A-B-C-C-B-A",
|
||||
},
|
||||
"results": results,
|
||||
}
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
output.write_text(json.dumps(payload, indent=2) + "\n")
|
||||
if markdown_output is not None:
|
||||
markdown_output.parent.mkdir(parents=True, exist_ok=True)
|
||||
markdown_output.write_text(render_markdown(payload))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark_command()
|
||||
@@ -1,650 +0,0 @@
|
||||
// BF16 decode-time GEMM primitive for linear layers: one entry point whose
|
||||
// internal path is sized by the decode batch M.
|
||||
//
|
||||
// M in [1, 8] — one CTA per weight row, the M rows held in registers
|
||||
// (any K).
|
||||
// M in (8, 32] — BM=16 tensor-core tiles; shape-driven configs (see
|
||||
// the dispatch table at the entry point) cover every
|
||||
// production shape (K % 8 == 0 and 16-byte-aligned
|
||||
// tensors).
|
||||
//
|
||||
// Both paths take row-major [N, K] weights, accumulate in FP32, fuse an
|
||||
// optional bias, and are inference-only.
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
|
||||
// Route kernel-launch failures through the torch error check instead of
|
||||
// common/launch.cuh's print+abort default. Must precede the kernel code.
|
||||
#define ASTRAI_LAUNCH_FAIL(err, what) C10_CUDA_CHECK(err)
|
||||
|
||||
#include "common/cp_async.cuh"
|
||||
#include "common/launch.cuh"
|
||||
#include "common/mma.cuh"
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kHalfCtaThreads = 128;
|
||||
constexpr int kWarpSize = 32;
|
||||
|
||||
__device__ __forceinline__ float warp_sum(float value) {
|
||||
#pragma unroll
|
||||
for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
|
||||
value += __shfl_down_sync(0xffffffff, value, offset);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
template <int Rows, int Threads>
|
||||
__global__ void skinny_gemm_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ weight,
|
||||
const __nv_bfloat16* __restrict__ bias,
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
int n,
|
||||
int k
|
||||
) {
|
||||
const int output_index = blockIdx.x;
|
||||
const int lane = threadIdx.x & (kWarpSize - 1);
|
||||
const int warp = threadIdx.x / kWarpSize;
|
||||
|
||||
float sums[Rows] = {};
|
||||
__shared__ float warp_sums[Rows][Threads / kWarpSize];
|
||||
// Weight row: scalar head/tail around a 16-byte-aligned uint4 middle so
|
||||
// any K is accepted while keeping 128-bit weight loads, which dominate
|
||||
// bandwidth on decode shapes. x pairs with scalar loads: it is a tiny
|
||||
// L1/L2-resident matrix, consecutive threads still touch contiguous
|
||||
// addresses, and no per-row alignment case analysis is needed.
|
||||
const __nv_bfloat16* __restrict__ wrow =
|
||||
weight + static_cast<int64_t>(output_index) * k;
|
||||
const unsigned whead_raw =
|
||||
((16u - (reinterpret_cast<uintptr_t>(wrow) & 15u)) & 15u) >> 1;
|
||||
const int whead = static_cast<int>(min(whead_raw, static_cast<unsigned>(k)));
|
||||
const int wvecs = (k - whead) / 8;
|
||||
const int wtail_start = whead + wvecs * 8;
|
||||
const uint4* __restrict__ w4 = reinterpret_cast<const uint4*>(wrow + whead);
|
||||
|
||||
// x chunks pair element-for-element with the aligned weight middle:
|
||||
// the uint4 view is rooted at ``x + whead`` (16-byte aligned by the
|
||||
// branch guard), and each row strides by ``k / 8`` vectors because its
|
||||
// first middle element sits ``whead`` scalars past ``row * k``. When
|
||||
// K % 8 == 0 and the weight row is already aligned (whead == 0, the
|
||||
// production case) this reduces to one pure uint4 loop with an empty
|
||||
// head/tail. Otherwise per-row uint4 loads are not 16-byte addressable,
|
||||
// and scalar x pairing keeps the kernel correct for any K while the
|
||||
// weight stream stays vectorized.
|
||||
if (k % 8 == 0 &&
|
||||
((reinterpret_cast<uintptr_t>(x) + 2u * static_cast<unsigned>(whead)) & 15u) == 0u) {
|
||||
const auto* x4 = reinterpret_cast<const uint4*>(x + whead);
|
||||
for (int v = threadIdx.x; v < wvecs; v += blockDim.x) {
|
||||
const uint4 wv_raw = w4[v];
|
||||
const auto* wv =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&wv_raw);
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
const uint4 xv_raw =
|
||||
x4[(static_cast<int64_t>(row) * (k / 8)) + v];
|
||||
const auto* xv =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&xv_raw);
|
||||
#pragma unroll
|
||||
for (int p = 0; p < 4; ++p) {
|
||||
sums[row] = fmaf(
|
||||
__bfloat162float(__low2bfloat16(xv[p])),
|
||||
__bfloat162float(__low2bfloat16(wv[p])),
|
||||
sums[row]
|
||||
);
|
||||
sums[row] = fmaf(
|
||||
__bfloat162float(__high2bfloat16(xv[p])),
|
||||
__bfloat162float(__high2bfloat16(wv[p])),
|
||||
sums[row]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for (int v = threadIdx.x; v < wvecs; v += blockDim.x) {
|
||||
const uint4 wv_raw = w4[v];
|
||||
const __nv_bfloat16* wv_s =
|
||||
reinterpret_cast<const __nv_bfloat16*>(&wv_raw);
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
const __nv_bfloat16* xv =
|
||||
x + static_cast<int64_t>(row) * k + whead + 8 * v;
|
||||
#pragma unroll
|
||||
for (int s = 0; s < 8; ++s) {
|
||||
sums[row] = fmaf(
|
||||
__bfloat162float(xv[s]),
|
||||
__bfloat162float(wv_s[s]),
|
||||
sums[row]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Head and tail remainders: plain scalar pairing, at most 14 elements.
|
||||
for (int i = threadIdx.x; i < whead; i += blockDim.x) {
|
||||
const float wv = __bfloat162float(wrow[i]);
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
sums[row] = fmaf(
|
||||
__bfloat162float(x[static_cast<int64_t>(row) * k + i]),
|
||||
wv,
|
||||
sums[row]
|
||||
);
|
||||
}
|
||||
}
|
||||
for (int i = wtail_start + threadIdx.x; i < k; i += blockDim.x) {
|
||||
const float wv = __bfloat162float(wrow[i]);
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
sums[row] = fmaf(
|
||||
__bfloat162float(x[static_cast<int64_t>(row) * k + i]),
|
||||
wv,
|
||||
sums[row]
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
sums[row] = warp_sum(sums[row]);
|
||||
}
|
||||
if (lane == 0) {
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
warp_sums[row][warp] = sums[row];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (warp == 0) {
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
float sum =
|
||||
lane < (Threads / kWarpSize) ? warp_sums[row][lane] : 0.0f;
|
||||
sum = warp_sum(sum);
|
||||
if (lane == 0) {
|
||||
if (bias != nullptr) {
|
||||
sum += __bfloat162float(bias[output_index]);
|
||||
}
|
||||
output[row * n + output_index] = __float2bfloat16_rn(sum);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int Rows, int Threads>
|
||||
void launch_skinny_gemm(
|
||||
const __nv_bfloat16* x,
|
||||
const __nv_bfloat16* weight,
|
||||
const __nv_bfloat16* bias,
|
||||
__nv_bfloat16* output,
|
||||
int n,
|
||||
int k,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
// Split-K was tried and rejected for the small-N shapes (GQA kv
|
||||
// projections): their ~48K uint4 loads already saturate thread-level
|
||||
// parallelism one load deep, so they sit on the launch+HBM latency
|
||||
// floor, and the fence+atomic+last-CTA partial round trip adds ~0.6us
|
||||
// of fixed sync cost (measured -27% at M=1, -113% at M=8 on L20).
|
||||
// The real fix for those shapes is fusing the QKV projections so the
|
||||
// tiny kv rows stop launching as standalone kernels at all.
|
||||
//
|
||||
// Decode is HBM weight-streaming bound. Every shape launches with a
|
||||
// single 128-thread CTA size: measured on L20 (sm_89) with L2-thrashing
|
||||
// weight rotation, flat 128t is within ~1% of a per-shape tuned
|
||||
// 128/256/512 mix at M=1 and M=8 and gives up at most ~2% at M=2-4
|
||||
// (512t down_proj, 256t gate/up/lm_head cells), while dodging the
|
||||
// Rows>=7 register cliff on the 307MB lm_head (+21% DRAM throughput
|
||||
// vs 256t at M=8). Re-measured
|
||||
// after the table refactor: 256t on wide-N gate/up wins only ~2.8% at
|
||||
// M=2-4 and ties at M=1/6, inside run-to-run drift. Simplicity keeps
|
||||
// winning over the last ~2%.
|
||||
skinny_gemm_kernel<Rows, Threads>
|
||||
<<<n, Threads, 0, stream>>>(
|
||||
x, weight, bias, output, n, k
|
||||
);
|
||||
ASTRAI_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
// Compile-time dispatch tables: replace a hand-written switch over M with
|
||||
// function-pointer tables indexed by M-1. Adding a Rows variant means adding
|
||||
// one table entry, not a new case block at the call site.
|
||||
using SkinnyGemmFn = void (*)(
|
||||
const __nv_bfloat16*, const __nv_bfloat16*, const __nv_bfloat16*,
|
||||
__nv_bfloat16*, int, int, cudaStream_t
|
||||
);
|
||||
|
||||
constexpr SkinnyGemmFn kSkinnyGemm[8] = {
|
||||
&launch_skinny_gemm<1, kHalfCtaThreads>,
|
||||
&launch_skinny_gemm<2, kHalfCtaThreads>,
|
||||
&launch_skinny_gemm<3, kHalfCtaThreads>,
|
||||
&launch_skinny_gemm<4, kHalfCtaThreads>,
|
||||
&launch_skinny_gemm<5, kHalfCtaThreads>,
|
||||
&launch_skinny_gemm<6, kHalfCtaThreads>,
|
||||
&launch_skinny_gemm<7, kHalfCtaThreads>,
|
||||
&launch_skinny_gemm<8, kHalfCtaThreads>,
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// M > 8 path: small-M (M <= 64) tiled GEMM.
|
||||
//
|
||||
// F.linear for decode batches in (8, 64]: cuBLAS tiles the small M as a
|
||||
// single tile row, which starves the grid (measured on L20, sm_89: 24-54
|
||||
// CTAs on 92 SMs, tensor pipe 24-42%, DRAM <= 72%). This kernel keeps the
|
||||
// M rows in one CTA tile and fills the SMs along N and K instead.
|
||||
//
|
||||
// CUTLASS-style configuration: the kernel is parameterized by template
|
||||
// parameters — CTA tile (BM x BN x BK), pipeline depth, thread count —
|
||||
// with the warp layout derived inside (kMt M fragments, kNt n16 tiles per
|
||||
// warp). Four families are instantiated (see the dispatch table at the
|
||||
// entry point): the default (BN=64, BK=64, 3 stages) and its BM=32 wide-N
|
||||
// variant, plus narrow-N deep-K rings (BK=256/128, BN=32, 64 threads); a
|
||||
// new shape is an instantiation, not a rewrite.
|
||||
//
|
||||
// Operand staging reuses the FP8 GEMM's scheme (see fp8/gemm/*.cuh and
|
||||
// docs/developer/cuda_kernels.md): 16B chunks XOR-swizzled with row & 7,
|
||||
// one barrier per k-tile, and a kStages+1 ring whose prefetch for tile
|
||||
// i+kStages lands in the slot tile i-1 released — no post-compute barrier.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
using bf16 = __nv_bfloat16;
|
||||
|
||||
// Logical (row, byte column) -> byte offset in one flat [rows * BK*2B]
|
||||
// staging tile. The 16B chunk index is XORed with row & (chunks - 1) so an
|
||||
// ldmatrix fragment load (8 consecutive rows x 16B) hits all 32 banks once.
|
||||
template <int BK>
|
||||
__device__ __forceinline__ int tile_off(int row, int byte_col) {
|
||||
constexpr int kRowBytes = BK * 2;
|
||||
constexpr int kChunks = kRowBytes / 16;
|
||||
static_assert(
|
||||
(kChunks & (kChunks - 1)) == 0, "swizzle needs a power-of-two chunk count"
|
||||
);
|
||||
return row * kRowBytes +
|
||||
(((byte_col >> 4) ^ (row & (kChunks - 1))) << 4) + (byte_col & 15);
|
||||
}
|
||||
|
||||
// Predicated staging of one [RowsTile x BK] operand slice from a
|
||||
// row-major [total_rows, K] tensor into its swizzled ring slot. Chunks past
|
||||
// the row count or past K zero-fill (the wrapper guarantees K % 8 == 0 and
|
||||
// 16B-aligned rows, so a misaligned *valid* chunk cannot occur).
|
||||
template <int RowsTile, int BK, int kThreads>
|
||||
__device__ __forceinline__ void stage_tile(
|
||||
char* slot, const bf16* __restrict__ src, int total_rows, int64_t k,
|
||||
int row0, int tid, int kt
|
||||
) {
|
||||
constexpr int kRowBytes = BK * 2;
|
||||
constexpr int kChunks = kRowBytes / 16;
|
||||
#pragma unroll
|
||||
for (int c = tid; c < RowsTile * kChunks; c += kThreads) {
|
||||
const int r = c / kChunks;
|
||||
const int c8 = c % kChunks;
|
||||
const int64_t kbase = (int64_t)kt * BK;
|
||||
const bool ok = row0 + r < total_rows && kbase + c8 * 8 + 8 <= k;
|
||||
char* dst = slot + tile_off<BK>(r, c8 * 16);
|
||||
if (ok) {
|
||||
astrai::cp_async_16(
|
||||
reinterpret_cast<bf16*>(dst),
|
||||
src + (int64_t)(row0 + r) * k + kbase + c8 * 8
|
||||
);
|
||||
} else {
|
||||
unsigned* w = reinterpret_cast<unsigned*>(dst);
|
||||
#pragma unroll
|
||||
for (int i = 0; i < 4; ++i)
|
||||
w[i] = 0u;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int BM, int BN, int BK, int kStages, int kThreads>
|
||||
__global__ __launch_bounds__(kThreads, 1) void tiled_gemm_kernel(
|
||||
const bf16* __restrict__ x,
|
||||
const bf16* __restrict__ w,
|
||||
const bf16* __restrict__ bias,
|
||||
bf16* __restrict__ out,
|
||||
int m,
|
||||
int n,
|
||||
int k
|
||||
) {
|
||||
constexpr int kMt = BM / 16; // m16 fragments
|
||||
constexpr int kWarps = kThreads / 32;
|
||||
constexpr int kNt = BN / (kWarps * 16); // n16 tiles per warp
|
||||
constexpr int kRowBytes = BK * 2;
|
||||
constexpr int kRing = kStages + 1; // ring slots per operand
|
||||
constexpr int kSegs = BK / 16; // m16k16 segments per tile
|
||||
constexpr int kSegXor = 32; // bytes: +2 chunks per segment
|
||||
static_assert(BM % 16 == 0, "BM must be a multiple of 16");
|
||||
static_assert(BN % (kWarps * 16) == 0, "warps must tile BN in n16 units");
|
||||
|
||||
extern __shared__ __align__(16) char smem[];
|
||||
char* const a_ring = smem;
|
||||
char* const b_ring = smem + kRing * BM * kRowBytes;
|
||||
|
||||
const int tid = threadIdx.x;
|
||||
const int warp = tid >> 5;
|
||||
const int lane = tid & 31;
|
||||
// Standard 2D grid mapping: grid.x covers M tiles, grid.y covers N
|
||||
// tiles. The grid fills the SMs along N; every block walks the whole
|
||||
// K range in a single pass.
|
||||
const int m0 = blockIdx.x * BM;
|
||||
const int n0 = blockIdx.y * BN;
|
||||
const int total_tiles = (k + BK - 1) / BK;
|
||||
|
||||
auto a_slot = [&](int t) {
|
||||
return a_ring + (t % kRing) * BM * kRowBytes;
|
||||
};
|
||||
auto b_slot = [&](int t) {
|
||||
return b_ring + (t % kRing) * BN * kRowBytes;
|
||||
};
|
||||
|
||||
#pragma unroll
|
||||
for (int s = 0; s < kStages; ++s) {
|
||||
if (s < total_tiles) {
|
||||
stage_tile<BM, BK, kThreads>(
|
||||
a_slot(s), x, m, k, m0, tid, s
|
||||
);
|
||||
stage_tile<BN, BK, kThreads>(b_slot(s), w, n, k, n0, tid, s);
|
||||
}
|
||||
astrai::cp_async_commit_group();
|
||||
}
|
||||
|
||||
// Per-lane ldmatrix fragment addresses (relative to each ring slot):
|
||||
// A x4: lanes 0-7 mat0 (rows 0-7, chunk k-lo), 8-15 mat1 (rows 8-15,
|
||||
// k-lo), 16-23 mat2 (rows 0-7, k-hi), 24-31 mat3 (rows 8-15, k-hi).
|
||||
// B x4 pair: lanes 0-7/8-15 the first n8 tile's k-lo/k-hi chunks,
|
||||
// 16-23/24-31 the second n8 tile's. Warp w owns the n16 tiles at
|
||||
// (w + j * kWarps) * 16 for j in [0, kNt).
|
||||
const int a_row = ((lane >> 3) & 1) * 8 + (lane & 7);
|
||||
const unsigned a_off = tile_off<BK>(a_row, (lane >> 4) * 16);
|
||||
unsigned b_off[kNt];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kNt; ++j) {
|
||||
const int b_row =
|
||||
(warp + j * kWarps) * 16 + (lane & 7) + (lane >> 4) * 8;
|
||||
b_off[j] = tile_off<BK>(b_row, ((lane >> 3) & 1) * 16);
|
||||
}
|
||||
const unsigned a_base0 = __cvta_generic_to_shared(a_ring) + a_off;
|
||||
unsigned b_base0[kNt];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kNt; ++j)
|
||||
b_base0[j] = __cvta_generic_to_shared(b_ring) + b_off[j];
|
||||
|
||||
float acc[kMt][kNt][2][4] = {};
|
||||
for (int i = 0; i < total_tiles; ++i) {
|
||||
astrai::cp_async_wait_group<kStages - 1>();
|
||||
__syncthreads();
|
||||
const unsigned a_base =
|
||||
a_base0 + (unsigned)((i % kRing) * BM * kRowBytes);
|
||||
unsigned b_base[kNt];
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kNt; ++j)
|
||||
b_base[j] = b_base0[j] + (unsigned)((i % kRing) * BN * kRowBytes);
|
||||
#pragma unroll
|
||||
for (int seg = 0; seg < kSegs; ++seg) {
|
||||
const unsigned a_seg = a_base ^ (unsigned)(seg * kSegXor);
|
||||
unsigned a4[kMt][4], b4[kNt][4];
|
||||
#pragma unroll
|
||||
for (int mt = 0; mt < kMt; ++mt)
|
||||
astrai::ldmatrix_x4_lane(
|
||||
a4[mt], a_seg + (unsigned)(mt * 16 * kRowBytes)
|
||||
);
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kNt; ++j)
|
||||
astrai::ldmatrix_x4_lane(
|
||||
b4[j], (b_base[j] ^ (unsigned)(seg * kSegXor))
|
||||
);
|
||||
#pragma unroll
|
||||
for (int mt = 0; mt < kMt; ++mt)
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kNt; ++j)
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt)
|
||||
astrai::mma_sync<bf16>(
|
||||
acc[mt][j][nt], a4[mt], b4[j] + nt * 2,
|
||||
acc[mt][j][nt]
|
||||
);
|
||||
}
|
||||
// Prefetch tile i+kStages into the slot tile i-1 released. The
|
||||
// barrier at the top of the next iteration separates every
|
||||
// thread's reads of that slot (iteration i-1) from these writes.
|
||||
const int pf = i + kStages;
|
||||
if (pf < total_tiles) {
|
||||
stage_tile<BM, BK, kThreads>(
|
||||
a_slot(pf), x, m, k, m0, tid, pf
|
||||
);
|
||||
stage_tile<BN, BK, kThreads>(b_slot(pf), w, n, k, n0, tid, pf);
|
||||
}
|
||||
astrai::cp_async_commit_group();
|
||||
}
|
||||
|
||||
const int row0 = lane >> 2;
|
||||
const int col0 = (lane & 3) * 2;
|
||||
#pragma unroll
|
||||
for (int mt = 0; mt < kMt; ++mt) {
|
||||
if (m0 + mt * 16 + row0 >= m)
|
||||
continue;
|
||||
const int64_t orow = (int64_t)(m0 + mt * 16 + row0) * n;
|
||||
#pragma unroll
|
||||
for (int j = 0; j < kNt; ++j) {
|
||||
#pragma unroll
|
||||
for (int nt = 0; nt < 2; ++nt) {
|
||||
const int col =
|
||||
n0 + (warp + j * kWarps) * 16 + col0 + nt * 8;
|
||||
if (col >= n)
|
||||
continue;
|
||||
float2 v, v8;
|
||||
v.x = acc[mt][j][nt][0];
|
||||
v.y = acc[mt][j][nt][1];
|
||||
v8.x = acc[mt][j][nt][2];
|
||||
v8.y = acc[mt][j][nt][3];
|
||||
if (bias != nullptr) {
|
||||
v.x += __bfloat162float(bias[col]);
|
||||
v.y += __bfloat162float(bias[col + 1]);
|
||||
v8.x += __bfloat162float(bias[col]);
|
||||
v8.y += __bfloat162float(bias[col + 1]);
|
||||
}
|
||||
if (col + 1 < n) {
|
||||
*reinterpret_cast<__nv_bfloat162*>(out + orow + col) =
|
||||
__floats2bfloat162_rn(v.x, v.y);
|
||||
if (m0 + mt * 16 + row0 + 8 < m)
|
||||
*reinterpret_cast<__nv_bfloat162*>(
|
||||
out + orow + 8 * n + col) =
|
||||
__floats2bfloat162_rn(v8.x, v8.y);
|
||||
} else {
|
||||
out[orow + col] = __float2bfloat16_rn(v.x);
|
||||
if (m0 + mt * 16 + row0 + 8 < m)
|
||||
out[orow + 8 * n + col] =
|
||||
__float2bfloat16_rn(v8.x);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int BM, int BN, int BK, int kStages, int kThreads>
|
||||
void launch_tiled_gemm(
|
||||
const bf16* x,
|
||||
const bf16* w,
|
||||
const bf16* bias,
|
||||
bf16* out,
|
||||
int m,
|
||||
int n,
|
||||
int k,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
constexpr int kRing = kStages + 1;
|
||||
constexpr int smem = kRing * (BM + BN) * BK * 2;
|
||||
// 99KB is the sm_86/89 per-block opt-in ceiling; configs above 48KB
|
||||
// (long-K BK=128) run one CTA per SM and pay a one-time attribute opt-in.
|
||||
static_assert(smem <= 99 * 1024, "family must fit the sm_86/89 opt-in ceiling");
|
||||
if constexpr (smem > 48 * 1024) {
|
||||
static const bool opted_in = [] {
|
||||
ASTRAI_CUDA_CHECK(cudaFuncSetAttribute(
|
||||
tiled_gemm_kernel<BM, BN, BK, kStages, kThreads>,
|
||||
cudaFuncAttributeMaxDynamicSharedMemorySize, smem
|
||||
));
|
||||
return true;
|
||||
}();
|
||||
(void)opted_in;
|
||||
}
|
||||
dim3 grid((m + BM - 1) / BM, (n + BN - 1) / BN);
|
||||
tiled_gemm_kernel<BM, BN, BK, kStages, kThreads>
|
||||
<<<grid, kThreads, smem, stream>>>(
|
||||
x, w, bias, out, m, n, k
|
||||
);
|
||||
ASTRAI_LAUNCH_CHECK();
|
||||
}
|
||||
|
||||
using TiledGemmFn = void (*)(
|
||||
const bf16*, const bf16*, const bf16*, bf16*, int, int, int,
|
||||
cudaStream_t
|
||||
);
|
||||
|
||||
// CTA tile configuration pairing the shape parameters with their matched
|
||||
// instantiation. Field order is canonical everywhere it appears — template
|
||||
// arguments, this struct, the dispatch table — as BM, BN, BK, then
|
||||
// pipeline depth (stages) and CTA size (threads). A new family is one
|
||||
// row here plus one select branch.
|
||||
struct TileConfig {
|
||||
int bm;
|
||||
int bn;
|
||||
int bk;
|
||||
int stages;
|
||||
int threads;
|
||||
TiledGemmFn launch;
|
||||
};
|
||||
|
||||
// Shape -> tile config, measured on L20 with L2-thrashing weight rotation.
|
||||
// The selector trades grid fill against K-loop serial latency:
|
||||
// - Wide N (n >= 4096): ceil(n/64) tiles already cover the SMs, the GEMM
|
||||
// is HBM-bound and the default config wins; BM widens to 32 at M>16 to
|
||||
// halve the re-staged weight stream.
|
||||
// - Narrow N: few N tiles leave the grid K-serial — widening the grid
|
||||
// does not help (measured: BN 64->32 ties, doubled m_tiles tie, kv at
|
||||
// 4 blocks ties q/o at 24); fewer, deeper K chunks do. BK=256 with a
|
||||
// 72KB two-stage ring is the winner while the grid fits one wave
|
||||
// (72KB smem means one CTA per SM); past one wave its 2-wave
|
||||
// quantization loses to BK=128's 36.9KB two-CTA ring.
|
||||
inline TileConfig select_tile_config(int m, int n, int k) {
|
||||
if (n >= 4096) {
|
||||
if (m > 16) {
|
||||
return {32, 64, 64, 3, 128,
|
||||
&launch_tiled_gemm<32, 64, 64, 3, 128>};
|
||||
}
|
||||
return {16, 64, 64, 3, 128, &launch_tiled_gemm<16, 64, 64, 3, 128>};
|
||||
}
|
||||
const int n_tiles = (n + 31) / 32;
|
||||
const int m_tiles = (m + 15) / 16;
|
||||
if (n_tiles * m_tiles <= 92) { // one wave on the 92-SM L20
|
||||
return {16, 32, 256, 2, 64, &launch_tiled_gemm<16, 32, 256, 2, 64>};
|
||||
}
|
||||
return {16, 32, 128, 2, 64, &launch_tiled_gemm<16, 32, 128, 2, 64>};
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
// Single entry point: M in [1, 8] routes to the register-resident skinny
|
||||
// GEMM kernel (any K), M in (8, 64] to the tiled kernel (K % 8 == 0 and
|
||||
// 16-byte-aligned tensors, checked at the branch).
|
||||
torch::Tensor bf16_gemm(
|
||||
torch::Tensor x,
|
||||
torch::Tensor weight,
|
||||
py::object bias_object
|
||||
) {
|
||||
TORCH_CHECK(x.is_cuda() && weight.is_cuda(), "x and weight must be CUDA tensors");
|
||||
TORCH_CHECK(x.device() == weight.device(), "x and weight must share device");
|
||||
TORCH_CHECK(
|
||||
x.scalar_type() == torch::kBFloat16 &&
|
||||
weight.scalar_type() == torch::kBFloat16,
|
||||
"x and weight must be bf16"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.dim() == 1 || x.dim() == 2,
|
||||
"x must have shape [K] or [M, K]"
|
||||
);
|
||||
TORCH_CHECK(weight.dim() == 2, "weight must have shape [N, K]");
|
||||
TORCH_CHECK(x.is_contiguous() && weight.is_contiguous(), "x and weight must be contiguous");
|
||||
TORCH_CHECK(
|
||||
!x.requires_grad() && !weight.requires_grad(),
|
||||
"bf16_gemm is inference-only and does not support autograd"
|
||||
);
|
||||
|
||||
const int64_t m = x.dim() == 1 ? 1 : x.size(0);
|
||||
const int64_t k = x.size(-1);
|
||||
const int64_t n = weight.size(0);
|
||||
TORCH_CHECK(weight.size(1) == k, "weight K must match x K");
|
||||
TORCH_CHECK(m >= 1 && m <= 64, "M must be in [1, 64]");
|
||||
TORCH_CHECK(k > 0 && n > 0, "N and K must be positive");
|
||||
TORCH_CHECK(
|
||||
k <= std::numeric_limits<int>::max() &&
|
||||
n <= std::numeric_limits<int>::max(),
|
||||
"N or K exceeds the CUDA launcher limit"
|
||||
);
|
||||
|
||||
torch::Tensor bias;
|
||||
const __nv_bfloat16* bias_ptr = nullptr;
|
||||
if (!bias_object.is_none()) {
|
||||
bias = bias_object.cast<torch::Tensor>();
|
||||
TORCH_CHECK(bias.is_cuda() && bias.device() == x.device(), "bias must share the CUDA device");
|
||||
TORCH_CHECK(bias.scalar_type() == torch::kBFloat16, "bias must be bf16");
|
||||
TORCH_CHECK(bias.dim() == 1 && bias.size(0) == n, "bias must have shape [N]");
|
||||
TORCH_CHECK(bias.is_contiguous(), "bias must be contiguous");
|
||||
TORCH_CHECK(!bias.requires_grad(), "bf16_gemm bias does not support autograd");
|
||||
bias_ptr = reinterpret_cast<const __nv_bfloat16*>(bias.data_ptr());
|
||||
}
|
||||
|
||||
const at::cuda::OptionalCUDAGuard guard(x.device());
|
||||
const auto* properties = at::cuda::getDeviceProperties(x.device().index());
|
||||
TORCH_CHECK(properties->major >= 8, "bf16_gemm requires compute capability 8.0+");
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
auto output = x.dim() == 1 ? torch::empty({n}, x.options())
|
||||
: torch::empty({m, n}, x.options());
|
||||
|
||||
const auto* x_ptr = reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
|
||||
const auto* weight_ptr =
|
||||
reinterpret_cast<const __nv_bfloat16*>(weight.data_ptr());
|
||||
auto* output_ptr = reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
|
||||
const int m_int = static_cast<int>(m);
|
||||
const int n_int = static_cast<int>(n);
|
||||
const int k_int = static_cast<int>(k);
|
||||
|
||||
if (m <= 8) {
|
||||
kSkinnyGemm[m_int - 1](
|
||||
x_ptr, weight_ptr, bias_ptr, output_ptr, n_int, k_int,
|
||||
stream.stream()
|
||||
);
|
||||
} else {
|
||||
TORCH_CHECK(
|
||||
k % 8 == 0 &&
|
||||
(reinterpret_cast<uintptr_t>(x.data_ptr()) & 15) == 0u &&
|
||||
(reinterpret_cast<uintptr_t>(weight.data_ptr()) & 15) == 0u,
|
||||
"M > 8 requires K to be a multiple of 8 and x/weight 16-byte aligned"
|
||||
);
|
||||
const TileConfig cfg = select_tile_config(m_int, n_int, k_int);
|
||||
cfg.launch(
|
||||
x_ptr, weight_ptr, bias_ptr, output_ptr, m_int, n_int, k_int,
|
||||
stream.stream()
|
||||
);
|
||||
}
|
||||
C10_CUDA_CHECK(cudaGetLastError());
|
||||
return output;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
|
||||
module.def(
|
||||
"bf16_gemm",
|
||||
&bf16_gemm,
|
||||
py::arg("x"),
|
||||
py::arg("weight"),
|
||||
py::arg("bias") = py::none(),
|
||||
"M in [1, 64] BF16 GEMM with optional fused bias "
|
||||
"(register-resident skinny GEMM path for M <= 8, tensor-core tiles above)"
|
||||
);
|
||||
}
|
||||
@@ -1,297 +0,0 @@
|
||||
// Fused small-M BF16 SwiGLU primitive for decode-time dense MLP layers.
|
||||
// One CTA per output column; each weight pair is read once and reused across
|
||||
// all decode rows. Bandwidth-bound in the cold-HBM decode regime, so variant
|
||||
// selection beyond the M=8 block-size rule is noise.
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include <cstdint>
|
||||
#include <limits>
|
||||
|
||||
namespace {
|
||||
|
||||
constexpr int kWarpSize = 32;
|
||||
|
||||
__device__ __forceinline__ float warp_sum(float value) {
|
||||
#pragma unroll
|
||||
for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) {
|
||||
value += __shfl_down_sync(0xffffffff, value, offset);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float round_bf16(float value) {
|
||||
return __bfloat162float(__float2bfloat16_rn(value));
|
||||
}
|
||||
|
||||
template <int Threads, int Rows>
|
||||
__global__ void bf16_swiglu_kernel(
|
||||
const __nv_bfloat16* __restrict__ x,
|
||||
const __nv_bfloat16* __restrict__ up_weight,
|
||||
const __nv_bfloat16* __restrict__ gate_weight,
|
||||
__nv_bfloat16* __restrict__ output,
|
||||
int n,
|
||||
int k
|
||||
) {
|
||||
constexpr int kWarps = Threads / kWarpSize;
|
||||
const int output_index = blockIdx.x;
|
||||
const int lane = threadIdx.x & (kWarpSize - 1);
|
||||
const int warp = threadIdx.x / kWarpSize;
|
||||
const int vector_count = k / 8;
|
||||
|
||||
float up_sums[Rows] = {};
|
||||
float gate_sums[Rows] = {};
|
||||
__shared__ float up_warp_sums[Rows][kWarps];
|
||||
__shared__ float gate_warp_sums[Rows][kWarps];
|
||||
|
||||
const auto* x4 = reinterpret_cast<const uint4*>(x);
|
||||
const auto* up4 = reinterpret_cast<const uint4*>(
|
||||
up_weight + static_cast<int64_t>(output_index) * k
|
||||
);
|
||||
const auto* gate4 = reinterpret_cast<const uint4*>(
|
||||
gate_weight + static_cast<int64_t>(output_index) * k
|
||||
);
|
||||
|
||||
// Read each pair of up/gate weight chunks once per CTA, then reuse it for
|
||||
// every active decode row. The fused epilogue removes two [M, N]
|
||||
// intermediates and the standalone SiLU and multiply launches.
|
||||
for (int vector_index = threadIdx.x;
|
||||
vector_index < vector_count;
|
||||
vector_index += blockDim.x) {
|
||||
const uint4 up_raw = up4[vector_index];
|
||||
const uint4 gate_raw = gate4[vector_index];
|
||||
const auto* up_values =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&up_raw);
|
||||
const auto* gate_values =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&gate_raw);
|
||||
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
const uint4 x_raw =
|
||||
x4[static_cast<int64_t>(row) * vector_count + vector_index];
|
||||
const auto* x_values =
|
||||
reinterpret_cast<const __nv_bfloat162*>(&x_raw);
|
||||
#pragma unroll
|
||||
for (int pair = 0; pair < 4; ++pair) {
|
||||
const float2 xv = __bfloat1622float2(x_values[pair]);
|
||||
const float2 uv = __bfloat1622float2(up_values[pair]);
|
||||
const float2 gv = __bfloat1622float2(gate_values[pair]);
|
||||
up_sums[row] = fmaf(xv.x, uv.x, up_sums[row]);
|
||||
up_sums[row] = fmaf(xv.y, uv.y, up_sums[row]);
|
||||
gate_sums[row] = fmaf(xv.x, gv.x, gate_sums[row]);
|
||||
gate_sums[row] = fmaf(xv.y, gv.y, gate_sums[row]);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
up_sums[row] = warp_sum(up_sums[row]);
|
||||
gate_sums[row] = warp_sum(gate_sums[row]);
|
||||
}
|
||||
if (lane == 0) {
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
up_warp_sums[row][warp] = up_sums[row];
|
||||
gate_warp_sums[row][warp] = gate_sums[row];
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
if (warp == 0) {
|
||||
#pragma unroll
|
||||
for (int row = 0; row < Rows; ++row) {
|
||||
float up = lane < kWarps ? up_warp_sums[row][lane] : 0.0f;
|
||||
float gate = lane < kWarps ? gate_warp_sums[row][lane] : 0.0f;
|
||||
up = warp_sum(up);
|
||||
gate = warp_sum(gate);
|
||||
if (lane == 0) {
|
||||
// Match the public composition's BF16 rounding boundaries:
|
||||
// BF16 linear outputs, BF16 SiLU output, then BF16 multiply.
|
||||
up = round_bf16(up);
|
||||
gate = round_bf16(gate);
|
||||
const float silu = round_bf16(gate / (1.0f + expf(-gate)));
|
||||
output[static_cast<int64_t>(row) * n + output_index] =
|
||||
__float2bfloat16_rn(up * silu);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int Threads, int Rows>
|
||||
void launch_bf16_swiglu(
|
||||
const __nv_bfloat16* x,
|
||||
const __nv_bfloat16* up_weight,
|
||||
const __nv_bfloat16* gate_weight,
|
||||
__nv_bfloat16* output,
|
||||
int n,
|
||||
int k,
|
||||
cudaStream_t stream
|
||||
) {
|
||||
bf16_swiglu_kernel<Threads, Rows><<<n, Threads, 0, stream>>>(
|
||||
x, up_weight, gate_weight, output, n, k
|
||||
);
|
||||
}
|
||||
|
||||
torch::Tensor bf16_swiglu(
|
||||
torch::Tensor x,
|
||||
torch::Tensor up_weight,
|
||||
torch::Tensor gate_weight
|
||||
) {
|
||||
TORCH_CHECK(
|
||||
x.is_cuda() && up_weight.is_cuda() && gate_weight.is_cuda(),
|
||||
"x, up_weight, and gate_weight must be CUDA tensors"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.device() == up_weight.device() && x.device() == gate_weight.device(),
|
||||
"x and weights must share a device"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.scalar_type() == torch::kBFloat16 &&
|
||||
up_weight.scalar_type() == torch::kBFloat16 &&
|
||||
gate_weight.scalar_type() == torch::kBFloat16,
|
||||
"x and weights must be bf16"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.dim() == 1 || x.dim() == 2,
|
||||
"x must have shape [K] or [M, K]"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
up_weight.dim() == 2 && gate_weight.dim() == 2,
|
||||
"weights must have shape [N, K]"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
x.is_contiguous() && up_weight.is_contiguous() &&
|
||||
gate_weight.is_contiguous(),
|
||||
"x and weights must be contiguous"
|
||||
);
|
||||
// The kernel loads all three streams as uint4; contiguous-but-offset
|
||||
// views would fault with an opaque "misaligned address" CUDA error, so
|
||||
// reject them here with an actionable message.
|
||||
TORCH_CHECK(
|
||||
(reinterpret_cast<uintptr_t>(x.data_ptr()) & 15u) == 0u,
|
||||
"bf16_swiglu requires 16-byte aligned x (storage_offset must keep "
|
||||
"data_ptr divisible by 16); clone the tensor or use the torch path"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
(reinterpret_cast<uintptr_t>(up_weight.data_ptr()) & 15u) == 0u,
|
||||
"bf16_swiglu requires 16-byte aligned up_weight (storage_offset "
|
||||
"must keep data_ptr divisible by 16); clone the tensor or use the "
|
||||
"torch path"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
(reinterpret_cast<uintptr_t>(gate_weight.data_ptr()) & 15u) == 0u,
|
||||
"bf16_swiglu requires 16-byte aligned gate_weight (storage_offset "
|
||||
"must keep data_ptr divisible by 16); clone the tensor or use the "
|
||||
"torch path"
|
||||
);
|
||||
TORCH_CHECK(
|
||||
!x.requires_grad() && !up_weight.requires_grad() &&
|
||||
!gate_weight.requires_grad(),
|
||||
"bf16_swiglu is inference-only and does not support autograd"
|
||||
);
|
||||
|
||||
const int64_t m = x.dim() == 1 ? 1 : x.size(0);
|
||||
const int64_t k = x.size(-1);
|
||||
const int64_t n = up_weight.size(0);
|
||||
TORCH_CHECK(m >= 1 && m <= 8, "M must be in [1, 8]");
|
||||
TORCH_CHECK(
|
||||
gate_weight.sizes() == up_weight.sizes(),
|
||||
"up_weight and gate_weight must have identical shapes"
|
||||
);
|
||||
TORCH_CHECK(up_weight.size(1) == k, "weight K must match x K");
|
||||
TORCH_CHECK(k > 0 && n > 0, "N and K must be positive");
|
||||
TORCH_CHECK(k % 8 == 0, "K must be divisible by 8");
|
||||
TORCH_CHECK(
|
||||
k <= std::numeric_limits<int>::max() &&
|
||||
n <= std::numeric_limits<int>::max(),
|
||||
"N or K exceeds the CUDA launcher limit"
|
||||
);
|
||||
|
||||
const at::cuda::OptionalCUDAGuard guard(x.device());
|
||||
const auto* properties = at::cuda::getDeviceProperties(x.device().index());
|
||||
TORCH_CHECK(
|
||||
properties->major >= 8,
|
||||
"bf16_swiglu requires compute capability 8.0+"
|
||||
);
|
||||
auto stream = at::cuda::getCurrentCUDAStream();
|
||||
auto output = x.dim() == 1 ? torch::empty({n}, x.options())
|
||||
: torch::empty({m, n}, x.options());
|
||||
|
||||
const auto* x_ptr =
|
||||
reinterpret_cast<const __nv_bfloat16*>(x.data_ptr());
|
||||
const auto* up_ptr =
|
||||
reinterpret_cast<const __nv_bfloat16*>(up_weight.data_ptr());
|
||||
const auto* gate_ptr =
|
||||
reinterpret_cast<const __nv_bfloat16*>(gate_weight.data_ptr());
|
||||
auto* output_ptr =
|
||||
reinterpret_cast<__nv_bfloat16*>(output.data_ptr());
|
||||
const int n_int = static_cast<int>(n);
|
||||
const int k_int = static_cast<int>(k);
|
||||
|
||||
// Block size 256 keeps the weight streams at the HBM bandwidth floor for
|
||||
// M in [1, 7]; M=8 halves the CTA so each thread owns more of the row
|
||||
// and the shared-memory reduction tree shrinks (measured on L20 with
|
||||
// rotated cold weights; larger CTAs only add idle warps).
|
||||
switch (m) {
|
||||
case 1:
|
||||
launch_bf16_swiglu<256, 1>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 2:
|
||||
launch_bf16_swiglu<256, 2>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 3:
|
||||
launch_bf16_swiglu<256, 3>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 4:
|
||||
launch_bf16_swiglu<256, 4>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 5:
|
||||
launch_bf16_swiglu<256, 5>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 6:
|
||||
launch_bf16_swiglu<256, 6>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 7:
|
||||
launch_bf16_swiglu<256, 7>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
case 8:
|
||||
launch_bf16_swiglu<128, 8>(
|
||||
x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream()
|
||||
);
|
||||
break;
|
||||
}
|
||||
C10_CUDA_CHECK(cudaGetLastError());
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
|
||||
module.def(
|
||||
"bf16_swiglu",
|
||||
&bf16_swiglu,
|
||||
py::arg("x"),
|
||||
py::arg("up_weight"),
|
||||
py::arg("gate_weight"),
|
||||
"M in [1, 8] fused BF16 up/gate projection and SwiGLU"
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user