ViperEkura
1798474316
- split-K removed entirely: tiled kernel walks K in one pass, no partials/semas workspace, no memset, single launch per call
- skinny GEMM (M<=8) dispatch table replaces the hand-written switch
- shape-driven four-family table replaces plan_gemm: wide-N (n>=4096) default {16,64,64,3,128} with BM=32 at M>16; narrow-N deep-K rings {16,32,256,2,64} while the grid fits one wave, {16,32,128,2,64} past it
- narrow-N is K-serial: widening the grid measurably does nothing (BN 64->32 ties, doubled m_tiles tie, kv at 4 blocks ties q/o at 24); deeper K chunks win until 72KB smem forces one CTA per SM and past one wave the 2-wave quantization loses to BK=128
- launch-check macros in common/launch.cuh; smem opt-in for the 72KB/60KB rings
- rename kernels/bf16_*.cu to gemm.cu/swiglu.cu; module names unchanged
- Python gate: lm_head (N>32768) falls back to cuBLAS, band narrows to M<=32
- drop the stale per-op benchmark narratives; fold the live numbers into cuda_kernels.md
Benchmark: NVIDIA L20 (sm_89, 92 SMs), CUDA 12.8, bf16, L2-thrash weight rotation, per-call medians at M=16: q/o 9.5us, kv 8.6us, gate/up 33.3us, down 33.7us (down -29% vs prior default). End-to-end 1B decode (gen 128, 3 trials, tokens/s vs cuBLAS): B=1 260 vs 252, B=8 1660 vs 1446, B=16 2464 vs 2437, B=32 3620 vs 3690. Prior split-K dispatch measured B=16 2243 / B=32 3393.
68 lines
2.7 KiB
Python
68 lines
2.7 KiB
Python
"""Inference-only fused SwiGLU selection for dense MLP layers."""
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch import Tensor
|
|
|
|
from astrai.extension.backend.linear import linear
|
|
from astrai.extension.dispatch import env_mode
|
|
from astrai.extension.loader import is_available
|
|
from astrai.extension.ops.swiglu import bf16_swiglu
|
|
|
|
|
|
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 batches 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())
|
|
|
|
|
|
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()
|
|
# The fused kernel reads all streams as uint4; contiguous-but-offset
|
|
# views are routed to the unfused chain instead of failing.
|
|
or (x.data_ptr() & 15) != 0
|
|
or (up_weight.data_ptr() & 15) != 0
|
|
or (gate_weight.data_ptr() & 15) != 0
|
|
or not is_available("bf16_swiglu")
|
|
)
|
|
|
|
|
|
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; ``auto`` (the default) uses
|
|
the fused primitive for decode batches with M in ``{1, ..., 8}``. The
|
|
fused kernel reads x once and covers both projections plus the SiLU
|
|
gate-multiply in a single launch, measured 9-15% faster than the
|
|
unfused chain per MLP call on L20 with L2-thrashing weight rotation.
|
|
"""
|
|
if env_mode("ASTRAI_SWIGLU") != "0" and _swiglu_capable(x, up_weight, gate_weight):
|
|
return _fused_swiglu(x, up_weight, gate_weight)
|
|
return _unfused_swiglu(x, up_weight, gate_weight)
|
|
|
|
|
|
__all__ = ["swiglu"]
|