Skip to content

Add standalone usage example and inference-only note #21

Description

@haix716

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions