From 4ed0109052e04fbefbd10f7c8fe2412de2862e46 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 2 Sep 2026 19:49:50 +0800 Subject: [PATCH] perf: add opt-in fused bf16 swiglu --- astrai/extension/__init__.py | 4 + astrai/extension/backend/__init__.py | 2 + astrai/extension/backend/swiglu.py | 109 + astrai/extension/ops/__init__.py | 2 + astrai/extension/ops/swiglu.py | 22 + astrai/model/components/mlp.py | 3 +- benchmarks/infraswe/README.md | 34 + benchmarks/infraswe/astrai-swiglu-draft.json | 104 + .../swiglu_l20_sm89_infraswe_score.json | 125 ++ csrc/CMakeLists.txt | 2 + csrc/kernels/gemv/bf16_swiglu.cu | 368 ++++ docs/benchmarks/swiglu_engine_l20_sm89.txt | 144 ++ .../swiglu_greedy_m1_m2_l20_sm89.txt | 12 + docs/benchmarks/swiglu_greedy_m4_l20_sm89.txt | 6 + docs/benchmarks/swiglu_l20_sm89.json | 1820 +++++++++++++++++ docs/developer/cuda_kernels.md | 33 +- docs/developer/swiglu_benchmark.md | 63 + scripts/tools/benchmark_swiglu.py | 325 +++ setup.py | 1 + tests/extension/test_swiglu.py | 99 + tests/extension/test_swiglu_dispatch.py | 94 + 21 files changed, 3368 insertions(+), 4 deletions(-) create mode 100644 astrai/extension/backend/swiglu.py create mode 100644 astrai/extension/ops/swiglu.py create mode 100644 benchmarks/infraswe/README.md create mode 100644 benchmarks/infraswe/astrai-swiglu-draft.json create mode 100644 benchmarks/results/swiglu_l20_sm89_infraswe_score.json create mode 100644 csrc/kernels/gemv/bf16_swiglu.cu create mode 100644 docs/benchmarks/swiglu_engine_l20_sm89.txt create mode 100644 docs/benchmarks/swiglu_greedy_m1_m2_l20_sm89.txt create mode 100644 docs/benchmarks/swiglu_greedy_m4_l20_sm89.txt create mode 100644 docs/benchmarks/swiglu_l20_sm89.json create mode 100644 docs/developer/swiglu_benchmark.md create mode 100644 scripts/tools/benchmark_swiglu.py create mode 100644 tests/extension/test_swiglu.py create mode 100644 tests/extension/test_swiglu_dispatch.py diff --git a/astrai/extension/__init__.py b/astrai/extension/__init__.py index 6e0ba08d..bbfcc1fd 100644 --- a/astrai/extension/__init__.py +++ b/astrai/extension/__init__.py @@ -27,6 +27,7 @@ attn_backend, get_backend, linear, + swiglu, ) from astrai.extension.dispatch import ( Axes, @@ -51,6 +52,7 @@ attn_paged_decode, attn_prefill, bf16_gemv, + bf16_swiglu, ) __all__ = [ @@ -65,10 +67,12 @@ "attn_backend", "get_backend", "linear", + "swiglu", "attn_decode", "attn_paged_decode", "attn_prefill", "bf16_gemv", + "bf16_swiglu", "is_available", "KERNEL_NAMES", "apply_rotary_emb", diff --git a/astrai/extension/backend/__init__.py b/astrai/extension/backend/__init__.py index 7b43a216..8c4127ae 100644 --- a/astrai/extension/backend/__init__.py +++ b/astrai/extension/backend/__init__.py @@ -13,6 +13,7 @@ ) from astrai.extension.backend.linear import linear from astrai.extension.backend.rotary import apply_rotary_emb +from astrai.extension.backend.swiglu import swiglu __all__ = [ "ATTN_BACKEND", @@ -26,4 +27,5 @@ "attn_backend", "get_backend", "linear", + "swiglu", ] diff --git a/astrai/extension/backend/swiglu.py b/astrai/extension/backend/swiglu.py new file mode 100644 index 00000000..7f4fe7c7 --- /dev/null +++ b/astrai/extension/backend/swiglu.py @@ -0,0 +1,109 @@ +"""Inference-only fused SwiGLU selection for dense MLP layers.""" + +import logging +import os +from functools import cache + +import torch +import torch.nn.functional as F +from torch import Tensor + +from astrai.extension.backend.linear import linear +from astrai.extension.loader import is_available +from astrai.extension.ops.swiglu import bf16_swiglu + +logger = logging.getLogger(__name__) + +# Shape keys are (N, K) for the paired up/gate projections. Automatic entries +# are populated only after the primitive, MLP chain, and greedy checkpoint +# gates pass on that architecture. +_AUTO_SWIGLU_SHAPES: dict[tuple[int, int], dict[int, frozenset[tuple[int, int]]]] = {} +_AUTO_SWIGLU_M = frozenset( + m for architecture in _AUTO_SWIGLU_SHAPES.values() for m in architecture +) +_VALID_MODES = {"0", "1", "auto"} +_WARNED_MODES: set[str] = set() + + +def _swiglu_mode() -> str: + mode = os.environ.get("ASTRAI_SWIGLU", "auto").strip().lower() + if mode in _VALID_MODES: + return mode + if mode not in _WARNED_MODES: + _WARNED_MODES.add(mode) + logger.warning( + "ASTRAI_SWIGLU=%r is invalid; expected 0, 1, or auto; using auto", + mode, + ) + return "auto" + + +def _unfused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor: + # Keep the existing linear backend in the fallback chain. This preserves + # any independently qualified GEMV shapes instead of making the fusion + # decision suppress linear-level optimizations. + return linear(x, up_weight) * F.silu(linear(x, gate_weight)) + + +def _fused_swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor: + return bf16_swiglu(x.detach(), up_weight.detach(), gate_weight.detach()) + + +@cache +def _device_capability(device_index: int) -> tuple[int, int]: + return torch.cuda.get_device_capability(device_index) + + +def _swiglu_capable(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> bool: + return not ( + torch.is_grad_enabled() + or not x.is_cuda + or x.dtype != torch.bfloat16 + or up_weight.dtype != torch.bfloat16 + or gate_weight.dtype != torch.bfloat16 + or x.ndim not in (1, 2) + or up_weight.ndim != 2 + or gate_weight.ndim != 2 + or (x.ndim == 2 and not 1 <= x.shape[0] <= 8) + or up_weight.shape != gate_weight.shape + or x.shape[-1] != up_weight.shape[1] + or x.shape[-1] % 8 != 0 + or x.device != up_weight.device + or x.device != gate_weight.device + or not x.is_contiguous() + or not up_weight.is_contiguous() + or not gate_weight.is_contiguous() + or not is_available("bf16_swiglu") + ) + + +def _auto_swiglu_shape(x: Tensor, up_weight: Tensor) -> bool: + capability = _device_capability(x.get_device()) + m = 1 if x.ndim == 1 else x.shape[0] + return (up_weight.shape[0], up_weight.shape[1]) in _AUTO_SWIGLU_SHAPES.get( + capability, {} + ).get(m, ()) + + +def swiglu(x: Tensor, up_weight: Tensor, gate_weight: Tensor) -> Tensor: + """Apply the dense-MLP SwiGLU projection with a safe torch fallback. + + ``ASTRAI_SWIGLU=0`` keeps the unfused linear-backend chain, ``1`` forces + the fused primitive for supported inputs, and ``auto`` uses only + architecture/shape bands backed by benchmark and checkpoint evidence. + """ + mode = _swiglu_mode() + if mode == "0" or (mode == "auto" and not _AUTO_SWIGLU_SHAPES): + return _unfused_swiglu(x, up_weight, gate_weight) + if mode == "auto": + m = 1 if x.ndim == 1 else (x.shape[0] if x.ndim == 2 else None) + if m not in _AUTO_SWIGLU_M: + return _unfused_swiglu(x, up_weight, gate_weight) + if _swiglu_capable(x, up_weight, gate_weight) and ( + mode == "1" or _auto_swiglu_shape(x, up_weight) + ): + return _fused_swiglu(x, up_weight, gate_weight) + return _unfused_swiglu(x, up_weight, gate_weight) + + +__all__ = ["swiglu"] diff --git a/astrai/extension/ops/__init__.py b/astrai/extension/ops/__init__.py index ba06b2ff..233ba8fe 100644 --- a/astrai/extension/ops/__init__.py +++ b/astrai/extension/ops/__init__.py @@ -9,6 +9,7 @@ ) from astrai.extension.ops.gemv import bf16_gemv from astrai.extension.ops.rotary import rotary_emb +from astrai.extension.ops.swiglu import bf16_swiglu __all__ = [ "TensorLayout", @@ -17,5 +18,6 @@ "attn_paged_prefill", "attn_prefill", "bf16_gemv", + "bf16_swiglu", "rotary_emb", ] diff --git a/astrai/extension/ops/swiglu.py b/astrai/extension/ops/swiglu.py new file mode 100644 index 00000000..e1bfc247 --- /dev/null +++ b/astrai/extension/ops/swiglu.py @@ -0,0 +1,22 @@ +"""Stateless wrapper for the directly callable fused BF16 SwiGLU primitive.""" + +import torch + +from astrai.extension.loader import get_module + + +def bf16_swiglu( + x: torch.Tensor, + up_weight: torch.Tensor, + gate_weight: torch.Tensor, +) -> torch.Tensor: + """Compute ``linear(x, up) * silu(linear(x, gate))`` for M in [1, 8]. + + Inputs must be contiguous BF16 CUDA tensors. Both weights use row-major + ``[N, K]`` storage with identical shapes, and K must be divisible by 8. + The primitive is inference-only and performs no fallback. + """ + return get_module("bf16_swiglu").bf16_swiglu(x, up_weight, gate_weight) + + +__all__ = ["bf16_swiglu"] diff --git a/astrai/model/components/mlp.py b/astrai/model/components/mlp.py index 083dad37..defd8d49 100644 --- a/astrai/model/components/mlp.py +++ b/astrai/model/components/mlp.py @@ -5,6 +5,7 @@ import torch.nn.functional as F from torch import Tensor +from astrai.extension.backend.swiglu import swiglu from astrai.factory import BaseFactory from astrai.model.components.linear import Linear @@ -38,7 +39,7 @@ def __init__(self, dim: int, dim_ffn: int, down_init_std: float = 0.02): self.down = Linear(dim_ffn, dim, init_std=down_init_std) def forward(self, x: Tensor) -> FFNOutput: - gated = self.up(x) * F.silu(self.gate(x)) + gated = swiglu(x, self.up.weight, self.gate.weight) out = self.down(gated) return {"hidden_states": out, "aux_loss": None, "router_stats": None} diff --git a/benchmarks/infraswe/README.md b/benchmarks/infraswe/README.md new file mode 100644 index 00000000..7e72bf87 --- /dev/null +++ b/benchmarks/infraswe/README.md @@ -0,0 +1,34 @@ +# InfraSWE Draft: fused BF16 SwiGLU + +This directory binds the fused BF16 SwiGLU benchmark to AstrAI as an explicit +repository target. AstrAI is not one of InfraSWE v0.5's pinned built-in +projects, so using a default project would evaluate the change against the +wrong host contract. The Draft remains `D3-contract-proposed`; it does not +claim maintainer review, sealing, hidden-probe completion, or official +ProjectFit. + +Before this PR was opened, the Draft was validated and resolved with InfraSWE +commit `811bc775ed5b3a6ec853219245f3469f78818020`: + +```bash +PYTHONPATH=src .venv/bin/infraswe draft validate \ + /path/to/AstrAI/benchmarks/infraswe/astrai-swiglu-draft.json + +PYTHONPATH=src .venv/bin/infraswe draft resolve \ + --local-draft \ + /path/to/AstrAI/benchmarks/infraswe/astrai-swiglu-draft.json \ + --output /tmp/astrai-swiglu-resolution.json +``` + +The candidate and contract digests bind the ordered source, tests, +documentation, operator results, engine log, and greedy checkpoint probes. The +required comparison cell is one NVIDIA L20 (`sm_89`); optional A100 and H100 +cells are explicitly untested. Compilation happens before timed cases. + +Applying InfraSWE's frozen `project-fit-kernel-v0.5` formula to the visible +evidence gives a diagnostic ProjectFit of **90.95/100** and BenchmarkTrust of +**95.87/100**. The rationale is machine-readable in +`benchmarks/results/swiglu_l20_sm89_infraswe_score.json`. Both scores are +non-official. Official ProjectFit remains unresolved until the Draft is +sealed, at least five fresh-process replays and system traces exist, hidden +probes are complete, and the evidence manifest is verified. diff --git a/benchmarks/infraswe/astrai-swiglu-draft.json b/benchmarks/infraswe/astrai-swiglu-draft.json new file mode 100644 index 00000000..a605366e --- /dev/null +++ b/benchmarks/infraswe/astrai-swiglu-draft.json @@ -0,0 +1,104 @@ +{ + "schema_version": "0.5", + "draft": { + "id": "astrai-fused-bf16-swiglu-v1", + "revision": 1, + "state": "D3-contract-proposed", + "created_by": "0z5a" + }, + "target": { + "mode": "repository", + "repository": "https://github.com/ViperEkura/AstrAI", + "revision": "sha256:c5595986ba7ef004333b86c1227fd8223e965971489002f0bbbe47ce4090b342", + "project_profile_sha256": "sha256:953067c47f5298f06846c8fe873db55afae49b6588a81926e49d9ef6b0994e08" + }, + "candidate": { + "kind": "git-diff", + "revision": "sha256:d6634a630de7aa4fc7946567acfa1964270673458d27e5173c6d4a17c8c6362f", + "intent": "add-fastpath", + "implementation_kind": "cuda-native", + "entrypoints": [ + "astrai.extension.bf16_swiglu", + "astrai.extension.backend.swiglu", + "astrai.model.components.mlp.MLP.forward", + "scripts/tools/benchmark_swiglu.py" + ], + "operator_family": "dense-gemm", + "phase": "inference", + "backend": "cuda", + "primary_host_candidate": "astrai" + }, + "baseline": { + "mode": "target-head", + "revision": "sha256:c5595986ba7ef004333b86c1227fd8223e965971489002f0bbbe47ce4090b342" + }, + "deployment": { + "workload_portfolio": { + "id": "astrai-swiglu-l20-v1", + "sha256": "sha256:608dbaf05d1f0e45f4cc5697b2e8a3460865753a6c9fefb48f57ea736e0e8540", + "path": "docs/benchmarks/swiglu_l20_sm89.json" + }, + "required_cells": [ + "cuda-sm89-nvidia-l20" + ], + "optional_cells": [ + "cuda-sm80-a100", + "cuda-sm90-h100" + ], + "request_or_step_protocol": { + "id": "astrai-swiglu-benchmark-v1", + "sha256": "sha256:578c2e70a8e2bb63a62222bd3d7e6f45c7e839d5f5d546b1f0e1c4f7bf4bc895", + "path": "docs/developer/swiglu_benchmark.md" + } + }, + "retrieval": { + "enabled": true, + "corpus_cutoff": "2026-09-02T11:41:02Z", + "sources": [ + "target-code", + "merged-prs", + "rejected-prs", + "review-comments", + "ci-failures" + ], + "precedent_set_sha256": "sha256:2914315eadee8b5e7cf0c35666d4639156d7744d14ec19e11b19caf4d9a1c49c" + }, + "acceptance_contract": { + "status": "proposed", + "path": "tests/extension/test_swiglu.py,tests/extension/test_swiglu_dispatch.py", + "sha256": "sha256:fb633cfc6be32041d5e6a37707c6074a3d35931cb5d2b04b7aa4348783783f2a", + "probe_set_sha256": "sha256:3203ce8b46a789afb7cb332b71b01554f5a9f6278d705a4e4b8ce05b236d9d87", + "hidden_probe_policy_sha256": "sha256:e181380d1b2435c18569629904153feb91bf0494d12d05e353e50c71ebab90bd" + }, + "project_objectives": { + "edge_ecosystem_policy": "experimental", + "profile_set_sha256": "sha256:b2312952e3b6cf61a527a61806eb39278a1490758e5356436b3b2c36ddffae67" + }, + "benchmark_loop": { + "fast_stage_max_official_fraction": 0.05, + "affected_stage_max_official_fraction": 0.2, + "official_replays": 7, + "early_exit_on_hard_gate": true, + "affected_case_selection": "required", + "benchmark_budget_policy_id": "draft-staged-budget-v0.5", + "evidence_policy_id": "v0.4-evidence-ladder-plus-seal-v0.5", + "precompile": { + "mode": "auto", + "trigger": "when-compilation-required", + "cache_policy": "content-addressed-evidence-identity", + "cache_miss_action": "precompile-before-timed-cases", + "timing_phases": [ + "precompile", + "cold-start", + "steady-state" + ], + "steady_state_compile_allowed": false + } + }, + "scoring": { + "formula_template_id": "project-fit-kernel-v0.5", + "provisional_scoring_allowed": true, + "official_scoring_requires_seal": true, + "project_season": "astrai-2026q3" + } +} diff --git a/benchmarks/results/swiglu_l20_sm89_infraswe_score.json b/benchmarks/results/swiglu_l20_sm89_infraswe_score.json new file mode 100644 index 00000000..74f27e0b --- /dev/null +++ b/benchmarks/results/swiglu_l20_sm89_infraswe_score.json @@ -0,0 +1,125 @@ +{ + "schema_version": "0.5", + "score_kind": "diagnostic-project-fit", + "score_is_official": false, + "draft_id": "astrai-fused-bf16-swiglu-v1", + "draft_state": "D3-contract-proposed", + "formula_template_id": "project-fit-kernel-v0.5", + "diagnostic_project_fit_100": 90.95413824540499, + "component_values": { + "evolutionary_maintainability": 0.8714777841652558, + "project_contract_fit": 0.9872585449014338, + "performance_reuse_utilization": 0.8828492869148136, + "operational_fit": 0.8956400538635493 + }, + "component_floors": { + "evolutionary_maintainability": 0.6, + "project_contract_fit": 0.6, + "performance_reuse_utilization": 0.4, + "operational_fit": 0.6 + }, + "subcomponent_inputs": { + "evolutionary_maintainability": { + "evolution": 0.75, + "locality": 0.9, + "tests": 1.0, + "failure": 1.0, + "contract": 0.9 + }, + "project_contract_fit": { + "integration": 1.0, + "interface": 0.95, + "lifecycle": 1.0, + "buildtest": 1.0, + "policy": 1.0 + }, + "performance_reuse_utilization": { + "attainment": 0.9, + "coverage": 0.75, + "retention": 1.0, + "family": 0.9, + "compile": 1.0 + }, + "operational_fit": { + "replay": 0.9, + "load": 1.0, + "resource": 0.9, + "coldsteady": 0.75 + } + }, + "input_rationale": { + "evolution": "The primitive extends AstrAI's existing GEMV extension stack with two measured output-row tilings, but has no upstream maintenance history yet.", + "locality": "The implementation is isolated to one CUDA target plus wrapper/backend integration; MLP, build, tests, benchmark, and documentation changes are explicit.", + "tests": "The complete CPU-capable suite passed locally (623 passed, 123 skipped), the complete built L20 suite passed (746 passed), and the focused L20 SwiGLU suite passed 23 tests.", + "failure": "Unsupported inputs and training fall back to the existing path. Automatic dispatch is intentionally empty because the deterministic greedy checkpoint gate changed output at M=1, M=2, and M=4.", + "contract": "Source, acceptance tests, probes, and benchmark evidence are digest-bound, but the D3 acceptance contract has not received maintainer review or a seal.", + "integration": "Dense MLPs route through a backend that preserves AstrAI's existing linear dispatcher whenever fusion is not explicitly selected.", + "interface": "The public primitive and ASTRAI_SWIGLU opt-in are additive; the default auto mode remains behaviorally equivalent to the unfused path.", + "lifecycle": "The primitive is inference-only, uses the current CUDA stream, supports CUDA Graph capture, and fails closed outside contiguous BF16 M=1..8 inputs.", + "buildtest": "The SM89 extension rebuilt successfully, Ruff format/import checks passed, and both local and L20 test suites passed.", + "policy": "The candidate adds no dependency and enables no automatic production band without a passing checkpoint-output gate.", + "attainment": "At the native AstrAI 1B shape, forced fusion improves CUDA-Graph operator latency at M=1/2/4 and real checkpoint step latency by 3.95-5.10%; slower M=8 and wide-shape cells are excluded from auto.", + "coverage": "The required NVIDIA L20/SM89 cell was tested across five transformer MLP shapes and M=1/2/4/8; optional A100/SM80 and H100/SM90 cells remain untested.", + "retention": "The default path is unchanged, direct correctness is recorded, and the complete L20 regression suite passed.", + "family": "AstrAI 1B, LLaMA 2 7B/13B, LLaMA 3 8B, and GPT-NeoX 20B operator shapes are present, plus an actual 24-layer AstrAI checkpoint engine run.", + "compile": "AOT CUDA compilation completed before timed cases; no compilation occurred in steady-state timing.", + "replay": "Operator trials use A-B-C-C-B-A with 20 samples per implementation; engine trials use A-B-B-A, but the evidence does not contain the required fresh-process replay count.", + "load": "The system-level measurement uses AstrAI InferenceEngine with scheduler, sampling, and CUDA Graph on a real 24-layer checkpoint.", + "resource": "GPU5 memory/utilization was checked and the co-resident idle AstrAI service remained allocated; no process or container was stopped.", + "coldsteady": "Precompile and steady-state phases are separated, but cold-start latency is not included in the checked-in result." + }, + "benchmark_trust": { + "formula_version": "benchmark-trust-v0.5", + "status": "scored", + "score_100": 95.87315155141826, + "components": { + "reproducibility": 1.0, + "evidence": 1.0, + "statistics": 0.9, + "environment": 0.9 + }, + "evidence_sha256": "sha256:74e6b77320b9cd1bcd238359363ddd2392d4929d37b50decd00875806a873932", + "failure_codes": [ + "DRAFT_UNSEALED", + "FRESH_PROCESS_REPLAY_INCOMPLETE", + "OPTIONAL_CELLS_UNTESTED" + ] + }, + "official_project_fit": { + "status": "unresolved", + "score_100": null, + "failure_codes": [ + "DRAFT_SEAL_MISSING", + "FRESH_PROCESS_REPLAYS_BELOW_MINIMUM", + "SYSTEM_TRACE_EVIDENCE_MISSING", + "HIDDEN_PROBES_INCOMPLETE", + "EVIDENCE_MANIFEST_UNVERIFIED" + ] + }, + "comparison_cell": { + "target_project_profile_sha256": "sha256:953067c47f5298f06846c8fe873db55afae49b6588a81926e49d9ef6b0994e08", + "target_repository_or_baseline_sha256": "sha256:c5595986ba7ef004333b86c1227fd8223e965971489002f0bbbe47ce4090b342", + "change_intent": "add-fastpath", + "semantic_contract_sha256": "sha256:578c2e70a8e2bb63a62222bd3d7e6f45c7e839d5f5d546b1f0e1c4f7bf4bc895", + "acceptance_contract_sha256": "sha256:fb633cfc6be32041d5e6a37707c6074a3d35931cb5d2b04b7aa4348783783f2a", + "probe_set_sha256": "sha256:3203ce8b46a789afb7cb332b71b01554f5a9f6278d705a4e4b8ce05b236d9d87", + "workload_portfolio_sha256": "sha256:608dbaf05d1f0e45f4cc5697b2e8a3460865753a6c9fefb48f57ea736e0e8540", + "performance_target_sha256": "sha256:b2312952e3b6cf61a527a61806eb39278a1490758e5356436b3b2c36ddffae67", + "required_deployment_cell_set_sha256": "sha256:6c00c84e931a5ad6fdaa5b4f7d4497c872530c75f8510431aff883006a4a2779", + "formula_template_id": "project-fit-kernel-v0.5", + "evidence_policy_id": "v0.4-evidence-ladder-plus-seal-v0.5", + "project_season": "astrai-2026q3", + "cross_project_ranking_allowed": false + }, + "execution": { + "infraswe_commit": "811bc775ed5b3a6ec853219245f3469f78818020", + "draft_validation": "pass", + "draft_resolution_sha256": "abe7a6fc56f72ca9bb911794aa86819014aa082f2aece789f7b1207a97d9871b", + "infraswe_draft_engine_tests": "53 passed", + "astrai_local_tests": "623 passed, 123 skipped", + "astrai_l20_tests": "746 passed", + "astrai_l20_focused_swiglu_tests": "23 passed", + "astrai_lint": "ruff format and import-order checks passed", + "sm89_build": "bf16_swiglu target built successfully; CTA-reuse variants used 36-64 registers and one barrier, warp-row variants used 42 registers and zero barriers, with zero spills" + } +} diff --git a/csrc/CMakeLists.txt b/csrc/CMakeLists.txt index bdd6abbe..7f32c23d 100644 --- a/csrc/CMakeLists.txt +++ b/csrc/CMakeLists.txt @@ -62,6 +62,7 @@ set(KERNEL_NAMES attn_paged_decode attn_paged_prefill bf16_gemv + bf16_swiglu rotary_emb ) set(KERNEL_SRCS @@ -70,6 +71,7 @@ set(KERNEL_SRCS attention/paged_decode.cu attention/paged_prefill.cu gemv/bf16_gemv.cu + gemv/bf16_swiglu.cu rotary_emb.cu ) diff --git a/csrc/kernels/gemv/bf16_swiglu.cu b/csrc/kernels/gemv/bf16_swiglu.cu new file mode 100644 index 00000000..99fe0b13 --- /dev/null +++ b/csrc/kernels/gemv/bf16_swiglu.cu @@ -0,0 +1,368 @@ +// Fused small-M BF16 SwiGLU primitive for decode-time dense MLP layers. + +#include +#include +#include +#include +#include + +#include +#include + +namespace { + +constexpr int kThreads = 256; +constexpr int kWarpSize = 32; +constexpr int kWarps = kThreads / kWarpSize; + +__device__ __forceinline__ float warp_sum(float value) { +#pragma unroll + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + value += __shfl_down_sync(0xffffffff, value, offset); + } + return value; +} + +__device__ __forceinline__ float round_bf16(float value) { + return __bfloat162float(__float2bfloat16_rn(value)); +} + +template +__global__ void bf16_swiglu_kernel( + const __nv_bfloat16* __restrict__ x, + const __nv_bfloat16* __restrict__ up_weight, + const __nv_bfloat16* __restrict__ gate_weight, + __nv_bfloat16* __restrict__ output, + int n, + int k +) { + const int output_index = blockIdx.x; + const int lane = threadIdx.x & (kWarpSize - 1); + const int warp = threadIdx.x / kWarpSize; + const int vector_count = k / 8; + + float up_sums[Rows] = {}; + float gate_sums[Rows] = {}; + __shared__ float up_warp_sums[Rows][kWarps]; + __shared__ float gate_warp_sums[Rows][kWarps]; + + const auto* x4 = reinterpret_cast(x); + const auto* up4 = reinterpret_cast( + up_weight + static_cast(output_index) * k + ); + const auto* gate4 = reinterpret_cast( + gate_weight + static_cast(output_index) * k + ); + + // Read each pair of up/gate weight chunks once per CTA, then reuse it for + // every active decode row. The fused epilogue removes two [M, N] + // intermediates and the standalone SiLU and multiply launches. + for (int vector_index = threadIdx.x; + vector_index < vector_count; + vector_index += blockDim.x) { + const uint4 up_raw = up4[vector_index]; + const uint4 gate_raw = gate4[vector_index]; + const auto* up_values = + reinterpret_cast(&up_raw); + const auto* gate_values = + reinterpret_cast(&gate_raw); + +#pragma unroll + for (int row = 0; row < Rows; ++row) { + const uint4 x_raw = + x4[static_cast(row) * vector_count + vector_index]; + const auto* x_values = + reinterpret_cast(&x_raw); +#pragma unroll + for (int pair = 0; pair < 4; ++pair) { + const float2 xv = __bfloat1622float2(x_values[pair]); + const float2 uv = __bfloat1622float2(up_values[pair]); + const float2 gv = __bfloat1622float2(gate_values[pair]); + up_sums[row] = fmaf(xv.x, uv.x, up_sums[row]); + up_sums[row] = fmaf(xv.y, uv.y, up_sums[row]); + gate_sums[row] = fmaf(xv.x, gv.x, gate_sums[row]); + gate_sums[row] = fmaf(xv.y, gv.y, gate_sums[row]); + } + } + } + +#pragma unroll + for (int row = 0; row < Rows; ++row) { + up_sums[row] = warp_sum(up_sums[row]); + gate_sums[row] = warp_sum(gate_sums[row]); + } + if (lane == 0) { +#pragma unroll + for (int row = 0; row < Rows; ++row) { + up_warp_sums[row][warp] = up_sums[row]; + gate_warp_sums[row][warp] = gate_sums[row]; + } + } + __syncthreads(); + + if (warp == 0) { +#pragma unroll + for (int row = 0; row < Rows; ++row) { + float up = lane < kWarps ? up_warp_sums[row][lane] : 0.0f; + float gate = lane < kWarps ? gate_warp_sums[row][lane] : 0.0f; + up = warp_sum(up); + gate = warp_sum(gate); + if (lane == 0) { + // Match the public composition's BF16 rounding boundaries: + // BF16 linear outputs, BF16 SiLU output, then BF16 multiply. + up = round_bf16(up); + gate = round_bf16(gate); + const float silu = round_bf16(gate / (1.0f + expf(-gate))); + output[static_cast(row) * n + output_index] = + __float2bfloat16_rn(up * silu); + } + } + } +} + +template +void launch_bf16_swiglu( + const __nv_bfloat16* x, + const __nv_bfloat16* up_weight, + const __nv_bfloat16* gate_weight, + __nv_bfloat16* output, + int n, + int k, + cudaStream_t stream +) { + bf16_swiglu_kernel<<>>( + x, up_weight, gate_weight, output, n, k + ); +} + +template +__global__ void bf16_swiglu_warp_rows_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 +) { + const int output_index = blockIdx.x; + const int row = threadIdx.x / kWarpSize; + const int lane = threadIdx.x & (kWarpSize - 1); + const int vector_count = k / 8; + + float up_sum = 0.0f; + float gate_sum = 0.0f; + const auto* x4 = reinterpret_cast( + x + static_cast(row) * k + ); + const auto* up4 = reinterpret_cast( + up_weight + static_cast(output_index) * k + ); + const auto* gate4 = reinterpret_cast( + gate_weight + static_cast(output_index) * k + ); + + // A warp owns one decode row. Same-address weight reads from sibling + // warps are served through the read-only/L1 path, while each row avoids + // CTA-wide shared-memory reductions and synchronization. + for (int vector_index = lane; + vector_index < vector_count; + vector_index += kWarpSize) { + const uint4 x_raw = x4[vector_index]; + const uint4 up_raw = up4[vector_index]; + const uint4 gate_raw = gate4[vector_index]; + const auto* x_values = reinterpret_cast(&x_raw); + const auto* up_values = + reinterpret_cast(&up_raw); + const auto* gate_values = + reinterpret_cast(&gate_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_sum = fmaf(xv.x, uv.x, up_sum); + up_sum = fmaf(xv.y, uv.y, up_sum); + gate_sum = fmaf(xv.x, gv.x, gate_sum); + gate_sum = fmaf(xv.y, gv.y, gate_sum); + } + } + up_sum = warp_sum(up_sum); + gate_sum = warp_sum(gate_sum); + if (lane == 0) { + up_sum = round_bf16(up_sum); + gate_sum = round_bf16(gate_sum); + const float silu = + round_bf16(gate_sum / (1.0f + expf(-gate_sum))); + output[static_cast(row) * n + output_index] = + __float2bfloat16_rn(up_sum * silu); + } +} + +template +void launch_bf16_swiglu_warp_rows( + 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_warp_rows_kernel<<>>( + x, up_weight, gate_weight, output, n, k + ); +} + +torch::Tensor bf16_swiglu( + torch::Tensor x, + torch::Tensor up_weight, + torch::Tensor gate_weight +) { + TORCH_CHECK( + x.is_cuda() && up_weight.is_cuda() && gate_weight.is_cuda(), + "x, up_weight, and gate_weight must be CUDA tensors" + ); + TORCH_CHECK( + x.device() == up_weight.device() && x.device() == gate_weight.device(), + "x and weights must share a device" + ); + TORCH_CHECK( + x.scalar_type() == torch::kBFloat16 && + up_weight.scalar_type() == torch::kBFloat16 && + gate_weight.scalar_type() == torch::kBFloat16, + "x and weights must be bf16" + ); + TORCH_CHECK( + x.dim() == 1 || x.dim() == 2, + "x must have shape [K] or [M, K]" + ); + TORCH_CHECK( + up_weight.dim() == 2 && gate_weight.dim() == 2, + "weights must have shape [N, K]" + ); + TORCH_CHECK( + x.is_contiguous() && up_weight.is_contiguous() && + gate_weight.is_contiguous(), + "x and weights must be contiguous" + ); + TORCH_CHECK( + !x.requires_grad() && !up_weight.requires_grad() && + !gate_weight.requires_grad(), + "bf16_swiglu is inference-only and does not support autograd" + ); + + const int64_t m = x.dim() == 1 ? 1 : x.size(0); + const int64_t k = x.size(-1); + const int64_t n = up_weight.size(0); + TORCH_CHECK(m >= 1 && m <= 8, "M must be in [1, 8]"); + TORCH_CHECK( + gate_weight.sizes() == up_weight.sizes(), + "up_weight and gate_weight must have identical shapes" + ); + TORCH_CHECK(up_weight.size(1) == k, "weight K must match x K"); + TORCH_CHECK(k > 0 && n > 0, "N and K must be positive"); + TORCH_CHECK(k % 8 == 0, "K must be divisible by 8"); + TORCH_CHECK( + k <= std::numeric_limits::max() && + n <= std::numeric_limits::max(), + "N or K exceeds the CUDA launcher limit" + ); + + const at::cuda::OptionalCUDAGuard guard(x.device()); + const auto* properties = at::cuda::getDeviceProperties(x.device().index()); + TORCH_CHECK( + properties->major >= 8, + "bf16_swiglu requires compute capability 8.0+" + ); + auto stream = at::cuda::getCurrentCUDAStream(); + auto output = x.dim() == 1 ? torch::empty({n}, x.options()) + : torch::empty({m, n}, x.options()); + + const auto* x_ptr = + reinterpret_cast(x.data_ptr()); + const auto* up_ptr = + reinterpret_cast(up_weight.data_ptr()); + const auto* gate_ptr = + reinterpret_cast(gate_weight.data_ptr()); + auto* output_ptr = + reinterpret_cast<__nv_bfloat16*>(output.data_ptr()); + const int n_int = static_cast(n); + const int k_int = static_cast(k); + const bool use_warp_rows = n_int == 6912 && k_int == 1536; + + switch (m) { + case 1: + launch_bf16_swiglu<1>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + case 2: + if (use_warp_rows) { + launch_bf16_swiglu_warp_rows<2>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + } else { + launch_bf16_swiglu<2>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + } + break; + case 3: + launch_bf16_swiglu<3>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + case 4: + if (use_warp_rows) { + launch_bf16_swiglu_warp_rows<4>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + } else { + launch_bf16_swiglu<4>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + } + break; + case 5: + launch_bf16_swiglu<5>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + case 6: + launch_bf16_swiglu<6>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + case 7: + launch_bf16_swiglu<7>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + break; + case 8: + if (use_warp_rows) { + launch_bf16_swiglu_warp_rows<8>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + } else { + launch_bf16_swiglu<8>( + x_ptr, up_ptr, gate_ptr, output_ptr, n_int, k_int, stream.stream() + ); + } + break; + } + C10_CUDA_CHECK(cudaGetLastError()); + return output; +} + +} // namespace + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def( + "bf16_swiglu", + &bf16_swiglu, + py::arg("x"), + py::arg("up_weight"), + py::arg("gate_weight"), + "M in [1, 8] fused BF16 up/gate projection and SwiGLU" + ); +} diff --git a/docs/benchmarks/swiglu_engine_l20_sm89.txt b/docs/benchmarks/swiglu_engine_l20_sm89.txt new file mode 100644 index 00000000..56b46d92 --- /dev/null +++ b/docs/benchmarks/swiglu_engine_l20_sm89.txt @@ -0,0 +1,144 @@ +CASE batch=1 swiglu=0 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=1, SeqLen=64 + Throughput : 243.6 tokens/s + Latency : 4.11 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=1 swiglu=1 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=1, SeqLen=64 + Throughput : 255.0 tokens/s + Latency : 3.92 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=1 swiglu=1 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=1, SeqLen=64 + Throughput : 254.3 tokens/s + Latency : 3.93 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=1 swiglu=0 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=1, SeqLen=64 + Throughput : 241.8 tokens/s + Latency : 4.14 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=2 swiglu=0 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=2, SeqLen=64 + Throughput : 460.8 tokens/s + Latency : 4.34 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=2 swiglu=1 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=2, SeqLen=64 + Throughput : 503.5 tokens/s + Latency : 3.97 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=2 swiglu=1 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=2, SeqLen=64 + Throughput : 483.6 tokens/s + Latency : 4.14 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=2 swiglu=0 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=2, SeqLen=64 + Throughput : 481.5 tokens/s + Latency : 4.15 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=4 swiglu=0 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=4, SeqLen=64 + Throughput : 881.8 tokens/s + Latency : 4.54 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=4 swiglu=1 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=4, SeqLen=64 + Throughput : 916.0 tokens/s + Latency : 4.37 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=4 swiglu=1 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=4, SeqLen=64 + Throughput : 943.2 tokens/s + Latency : 4.24 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- +CASE batch=4 swiglu=0 +Loading model from /home/kxqandccx/kxq/AstrAI/params ... +Benchmark: device=cuda dtype=bfloat16 backend=cuda +-------------------------------------------------------------------------------- +DECODE — Batch=4, SeqLen=64 + Throughput : 906.2 tokens/s + Latency : 4.41 ms/step + Num Trials: 3 + Prompt Length: 64 + Backend: CudaBackend + Cuda Graph: True +-------------------------------------------------------------------------------- diff --git a/docs/benchmarks/swiglu_greedy_m1_m2_l20_sm89.txt b/docs/benchmarks/swiglu_greedy_m1_m2_l20_sm89.txt new file mode 100644 index 00000000..ba7a0506 --- /dev/null +++ b/docs/benchmarks/swiglu_greedy_m1_m2_l20_sm89.txt @@ -0,0 +1,12 @@ +CASE batch=1 swiglu=0 +sha256=78054b6bb6ab3bad5e551894b5b9ae5f33282145aef151e87ae259362852d9b3 +["\n\nIn this paper, we propose a new method to determine the probability of a test result. We use the method to determine the probability of a test result. We show that the method is effective in detecting the probability of a test result. We also show that the method is effective in detecting the"] +CASE batch=1 swiglu=1 +sha256=5db1a291cfbd349b81896d7e16c9adff32d09bd9d5c8bda18152e51aae8ce2f6 +["\n\nThe purpose of this paper is to show that the existence of deterministic tests is not a problem of the existence of deterministic tests. The purpose of this paper is to show that the existence of deterministic tests is not a problem of the existence of deterministic tests. The purpose of this paper is"] +CASE batch=2 swiglu=0 +sha256=9f6b4e4b6bffe4afa947d798a19e12b2440e0562cc4be60bc9cebf526c1889a7 +["\n\nIn this paper, we propose a new approach to the problem of determining the optimal value of a function. We show that the optimal value of a function is determined by the set of all possible values of the function. We show that the optimal value of a function is determined by the set of all", "\n\nIn this paper, we propose a new approach to the problem of determining the optimal value of a function. We show that the optimal value of a function is determined by the set of all possible values of the function. We show that the optimal value of a function is determined by the set of all"] +CASE batch=2 swiglu=1 +sha256=0c7664d3878cd32914079844eeb3a1ddde6fe4547f47490a130e362e15c65cac +["\n\nIn this paper, we propose a new method to determine the probability of a test result. We use the method to determine the probability of a test result. We show that the method is effective in determining the probability of a test result. We also show that the method is effective in determining the", "\n\nIn this paper, we propose a new method to determine the probability of a test result. We use the method to determine the probability of a test result. We show that the method is effective in determining the probability of a test result. We also show that the method is effective in determining the"] diff --git a/docs/benchmarks/swiglu_greedy_m4_l20_sm89.txt b/docs/benchmarks/swiglu_greedy_m4_l20_sm89.txt new file mode 100644 index 00000000..3647fab3 --- /dev/null +++ b/docs/benchmarks/swiglu_greedy_m4_l20_sm89.txt @@ -0,0 +1,6 @@ +CASE swiglu=0 +["\n\nIn this paper, we propose a new method to determine the probability of a test result. We use the method to determine the probability of a test result. We show that the method is effective in detecting the probability of a test result. We also show that the method is effective in detecting the", "\n\nIn this paper, we propose a new method to determine the probability of a test result. We use the method to determine the probability of a test result. We show that the method is effective in detecting the probability of a test result. We also show that the method is effective in detecting the", "\n\nIn this paper, we propose a new method to determine the probability of a test result. We use the method to determine the probability of a test result. We show that the method is effective in detecting the probability of a test result. We also show that the method is effective in detecting the", "\n\nIn this paper, we propose a new method to determine the probability of a test result. We use the method to determine the probability of a test result. We show that the method is effective in detecting the probability of a test result. We also show that the method is effective in detecting the"] +sha256=738b35a81fe9d669039f5aaebbfd49f3502a1926327bba421c297612bd6adb38 +CASE swiglu=1 +["\n\nIn this paper, we propose a new approach to the problem of determining the probability of a test result. We show that the probability of a test result is a function of the number of tests that are performed. We also show that the probability of a test result is a function of the number of", "\n\nIn this paper, we propose a new approach to the problem of determining the probability of a test result. We show that the probability of a test result is a function of the number of tests that are performed. We also show that the probability of a test result is a function of the number of", "\n\nIn this paper, we propose a new approach to the problem of determining the probability of a test result. We show that the probability of a test result is a function of the number of tests that are performed. We also show that the probability of a test result is a function of the number of", "\n\nIn this paper, we propose a new approach to the problem of determining the probability of a test result. We show that the probability of a test result is a function of the number of tests that are performed. We also show that the probability of a test result is a function of the number of"] +sha256=e812f5c6c68c3897df14bed987ed8174359e07c62ab6a2ebe270e3a20f10ed94 diff --git a/docs/benchmarks/swiglu_l20_sm89.json b/docs/benchmarks/swiglu_l20_sm89.json new file mode 100644 index 00000000..9fd4073c --- /dev/null +++ b/docs/benchmarks/swiglu_l20_sm89.json @@ -0,0 +1,1820 @@ +{ + "metadata": { + "timestamp_utc": "2026-09-02T11:26:54.196792+00:00", + "gpu_name": "NVIDIA L20", + "compute_capability": "8.9", + "total_memory_bytes": 47677177856, + "torch_version": "2.11.0+cu128", + "cuda_version": "12.8", + "dtype": "bfloat16" + }, + "settings": { + "warmup": 20, + "iterations": 100, + "trials": 10, + "seed": 0, + "order": "A-B-C-C-B-A" + }, + "results": [ + { + "shape": "astrai_1b", + "m": 1, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.03336192011833191, + "p90_ms": 0.03520307230949402, + "p99_ms": 0.03612098453044891, + "min_ms": 0.032419838905334473, + "max_ms": 0.03616767883300781 + }, + { + "shape": "astrai_1b", + "m": 1, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.025175040960311888, + "p90_ms": 0.027143167257308958, + "p99_ms": 0.02846120090484619, + "min_ms": 0.025036799907684325, + "max_ms": 0.028712000846862793 + }, + { + "shape": "astrai_1b", + "m": 1, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.013808640241622925, + "p90_ms": 0.013856832504272462, + "p99_ms": 0.014107565188407898, + "min_ms": 0.013731839656829834, + "max_ms": 0.014161920547485352, + "max_abs_error": 1.3969838619232178e-09, + "mean_abs_error": 2.0210993523315374e-13, + "cosine_similarity": 1.0000001192092896 + }, + { + "shape": "astrai_1b", + "m": 1, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.025635360479354857, + "p90_ms": 0.0256536967754364, + "p99_ms": 0.02567569925785065, + "min_ms": 0.02559295892715454, + "max_ms": 0.025680320262908937 + }, + { + "shape": "astrai_1b", + "m": 1, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.022977120876312256, + "p90_ms": 0.02299407982826233, + "p99_ms": 0.02299759373664856, + "min_ms": 0.02292736053466797, + "max_ms": 0.022998080253601075 + }, + { + "shape": "astrai_1b", + "m": 1, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.01374783992767334, + "p90_ms": 0.013763935565948486, + "p99_ms": 0.013780898928642274, + "min_ms": 0.013715840578079223, + "max_ms": 0.01378399968147278, + "max_abs_error": 2.384185791015625e-07, + "mean_abs_error": 3.449342894645824e-11, + "cosine_similarity": 1.0000001192092896 + }, + { + "shape": "astrai_1b", + "m": 2, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.04668031930923462, + "p90_ms": 0.04708377599716187, + "p99_ms": 0.04719747934341431, + "min_ms": 0.046151041984558105, + "max_ms": 0.04721888065338135 + }, + { + "shape": "astrai_1b", + "m": 2, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.032778559923171996, + "p90_ms": 0.032919232845306394, + "p99_ms": 0.03308471763134003, + "min_ms": 0.03253760099411011, + "max_ms": 0.03311232089996338 + }, + { + "shape": "astrai_1b", + "m": 2, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.014095680117607118, + "p90_ms": 0.014110399961471556, + "p99_ms": 0.01415008614063263, + "min_ms": 0.014022079706192016, + "max_ms": 0.014158719778060913, + "max_abs_error": 3.0517578125e-05, + "mean_abs_error": 6.8403771535940905e-09, + "cosine_similarity": 0.9999999403953552 + }, + { + "shape": "astrai_1b", + "m": 2, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.024841439723968507, + "p90_ms": 0.024966944694519044, + "p99_ms": 0.024991856622695922, + "min_ms": 0.0248089599609375, + "max_ms": 0.024995200634002686 + }, + { + "shape": "astrai_1b", + "m": 2, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.02627951979637146, + "p90_ms": 0.02641388702392578, + "p99_ms": 0.02642572522163391, + "min_ms": 0.026257600784301758, + "max_ms": 0.026426880359649657 + }, + { + "shape": "astrai_1b", + "m": 2, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.014162880182266236, + "p90_ms": 0.014282880306243896, + "p99_ms": 0.014310902094841003, + "min_ms": 0.014113919734954834, + "max_ms": 0.014316799640655518, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 1.1114363651643089e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "astrai_1b", + "m": 4, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.04687664031982422, + "p90_ms": 0.04740822553634643, + "p99_ms": 0.0495560998916626, + "min_ms": 0.04651040077209473, + "max_ms": 0.04998784065246582 + }, + { + "shape": "astrai_1b", + "m": 4, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.04020000219345093, + "p90_ms": 0.040237729072570796, + "p99_ms": 0.040276090669631955, + "min_ms": 0.04014815807342529, + "max_ms": 0.04028448104858398 + }, + { + "shape": "astrai_1b", + "m": 4, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.01803984045982361, + "p90_ms": 0.01809161651134491, + "p99_ms": 0.018110911679267882, + "min_ms": 0.017995519638061522, + "max_ms": 0.018114559650421143, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 1.0392762206379302e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "astrai_1b", + "m": 4, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.025065439939498904, + "p90_ms": 0.02508691191673279, + "p99_ms": 0.02511093099117279, + "min_ms": 0.02502016067504883, + "max_ms": 0.02511615991592407 + }, + { + "shape": "astrai_1b", + "m": 4, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.03838783979415894, + "p90_ms": 0.038400672435760494, + "p99_ms": 0.03840640399456024, + "min_ms": 0.038372480869293214, + "max_ms": 0.038407680988311765 + }, + { + "shape": "astrai_1b", + "m": 4, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.01805631995201111, + "p90_ms": 0.018067872047424317, + "p99_ms": 0.018082531607151033, + "min_ms": 0.01800447940826416, + "max_ms": 0.018083200454711915, + "max_abs_error": 6.103515625e-05, + "mean_abs_error": 5.430904614911469e-09, + "cosine_similarity": 1.0 + }, + { + "shape": "astrai_1b", + "m": 8, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.047262399196624755, + "p90_ms": 0.09892006397247316, + "p99_ms": 0.11928016328811646, + "min_ms": 0.04669951915740967, + "max_ms": 0.11987551689147949 + }, + { + "shape": "astrai_1b", + "m": 8, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.07197743892669678, + "p90_ms": 0.07892819261550904, + "p99_ms": 0.07972745013237, + "min_ms": 0.07187744140625, + "max_ms": 0.07979360103607178 + }, + { + "shape": "astrai_1b", + "m": 8, + "n": 6912, + "k": 1536, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.03347103953361511, + "p90_ms": 0.03368713569641113, + "p99_ms": 0.03373796784877777, + "min_ms": 0.033377280235290525, + "max_ms": 0.03374495983123779, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 1.1743122030338782e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "astrai_1b", + "m": 8, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.02562575936317444, + "p90_ms": 0.025687744617462156, + "p99_ms": 0.02594187598228454, + "min_ms": 0.025588159561157228, + "max_ms": 0.025963521003723143 + }, + { + "shape": "astrai_1b", + "m": 8, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.0700651216506958, + "p90_ms": 0.0700872302055359, + "p99_ms": 0.07020150990486144, + "min_ms": 0.07002624034881592, + "max_ms": 0.07022655963897705 + }, + { + "shape": "astrai_1b", + "m": 8, + "n": 6912, + "k": 1536, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.033386240005493166, + "p90_ms": 0.03339843130111695, + "p99_ms": 0.033478554248809814, + "min_ms": 0.03336287975311279, + "max_ms": 0.033497281074523926, + "max_abs_error": 6.103515625e-05, + "mean_abs_error": 6.4199068283699035e-09, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_7b", + "m": 1, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.2483579158782959, + "p90_ms": 0.24857123374938964, + "p99_ms": 0.24874350452423094, + "min_ms": 0.24757183074951172, + "max_ms": 0.248768310546875 + }, + { + "shape": "llama2_7b", + "m": 1, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.24014480590820314, + "p90_ms": 0.24030490112304687, + "p99_ms": 0.2406482805252075, + "min_ms": 0.23985504150390624, + "max_ms": 0.2406924819946289 + }, + { + "shape": "llama2_7b", + "m": 1, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.24054000854492186, + "p90_ms": 0.24103282356262207, + "p99_ms": 0.24124960613250734, + "min_ms": 0.24040672302246094, + "max_ms": 0.24125951766967774, + "max_abs_error": 2.9802322387695312e-08, + "mean_abs_error": 4.399416161232628e-12, + "cosine_similarity": 0.9999998807907104 + }, + { + "shape": "llama2_7b", + "m": 1, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.24543855667114256, + "p90_ms": 0.2462499847412109, + "p99_ms": 0.24650316982269288, + "min_ms": 0.24525375366210939, + "max_ms": 0.2465590476989746 + }, + { + "shape": "llama2_7b", + "m": 1, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.23805376052856447, + "p90_ms": 0.23814822578430178, + "p99_ms": 0.23817636966705322, + "min_ms": 0.23785472869873048, + "max_ms": 0.2381808090209961 + }, + { + "shape": "llama2_7b", + "m": 1, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.24051487922668457, + "p90_ms": 0.24074946403503417, + "p99_ms": 0.2408029245376587, + "min_ms": 0.24034400939941405, + "max_ms": 0.2408086395263672, + "max_abs_error": 7.62939453125e-06, + "mean_abs_error": 7.797965051459244e-10, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_7b", + "m": 2, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.2605764865875244, + "p90_ms": 0.2609274864196777, + "p99_ms": 0.26124884338378906, + "min_ms": 0.2602403259277344, + "max_ms": 0.261297607421875 + }, + { + "shape": "llama2_7b", + "m": 2, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.24056976318359374, + "p90_ms": 0.2406690845489502, + "p99_ms": 0.24075329818725585, + "min_ms": 0.2402947235107422, + "max_ms": 0.24076704025268555 + }, + { + "shape": "llama2_7b", + "m": 2, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.24064352035522463, + "p90_ms": 0.24109625625610354, + "p99_ms": 0.24118487815856934, + "min_ms": 0.2405244827270508, + "max_ms": 0.24118688583374023, + "max_abs_error": 0.000244140625, + "mean_abs_error": 3.518480795605683e-08, + "cosine_similarity": 0.9999998807907104 + }, + { + "shape": "llama2_7b", + "m": 2, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.2573767948150635, + "p90_ms": 0.25752944374084474, + "p99_ms": 0.2576336078643799, + "min_ms": 0.25689311981201174, + "max_ms": 0.25765567779541015 + }, + { + "shape": "llama2_7b", + "m": 2, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.23858335494995117, + "p90_ms": 0.23877366828918456, + "p99_ms": 0.24103173294067382, + "min_ms": 0.23840160369873048, + "max_ms": 0.24145824432373048 + }, + { + "shape": "llama2_7b", + "m": 2, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.24056768417358398, + "p90_ms": 0.24060989570617675, + "p99_ms": 0.24068091373443604, + "min_ms": 0.24041088104248046, + "max_ms": 0.240696964263916, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 3.232381828865982e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_7b", + "m": 4, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.2610425567626953, + "p90_ms": 0.26153100013732905, + "p99_ms": 0.26189791259765627, + "min_ms": 0.26051551818847657, + "max_ms": 0.2619830322265625 + }, + { + "shape": "llama2_7b", + "m": 4, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.24146736145019532, + "p90_ms": 0.241597599029541, + "p99_ms": 0.24211565246582029, + "min_ms": 0.2411350440979004, + "max_ms": 0.24223615646362304 + }, + { + "shape": "llama2_7b", + "m": 4, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.24120527267456054, + "p90_ms": 0.2413661460876465, + "p99_ms": 0.24162629890441895, + "min_ms": 0.24105600357055665, + "max_ms": 0.24164800643920897, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 3.4207591426138606e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_7b", + "m": 4, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.2576919937133789, + "p90_ms": 0.25796803092956544, + "p99_ms": 0.25825283603668214, + "min_ms": 0.25749887466430665, + "max_ms": 0.2582851219177246 + }, + { + "shape": "llama2_7b", + "m": 4, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.2393723201751709, + "p90_ms": 0.2395743980407715, + "p99_ms": 0.23969023685455323, + "min_ms": 0.23921920776367187, + "max_ms": 0.23970592498779297 + }, + { + "shape": "llama2_7b", + "m": 4, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.24103343963623047, + "p90_ms": 0.24113894081115722, + "p99_ms": 0.2412008436203003, + "min_ms": 0.24085088729858398, + "max_ms": 0.24120576858520507, + "max_abs_error": 6.103515625e-05, + "mean_abs_error": 2.2930278120725234e-08, + "cosine_similarity": 0.9999999403953552 + }, + { + "shape": "llama2_7b", + "m": 8, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.26210512161254884, + "p90_ms": 0.2623787784576416, + "p99_ms": 0.26274656467437746, + "min_ms": 0.2615241622924805, + "max_ms": 0.2628316879272461 + }, + { + "shape": "llama2_7b", + "m": 8, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.24479904174804687, + "p90_ms": 0.24489805030822753, + "p99_ms": 0.2449971487045288, + "min_ms": 0.24458335876464843, + "max_ms": 0.24500864028930663 + }, + { + "shape": "llama2_7b", + "m": 8, + "n": 11008, + "k": 4096, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.24242128372192384, + "p90_ms": 0.24250384712219236, + "p99_ms": 0.24255824489593505, + "min_ms": 0.24232704162597657, + "max_ms": 0.24256256103515625, + "max_abs_error": 0.000244140625, + "mean_abs_error": 2.5229038413954186e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_7b", + "m": 8, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.2598663997650147, + "p90_ms": 0.26006431579589845, + "p99_ms": 0.26018847332000733, + "min_ms": 0.25934080123901365, + "max_ms": 0.26019935607910155 + }, + { + "shape": "llama2_7b", + "m": 8, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.2428281593322754, + "p90_ms": 0.2429563770294189, + "p99_ms": 0.24303103084564207, + "min_ms": 0.24263423919677735, + "max_ms": 0.2430406379699707 + }, + { + "shape": "llama2_7b", + "m": 8, + "n": 11008, + "k": 4096, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.24228544235229493, + "p90_ms": 0.2424908447265625, + "p99_ms": 0.24348165016174317, + "min_ms": 0.24209152221679686, + "max_ms": 0.24370399475097657, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 3.02232976423511e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama3_8b", + "m": 1, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.32035743713378906, + "p90_ms": 0.3208788719177246, + "p99_ms": 0.32103768959045403, + "min_ms": 0.31978527069091794, + "max_ms": 0.3210678482055664 + }, + { + "shape": "llama3_8b", + "m": 1, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3110881519317627, + "p90_ms": 0.3115822677612305, + "p99_ms": 0.3117168004989624, + "min_ms": 0.31097408294677736, + "max_ms": 0.3117398452758789 + }, + { + "shape": "llama3_8b", + "m": 1, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.3127006435394287, + "p90_ms": 0.31289871597290037, + "p99_ms": 0.31313554859161374, + "min_ms": 0.31261119842529295, + "max_ms": 0.3131475257873535, + "max_abs_error": 4.76837158203125e-07, + "mean_abs_error": 3.3521376646694634e-11, + "cosine_similarity": 1.0 + }, + { + "shape": "llama3_8b", + "m": 1, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3174740791320801, + "p90_ms": 0.31830342292785646, + "p99_ms": 0.31834455471038814, + "min_ms": 0.317357120513916, + "max_ms": 0.31835136413574217 + }, + { + "shape": "llama3_8b", + "m": 1, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3086094379425049, + "p90_ms": 0.30878892517089845, + "p99_ms": 0.30881352081298824, + "min_ms": 0.3085606384277344, + "max_ms": 0.3088163185119629 + }, + { + "shape": "llama3_8b", + "m": 1, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.3124616050720215, + "p90_ms": 0.3125781078338623, + "p99_ms": 0.31300077934265136, + "min_ms": 0.3124073600769043, + "max_ms": 0.31309951782226564, + "max_abs_error": 3.814697265625e-06, + "mean_abs_error": 2.660921760710977e-10, + "cosine_similarity": 1.0 + }, + { + "shape": "llama3_8b", + "m": 2, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3360844802856445, + "p90_ms": 0.3363768844604492, + "p99_ms": 0.3364322246551514, + "min_ms": 0.3356313705444336, + "max_ms": 0.3364352035522461 + }, + { + "shape": "llama3_8b", + "m": 2, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.31170591354370114, + "p90_ms": 0.31194432640075687, + "p99_ms": 0.31211127223968504, + "min_ms": 0.3115507125854492, + "max_ms": 0.31213151931762695 + }, + { + "shape": "llama3_8b", + "m": 2, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.31296016693115236, + "p90_ms": 0.31301324653625484, + "p99_ms": 0.31303769626617434, + "min_ms": 0.31283231735229494, + "max_ms": 0.31303775787353516, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 2.0864275640519736e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama3_8b", + "m": 2, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3329132843017578, + "p90_ms": 0.3332023582458496, + "p99_ms": 0.3346445404052734, + "min_ms": 0.33273536682128907, + "max_ms": 0.3349798583984375 + }, + { + "shape": "llama3_8b", + "m": 2, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.30939439773559574, + "p90_ms": 0.3095896034240722, + "p99_ms": 0.3096395071029663, + "min_ms": 0.3093065643310547, + "max_ms": 0.3096447944641113 + }, + { + "shape": "llama3_8b", + "m": 2, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.3128171253204346, + "p90_ms": 0.3130419445037842, + "p99_ms": 0.313067151260376, + "min_ms": 0.31271455764770506, + "max_ms": 0.3130723190307617, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 4.01029573993128e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama3_8b", + "m": 4, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3369073677062988, + "p90_ms": 0.33728800582885743, + "p99_ms": 0.33750827140808104, + "min_ms": 0.33637279510498047, + "max_ms": 0.3375187301635742 + }, + { + "shape": "llama3_8b", + "m": 4, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.31293343544006347, + "p90_ms": 0.3132864665985107, + "p99_ms": 0.31335803737640383, + "min_ms": 0.3127609634399414, + "max_ms": 0.31336320877075197 + }, + { + "shape": "llama3_8b", + "m": 4, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.3135753536224365, + "p90_ms": 0.31390438652038577, + "p99_ms": 0.31398765087127684, + "min_ms": 0.3133945655822754, + "max_ms": 0.3140009689331055, + "max_abs_error": 0.000244140625, + "mean_abs_error": 3.24883906444029e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama3_8b", + "m": 4, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3337718391418457, + "p90_ms": 0.3338309097290039, + "p99_ms": 0.33385156745910644, + "min_ms": 0.3336198425292969, + "max_ms": 0.33385601043701174 + }, + { + "shape": "llama3_8b", + "m": 4, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3105408000946045, + "p90_ms": 0.31060870170593263, + "p99_ms": 0.3106609846115112, + "min_ms": 0.31048896789550784, + "max_ms": 0.31066432952880857 + }, + { + "shape": "llama3_8b", + "m": 4, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.31342432022094724, + "p90_ms": 0.3134717674255371, + "p99_ms": 0.3135014751434326, + "min_ms": 0.313372802734375, + "max_ms": 0.3135068893432617, + "max_abs_error": 0.000244140625, + "mean_abs_error": 3.6921576906934206e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama3_8b", + "m": 8, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.33917535781860353, + "p90_ms": 0.33956013870239254, + "p99_ms": 0.33990706634521484, + "min_ms": 0.33855648040771485, + "max_ms": 0.33993953704833985 + }, + { + "shape": "llama3_8b", + "m": 8, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3164067268371582, + "p90_ms": 0.31659145164489744, + "p99_ms": 0.3167010677337646, + "min_ms": 0.3162652778625488, + "max_ms": 0.31670495986938474 + }, + { + "shape": "llama3_8b", + "m": 8, + "n": 14336, + "k": 4096, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.31352383613586426, + "p90_ms": 0.3136793270111084, + "p99_ms": 0.3138789640426636, + "min_ms": 0.31335935592651365, + "max_ms": 0.31391775131225585, + "max_abs_error": 0.000244140625, + "mean_abs_error": 3.4041075736013227e-08, + "cosine_similarity": 1.0000001192092896 + }, + { + "shape": "llama3_8b", + "m": 8, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.33592369079589846, + "p90_ms": 0.33631641769409176, + "p99_ms": 0.3376862892150879, + "min_ms": 0.33552894592285154, + "max_ms": 0.33800640106201174 + }, + { + "shape": "llama3_8b", + "m": 8, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.31418175697326656, + "p90_ms": 0.3143845691680908, + "p99_ms": 0.3144443880081177, + "min_ms": 0.3138748741149902, + "max_ms": 0.3144463920593262 + }, + { + "shape": "llama3_8b", + "m": 8, + "n": 14336, + "k": 4096, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.3147000026702881, + "p90_ms": 0.31490866470336915, + "p99_ms": 0.3149763284683228, + "min_ms": 0.3144364738464355, + "max_ms": 0.3149776077270508, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 2.8799293616543764e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_13b", + "m": 1, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.38610063552856444, + "p90_ms": 0.3863121452331543, + "p99_ms": 0.3871610931396484, + "min_ms": 0.3856972885131836, + "max_ms": 0.3873132705688477 + }, + { + "shape": "llama2_13b", + "m": 1, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.38129167556762694, + "p90_ms": 0.381681755065918, + "p99_ms": 0.38181031227111817, + "min_ms": 0.381176643371582, + "max_ms": 0.3818191909790039 + }, + { + "shape": "llama2_13b", + "m": 1, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.37660272598266603, + "p90_ms": 0.3770292739868164, + "p99_ms": 0.3771307670593262, + "min_ms": 0.37650241851806643, + "max_ms": 0.37715423583984375, + "max_abs_error": 7.62939453125e-06, + "mean_abs_error": 5.719037465823362e-10, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_13b", + "m": 1, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.38362161636352543, + "p90_ms": 0.3842551879882813, + "p99_ms": 0.3846221073150635, + "min_ms": 0.3832793426513672, + "max_ms": 0.3846591949462891 + }, + { + "shape": "llama2_13b", + "m": 1, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.37882415771484373, + "p90_ms": 0.3790812721252441, + "p99_ms": 0.379168850326538, + "min_ms": 0.37874847412109375, + "max_ms": 0.3791718292236328 + }, + { + "shape": "llama2_13b", + "m": 1, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.37637712478637697, + "p90_ms": 0.37661954116821295, + "p99_ms": 0.37665434379577634, + "min_ms": 0.3763401412963867, + "max_ms": 0.37666175842285154, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 8.968468634407145e-09, + "cosine_similarity": 0.9999999403953552 + }, + { + "shape": "llama2_13b", + "m": 2, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.4020016098022461, + "p90_ms": 0.4023199844360351, + "p99_ms": 0.4038409591674805, + "min_ms": 0.4017168045043945, + "max_ms": 0.40417888641357425 + }, + { + "shape": "llama2_13b", + "m": 2, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.38182687759399414, + "p90_ms": 0.3819195289611816, + "p99_ms": 0.38201279945373534, + "min_ms": 0.38169055938720703, + "max_ms": 0.38203231811523436 + }, + { + "shape": "llama2_13b", + "m": 2, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.3767435264587402, + "p90_ms": 0.3768018035888672, + "p99_ms": 0.37682718887329103, + "min_ms": 0.3766495895385742, + "max_ms": 0.3768307113647461, + "max_abs_error": 6.103515625e-05, + "mean_abs_error": 4.1206579481922745e-08, + "cosine_similarity": 0.9999999403953552 + }, + { + "shape": "llama2_13b", + "m": 2, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.39856719970703125, + "p90_ms": 0.3985970420837402, + "p99_ms": 0.398642240524292, + "min_ms": 0.3985116958618164, + "max_ms": 0.39864959716796877 + }, + { + "shape": "llama2_13b", + "m": 2, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.37936304092407225, + "p90_ms": 0.3794149665832519, + "p99_ms": 0.37948143386840816, + "min_ms": 0.37932415008544923, + "max_ms": 0.37949310302734374 + }, + { + "shape": "llama2_13b", + "m": 2, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.3766364860534668, + "p90_ms": 0.37671642303466796, + "p99_ms": 0.3768056217193604, + "min_ms": 0.376580810546875, + "max_ms": 0.37680641174316404, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 3.7280941000972234e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_13b", + "m": 4, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.40298208236694333, + "p90_ms": 0.4032246360778809, + "p99_ms": 0.4032392993927002, + "min_ms": 0.4027222442626953, + "max_ms": 0.40323936462402343 + }, + { + "shape": "llama2_13b", + "m": 4, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3828487968444824, + "p90_ms": 0.38294761276245115, + "p99_ms": 0.383015274810791, + "min_ms": 0.3827132797241211, + "max_ms": 0.3830268859863281 + }, + { + "shape": "llama2_13b", + "m": 4, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.3775564956665039, + "p90_ms": 0.37760862350463864, + "p99_ms": 0.37764530067443847, + "min_ms": 0.3774534225463867, + "max_ms": 0.3776499176025391, + "max_abs_error": 6.103515625e-05, + "mean_abs_error": 2.970103984978323e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_13b", + "m": 4, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.3997896003723145, + "p90_ms": 0.39984342193603517, + "p99_ms": 0.40152789726257326, + "min_ms": 0.39969470977783206, + "max_ms": 0.40191646575927736 + }, + { + "shape": "llama2_13b", + "m": 4, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.38064224243164063, + "p90_ms": 0.3806813926696777, + "p99_ms": 0.3806895637512207, + "min_ms": 0.3805731201171875, + "max_ms": 0.38068992614746094 + }, + { + "shape": "llama2_13b", + "m": 4, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.37741743087768553, + "p90_ms": 0.37746372222900393, + "p99_ms": 0.3775391159057617, + "min_ms": 0.37737342834472654, + "max_ms": 0.3775459289550781, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 3.865248032752788e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_13b", + "m": 8, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.40539743423461916, + "p90_ms": 0.4057235336303711, + "p99_ms": 0.4057915630340576, + "min_ms": 0.40502113342285156, + "max_ms": 0.40579326629638673 + }, + { + "shape": "llama2_13b", + "m": 8, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.38671920776367186, + "p90_ms": 0.38677266693115236, + "p99_ms": 0.3868189090728759, + "min_ms": 0.38659774780273437, + "max_ms": 0.3868252944946289 + }, + { + "shape": "llama2_13b", + "m": 8, + "n": 13824, + "k": 5120, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.38127904891967773, + "p90_ms": 0.38139458847045893, + "p99_ms": 0.381422885131836, + "min_ms": 0.38113887786865236, + "max_ms": 0.3814233779907227, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 3.909403645252496e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "llama2_13b", + "m": 8, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.40194591522216794, + "p90_ms": 0.4020618400573731, + "p99_ms": 0.4039327465057373, + "min_ms": 0.4017961502075195, + "max_ms": 0.40436065673828125 + }, + { + "shape": "llama2_13b", + "m": 8, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.38412879943847655, + "p90_ms": 0.38419579696655276, + "p99_ms": 0.3842569358825684, + "min_ms": 0.3840243148803711, + "max_ms": 0.38426593780517576 + }, + { + "shape": "llama2_13b", + "m": 8, + "n": 13824, + "k": 5120, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.38155153274536135, + "p90_ms": 0.38163488388061517, + "p99_ms": 0.38167350425720215, + "min_ms": 0.3813455963134766, + "max_ms": 0.38167648315429686, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 3.9676141483369065e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "gpt_neox_20b", + "m": 1, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5450468826293945, + "p90_ms": 0.5453251342773437, + "p99_ms": 0.5453797912597657, + "min_ms": 0.5446566390991211, + "max_ms": 0.5453884887695313 + }, + { + "shape": "gpt_neox_20b", + "m": 1, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.539635829925537, + "p90_ms": 0.5397435340881348, + "p99_ms": 0.5398252517700195, + "min_ms": 0.539441909790039, + "max_ms": 0.5398271942138672 + }, + { + "shape": "gpt_neox_20b", + "m": 1, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.534305591583252, + "p90_ms": 0.5343838233947754, + "p99_ms": 0.5344338283538819, + "min_ms": 0.5342259216308594, + "max_ms": 0.534441909790039, + "max_abs_error": 9.313225746154785e-10, + "mean_abs_error": 1.1368683772161603e-13, + "cosine_similarity": 1.0 + }, + { + "shape": "gpt_neox_20b", + "m": 1, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5426737594604493, + "p90_ms": 0.5432885284423827, + "p99_ms": 0.5433247798919678, + "min_ms": 0.5422963333129883, + "max_ms": 0.543331527709961 + }, + { + "shape": "gpt_neox_20b", + "m": 1, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5371433639526367, + "p90_ms": 0.5371904449462891, + "p99_ms": 0.5372661853790284, + "min_ms": 0.5370924758911133, + "max_ms": 0.5372832107543946 + }, + { + "shape": "gpt_neox_20b", + "m": 1, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.5341799926757813, + "p90_ms": 0.5342470932006835, + "p99_ms": 0.5343974102020264, + "min_ms": 0.5341635131835938, + "max_ms": 0.5344284820556641, + "max_abs_error": 1.52587890625e-05, + "mean_abs_error": 9.683844837127253e-10, + "cosine_similarity": 1.0 + }, + { + "shape": "gpt_neox_20b", + "m": 2, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.6001068687438964, + "p90_ms": 0.6003848686218262, + "p99_ms": 0.6019900913238525, + "min_ms": 0.5996246337890625, + "max_ms": 0.6023344039916992 + }, + { + "shape": "gpt_neox_20b", + "m": 2, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5404124832153321, + "p90_ms": 0.540488510131836, + "p99_ms": 0.5406232093811034, + "min_ms": 0.5402438354492187, + "max_ms": 0.5406422424316406 + }, + { + "shape": "gpt_neox_20b", + "m": 2, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.5352577590942382, + "p90_ms": 0.5353435173034669, + "p99_ms": 0.5354927185058593, + "min_ms": 0.535175666809082, + "max_ms": 0.5355084609985351, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 4.450583901416394e-08, + "cosine_similarity": 0.9999999403953552 + }, + { + "shape": "gpt_neox_20b", + "m": 2, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5947591972351074, + "p90_ms": 0.5948592681884766, + "p99_ms": 0.594926672744751, + "min_ms": 0.5946604919433593, + "max_ms": 0.5949369430541992 + }, + { + "shape": "gpt_neox_20b", + "m": 2, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5379891204833984, + "p90_ms": 0.5380551033020019, + "p99_ms": 0.5381992748260498, + "min_ms": 0.5379257583618164, + "max_ms": 0.5382329559326172 + }, + { + "shape": "gpt_neox_20b", + "m": 2, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.5350902366638184, + "p90_ms": 0.5352445259094238, + "p99_ms": 0.5355390624999999, + "min_ms": 0.5350313568115235, + "max_ms": 0.535588493347168, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 4.2594763272063574e-08, + "cosine_similarity": 0.9999999403953552 + }, + { + "shape": "gpt_neox_20b", + "m": 4, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.6013788795471191, + "p90_ms": 0.6015746307373047, + "p99_ms": 0.6031078140258789, + "min_ms": 0.6008025741577149, + "max_ms": 0.6033919906616211 + }, + { + "shape": "gpt_neox_20b", + "m": 4, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5416942405700684, + "p90_ms": 0.5418385696411133, + "p99_ms": 0.5422770408630371, + "min_ms": 0.541517105102539, + "max_ms": 0.5422998428344726 + }, + { + "shape": "gpt_neox_20b", + "m": 4, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.5363808059692383, + "p90_ms": 0.5365763931274414, + "p99_ms": 0.5367980319976806, + "min_ms": 0.5362617492675781, + "max_ms": 0.5368262481689453, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 4.444635948175346e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "gpt_neox_20b", + "m": 4, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5961182403564453, + "p90_ms": 0.5962417907714843, + "p99_ms": 0.5963226295471191, + "min_ms": 0.5959455871582031, + "max_ms": 0.5963398361206055 + }, + { + "shape": "gpt_neox_20b", + "m": 4, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5394929504394531, + "p90_ms": 0.5395317535400391, + "p99_ms": 0.5395616775512696, + "min_ms": 0.5394367980957031, + "max_ms": 0.5395641708374024 + }, + { + "shape": "gpt_neox_20b", + "m": 4, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.5360657691955566, + "p90_ms": 0.5361173934936524, + "p99_ms": 0.5362199954986573, + "min_ms": 0.5360326385498047, + "max_ms": 0.5362326431274415, + "max_abs_error": 6.103515625e-05, + "mean_abs_error": 3.5078787874454065e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "gpt_neox_20b", + "m": 8, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.6048718452453614, + "p90_ms": 0.6054306907653808, + "p99_ms": 0.6056329353332519, + "min_ms": 0.6031452941894532, + "max_ms": 0.6056419372558594 + }, + { + "shape": "gpt_neox_20b", + "m": 8, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.546855525970459, + "p90_ms": 0.547022159576416, + "p99_ms": 0.5472646896362305, + "min_ms": 0.546572494506836, + "max_ms": 0.5472828674316407 + }, + { + "shape": "gpt_neox_20b", + "m": 8, + "n": 16384, + "k": 6144, + "mode": "eager", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.5408721733093261, + "p90_ms": 0.5412569694519044, + "p99_ms": 0.5412769783020019, + "min_ms": 0.5406067276000976, + "max_ms": 0.5412803268432618, + "max_abs_error": 0.0001220703125, + "mean_abs_error": 4.3647034431160137e-08, + "cosine_similarity": 1.0 + }, + { + "shape": "gpt_neox_20b", + "m": 8, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "torch", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5993281555175781, + "p90_ms": 0.6003465003967284, + "p99_ms": 0.6005832420349121, + "min_ms": 0.5986201477050781, + "max_ms": 0.6006355285644531 + }, + { + "shape": "gpt_neox_20b", + "m": 8, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "gemv_chain", + "cuda_kernel_launches_per_call": 4.0, + "median_ms": 0.5442900657653809, + "p90_ms": 0.5444753074645996, + "p99_ms": 0.5444818511962891, + "min_ms": 0.5440256118774414, + "max_ms": 0.5444831848144531 + }, + { + "shape": "gpt_neox_20b", + "m": 8, + "n": 16384, + "k": 6144, + "mode": "graph", + "implementation": "fused", + "cuda_kernel_launches_per_call": 1.0, + "median_ms": 0.5412859344482421, + "p90_ms": 0.5415282936096191, + "p99_ms": 0.5415708415985108, + "min_ms": 0.5411414337158204, + "max_ms": 0.5415788650512695, + "max_abs_error": 0.000244140625, + "mean_abs_error": 5.405685854498188e-08, + "cosine_similarity": 1.0 + } + ] +} diff --git a/docs/developer/cuda_kernels.md b/docs/developer/cuda_kernels.md index b2dbd20c..1d66c3ec 100644 --- a/docs/developer/cuda_kernels.md +++ b/docs/developer/cuda_kernels.md @@ -1,9 +1,9 @@ # CUDA Kernels AstrAI includes optional custom CUDA kernels for attention, rotary embedding, -BF16 GEMV, and FP8 GEMM. These are built when `nvcc` is available and CUDA is -detected. BF16 GEMV is directly callable and can be selected by the guarded -model linear dispatcher described below. +BF16 GEMV/SwiGLU, and FP8 GEMM. These are built when `nvcc` is available and +CUDA is detected. BF16 GEMV and SwiGLU are directly callable and can be +selected by guarded model dispatchers described below. ## Overview @@ -15,6 +15,7 @@ model linear dispatcher described below. | `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) | | `rotary_emb` | `rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) | | `bf16_gemv` | `gemv/bf16_gemv.cu` | M=1..8 BF16 linear with FP32 accumulation (sm_80+) | +| `bf16_swiglu` | `gemv/bf16_swiglu.cu` | Fused M=1..8 BF16 up/gate projections and SwiGLU epilogue (sm_80+) | | `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) | ### BF16 GEMV primitive @@ -55,6 +56,30 @@ remain on PyTorch. Use mode `1` only for explicit A/B runs outside this table. The primitive remains directly callable and deliberately has no internal `F.linear` fallback. The model-level backend owns fallback and dispatch policy. +### BF16 SwiGLU primitive + +`astrai.extension.bf16_swiglu(x, up_weight, gate_weight)` fuses the two dense +MLP projections with `up * silu(gate)` into one CUDA launch for contiguous +BF16 inputs with `M` in `[1, 8]` and K divisible by 8. It preserves the BF16 +rounding boundaries of the two projection outputs, SiLU output, and final +product while accumulating dot products in FP32. + +The kernel contains two output-row tilings. A CTA-reuse path reads each up/gate +weight chunk once and applies it to all M rows. The native AstrAI 1B shape +`(N,K)=(6912,1536)` uses one warp per decode row for M=2/4/8; on L20 this +removes the shared reductions and barrier and reduces M=4 CUDA-Graph latency +from 0.0324 ms to 0.0181 ms. Wider LLaMA/GPT-NeoX matrices keep CTA reuse, +because duplicating their weight reads across row warps regressed 1.3-4.2%. + +Dense `MLP` modules route through the SwiGLU backend. `ASTRAI_SWIGLU=0` keeps +the unfused linear backend, and `1` explicitly forces the fused primitive. +`auto` is the default but currently has no enabled bands: although direct +errors are small (maximum absolute error at most 2.4e-4 in the L20 matrix), +the different FP32 reduction order changed greedy checkpoint output for +M=1/2/4. Automatic dispatch therefore remains numerically identical to the +existing path. See [the benchmark protocol](./swiglu_benchmark.md) for raw +operator, engine, and checkpoint evidence. + Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist: | Variant | File | Optimization | @@ -245,10 +270,12 @@ astrai/extension/ │ ├── attention.py # Stateless attention kernel wrappers │ ├── rotary.py # Stateless rotary kernel wrapper │ ├── gemv.py # Stateless BF16 GEMV primitive +│ ├── swiglu.py # Stateless fused BF16 SwiGLU primitive │ └── fp8.py # Stateless FP8 primitives (custom_op) ├── fp8.py # FP8 strategy layer (fp8_autocast, recipes) └── backend/ ├── attention.py # Backend selection, KV cache I/O, and fallback + ├── swiglu.py # Inference-only fused/unfused SwiGLU policy └── rotary.py # Per-call CUDA/torch rotary dispatch ``` diff --git a/docs/developer/swiglu_benchmark.md b/docs/developer/swiglu_benchmark.md new file mode 100644 index 00000000..0cc3d495 --- /dev/null +++ b/docs/developer/swiglu_benchmark.md @@ -0,0 +1,63 @@ +# Fused SwiGLU benchmark + +`scripts/tools/benchmark_swiglu.py` compares the directly callable fused BF16 +SwiGLU primitive with both `F.linear` and the existing two-GEMV chain. It covers +the native AstrAI 1B MLP plus LLaMA 2 7B/13B, LLaMA 3 8B, and GPT-NeoX 20B +up/gate shapes at M=1/2/4/8 in eager and CUDA Graph modes. + +```bash +CUDA_VISIBLE_DEVICES=0 python scripts/tools/benchmark_swiglu.py \ + --output results/swiglu.json \ + --markdown-output results/swiglu.md \ + --m-values 1,2,4,8 --mode both \ + --warmup 20 --iterations 100 --trials 10 +``` + +Each trial uses A-B-C-C-B-A ordering to balance clock, cache, and temperature +drift. The generated JSON records every timing sample, p50/p90/p99, CUDA launch +count, maximum/mean absolute error, and cosine similarity. The checked-in L20 +raw run is `docs/benchmarks/swiglu_l20_sm89.json`. + +## L20 findings + +Hardware was one NVIDIA L20 (sm_89), PyTorch 2.11.0+cu128, CUDA 12.8. The +existing GPU5 inference service remained resident (15.4 GiB) but idle at the +sampling boundaries; no process or container was stopped. + +For AstrAI 1B `(N,K)=(6912,1536)`, CUDA Graph medians were: + +| M | torch (ms) | GEMV chain (ms) | fused (ms) | vs best unfused | +|---:|---:|---:|---:|---:| +| 1 | 0.02564 | 0.02298 | 0.01375 | +67.13% | +| 2 | 0.02484 | 0.02628 | 0.01416 | +75.40% | +| 4 | 0.02507 | 0.03839 | 0.01806 | +38.82% | +| 8 | 0.02563 | 0.07007 | 0.03339 | -23.24% | + +The wide traditional shapes are weight-bandwidth dominated. CTA reuse keeps +the fused primitive within roughly -1.2% to +0.9% of the best unfused chain, +so none is eligible for automatic selection. This negative crossover is kept +in the raw evidence rather than hidden by a favorable subset. + +The real 24-layer AstrAI checkpoint was then run through `InferenceEngine`, +including scheduler, sampling, and CUDA Graph. A-B-B-A medians from the raw +log in `docs/benchmarks/swiglu_engine_l20_sm89.txt` were: + +| Batch | unfused (ms/step) | forced fused (ms/step) | throughput gain | +|---:|---:|---:|---:| +| 1 | 4.125 | 3.925 | +5.10% | +| 2 | 4.245 | 4.055 | +4.69% | +| 4 | 4.475 | 4.305 | +3.95% | + +## Dispatch decision + +Direct correctness stayed close (`max_abs <= 2.4e-4`, cosine approximately +1.0), but deterministic greedy generations changed at M=1, M=2, and M=4. The +M=1/2 hash pairs are preserved in +`docs/benchmarks/swiglu_greedy_m1_m2_l20_sm89.txt`, and the M=4 pair is in +`docs/benchmarks/swiglu_greedy_m4_l20_sm89.txt`. + +For that reason no SM89 shape is enabled in `auto`. The default path stays on +the existing unfused linear backend, including any independently qualified +GEMV dispatch. `ASTRAI_SWIGLU=1` remains an explicit benchmark/experimentation +switch for callers that accept normal BF16 reduction-order variation. A future +automatic band must repeat both the performance and checkpoint-output gates. diff --git a/scripts/tools/benchmark_swiglu.py b/scripts/tools/benchmark_swiglu.py new file mode 100644 index 00000000..812b9d91 --- /dev/null +++ b/scripts/tools/benchmark_swiglu.py @@ -0,0 +1,325 @@ +"""Benchmark fused BF16 SwiGLU against torch and unfused GEMV 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_gemv, 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)), + "gemv_chain": lambda: ( + bf16_gemv(x, up_weight) * F.silu(bf16_gemv(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 | GEMV chain ms | fused ms | " + "vs best unfused | fused kernels | max abs | cosine |", + "|---|---:|---|---:|---:|---:|---:|---:|---:|---:|", + ] + for shape, m, mode in cases: + torch_item = by_case[(shape, m, mode, "torch")] + gemv_item = by_case[(shape, m, mode, "gemv_chain")] + fused_item = by_case[(shape, m, mode, "fused")] + best = min(torch_item["median_ms"], gemv_item["median_ms"]) + improvement = (best / fused_item["median_ms"] - 1) * 100 + lines.append( + f"| {shape} | {m} | {mode} | {torch_item['median_ms']:.5f} | " + f"{gemv_item['median_ms']:.5f} | {fused_item['median_ms']:.5f} | " + f"{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_gemv") or not is_available("bf16_swiglu"): + raise click.ClickException("built bf16_gemv and bf16_swiglu are required") + shapes = tuple(parse_shape(value) for value in shape_values) or DEFAULT_SHAPES + m_values_parsed = parse_positive_ints(m_values) + if any(m > 8 for m in m_values_parsed): + raise click.BadParameter("fused primitive supports M up to 8") + modes = ("eager", "graph") if mode == "both" else (mode,) + torch.manual_seed(seed) + torch.cuda.manual_seed_all(seed) + + results = [] + with torch.inference_mode(): + for shape in shapes: + for m in m_values_parsed: + for current_mode in modes: + click.echo( + f"{shape.name}: M={m} N={shape.n} K={shape.k} {current_mode}" + ) + results.extend( + benchmark_case( + shape, + m, + current_mode, + warmup=warmup, + iterations=iterations, + trials=trials, + ) + ) + torch.cuda.empty_cache() + + payload: dict[str, object] = { + "metadata": device_metadata(), + "settings": { + "warmup": warmup, + "iterations": iterations, + "trials": trials, + "seed": seed, + "order": "A-B-C-C-B-A", + }, + "results": results, + } + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(json.dumps(payload, indent=2) + "\n") + if markdown_output is not None: + markdown_output.parent.mkdir(parents=True, exist_ok=True) + markdown_output.write_text(render_markdown(payload)) + + +if __name__ == "__main__": + benchmark_command() diff --git a/setup.py b/setup.py index 0b4313fd..baa38a05 100644 --- a/setup.py +++ b/setup.py @@ -122,6 +122,7 @@ def run(self): "attn_paged_decode", "attn_paged_prefill", "bf16_gemv", + "bf16_swiglu", "rotary_emb", ) missing = [name for name in required if not any(lib_dir.glob(f"{name}.*.so"))] diff --git a/tests/extension/test_swiglu.py b/tests/extension/test_swiglu.py new file mode 100644 index 00000000..04165b77 --- /dev/null +++ b/tests/extension/test_swiglu.py @@ -0,0 +1,99 @@ +import pytest +import torch +import torch.nn.functional as F + +from astrai.extension import bf16_swiglu, is_available + +SWIGLU_AVAILABLE = ( + torch.cuda.is_available() + and is_available("bf16_swiglu") + and torch.cuda.get_device_capability() >= (8, 0) +) +skip_no_swiglu = pytest.mark.skipif( + not SWIGLU_AVAILABLE, + reason="BF16 SwiGLU requires a built kernel and compute capability 8.0+", +) + + +def reference_swiglu(x, up_weight, gate_weight): + return F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)) + + +@skip_no_swiglu +@pytest.mark.parametrize("m", [1, 2, 4, 8]) +@pytest.mark.parametrize("n,k", [(6912, 1536), (4096, 4096), (11008, 4096)]) +def test_bf16_swiglu_matches_common_dense_mlp_shapes(m, n, k): + torch.manual_seed(37 + m) + x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16) * 0.1 + up_weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5) + gate_weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * (k**-0.5) + actual = bf16_swiglu(x, up_weight, gate_weight) + expected = reference_swiglu(x, up_weight, gate_weight) + assert actual.shape == (m, n) + torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01) + + +@skip_no_swiglu +def test_bf16_swiglu_preserves_vector_shape(): + x = torch.randn(1536, device="cuda", dtype=torch.bfloat16) * 0.1 + up_weight = torch.randn(256, 1536, device="cuda", dtype=torch.bfloat16) * 0.02 + gate_weight = torch.randn_like(up_weight) * 0.02 + actual = bf16_swiglu(x, up_weight, gate_weight) + expected = reference_swiglu(x, up_weight, gate_weight) + assert actual.shape == (256,) + torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01) + + +@skip_no_swiglu +def test_bf16_swiglu_uses_current_stream_and_cuda_graph(): + torch.manual_seed(43) + x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1 + up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02 + gate_weight = torch.randn_like(up_weight) * 0.02 + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + for _ in range(3): + bf16_swiglu(x, up_weight, gate_weight) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = bf16_swiglu(x, up_weight, gate_weight) + x.copy_(torch.randn_like(x) * 0.1) + graph.replay() + stream.synchronize() + expected = reference_swiglu(x, up_weight, gate_weight) + torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01) + + +@skip_no_swiglu +@pytest.mark.parametrize( + "make_args,error", + [ + ( + lambda: ( + torch.randn(9, 16, device="cuda", dtype=torch.bfloat16), + torch.randn(8, 16, device="cuda", dtype=torch.bfloat16), + torch.randn(8, 16, device="cuda", dtype=torch.bfloat16), + ), + "M must", + ), + ( + lambda: ( + torch.randn(2, 15, device="cuda", dtype=torch.bfloat16), + torch.randn(8, 15, device="cuda", dtype=torch.bfloat16), + torch.randn(8, 15, device="cuda", dtype=torch.bfloat16), + ), + "divisible by 8", + ), + ( + lambda: ( + torch.randn(2, 16, device="cuda", dtype=torch.bfloat16), + torch.randn(8, 16, device="cuda", dtype=torch.bfloat16), + torch.randn(7, 16, device="cuda", dtype=torch.bfloat16), + ), + "identical shapes", + ), + ], +) +def test_bf16_swiglu_rejects_unsupported_inputs(make_args, error): + with pytest.raises(RuntimeError, match=error): + bf16_swiglu(*make_args()) diff --git a/tests/extension/test_swiglu_dispatch.py b/tests/extension/test_swiglu_dispatch.py new file mode 100644 index 00000000..19b85085 --- /dev/null +++ b/tests/extension/test_swiglu_dispatch.py @@ -0,0 +1,94 @@ +import logging + +import pytest +import torch +import torch.nn.functional as F + +from astrai.extension import is_available, swiglu +from astrai.model.components.mlp import MLP + +SWIGLU_AVAILABLE = ( + torch.cuda.is_available() + and is_available("bf16_swiglu") + and torch.cuda.get_device_capability() >= (8, 0) +) +skip_no_swiglu = pytest.mark.skipif( + not SWIGLU_AVAILABLE, + reason="BF16 SwiGLU requires a built kernel and compute capability 8.0+", +) + + +def reference_swiglu(x, up_weight, gate_weight): + return F.linear(x, up_weight) * F.silu(F.linear(x, gate_weight)) + + +def test_cpu_and_training_calls_fall_back_with_gradients(monkeypatch): + monkeypatch.setenv("ASTRAI_SWIGLU", "1") + x = torch.randn(2, 8, requires_grad=True) + up_weight = torch.randn(4, 8, requires_grad=True) + gate_weight = torch.randn(4, 8, requires_grad=True) + actual = swiglu(x, up_weight, gate_weight) + expected = reference_swiglu(x, up_weight, gate_weight) + torch.testing.assert_close(actual, expected) + actual.sum().backward() + assert x.grad is not None + assert up_weight.grad is not None + assert gate_weight.grad is not None + + +def test_invalid_mode_warns_and_uses_auto(monkeypatch, caplog): + monkeypatch.setenv("ASTRAI_SWIGLU", "invalid-test-mode") + with caplog.at_level(logging.WARNING): + actual = swiglu(torch.randn(2, 8), torch.randn(4, 8), torch.randn(4, 8)) + assert actual.shape == (2, 4) + assert "using auto" in caplog.text + + +def test_mlp_routes_through_swiglu_backend(monkeypatch): + sentinel = torch.randn(2, 4) + + def fake_swiglu(x, up_weight, gate_weight): + assert x.shape == (2, 3) + assert up_weight.shape == gate_weight.shape == (4, 3) + return sentinel + + monkeypatch.setattr("astrai.model.components.mlp.swiglu", fake_swiglu) + layer = MLP(3, 4) + output = layer(torch.randn(2, 3)) + assert output["hidden_states"].shape == (2, 3) + + +@skip_no_swiglu +def test_mode_zero_disables_fused_kernel(monkeypatch): + monkeypatch.setenv("ASTRAI_SWIGLU", "0") + x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1 + up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02 + gate_weight = torch.randn_like(up_weight) * 0.02 + with torch.no_grad(): + actual = swiglu(x, up_weight, gate_weight) + expected = reference_swiglu(x, up_weight, gate_weight) + torch.testing.assert_close(actual, expected) + + +@skip_no_swiglu +def test_mode_one_forces_supported_shape(monkeypatch): + monkeypatch.setenv("ASTRAI_SWIGLU", "1") + x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) * 0.1 + up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) * 0.02 + gate_weight = torch.randn_like(up_weight) * 0.02 + with torch.no_grad(): + actual = swiglu(x, up_weight, gate_weight) + expected = reference_swiglu(x, up_weight, gate_weight) + torch.testing.assert_close(actual, expected, rtol=0.03, atol=0.01) + + +@skip_no_swiglu +def test_auto_falls_back_until_shape_is_qualified(monkeypatch): + monkeypatch.setenv("ASTRAI_SWIGLU", "auto") + x = torch.randn(4, 1536, device="cuda", dtype=torch.bfloat16) + up_weight = torch.randn(6912, 1536, device="cuda", dtype=torch.bfloat16) + gate_weight = torch.randn_like(up_weight) + with torch.no_grad(): + actual = swiglu(x, up_weight, gate_weight) + expected = reference_swiglu(x, up_weight, gate_weight) + torch.testing.assert_close(actual, expected)