biondizzle
2649488d13
- Kernel loops over KV tiles internally with running max/sum rescale - SMEM accumulator sOacc[hd] replaces TMEM accumulation across tiles - P is UN-NORMALIZED for multi-tile (exp(s-max), not /sum) - Per KV tile: QK→softmax→PV→TMEM→read→add to sOacc - Final: O = sOacc / running_sum - Single tile (n_kv_tiles=1): same as before, no rescale - Updated CAPI, Python loader, production.py fast path - Added multi-tile test cases (N=256, 512)
268 lines
8.0 KiB
Python
268 lines
8.0 KiB
Python
"""DSV4 FMHA — 6-warp multi-head decode kernel loader.
|
|
|
|
Precompiles the raw CUDA kernel with nvcc (sm_100a) on first use,
|
|
then loads the .so via ctypes. This bypasses torch.utils.cpp_extension
|
|
which compiles with -arch=sm_100 (missing tcgen05 support).
|
|
|
|
Decode-only: T=1, single KV segment (N <= 128).
|
|
Supports MHA, MQA, and GQA attention patterns.
|
|
"""
|
|
|
|
import torch
|
|
import logging
|
|
import os
|
|
import subprocess
|
|
import ctypes
|
|
from typing import Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
KERNEL_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
REPO_ROOT = os.path.normpath(os.path.join(KERNEL_DIR, "..", ".."))
|
|
SOURCE = os.path.join(KERNEL_DIR, "fmha_multihead_capi.cu")
|
|
BUILD_DIR = os.path.join(REPO_ROOT, "build", "fmha_multihead")
|
|
SO_NAME = "libfmha_multihead_decode.so"
|
|
|
|
_lib = None
|
|
_lib_lock = False
|
|
|
|
|
|
def _find_nvcc():
|
|
"""Find nvcc on the system."""
|
|
for c in ["/usr/local/cuda-13.2/bin/nvcc", "/usr/local/cuda/bin/nvcc"]:
|
|
if os.path.isfile(c):
|
|
return c
|
|
# Try PATH
|
|
import shutil
|
|
nvcc = shutil.which("nvcc")
|
|
if nvcc:
|
|
return nvcc
|
|
raise RuntimeError("nvcc not found — required for tcgen05 kernel compilation")
|
|
|
|
|
|
def _ensure_built():
|
|
"""Build the shared library with nvcc if needed. Returns .so path."""
|
|
global _lib
|
|
if _lib is not None:
|
|
return _lib
|
|
|
|
so_path = os.path.join(BUILD_DIR, SO_NAME)
|
|
|
|
# Check if rebuild needed
|
|
need_build = True
|
|
if os.path.isfile(so_path):
|
|
src_mtime = os.path.getmtime(SOURCE)
|
|
for dep in ["fmha_common.cuh", "fmha_umma_desc.cuh",
|
|
"fmha_6warp_multihead.cuh", "fmha_multihead_capi.cu"]:
|
|
dep_path = os.path.join(KERNEL_DIR, dep)
|
|
if os.path.isfile(dep_path):
|
|
src_mtime = max(src_mtime, os.path.getmtime(dep_path))
|
|
need_build = src_mtime > os.path.getmtime(so_path)
|
|
|
|
if not need_build:
|
|
logger.info(f"Using cached {so_path}")
|
|
_lib = ctypes.CDLL(so_path)
|
|
return _lib
|
|
|
|
logger.info(f"Building {SO_NAME} with nvcc (sm_100a)...")
|
|
os.makedirs(BUILD_DIR, exist_ok=True)
|
|
|
|
nvcc = _find_nvcc()
|
|
cmd = [
|
|
nvcc,
|
|
"-std=c++20",
|
|
"-shared",
|
|
"-Xcompiler", "-fPIC",
|
|
"-gencode=arch=compute_100a,code=sm_100a",
|
|
"-gencode=arch=compute_100a,code=compute_100a",
|
|
f"-I{KERNEL_DIR}",
|
|
f"-I{REPO_ROOT}",
|
|
"-O3",
|
|
"--expt-relaxed-constexpr",
|
|
SOURCE,
|
|
"-o", so_path,
|
|
"-lcudart",
|
|
"-lcuda",
|
|
]
|
|
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(f"nvcc compilation failed:\n{result.stderr}")
|
|
if result.stderr:
|
|
logger.debug(f"nvcc warnings:\n{result.stderr}")
|
|
|
|
_lib = ctypes.CDLL(so_path)
|
|
logger.info(f"Built and loaded {so_path}")
|
|
return _lib
|
|
|
|
|
|
def _get_lib():
|
|
"""Get or build the shared library."""
|
|
global _lib_lock
|
|
if _lib is not None:
|
|
return _lib
|
|
if _lib_lock:
|
|
raise RuntimeError("Recursive build")
|
|
_lib_lock = True
|
|
try:
|
|
return _ensure_built()
|
|
finally:
|
|
_lib_lock = False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Kernel launch via ctypes
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def fmha_multihead_decode_raw(
|
|
q: torch.Tensor, # (batch, n_h, 1, hd) BF16, contiguous
|
|
k: torch.Tensor, # (batch, n_kv, N, hd) BF16, contiguous
|
|
v: torch.Tensor, # (batch, n_kv, hd, N) BF16, contiguous
|
|
scale: float,
|
|
n_comp: int,
|
|
swa_len: int,
|
|
is_causal: bool,
|
|
attn_sink: torch.Tensor, # (batch, n_h) FP32 — unused by kernel currently
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Launch the 6-warp multi-head FMHA kernel. Returns (O, LSE).
|
|
|
|
O: (batch, n_h, 1, hd) BF16
|
|
LSE: (batch, n_h, 1) FP32
|
|
"""
|
|
lib = _get_lib()
|
|
|
|
B = q.shape[0]
|
|
n_h = q.shape[1]
|
|
hd = q.shape[3]
|
|
n_kv = k.shape[1]
|
|
N = k.shape[2]
|
|
|
|
assert q.shape[2] == 1, f"Decode requires T=1, got T={q.shape[2]}"
|
|
assert hd in (64, 128, 256), f"Unsupported hd={hd}"
|
|
assert N > 0, f"N must be positive, got N={N}"
|
|
|
|
q_per_kv = n_h // n_kv
|
|
|
|
# GQA: expand K/V to (1, n_h, ...) so head_idx * stride = correct data
|
|
if n_kv < n_h:
|
|
k = k.repeat_interleave(q_per_kv, dim=1)
|
|
v = v.repeat_interleave(q_per_kv, dim=1)
|
|
|
|
# The kernel template has SK_TILE=128 hardcoded in the softmax loop.
|
|
# When N < 128, pad K and V to 128 so the kernel processes zeros for
|
|
# the extra positions (correctly gets zero attention weight after softmax).
|
|
# When N > 128, the kernel loops over KV tiles internally (FA2 online softmax).
|
|
if N < 128:
|
|
pad_len = 128 - N
|
|
k = torch.cat([k,
|
|
torch.zeros(k.shape[0], k.shape[1], pad_len, hd, dtype=torch.bfloat16, device=k.device)], dim=2)
|
|
v = torch.cat([v,
|
|
torch.zeros(v.shape[0], v.shape[1], hd, pad_len, dtype=torch.bfloat16, device=v.device)], dim=3)
|
|
k = k.contiguous()
|
|
v = v.contiguous()
|
|
N = 128
|
|
else:
|
|
k = k.contiguous()
|
|
v = v.contiguous()
|
|
|
|
# Pad N to multiple of 128 for the KV tile loop
|
|
n_kv_tiles = (N + 127) // 128
|
|
N_padded = n_kv_tiles * 128
|
|
if N < N_padded:
|
|
pad_len = N_padded - N
|
|
k = torch.cat([k,
|
|
torch.zeros(k.shape[0], k.shape[1], pad_len, hd, dtype=torch.bfloat16, device=k.device)], dim=2)
|
|
v = torch.cat([v,
|
|
torch.zeros(v.shape[0], v.shape[1], hd, pad_len, dtype=torch.bfloat16, device=v.device)], dim=3)
|
|
k = k.contiguous()
|
|
v = v.contiguous()
|
|
N = N_padded
|
|
|
|
q = q.contiguous()
|
|
|
|
o = torch.zeros(B, n_h, 1, hd, dtype=torch.bfloat16, device=q.device)
|
|
lse = torch.zeros(B, n_h, 1, dtype=torch.float32, device=q.device)
|
|
|
|
# Compute strides in BF16 elements
|
|
q_hs = q.stride(1)
|
|
q_bs = q.stride(0)
|
|
k_hs = k.stride(1)
|
|
k_bs = k.stride(0)
|
|
v_hs = v.stride(1)
|
|
v_bs = v.stride(0)
|
|
o_hs = o.stride(1)
|
|
o_bs = o.stride(0)
|
|
lse_hs = lse.stride(1)
|
|
lse_bs = lse.stride(0)
|
|
|
|
# Call the C API
|
|
ret = lib.fmha_multihead_decode_launch(
|
|
ctypes.c_void_p(q.data_ptr()),
|
|
ctypes.c_void_p(k.data_ptr()),
|
|
ctypes.c_void_p(v.data_ptr()),
|
|
ctypes.c_void_p(o.data_ptr()),
|
|
ctypes.c_void_p(lse.data_ptr()),
|
|
ctypes.c_int(B),
|
|
ctypes.c_int(n_h),
|
|
ctypes.c_int(n_kv),
|
|
ctypes.c_int(N),
|
|
ctypes.c_int(hd),
|
|
ctypes.c_int(q_hs),
|
|
ctypes.c_int(q_bs),
|
|
ctypes.c_int(k_hs),
|
|
ctypes.c_int(k_bs),
|
|
ctypes.c_int(v_hs),
|
|
ctypes.c_int(v_bs),
|
|
ctypes.c_int(o_hs),
|
|
ctypes.c_int(o_bs),
|
|
ctypes.c_int(lse_hs),
|
|
ctypes.c_int(lse_bs),
|
|
ctypes.c_float(scale),
|
|
)
|
|
|
|
if ret != 0:
|
|
raise RuntimeError(f"Kernel launch failed with code {ret}")
|
|
|
|
return o, lse
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Custom op registration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
@torch.library.custom_op("dsv4::fmha_multihead_decode", mutates_args=())
|
|
def fmha_multihead_decode(
|
|
q: torch.Tensor,
|
|
k: torch.Tensor,
|
|
v: torch.Tensor,
|
|
scale: float,
|
|
n_comp: int,
|
|
swa_len: int,
|
|
is_causal: bool,
|
|
attn_sink: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
o, _ = fmha_multihead_decode_raw(
|
|
q, k, v, scale, n_comp, swa_len, is_causal, attn_sink
|
|
)
|
|
return o
|
|
|
|
|
|
@fmha_multihead_decode.register_fake
|
|
def _(q, k, v, scale, n_comp, swa_len, is_causal, attn_sink):
|
|
return torch.empty_like(q)
|
|
|
|
|
|
def fmha_multihead_decode_with_lse(
|
|
q, k, v, scale, n_comp=0, swa_len=0, is_causal=False, attn_sink=None,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
if attn_sink is None:
|
|
attn_sink = torch.zeros(q.shape[0], q.shape[1],
|
|
dtype=torch.float32, device=q.device)
|
|
return fmha_multihead_decode_raw(
|
|
q, k, v, scale, n_comp, swa_len, is_causal, attn_sink
|
|
)
|
|
|
|
|
|
def can_use_6warp_decode(T: int, N: int, hd: int, n_segments: int) -> bool:
|
|
return T == 1 and hd in (64, 128, 256)
|