396 lines
18 KiB
Markdown
396 lines
18 KiB
Markdown
# CUDA Kernels
|
||
|
||
AstrAI includes optional custom CUDA kernels for attention, rotary embedding, and FP8 GEMM. These are built when `nvcc` is available and CUDA is detected, and are dispatched via the `CudaBackend` attention backend, auto-dispatched for rotary, or invoked through the FP8 linear primitives.
|
||
|
||
## Overview
|
||
|
||
| Kernel | File | Description |
|
||
|--------|------|-------------|
|
||
| `attn_decode` | `attention/decode.cu` | GQA decode attention (split-KV) |
|
||
| `attn_prefill` | `attention/prefill.cu` | GQA prefill attention (split-Q) |
|
||
| `attn_paged_decode` | `attention/paged_decode.cu` | Paged KV cache decode attention |
|
||
| `attn_paged_prefill` | `attention/paged_prefill.cu` | Paged KV cache prefill attention (ragged batch) |
|
||
| `rotary_emb` | `rotary/rotary_emb.cu` | Fused rotary embedding (cos/sin lookup + rotation) |
|
||
| `fp8_ops` | `fp8/ops.cu` | FP8 quantization + tensor-core GEMM (sm_89+) |
|
||
|
||
Additionally, optimized `.cuh` variants with tensor-core MMA (Matrix Multiply-Accumulate) exist:
|
||
|
||
| Variant | File | Optimization |
|
||
|---------|------|--------------|
|
||
| Split-KV MMA decode | `attention/decode_split_kv_mma.cuh` | Split KV across warps + MMA (sm_80+) |
|
||
| Split-Q MMA prefill | `attention/prefill_split_q_mma.cuh` | Split Q across warps + MMA (sm_80+) |
|
||
|
||
> The paged and non-paged paths share one kernel body. Prefill is templated on
|
||
> an independent Q schedule (`DenseQSchedule` / `PackedQSchedule`) and KV
|
||
> source (`ContigKV` / `PagedKV`); decode only needs the KV source. There are
|
||
> no separate `attn_paged_*.cuh` files.
|
||
|
||
### Rotary Embedding Kernel
|
||
|
||
The `rotary_emb` kernel (`csrc/kernels/rotary/rotary_emb.cu`) fuses cos/sin lookup and rotation into a single kernel:
|
||
|
||
- One thread per (head, dim-pair), vectorized `__nv_bfloat162` load/store
|
||
- f32 cos/sin input, bf16 compute and output
|
||
- 256-thread blocks, grid-stride loop
|
||
- Auto-dispatched via `apply_rotary_emb` in `astrai/extension/backend/rotary.py` (CUDA when available + inference mode, else torch complex-multiply fallback)
|
||
- No context-manager backend needed — rotary is backend-agnostic, both attention backends benefit
|
||
|
||
Standalone benchmark vs torch complex-multiply (48 calls = 24 layers × q+k): 6-9x faster, max diff 0 (decode) to 3e-2 (large prefill, bf16).
|
||
|
||
### FP8 GEMM / Linear Kernel
|
||
|
||
The `fp8_ops` family (`csrc/kernels/fp8/`) accelerates bf16 linear layers by
|
||
quantizing to FP8 and running tensor-core GEMMs (**requires sm_89+**; fp8
|
||
`mma.sync.m16n8k32` only exists on Ada/Hopper). It follows the same three-layer
|
||
style as attention, but split into **three** files:
|
||
|
||
| File | Role |
|
||
|------|------|
|
||
| `fp8/common.h` | `FP8Format` enum (E4M3/E5M2), `Fp8GemmTraits<Fmt, BlockM, BlockN, K, Stages>`, `FP8Params` POD — no torch |
|
||
| `fp8/gemm.cuh` | pure-CUDA device code: `fp8_quantize_kernel` (BF16→FP8 + amax), `fp8_gemm_kernel` (pre-quantized GEMM, 128×64 CTA / 64×16 warp / 3-stage cp.async) — no torch |
|
||
| `fp8/ops.cu` | binding only: `check_fp8_device` (sm_89+), param packing, launch dispatch, pybind → module `fp8_ops` |
|
||
|
||
Scale semantics follow `torch._scaled_mm` (quantization step size: divide by
|
||
`scale`; the kernel computes the reciprocal internally — the interface never
|
||
takes `*_inv`). `amax` is always returned in the original bf16 domain.
|
||
|
||
Python layer (two levels): `astrai/extension/ops/fp8.py` provides stateless
|
||
primitives (`quantize_bf16` / `mm_fp8` / `linear_forward_fp8` /
|
||
`linear_backward_fp8`) via `torch.library.custom_op`, and
|
||
`astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed /
|
||
dynamic scaling recipes, `fp8_linear_forward/backward` wiring `aten::linear`
|
||
on CUDA). See the FP8 section in `AGENTS.md` for full detail.
|
||
|
||
## Build System
|
||
|
||
### Auto-detection
|
||
|
||
Kernels are built when **both** of these conditions are met:
|
||
1. `nvcc` is available on `PATH`
|
||
2. `torch.cuda.is_available()` returns `True`
|
||
|
||
Unless `CSRC_KERNELS=false` is set explicitly.
|
||
|
||
### Manual build
|
||
|
||
```bash
|
||
# During install
|
||
CSRC_KERNELS=true pip install -e . --no-build-isolation
|
||
|
||
# Rebuild after editing .cu/.cuh files
|
||
CSRC_KERNELS=true python setup.py build_ext --inplace
|
||
# Output: astrai/extension/lib/*.so
|
||
|
||
# Or invoke CMake directly
|
||
cmake -S csrc -B build/cmake \
|
||
-DTORCH_HOME=<site-packages>/torch \
|
||
-DPYTHON_INCLUDE_DIR=<python include> \
|
||
-DPY_SOABI=cpython-312-x86_64-linux-gnu
|
||
cmake --build build/cmake -j 16
|
||
```
|
||
|
||
### Architecture flags
|
||
|
||
`setup.py` passes the GPU compute capability to CMake via `ASTRAI_CUDA_ARCH`. When
|
||
unset, `setup.py` auto-detects the real GPU capability through
|
||
`torch.cuda.get_device_capability()`; the CMake fallback default is `80` (sm_80):
|
||
|
||
- **sm_80+** (Ampere and later): enables the tensor-core MMA path
|
||
(`mma.sync.m16n8k16.bf16` for bf16 attention, `mma.sync.m16n8k32` for FP8).
|
||
- **sm_89+**: required for the FP8 family (`fp8_ops`) — FP8 tensor-core
|
||
instructions only exist on Ada/Hopper and newer. On older architectures,
|
||
CMake emits a warning and skips the `fp8_ops` target so the remaining CUDA
|
||
kernels still build successfully.
|
||
- **`-DASTRAI_NO_MMA`** is a manual escape hatch only — the build never defines
|
||
it automatically. To disable the MMA path, add it to `NVCC_FLAGS` yourself;
|
||
all supported build targets are sm_80+.
|
||
|
||
### Build configuration
|
||
|
||
`csrc/CMakeLists.txt` defines the CUDA extension build:
|
||
|
||
```
|
||
NVCC_FLAGS = -O3 --expt-relaxed-constexpr --use_fast_math
|
||
--ptxas-options=-O3,-v --extra-device-vectorization --threads=16
|
||
```
|
||
|
||
Each kernel in `astrai/extension/lib` is compiled as an independent pybind11 module (one `.so` per kernel, named `<kernel>.cpython-*-x86_64-linux-gnu.so`). CMake builds all six kernel targets in parallel via `cmake --build -j N`. The target list is the **single source of truth**: `KERNEL_NAMES` and the parallel `KERNEL_SRCS` list in `csrc/CMakeLists.txt`; `astrai/extension/loader.py` auto-discovers the compiled `.so` files.
|
||
|
||
## Python Extension Architecture
|
||
|
||
The Python extension package separates low-level kernel bindings from execution
|
||
policy:
|
||
|
||
```text
|
||
astrai/extension/
|
||
├── __init__.py # Stable public API
|
||
├── loader.py # Optional compiled-module discovery and loading
|
||
├── ops/
|
||
│ ├── attention.py # Stateless attention kernel wrappers
|
||
│ ├── rotary.py # Stateless rotary kernel wrapper
|
||
│ └── fp8.py # Stateless FP8 primitives (custom_op)
|
||
├── fp8.py # FP8 strategy layer (fp8_autocast, recipes)
|
||
└── backend/
|
||
├── attention.py # Backend selection, KV cache I/O, and fallback
|
||
└── rotary.py # Per-call CUDA/torch rotary dispatch
|
||
```
|
||
|
||
The dependency direction is one-way:
|
||
|
||
```text
|
||
model / inference
|
||
|
|
||
v
|
||
extension public API
|
||
|
|
||
v
|
||
backend policy ---> ops wrappers ---> loader ---> compiled .so
|
||
|
|
||
+-----------> torch / flash-attn fallback
|
||
```
|
||
|
||
`ops` must not import `backend`. This keeps direct kernel bindings independent
|
||
of model, cache, fallback, and backend-selection policy.
|
||
|
||
### Ops Layer
|
||
|
||
`astrai.extension.ops` is the low-level boundary around compiled extensions:
|
||
|
||
- Wrappers are stateless and map Python arguments to pybind or
|
||
`torch.library.custom_op` calls.
|
||
- Wrappers validate kernel availability and raise `RuntimeError` when a
|
||
requested extension was not built.
|
||
- Wrappers do not choose another implementation, gather KV cache entries, or
|
||
decide whether an input is supported by a backend.
|
||
- Tests that specifically exercise a compiled kernel may import from
|
||
`astrai.extension.ops`.
|
||
|
||
For example, `attn_prefill(...)` means "run this CUDA kernel" rather than "run
|
||
attention using the best available implementation":
|
||
|
||
```python
|
||
from astrai.extension.ops import attn_prefill
|
||
|
||
output = attn_prefill(q, k, v, mask=mask, is_causal=True)
|
||
```
|
||
|
||
If the kernel is unavailable, this call fails. Callers that need fallback and
|
||
capability dispatch must use the public `attention(...)` entry point instead.
|
||
|
||
### Backend Layer
|
||
|
||
`astrai.extension.backend` owns execution policy:
|
||
|
||
- It selects CUDA, FlashAttention, or torch-native attention.
|
||
- It checks per-call constraints such as dtype, shape, head dimension, cache
|
||
availability, and installed optional dependencies.
|
||
- It owns KV cache writes and reads because those operations differ by backend.
|
||
- It provides torch fallbacks and raises when an explicitly requested backend
|
||
cannot handle a call.
|
||
- Rotary dispatch follows the same boundary without a backend class: the
|
||
policy layer chooses the fused op for supported inference calls and otherwise
|
||
uses the autograd-compatible torch implementation.
|
||
|
||
Normal model and inference code should import the stable API from
|
||
`astrai.extension`:
|
||
|
||
```python
|
||
from astrai.extension import ATTN_BACKEND, attention, attn_backend
|
||
|
||
output = attention(q, k, v, kv_cache=cache, layer_id=layer_id, fwd="decode")
|
||
|
||
with attn_backend(ATTN_BACKEND.TORCH_NATIVE):
|
||
output = attention(q, k, v)
|
||
```
|
||
|
||
The package root re-exports the supported high-level API and selected direct
|
||
kernel wrappers. Internal code should use `astrai.extension.backend` only when
|
||
it needs a backend type or policy implementation, and `astrai.extension.ops`
|
||
only when it deliberately requires one exact kernel.
|
||
|
||
### Placement Rules
|
||
|
||
When extending this package:
|
||
|
||
| Change | Location |
|
||
|--------|----------|
|
||
| Add a pybind call for a compiled kernel | `astrai/extension/ops/` |
|
||
| Add argument translation required by the compiled ABI | `astrai/extension/ops/` |
|
||
| Add capability checks or implementation selection | `astrai/extension/backend/` |
|
||
| Add a torch or third-party fallback | `astrai/extension/backend/` |
|
||
| Add attention KV cache behavior | `astrai/extension/backend/attention.py` |
|
||
| Expose a supported user-facing symbol | `astrai/extension/__init__.py` |
|
||
|
||
Imports belong at module scope. Optional dependencies such as `flash_attn` may
|
||
use a module-level guarded import. Type-only imports that would create a runtime
|
||
cycle belong under `TYPE_CHECKING`.
|
||
|
||
## Attention Backend
|
||
|
||
`astrai/extension/backend/attention.py` provides the backend abstraction:
|
||
|
||
- **`AttentionBackend`** (ABC): `fwd_decode` / `fwd_prefill` abstract methods, `forward` dispatches by q_len
|
||
- **`CudaBackend`**: CUDA kernel dispatch — decode via `attn_paged_decode` (page_size=1), prefill via `attn_paged_prefill` (ragged batch, `qo_indptr` + `kv_indptr`). Default on GPU.
|
||
- **`FlashAttnBackend`**: Optional flash-attn dispatch with `flash_attn_with_kvcache` fast path.
|
||
- **`TorchNativeBackend`**: SDPA with indirect KV cache gather (always-available fallback)
|
||
|
||
Default priority: cuda > flash > torch. Set ``ASTR_BACKEND=cuda|torch_native|flash``
|
||
to override the default.
|
||
|
||
Select a backend via context manager (mirrors `torch.nn.attention.sdpa_kernel`):
|
||
|
||
```python
|
||
from astrai.extension import attn_backend, ATTN_BACKEND
|
||
|
||
with attn_backend(ATTN_BACKEND.CUDA):
|
||
engine.generate("hello")
|
||
```
|
||
|
||
The `attention(...)` policy entry point falls back to `FlashAttnBackend` (when
|
||
flash-attn is installed and supports the call) or `TorchNativeBackend` when the
|
||
automatically selected CUDA backend cannot handle an input. Resolution
|
||
precedence is: explicit `attn_backend(...)` context > `ASTR_BACKEND` env >
|
||
default. An explicit `attn_backend(...)` selection is strict and raises instead
|
||
of silently switching implementations; the env override (and the implicit
|
||
default) fall back to the first compatible backend when incapable. Training
|
||
calls (`fwd=None`, no KV cache) resolve by capability: the CUDA cache kernels
|
||
cannot run without a cache, so they fall back to flash (mask-free/causal calls
|
||
only) and finally to torch SDPA.
|
||
|
||
### Rotary Backend
|
||
|
||
`astrai/extension/backend/rotary.py` provides `apply_rotary_emb(x, (cos, sin))` with auto-dispatch:
|
||
|
||
- **CUDA path**: calls `rotary_emb` kernel directly when available, input is bf16 on CUDA, and `torch.is_grad_enabled()` is `False` (inference)
|
||
- **Torch fallback**: complex multiply (`torch.view_as_complex` → `torch.complex` multiply → `torch.view_as_real`), used during training (supports autograd) or when kernel unavailable
|
||
|
||
No context-manager switching needed — the dispatch is automatic per call.
|
||
|
||
## Python Wrappers
|
||
|
||
`astrai/extension/ops/attention.py` provides Python wrappers for each compiled attention kernel. Each wrapper calls its CUDA kernel directly and raises `RuntimeError` if the `.so` is not available. Fallback to torch SDPA is handled by the attention backend, not the wrapper functions.
|
||
|
||
`astrai/extension/ops/rotary.py` provides the wrapper for the rotary embedding kernel. Fallback to torch complex multiply is handled by `backend/rotary.py`.
|
||
|
||
Interface (all functions):
|
||
```
|
||
is_causal: True = causal mask; False = non-causal
|
||
mask: 2D [batch, kv_len] or 3D [batch, q_len, kv_len] (bool, True=keep)
|
||
```
|
||
|
||
Layout convention: all q/k/v are `[batch, seq_len, n_heads, head_dim]` (blhd). Scale is always `1/sqrt(head_dim)`.
|
||
|
||
### Q Scheduling and KV Addressing
|
||
|
||
Prefill separates Q work scheduling from KV storage:
|
||
|
||
- `DenseQSchedule` maps a rectangular grid directly with
|
||
`batch = blockIdx.z` and `q_tile = blockIdx.x`.
|
||
- `PackedQSchedule` consumes a compact work map for a packed
|
||
`[total_q, q_heads, head_dim]` tensor.
|
||
- `ContigKV` and `PagedKV` only provide KV lengths and translate logical KV
|
||
positions into physical addresses. They do not schedule Q blocks.
|
||
|
||
For ragged Q lengths `[70, 10, 130]` and 64 rows per Q tile, cache binding
|
||
builds:
|
||
|
||
```text
|
||
qo_indptr = [0, 70, 80, 210]
|
||
q_tile_to_batch = [0, 0, 1, 2, 2, 2]
|
||
q_tile_to_index = [0, 1, 0, 0, 1, 2]
|
||
```
|
||
|
||
Paged prefill launches:
|
||
|
||
```text
|
||
grid.x = num_q_tiles # 6, exactly the valid ragged work items
|
||
grid.y = q_heads
|
||
grid.z = 1
|
||
```
|
||
|
||
Each block resolves its request and request-local tile in O(1):
|
||
|
||
```cpp
|
||
batch = q_tile_to_batch[blockIdx.x];
|
||
q_tile = q_tile_to_index[blockIdx.x];
|
||
```
|
||
|
||
The kernel then uses `qo_indptr[batch]` for the packed Q base and adjacent
|
||
`qo_indptr` / `kv_indptr` entries for that request's Q and KV lengths. This
|
||
avoids the previous per-block linear scan over the batch, shared-memory
|
||
broadcast, mapping barrier, and upper-bound grid with potentially invalid
|
||
blocks.
|
||
|
||
## Standalone Testing
|
||
|
||
Each `csrc/tests/*.cu` file has the `nvcc` compile command in its header comment. Example:
|
||
|
||
```bash
|
||
nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
|
||
--ptxas-options=-O3,-v --extra-device-vectorization \
|
||
-Xcompiler -fopenmp csrc/tests/attn_test.cu -o /tmp/test && /tmp/test
|
||
```
|
||
|
||
Test files:
|
||
- `attn_test.cu` — decode + prefill kernels (correctness tables + benchmarks)
|
||
- `attn_paged_test.cu` — paged decode/prefill kernels
|
||
- `fp8_mma_test.cu` — BF16→FP8→BF16 MMA demo (sm_89)
|
||
|
||
## Benchmarks
|
||
|
||
Hardware: NVIDIA L20 (sm_89, 46 GB), CUDA 12.8, driver 570.86.
|
||
|
||
Reproduce (decode + prefill in `attn_test.cu`, paged in `attn_paged_test.cu`):
|
||
```bash
|
||
nvcc -I csrc -arch=sm_89 -O3 --use_fast_math \
|
||
--ptxas-options=-O3,-v --extra-device-vectorization \
|
||
-Xcompiler -fopenmp csrc/tests/attn_test.cu -o /tmp/test && /tmp/test
|
||
```
|
||
|
||
## Known Optimization Targets
|
||
|
||
- **Decode D=256**: spill eliminated (BC=16 + STAGES=2), but still 248 regs — further tiling could help.
|
||
- **Prefill single-batch**: bandwidth low (22 GB/s at q=kv=2048) — compute-bound at ~94 TFLOP/s (near L20 bf16 ceiling ~193 TFLOP/s for non-causal).
|
||
- **Decode single-batch**: bandwidth low (113 GB/s at kv=512, 13% of 864 GB/s theoretical) — small kv underutilizes SMs despite split-KV; scales to 757 GB/s (88%) at B=16+.
|
||
|
||
## File Layout
|
||
|
||
```
|
||
csrc/
|
||
├── CMakeLists.txt # CMake build: kernel registry (KERNEL_NAMES / KERNEL_SRCS), torch/pybind11 linking
|
||
├── kernels/
|
||
│ ├── common/ # cross-family pure-CUDA helpers (no torch)
|
||
│ │ ├── device.cuh # sm_at_least(), kMinSmForFp8* constants
|
||
│ │ └── mma.cuh # shared mma_sync<InT> + mma_shape<InT> (bf16 m16n8k16 / fp8 m16n8k32) + ldmatrix_x2/x4<T>
|
||
│ ├── attention/ # attention family (module names keep the attn_* prefix)
|
||
│ │ ├── common.h # AttentionParams POD, TensorLayout enum (BHLD/BLHD)
|
||
│ │ ├── warp_utils.cuh # warp reduction helpers
|
||
│ │ ├── layout_policies.cuh # KV addressing policies: DenseQSchedule/PackedQSchedule, ContigKV/PagedKV
|
||
│ │ ├── mma_utils.cuh # ldmatrix/pack helpers + online-softmax (bf16 mma via common/mma.cuh)
|
||
│ │ ├── entry_utils.cuh # torch binding helpers: DISPATCH_HEAD_DIM, pack_*_params
|
||
│ │ ├── dispatchers.cuh # pure-CUDA launchers: dispatch_decode/prefill (+paged), split-K math
|
||
│ │ ├── decode_split_kv.cuh # decode kernel, scalar (split-KV)
|
||
│ │ ├── decode_split_kv_mma.cuh # decode kernel, MMA + split-K
|
||
│ │ ├── prefill_split_q.cuh # prefill kernel, scalar (split-Q)
|
||
│ │ ├── prefill_split_q_mma.cuh # prefill kernel, MMA (split-Q, packed/ragged Q schedule)
|
||
│ │ ├── decode.cu # → module attn_decode
|
||
│ │ ├── prefill.cu # → module attn_prefill
|
||
│ │ ├── paged_decode.cu # → module attn_paged_decode
|
||
│ │ └── paged_prefill.cu # → module attn_paged_prefill
|
||
│ ├── rotary/
|
||
│ │ └── rotary_emb.cu # rotary embedding (kernel + binding in one file) → module rotary_emb
|
||
│ └── fp8/ # FP8 family (module name fp8_ops)
|
||
│ ├── common.h # FP8Format enum, Fp8GemmTraits, FP8Params POD (no torch)
|
||
│ ├── gemm.cuh # FP8 device code: quantize + pre-quantized GEMM kernels (no torch)
|
||
│ └── mm.cu # binding only: validation, param packing, launch dispatch, pybind
|
||
└── tests/
|
||
├── test_utils.cuh # Shared test utilities (now_ms, f2bf, bf2f, randf)
|
||
├── attn_test.cu # Decode + prefill kernels
|
||
├── attn_paged_test.cu # Paged decode/prefill kernels
|
||
└── fp8_mma_test.cu # BF16→FP8→BF16 MMA demo
|
||
```
|
||
|
||
Compiled `.so` files are placed in `astrai/extension/lib/`, separate from Python source files.
|
||
|
||
> Document Update Time: 2026-08-22
|