---
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.