perf: add opt-in fused BF16 SwiGLU

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",

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",
 ]

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"]

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",
 ]

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"]

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}
 

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.

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"
+  }
+}