---
library_name: kernels
license: bsd-3-clause
tags: [cuda, blackwell, flash-attention, cute-dsl, inference]
---

# flashrt/fa4-cute-runtime

Forward-only FlashAttention-4 CuTe DSL runtime used by FlashRT's GROOT N1.7
and PI0.5 Thor pipelines. The source is vendored under the private `flashrt_fa4`
namespace and does not shadow an installed `flash_attn` package.

This package adds the dedicated Blackwell D256 2CTA forward path required by
PI0.5's 8-Q/1-KV-head PaliGemma encoder. The community FlashAttention-4 package
already covers the D48/D72/D128 family; it does not currently expose D256 on
SM100/SM110.

## Functions

- `flash_attn_func`
- `flash_attn_varlen_func`
- `forward_static`

```python
from kernels import get_kernel

fa4 = get_kernel("flashrt/fa4-cute-runtime", version=1)
out = fa4.flash_attn_func(q, k, v, causal=False)
```

The vendored forward wrapper returns `(out, lse)`; use `result[0]` when only
the attention output is needed.

For a CUDA Graph hot path, preallocate the output and use the allocation-free
entry point:

```python
out = torch.empty_like(q)
fa4.forward_static(q, k, v, out, causal=False)
```

For a padded fixed-shape graph, pass the valid K/V length as a CUDA int32
tensor. PI0.5 uses this form for its encoder cache:

```python
seqused_k = torch.tensor([valid_k], device="cuda", dtype=torch.int32)
fa4.forward_static(
    q, k_padded, v_padded, out,
    causal=False,
    pack_gqa=True,
    seqused_k=seqused_k,
)
```

Inputs follow FlashAttention's `(batch, sequence, heads, head_dim)` contract.
Qualified model profiles include D72 MHA and D256 GQA (8 Q heads / 1 KV head),
with both dense and `seqused_k` execution. This package targets SM100-family
Blackwell forward inference and requires
CUDA 13 plus `nvidia-cutlass-dsl` 4.4.x, 4.5.x, or 4.6.x. The wrapper selects Thor's
accepted architecture alias according to the installed DSL version.

This is an execution backend rather than a universal SDPA replacement. Select
it with model-shape profiling; the GROOT causal GQA profile benefits while some
short vision profiles remain faster on PyTorch SDPA.