Problem
The README only shows integration via FLA's chunk_kda interface. There is no standalone example showing direct usage of the flash_kda.fwd API, and no minimal reproduction script.
Recommendation
Add examples/basic_usage.py:
import torch
from flash_kda import fwd
B, H, T, D = 2, 8, 1024, 128
q = torch.randn(B, H, T, D, dtype=torch.bfloat16, device='cuda')
k = torch.randn_like(q)
v = torch.randn_like(q)
g = torch.randn(B, H, T, dtype=torch.float32, device='cuda')
beta = torch.rand(B, H, T, dtype=torch.float32, device='cuda').sigmoid()
scale = D ** -0.5
out = fwd(q, k, v, g, beta, scale)
print(f"Output shape: {out.shape}") # [2, 8, 1024, 128]
Also add a note in README that this is inference-only (no backward pass) — users needing training should fall back to FLA's Triton path.
Impact
Documentation. Reduces friction for new users who want to try the kernel directly without understanding FLA's internals.
Problem
The README only shows integration via FLA's
chunk_kdainterface. There is no standalone example showing direct usage of theflash_kda.fwdAPI, and no minimal reproduction script.Recommendation
Add
examples/basic_usage.py:Also add a note in README that this is inference-only (no backward pass) — users needing training should fall back to FLA's Triton path.
Impact
Documentation. Reduces friction for new users who want to try the kernel directly without understanding FLA's internals.