perf: dispatch linear gemv by decode batch size and unify extension style
- replace the per-shape auto tables in the linear backend with an M-banded rule (M in [2,4] on compute capability 8.0+) that measured at the HBM bandwidth floor across every family, and fold the capability check into the capable guard - drop the unreachable swiglu auto shape-table machinery so both backends share one env-mode ladder via the new dispatch.env_mode helper - add __all__ across extension modules, name the rotary registration records, and unify typing to the typing-module style - rewrite test_linear_dispatch.py around behavioral routing assertions and document the M-banded policy in the developer docs - Benchmark: L20 SM89, Python dispatch overhead 2.9us to 1.5us, auto now covers every projection shape at M in [2,4].
This commit is contained in:
@@ -1,18 +1,19 @@
|
||||
"""CUDA attention kernel wrappers with torch fallback.
|
||||
"""CUDA kernel wrappers, operator dispatch, and backend selection.
|
||||
|
||||
Public API:
|
||||
- ``attn_decode`` — single-query decode attention
|
||||
- ``attn_prefill`` — multi-query prefill attention
|
||||
- ``attn_paged_decode`` — paged decode attention (direct page-table access)
|
||||
- ``AttentionBackend`` — ABC for attention computation strategies
|
||||
- ``TorchNativeBackend`` — default SDPA backend with KV cache I/O
|
||||
- ``CudaBackend`` — CUDA kernel backend with paged decode + prefill
|
||||
- ``attention``, ``linear``, ``swiglu``, ``apply_rotary_emb`` — op
|
||||
families with safe torch fallbacks (see ``astrai.extension.backend``)
|
||||
- ``attn_decode`` / ``attn_prefill`` / ``attn_paged_decode`` /
|
||||
``attn_paged_prefill`` — direct attention kernel wrappers
|
||||
- ``bf16_gemv`` / ``bf16_swiglu`` — directly callable linear/MLP kernels
|
||||
- ``AttentionBackend`` / ``TorchNativeBackend`` / ``CudaBackend`` /
|
||||
``FlashAttnBackend`` — attention backend strategies
|
||||
- ``resolve`` / ``explain`` / ``op_backend`` / ``env_mode`` — the shared
|
||||
operator dispatcher (see ``astrai.extension.dispatch``)
|
||||
|
||||
Layout convention: all q/k/v are ``[batch, seq_len, n_heads, head_dim]``
|
||||
(blhd). Scale is always ``1/sqrt(head_dim)``.
|
||||
|
||||
Each wrapper calls its compiled CUDA kernel directly. Fallback to torch
|
||||
SDPA is handled by the attention backend, not the wrapper functions.
|
||||
(blhd). Scale is always ``1/sqrt(head_dim)``. Wrapper functions call their
|
||||
compiled CUDA kernels directly; fallback is the backend's responsibility.
|
||||
"""
|
||||
|
||||
from astrai.extension.backend import (
|
||||
@@ -36,6 +37,7 @@ from astrai.extension.dispatch import (
|
||||
Resolution,
|
||||
Spec,
|
||||
axis,
|
||||
env_mode,
|
||||
explain,
|
||||
explain_plan,
|
||||
op_backend,
|
||||
@@ -82,6 +84,7 @@ __all__ = [
|
||||
"Resolution",
|
||||
"Spec",
|
||||
"axis",
|
||||
"env_mode",
|
||||
"explain",
|
||||
"explain_plan",
|
||||
"op_backend",
|
||||
|
||||
Reference in New Issue
Block a user