ViperEkura
3e67b4f88d
- Single-kernel rotary embedding (cos/sin lookup + rotation) replaces PyTorch complex-multiply path (3 kernel launches + f32 upcast per call) - RotaryEmbedding now stores cos_table/sin_table and returns (cos, sin) f32 tuple instead of a complex tensor - apply_rotary_emb in rotary_backend.py auto-dispatches: CUDA kernel if available, else torch complex-multiply fallback; backend-agnostic (both attention backends benefit) - Kernel: 256-thread blocks, grid-stride loop, vectorized __nv_bfloat162 load/store, f32 compute, bf16 out - Standalone kernel 6-9x faster than torch across decode/prefill shapes, max diff 0 (decode) to 3e-2 (large prefill, bf16) - Benchmark (L20, bf16, CUDA backend): B=1 9.48->7.25ms (+31%), B=4 10.73->7.67ms (+40%), B=8 10.77->7.81ms (+38%), B=16 10.79->7.83ms (+38%)
76 lines
2.0 KiB
Python
76 lines
2.0 KiB
Python
from pathlib import Path
|
|
|
|
|
|
def cuda_toolkit_version() -> tuple[int, int] | None:
|
|
"""Return ``(major, minor)`` of the nvcc on PATH, or ``None``.
|
|
|
|
Used by ``setup.py`` to detect nvcc/torch CUDA version mismatches
|
|
(e.g. nvcc 13.0 with a cu128 torch wheel) which cause cryptic ABI errors.
|
|
"""
|
|
import shutil
|
|
import subprocess
|
|
|
|
nvcc = shutil.which("nvcc")
|
|
if nvcc is None:
|
|
return None
|
|
try:
|
|
out = subprocess.check_output(
|
|
[nvcc, "--version"], stderr=subprocess.STDOUT, text=True
|
|
)
|
|
for line in out.splitlines():
|
|
if "release" in line:
|
|
ver = line.split("release")[1].split(",")[0].strip()
|
|
major, minor = ver.split(".")
|
|
return (int(major), int(minor))
|
|
except Exception:
|
|
pass
|
|
return None
|
|
|
|
|
|
def _arch_flags() -> list[str]:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
cap = torch.cuda.get_device_capability()
|
|
else:
|
|
cap = (8, 0)
|
|
ver = f"{cap[0]}{cap[1]}"
|
|
flags = [f"-gencode=arch=compute_{ver},code=sm_{ver}"]
|
|
# tensor-core mma path (mma.sync.m16n8k16.bf16) requires sm_80+; decide the
|
|
# kernel dispatch at build time via this define rather than at runtime.
|
|
if cap[0] < 8:
|
|
flags.append("-DASTRAI_NO_MMA")
|
|
return flags
|
|
|
|
|
|
_kernels_dir = Path("csrc/kernels")
|
|
REGISTRY: dict[str, dict] = {}
|
|
|
|
CXX_FLAGS = ["-O3", "-funroll-loops"]
|
|
NVCC_FLAGS = [
|
|
"-O3",
|
|
"--expt-relaxed-constexpr",
|
|
"--use_fast_math",
|
|
"--ptxas-options=-O3,-v",
|
|
"--extra-device-vectorization",
|
|
"--threads=8",
|
|
]
|
|
|
|
|
|
def register(name: str, sources: list[str] | None = None, **kwargs):
|
|
if sources is None:
|
|
sources = [str(_kernels_dir / f"{name}.cu")]
|
|
REGISTRY[name] = {
|
|
"sources": sources,
|
|
"cxx_flags": [*CXX_FLAGS],
|
|
"nvcc_flags": [*NVCC_FLAGS, *_arch_flags()],
|
|
"extra_link_args": kwargs.pop("extra_link_args", []),
|
|
**kwargs,
|
|
}
|
|
|
|
|
|
register("attn_decode")
|
|
register("attn_prefill")
|
|
register("attn_paged_decode")
|
|
register("rotary_emb")
|