Kernels:
Trusted publisher
File size: 2,276 Bytes
964dbb0 a717a14 964dbb0 a717a14 2c1bc68 a717a14 2c1bc68 a717a14 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | ---
library_name: kernels
license: apache-2.0
tags:
- kernels
- helion
- attention
---
# attention (Helion)
A scaled-dot-product / flash-attention kernel written in
[Helion](https://github.com/pytorch/helion) and packaged for the
[`kernels` library](https://github.com/huggingface/kernels). It is a
Triton-backend kernel that ships **pre-tuned configs** (via
`helion.aot_kernel`), so the first call selects a known-good configuration and
compiles it — no autotuning search at run time.
## How to use
```python
# make sure `kernels` is installed: `pip install -U kernels`
import torch
from kernels import get_kernel
kernel = get_kernel("HelionDSL/attention", version=1)
# flash-attn layout: (batch, seq, heads, head_dim), fp16/bf16
q = torch.randn(2, 512, 8, 64, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)
out = kernel.flash_attn_func(q, k, v, causal=False)
```
`get_kernel` requires a pinned version or revision — use `version=1` (the stable
kernel-API major, backed by the repo's `v1` branch) or `revision="main"` /
`revision="<commit>"` for an explicit Hub revision.
`HelionDSL` is a trusted kernel publisher on the Hub, so `trust_remote_code=True`
is **not** required. (On older `kernels` releases that predate publisher trust,
pass `trust_remote_code=True` to `get_kernel`.)
## Available functions
- `flash_attn_func(q, k, v, softmax_scale=None, causal=False)` — flash-attn
compatible entry point. Inputs and output are `(batch, seq, heads, head_dim)`.
- `attention(q, k, v, causal=False, out=None)` — SDPA-layout entry point. Inputs
and output are `(batch, heads, seq, head_dim)`.
Both accept `torch.float16` / `torch.bfloat16` and support `causal=True`.
## Available layers
- `Attention` — an `nn.Module` wrapper (`layers.Attention(causal=False)`) whose
`forward(q, k, v)` calls the kernel.
## Notes
- **Backend:** Helion → Triton. Correctness is verified against
`torch.nn.functional.scaled_dot_product_attention`.
- **Pre-tuned:** ships `_helion_aot_attention_cuda_sm100.py` and
`_helion_aot_attention_cuda_sm90.py`; the right one is selected for your GPU
(with compute-capability fallback). On a shape/GPU without a shipped config,
the kernel falls back to normal Helion autotuning.
|