0z5a
d4a292b36b
- deepen common-shape BF16 GEMV tuning with warp-row tiling for LLaMA/Qwen2/GPT-NeoX/OPT decode projections - add fused BF16 up/gate SwiGLU CUDA primitive with ASTRAI_SWIGLU=0/1/auto dispatch - keep the unfused linear backend as the default path; auto enables no shape until per-architecture checkpoint gates pass - fall back to the linear/torch chain when kernels are absent, on CPU, in training, or outside supported M/K/dtype shapes - add gemv/swiglu benchmark scripts, dispatch and parity tests, and kernel documentation Benchmark: NVIDIA L20 (sm_89), CUDA 12.8, PyTorch 2.11.0+cu128, idle GPU. AstrAI 1B config (24 layers, hidden 1536, vocab 100000), BF16, prompt 128, 32 greedy decode tokens, CUDA graphs enabled, A/B in separate interleaved processes (3 rounds, 8 trials each, medians). Default vs ASTRAI_SWIGLU=1 per generate call: batch 1 134.8->129.1 ms (+4.44%), batch 2 136.2->130.9 ms (+4.06%), batch 4 145.5->140.3 ms (+3.66%). Greedy output identical at batch 1, differs at batch 2/4, so auto stays unfused by default; kernelless fallback verified bit-identical greedy.
110 lines
3.8 KiB
Python
110 lines
3.8 KiB
Python
"""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"]
|