### Setup Clangd for Development

Source: https://github.com/moonshotai/flashkda/blob/master/README.md

Run the setup script to configure IntelliSense for CUDA/C++ sources by generating a `.clangd` file and installing global clangd configuration.

```bash
bash setup_clangd.sh
```

--------------------------------

### Install FlashKDA

Source: https://github.com/moonshotai/flashkda/blob/master/README.md

Clone the repository, update submodules, and install the package. This is the standard installation procedure.

```bash
git clone https://github.com/MoonshotAI/FlashKDA.git flash-kda
cd flash-kda
git submodule update --init --recursive
pip install -v --no-build-isolation .
```

--------------------------------

### Python Example for Varlen Batching

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/architecture.md

This Python snippet demonstrates the setup for variable-length batching using cu_seqlens on CUDA. It shows how to define sequence boundaries and the expected shape for state tensors when handling multiple sequences of different lengths.

```python
# Two sequences: lengths 2048, 1024
cu_seqlens = torch.tensor([0, 2048, 3072], dtype=torch.int64, device='cuda')
# State shape: [2, H, 128, 128] (one state per sequence)
```

--------------------------------

### Install FlashKDA for All Architectures

Source: https://github.com/moonshotai/flashkda/blob/master/README.md

Compile FlashKDA for all supported CUDA architectures during installation. Use this for wheel or CI builds.

```bash
FLASH_KDA_CUDA_ARCHS=all pip install -v --no-build-isolation .
```

--------------------------------

### FlashKDA Forward Example Usage

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/api-reference-fwd.md

An example demonstrating how to set up dimensions and create input tensors for the flash_kda.fwd function. Ensure tensors are on the CUDA device and have the correct dtype.

```python
import torch
import flash_kda
import math

# Set up dimensions
batch_size, seq_len, num_heads, dim = 1, 4096, 96, 128
scale = 1.0 / math.sqrt(dim)

# Create input tensors
q = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device='cuda')
k = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device='cuda')
v = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device='cuda')
g = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device='cuda')
beta = torch.randn(batch_size, seq_len, num_heads, dtype=torch.bfloat16, device='cuda')

```

--------------------------------

### Python Setup Script Defaults

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/configuration.md

Defines supported CUDA architectures and default NVCC compilation flags within the setup.py script.

```python
SUPPORTED_CUDA_ARCHS = ["90a", "100a", "103a", "120a"]

```

--------------------------------

### Install with Specific CUDA Archs

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/configuration.md

Installs FlashKDA by compiling for a specified list of CUDA architectures. Use this for targeted builds.

```bash
FLASH_KDA_CUDA_ARCHS=90a,100a pip install -v --no-build-isolation .
```

--------------------------------

### Install with Custom NVCC Threads

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/configuration.md

Installs FlashKDA while setting the number of parallel threads for NVCC compilation. Higher values can speed up compilation at the cost of memory.

```bash
NVCC_THREADS=16 pip install -v --no-build-isolation .
```

--------------------------------

### Install with Auto CUDA Archs

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/configuration.md

Installs FlashKDA by automatically detecting the current CUDA device capability. Requires a visible CUDA device.

```bash
pip install -v --no-build-isolation .
```

--------------------------------

### Get FlashKDA Module Version

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/python-module-reference.md

Inspect the installed version of the flash_kda module. This version is managed by setup.py and may include a git hash or custom suffix.

```python
import flash_kda
print(flash_kda.__version__)  # '0.0.1+<git-hash>' or custom suffix
```

--------------------------------

### Documentation Structure Diagram

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/MANIFEST.md

Illustrates the hierarchical and interconnected structure of the FlashKDA documentation, starting from the README.md entry point.

```text
README.md (entry point)
    ↓
INDEX.md (navigation guide)
    ├→ api-reference-fwd.md
    ├→ api-reference-workspace.md
    ├→ python-module-reference.md
    ├→ types.md
    ├→ errors.md
    ├→ configuration.md
    ├→ architecture.md
    ├→ performance-notes.md
    └→ usage-examples.md
```

--------------------------------

### Install with Custom Version Suffix

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/configuration.md

Installs FlashKDA with a custom version suffix appended to the package version. This is useful for development builds.

```bash
FLASH_KDA_VERSION_SUFFIX=+custom pip install -v --no-build-isolation .
```

--------------------------------

### FlashKDA Type Conversions with Mixed Precision

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/usage-examples.md

This example demonstrates how to handle inputs in different precisions (e.g., float32 and bfloat16) when using `flash_kda.fwd`. Inputs can be converted to bfloat16 before processing and back to float32 if needed.

```python
import torch
import flash_kda
import math

# Inputs may come in float32 or float16
q_fp32 = torch.randn(1, 4096, 96, 128, dtype=torch.float32, device='cuda')
k_fp32 = torch.randn(1, 4096, 96, 128, dtype=torch.float32, device='cuda')
v_fp32 = torch.randn(1, 4096, 96, 128, dtype=torch.float32, device='cuda')

# Convert to bfloat16
q = q_fp32.to(torch.bfloat16)
k = k_fp32.to(torch.bfloat16)
v = v_fp32.to(torch.bfloat16)

# Prepare remaining inputs
g = torch.randn(1, 4096, 96, 128, dtype=torch.bfloat16, device='cuda')
beta = torch.randn(1, 4096, 96, dtype=torch.bfloat16, device='cuda')
A_log = torch.rand(96, dtype=torch.float32, device='cuda')
dt_bias = torch.randn(96, 128, dtype=torch.float32, device='cuda')

out = torch.empty_like(q)

flash_kda.fwd(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=1.0 / math.sqrt(128),
    out=out,
    A_log=A_log,
    dt_bias=dt_bias,
    lower_bound=-5.0
)

# Convert back to float32 if needed
out_fp32 = out.to(torch.float32)
```

--------------------------------

### FlashKDA Recurrent State with Float32 Precision

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/usage-examples.md

This example shows how to use `flash_kda.fwd` when higher precision is required for the recurrent state. The `initial_state` and `final_state` can be explicitly set to `torch.float32` for more accurate state accumulation.

```python
import torch
import flash_kda

batch_size, seq_len, num_heads, dim = 1, 4096, 96, 128
device = 'cuda'

# Input tensors in bfloat16
q = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
k = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
v = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
g = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
beta = torch.randn(batch_size, seq_len, num_heads, dtype=torch.bfloat16, device=device)

A_log = torch.rand(num_heads, dtype=torch.float32, device=device)
dt_bias = torch.randn(num_heads, dim, dtype=torch.float32, device=device)

out = torch.empty(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)

# Use float32 for state (higher precision accumulation)
initial_state = torch.zeros(batch_size, num_heads, dim, dim, dtype=torch.float32, device=device)
final_state = torch.empty(batch_size, num_heads, dim, dim, dtype=torch.float32, device=device)

flash_kda.fwd(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=1.0 / 128,
    out=out,
    A_log=A_log,
    dt_bias=dt_bias,
    lower_bound=-5.0,
    initial_state=initial_state,
    final_state=final_state
)

# final_state is in float32
print(f"Final state dtype: {final_state.dtype}")  # torch.float32
```

--------------------------------

### Use FlashKDA as FLA Backend

Source: https://github.com/moonshotai/flashkda/blob/master/README.md

Integrate FlashKDA with flash-linear-attention by calling `chunk_kda` within `torch.inference_mode()`. Ensure `flash-linear-attention` is installed.

```python
import torch
from fla.ops.kda import chunk_kda

with torch.inference_mode():
    out, final_state = chunk_kda(
        q=q, k=k, v=v, g=g, beta=beta,
        scale=scale,
        initial_state=h0,
        output_final_state=True,
        use_gate_in_kernel=True,
        use_qk_l2norm_in_kernel=True,
        use_beta_sigmoid_in_kernel=True,
        safe_gate=True,
        A_log=A_log, dt_bias=dt_bias,
        lower_bound=lower_bound,
        transpose_state_layout=True,
        cu_seqlens=cu_seqlens,
    )
```

--------------------------------

### Import flash_kda Module Components

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/python-module-reference.md

Imports necessary components from the flash_kda package and its C++ extension. Ensure the C++ extension is correctly built and installed to avoid ImportErrors.

```python
import torch
from flash_kda_C import fwd as _fwd_raw, get_workspace_size
```

--------------------------------

### Integrate FlashKDA with flash-linear-attention

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/usage-examples.md

This snippet shows how to use `chunk_kda` from `fla.ops.kda` with FlashKDA as the backend. Ensure FlashKDA is installed for automatic dispatch. Set the environment variable `FLA_FLASH_KDA=0` to disable.

```python
import torch
from fla.ops.kda import chunk_kda

# Enable FlashKDA backend (default if installed)
# Set FLA_FLASH_KDA=0 environment variable to disable

with torch.inference_mode():
    out, final_state = chunk_kda(
        q=q,  # [B, T, H, D]
        k=k,  # [B, T, H, D]
        v=v,  # [B, T, H, D]
        g=g,  # [B, T, H, D] gate logits
        beta=beta,  # [B, T, H] beta logits (pre-sigmoid)
        scale=scale,  # scalar
        initial_state=h0,
        output_final_state=True,
        use_gate_in_kernel=True,
        use_qk_l2norm_in_kernel=True,
        use_beta_sigmoid_in_kernel=True,
        safe_gate=True,
        A_log=A_log,  # [H]
        dt_bias=dt_bias,  # [H, D]
        lower_bound=lower_bound,
        transpose_state_layout=True,
        cu_seqlens=cu_seqlens  # optional
    )

# On success, this dispatches to FlashKDA internally
# Check logs: set logging.basicConfig(level=logging.INFO) to see "[FLA Backend] kda.chunk_kda -> flashkda"
```

--------------------------------

### FlashKDA Forward Pass (Stateless)

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/MANIFEST.md

Example of a basic stateless forward pass using the flash_kda.fwd() function. This is suitable for simple, non-recurrent computations.

```python
import flash_kda
import torch

# Example usage (stateless)
# Assume q, k, v, and other parameters are already defined Tensors
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)

# Call the forward pass
output = flash_kda.fwd(q, k, v)

```

--------------------------------

### FlashKDA Class-based Inference Loop

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/MANIFEST.md

Provides an example of structuring inference using a class, which can manage state and workspace allocation across multiple forward passes.

```python
import flash_kda
import torch

class FlashKDAModel:
    def __init__(self, batch_size, seq_len, num_heads, head_dim):
        self.batch_size = batch_size
        self.seq_len = seq_len
        self.num_heads = num_heads
        self.head_dim = head_dim
        
        # Pre-allocate workspace
        workspace_size = flash_kda.get_workspace_size(
            batch_size=self.batch_size,
            seq_len=self.seq_len,
            num_heads=self.num_heads,
            head_dim=self.head_dim
        )
        self.workspace = torch.empty(workspace_size, device='cuda', dtype=torch.uint8)
        self.state = None

    def forward(self, q, k, v):
        # Assume q, k, v are already prepared tensors
        output, self.state = flash_kda.fwd(
            q, k, v, 
            workspace=self.workspace,
            state=self.state
        )
        return output

# Example usage:
# model = FlashKDAModel(batch_size, seq_len, num_heads, head_dim)
# q_batch, k_batch, v_batch = ... # Prepare input tensors
# result = model.forward(q_batch, k_batch, v_batch)

```

--------------------------------

### Profile Kernel Execution with CUDA Events

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/performance-notes.md

Use CUDA events to accurately measure the execution time of the FlashKDA forward pass kernel. Ensure synchronization after recording events to get the correct elapsed time.

```python
import torch
import flash_kda

# Profile kernel execution
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()
flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound)
end.record()
torch.cuda.synchronize()

elapsed_ms = start.elapsed_time(end)
print(f"Kernel time: {elapsed_ms:.3f} ms")

```

--------------------------------

### Entry Point via flash-linear-attention (FLA)

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/python-module-reference.md

Utilize FlashKDA through the flash-linear-attention (FLA) library. FLA internally dispatches to flash_kda.fwd when available, providing a unified interface.

```python
from fla.ops.kda import chunk_kda
# Internally dispatches to flash_kda.fwd if available
chunk_kda(...)
```

--------------------------------

### Direct Import Entry Point

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/python-module-reference.md

Invoke FlashKDA functionality directly by importing the module. This is the primary method for using the library's core functions.

```python
import flash_kda
flash_kda.fwd(...)
```

--------------------------------

### Enable Debug Logging

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/configuration.md

Configure Python's logging module to INFO level to observe FlashKDA backend messages, including successful dispatches or rejection reasons.

```python
import logging
logging.basicConfig(level=logging.INFO)
# Now see "[FLA Backend] kda.chunk_kda -> flashkda" on success
# or "... rejected: <reason>" on dispatch failure
```

--------------------------------

### FlashKDA Computation Pipeline

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/architecture.md

Illustrates the data flow through the two main kernels (K1 and K2) and the intermediate workspace storage.

```text
Input Tensors
    ↓
[K1: Token-Parallel Stage]
    • g activation: exp(A_log) * sigmoid(g + dt_bias)
    • L2-norm: q, k → normalized q, k
    • Decay: exp(cumsum(g))
    • Matrix Inversion: 16×16 matrix (fp16)
    • Score Construction: L, Mqk
    ↓
Workspace Storage
    • k_decayed: [ht, CHUNK, 128]
    • q_decayed: [ht, CHUNK, 128]
    • k_restored: [ht, CHUNK, 128]
    • g_total: [ht, 128]
    • INV, Mqk: [ht, CHUNK, CHUNK]
    ↓
[K2: Recurrence Stage]
    • Chunk-by-chunk delta-rule:
      h_new = beta * h_old + (1 - beta) * contribution
    • Output accumulation
    • Final state: h_final → [N, H, 128, 128]
    ↓
Output Tensors
    ✓ out: [B, T, H, 128]
    ✓ final_state: [N, H, 128, 128] (optional)

```

--------------------------------

### Stateful Recurrent Computation with FlashKDA

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/usage-examples.md

Illustrates how to use FlashKDA for recurrent computations where state is maintained across forward passes. Initialize `initial_state` with zeros or the `final_state` from a previous computation.

```python
import torch
import flash_kda
import math

batch_size, seq_len, num_heads, dim = 2, 2048, 64, 128
scale = 1.0 / math.sqrt(dim)
device = 'cuda'

# Input tensors
q = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
k = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
v = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
g = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
beta = torch.randn(batch_size, seq_len, num_heads, dtype=torch.bfloat16, device=device)

A_log = torch.rand(num_heads, dtype=torch.float32, device=device)
dt_bias = torch.randn(num_heads, dim, dtype=torch.float32, device=device)

# Initialize recurrent state (or zero if starting fresh)
initial_state = torch.zeros(batch_size, num_heads, dim, dim, dtype=torch.bfloat16, device=device)
# Or load from previous computation:
# initial_state = previous_final_state

# Allocate output and final state
out = torch.empty(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
final_state = torch.empty(batch_size, num_heads, dim, dim, dtype=torch.bfloat16, device=device)

# Forward pass with state
flash_kda.fwd(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=scale,
    out=out,
    A_log=A_log,
    dt_bias=dt_bias,
    lower_bound=-5.0,
    initial_state=initial_state,
    final_state=final_state
)

# Use final_state as initial_state for next forward pass
print(f"Final state shape: {final_state.shape}")  # [2, 64, 128, 128]
```

--------------------------------

### Pre-allocate Workspace for FlashKDA

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/usage-examples.md

Shows how to calculate and pre-allocate the necessary workspace memory for FlashKDA operations, which is crucial for memory-constrained environments. This avoids dynamic memory allocation during computation.

```python
import torch
import flash_kda

batch_size, seq_len, num_heads, dim = 1, 8192, 96, 128
T_total = batch_size * seq_len
H = num_heads

# Pre-compute workspace size
workspace_size = flash_kda.get_workspace_size(T_total, H, N=batch_size)
print(f"Workspace: {workspace_size / 1e9:.2f} GB")

# Pre-allocate
workspace = torch.empty(workspace_size, dtype=torch.uint8, device='cuda')

# Now you know exactly how much memory is needed
# For memory-constrained scenarios, adjust batch_size or seq_len accordingly

```

--------------------------------

### Run FlashKDA Tests

Source: https://github.com/moonshotai/flashkda/blob/master/README.md

Execute the test suite for FlashKDA using the provided test script. This verifies the correctness of the kernels.

```bash
bash tests/test.sh
```

--------------------------------

### Clangd Configuration for CUDA

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/configuration.md

Configuration settings for clangd to enable IntelliSense support for CUDA files during development.

```yaml
-xcuda
--cuda-gpu-arch=sm_90
-D__CUDA_ARCH__=900
-DCUTLASS_ARCH_MMA_SM90_SUPPORTED=1

```

--------------------------------

### Ensuring Input Shape Compatibility

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/errors.md

Verify that the shapes of `k`, `v`, `g`, and `out` match `q`. Ensure `beta` has a shape compatible with `q`'s batch, time, and hidden dimensions.

```python
# Shape compatibility checks are implicit in the library's internal TORCH_CHECKs.
# Ensure tensors are created with compatible shapes before calling flash_kda.fwd.
```

--------------------------------

### Basic Stateless Forward Pass with FlashKDA

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/usage-examples.md

Demonstrates a simple forward pass using FlashKDA for stateless computation. Ensure all input tensors are on the CUDA device and have the correct data type (e.g., bfloat16).

```python
import torch
import flash_kda
import math

# Configuration
batch_size, seq_len, num_heads, dim = 1, 4096, 96, 128
scale = 1.0 / math.sqrt(dim)
device = 'cuda'

# Create input tensors
q = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
k = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
v = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
g = torch.randn(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)
beta = torch.randn(batch_size, seq_len, num_heads, dtype=torch.bfloat16, device=device)

# Gate parameters (typically learned during training)
A_log = torch.rand(num_heads, dtype=torch.float32, device=device)
dt_bias = torch.randn(num_heads, dim, dtype=torch.float32, device=device)
lower_bound = -5.0

# Allocate output
out = torch.empty(batch_size, seq_len, num_heads, dim, dtype=torch.bfloat16, device=device)

# Forward pass
flash_kda.fwd(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=scale,
    out=out,
    A_log=A_log,
    dt_bias=dt_bias,
    lower_bound=lower_bound
)

print(f"Output shape: {out.shape}")  # [1, 4096, 96, 128]
```

--------------------------------

### FlashKDA Workspace Pre-allocation

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/MANIFEST.md

Shows how to pre-allocate the workspace tensor and pass it to the flash_kda.fwd() function. This can be more efficient than letting the function allocate it internally on each call.

```python
import flash_kda
import torch

# Example: Pre-allocate workspace
workspace_size = flash_kda.get_workspace_size(batch_size, seq_len, num_heads, head_dim)
workspace = torch.empty(workspace_size, device='cuda', dtype=torch.uint8)

# Assume q, k, v are defined
output = flash_kda.fwd(q, k, v, workspace=workspace)

```

--------------------------------

### Fast Math Approximations (CUDA PTX)

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/architecture.md

Shows specific CUDA PTX instructions used for fast approximations in gating calculations, offering significant speedups.

```cuda
tanh.approx.f32    // Sigmoid via tanh approximation
ex2.approx.ftz.f32 // Base-2 exponential with flush-to-zero

```

--------------------------------

### Direct Usage of flash_kda Functions

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/python-module-reference.md

Demonstrates how to directly call the `fwd` function for the forward pass or `get_workspace_size` to pre-calculate buffer requirements. The `fwd` function is the primary interface for performing the attention computation.

```python
import flash_kda

# Call fwd directly
flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound)

# Or call get_workspace_size (rarely needed; done internally by fwd)
size = flash_kda.get_workspace_size(T_total, H, N)
```

--------------------------------

### Prepare for Variable-Length Batching by Concatenating Sequences

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/errors.md

Demonstrates how to prepare input tensors for variable-length batching by concatenating sequences into a single batch element (B=1) and creating the `cu_seqlens` tensor. This is required when `cu_seqlens` is provided.

```python
# Ensure B=1 for variable-length batching
# Concatenate multiple sequences into a single batch element
q = torch.cat(sequence_list, dim=1)  # Shape: [1, T_total, H, D]
cu_seqlens = torch.cumsum(torch.tensor([0] + seq_lengths), dim=0)
```

--------------------------------

### FlashKDA Workspace Size Calculation

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/MANIFEST.md

Demonstrates how to calculate the required workspace size using flash_kda.get_workspace_size(). This is crucial for pre-allocating memory, especially in memory-constrained environments.

```python
import flash_kda
import torch

# Example usage for get_workspace_size
batch_size = 32
seq_len = 1024
num_heads = 12
head_dim = 64

workspace_size = flash_kda.get_workspace_size(
    batch_size=batch_size,
    seq_len=seq_len,
    num_heads=num_heads,
    head_dim=head_dim
)
print(f"Required workspace size: {workspace_size} bytes")

# Pre-allocate workspace if needed
# workspace = torch.empty(workspace_size, device='cuda', dtype=torch.uint8)
# output = flash_kda.fwd(q, k, v, workspace=workspace)

```

--------------------------------

### C++ Kernel Dispatch Logic

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/architecture.md

This C++ code snippet from flash_kda.cpp shows the macro used to dispatch to different kernel instantiations based on state configuration (stateless, FP32 state, BF16 state). It determines the optimal code path at runtime.

```cpp
#define DISPATCH_STATE(VL) \
    if (!has_state_in && !has_state_out) { \
        LAUNCH(false, false, false, VL);     // Stateless \
    } else if (has_state_in && has_state_out && state_fp32) { \
        LAUNCH(true, true, true, VL);        // FP32 state \
    } else if (has_state_in && has_state_out && !state_fp32) { \
        LAUNCH(true, true, false, VL);       // BF16 state \
    } ...
```

--------------------------------

### FlashKDA Forward Pass (Stateful)

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/MANIFEST.md

Demonstrates stateful computation using flash_kda.fwd() for recurrent scenarios. This mode maintains and updates internal states across calls.

```python
import flash_kda
import torch

# Example usage (stateful)
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16)

# Initial state (can be None or a pre-allocated tensor)
initial_state = None

# Call the forward pass with state
output, final_state = flash_kda.fwd(q, k, v, state=initial_state)

```

--------------------------------

### Calculate and Print Throughput

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/performance-notes.md

Calculates theoretical FLOPs and TFLOPS based on model dimensions and elapsed time. Use this to estimate raw computational throughput.

```python
flops = 2 * B * T * H * D * D + 2 * B * T * H * D  # Simplified
tflops = flops / (elapsed_ms * 1e-3) / 1e12
print(f"Throughput: {tflops:.1f} TFLOPS")
```

--------------------------------

### FlashKDA Forward Pass (Variable-Length Sequences)

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/MANIFEST.md

Illustrates handling variable-length sequences using packed sequences with cu_seqlens. This is essential for efficient batching of sequences with differing lengths.

```python
import flash_kda
import torch

# Example usage (variable-length sequences)
# Assume q, k, v are concatenated tensors and cu_seqlens define sequence boundaries
# q = torch.cat([seq1_q, seq2_q, ...])
# k = torch.cat([seq1_k, seq2_k, ...])
# v = torch.cat([seq1_v, seq2_v, ...])
# cu_seqlens = torch.tensor([0, len1, len1+len2, ...], device='cuda')

# Call the forward pass with variable-length sequences
# output = flash_kda.fwd(q, k, v, cu_seqlens=cu_seqlens)

```

--------------------------------

### FlashKDA Inference Loop with State Management

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/usage-examples.md

Implements a KDAAttention class to perform forward passes using flash_kda.fwd, including handling optional initial states and returning the final state. This is useful for sequential processing or recurrent models.

```python
import torch
import flash_kda
import math

class KDAAttention:
    def __init__(self, num_heads, dim, device='cuda'):
        self.num_heads = num_heads
        self.dim = dim
        self.device = device
        self.scale = 1.0 / math.sqrt(dim)
        
        # Initialize gate parameters (normally trained)
        self.A_log = torch.rand(num_heads, dtype=torch.float32, device=device)
        self.dt_bias = torch.randn(num_heads, dim, dtype=torch.float32, device=device)
        self.lower_bound = -5.0
    
    def forward(self, q, k, v, g, beta, initial_state=None):
        """Forward pass with optional state."""
        batch_size, seq_len = q.shape[:2]
        out = torch.empty_like(q)
        
        if initial_state is None:
            initial_state = torch.zeros(
                batch_size, self.num_heads, self.dim, self.dim,
                dtype=torch.bfloat16, device=self.device
            )
        
        final_state = torch.empty_like(initial_state)
        
        flash_kda.fwd(
            q=q.to(torch.bfloat16),
            k=k.to(torch.bfloat16),
            v=v.to(torch.bfloat16),
            g=g.to(torch.bfloat16),
            beta=beta.to(torch.bfloat16),
            scale=self.scale,
            out=out,
            A_log=self.A_log,
            dt_bias=self.dt_bias,
            lower_bound=self.lower_bound,
            initial_state=initial_state,
            final_state=final_state
        )
        
        return out, final_state

# Usage
attn = KDAAttention(num_heads=96, dim=128)
q = torch.randn(2, 4096, 96, 128, dtype=torch.bfloat16, device='cuda')
k = torch.randn(2, 4096, 96, 128, dtype=torch.bfloat16, device='cuda')
v = torch.randn(2, 4096, 96, 128, dtype=torch.bfloat16, device='cuda')
g = torch.randn(2, 4096, 96, 128, dtype=torch.bfloat16, device='cuda')
beta = torch.randn(2, 4096, 96, dtype=torch.bfloat16, device='cuda')

out, final_state = attn.forward(q, k, v, g, beta)
print(f"Output shape: {out.shape}")
print(f"State shape: {final_state.shape}")

```

--------------------------------

### Variable-Length Sequence Processing with FlashKDA

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/usage-examples.md

Shows how to efficiently process sequences of varying lengths using FlashKDA by concatenating them and providing cumulative sequence lengths. Ensure `cu_seqlens` is correctly computed and on the CUDA device.

```python
import torch
import flash_kda
import math

# Two sequences with different lengths
seq_lens = [1024, 2048]
total_len = sum(seq_lens)
num_heads, dim = 96, 128
scale = 1.0 / math.sqrt(dim)
device = 'cuda'

# Concatenate sequences into [1, T_total, H, D] format
q_list = [torch.randn(1, l, num_heads, dim, dtype=torch.bfloat16, device=device) for l in seq_lens]
k_list = [torch.randn(1, l, num_heads, dim, dtype=torch.bfloat16, device=device) for l in seq_lens]
v_list = [torch.randn(1, l, num_heads, dim, dtype=torch.bfloat16, device=device) for l in seq_lens]
g_list = [torch.randn(1, l, num_heads, dim, dtype=torch.bfloat16, device=device) for l in seq_lens]
beta_list = [torch.randn(1, l, num_heads, dtype=torch.bfloat16, device=device) for l in seq_lens]

q = torch.cat(q_list, dim=1)  # [1, 3072, 96, 128]
k = torch.cat(k_list, dim=1)
v = torch.cat(v_list, dim=1)
g = torch.cat(g_list, dim=1)
beta = torch.cat(beta_list, dim=1)

# Cumulative sequence lengths
cu_seqlens = torch.cumsum(
    torch.tensor([0] + seq_lens, dtype=torch.int64),
    dim=0
).to(device)  # [0, 1024, 3072]

A_log = torch.rand(num_heads, dtype=torch.float32, device=device)
dt_bias = torch.randn(num_heads, dim, dtype=torch.float32, device=device)

# One state per sequence
out = torch.empty(1, total_len, num_heads, dim, dtype=torch.bfloat16, device=device)
initial_state = torch.zeros(len(seq_lens), num_heads, dim, dim, dtype=torch.bfloat16, device=device)
final_state = torch.empty(len(seq_lens), num_heads, dim, dim, dtype=torch.bfloat16, device=device)

# Forward pass
flash_kda.fwd(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=scale,
    out=out,
    A_log=A_log,
    dt_bias=dt_bias,
    lower_bound=-5.0,
    initial_state=initial_state,
    final_state=final_state,
    cu_seqlens=cu_seqlens
)

print(f"Output shape: {out.shape}")  # [1, 3072, 96, 128]
print(f"Final state shape: {final_state.shape}")  # [2, 96, 128, 128]
```

--------------------------------

### FlashKDA Type Conversion (float32 to bfloat16)

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/MANIFEST.md

Shows how to convert tensors between float32 and bfloat16 data types. This is useful for managing numerical precision and memory usage.

```python
import torch

# Example: Convert float32 tensor to bfloat16
tensor_fp32 = torch.randn(10, 10, device='cuda', dtype=torch.float32)
tensor_bf16 = tensor_fp32.to(torch.bfloat16)

# Example: Convert bfloat16 tensor to float32
tensor_bf16_2 = torch.randn(10, 10, device='cuda', dtype=torch.bfloat16)
tensor_fp32_2 = tensor_bf16_2.to(torch.float32)

```

--------------------------------

### flash_kda.fwd (Stateless)

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/api-reference-fwd.md

Basic stateless usage of the flash_kda.fwd function. This is the simplest way to call the function, without managing recurrent states.

```APIDOC
## flash_kda.fwd (Stateless)

### Description
Performs a stateless forward pass using the flash_kda kernel. This is suitable for standard attention computations where state is not maintained between calls.

### Method
`flash_kda.fwd`

### Parameters
- **q** (Tensor) - Query tensor.
- **k** (Tensor) - Key tensor.
- **v** (Tensor) - Value tensor.
- **g** (Tensor) - Gate tensor.
- **beta** (Tensor) - Beta parameter.
- **scale** (float) - Scaling factor.
- **out** (Tensor) - Output tensor to store results.
- **A_log** (Tensor) - Logarithm of the A matrix, per head.
- **dt_bias** (Tensor) - Delta-time bias, per head.
- **lower_bound** (float) - Lower bound for calculations.

### Request Example
```python
flash_kda.fwd(
    q=q,
    k=k,
    v=v,
    g=g,
    beta=beta,
    scale=scale,
    out=out,
    A_log=A_log,
    dt_bias=dt_bias,
    lower_bound=lower_bound
)
```

### Response
- **out** (Tensor) - The computed output tensor.

### Errors
- **TORCH_CHECK(all tensors on CUDA)**: All input tensors must be on the CUDA device.
- **TORCH_CHECK(all tensors contiguous)**: All input tensors must be memory-contiguous.
- **TORCH_CHECK(q/k/v/g/beta/out bfloat16)**: Input dtype mismatch for q, k, v, g, beta, or out.
- **TORCH_CHECK(A_log/dt_bias float32)**: A_log or dt_bias must be float32.
- **TORCH_CHECK(input dimensions correct)**: Input tensors must have the expected rank (4D for q, k, v, g, out; 3D for beta).
- **TORCH_CHECK(D == 128)**: Dimension K or V must be 128.
```

--------------------------------

### BibTeX Citation for FlashKDA

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/README.md

This is the BibTeX entry for citing the FlashKDA project. It includes title, authors, year, publisher, and a URL.

```bibtex
@misc{flashkda2026,
      title={FlashKDA: Flash Kimi Delta Attention},
      author={Yutian Chen, Zhiyuan Li, Yucheng Wang, Ming Wei},
      year={2026},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/MoonshotAI/FlashKDA}},
}

```

--------------------------------

### Profile with Nsys

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/performance-notes.md

Profiles the Python script using NVIDIA Nsys to capture CUDA events and performance metrics. The output file can be analyzed for detailed performance insights.

```bash
# Profile with NVIDIA Nsys
nsys profile -o profile_flashkda \
  -c cudaImportSynchronousMemOps=on \
  python benchmark_script.py

# Analyze
nsys stats profile_flashkda.qdrep
```

--------------------------------

### Initialize Recurrent States for Batched or Variable-Length Modes

Source: https://github.com/moonshotai/flashkda/blob/master/_autodocs/errors.md

Provides code to correctly initialize `initial_state` and `final_state` tensors based on whether the model is in batched mode (B) or variable-length mode (using `cu_seqlens`). Ensures states match the required [N, H, D, D] shape.

```python
# Batched mode
initial_state = torch.zeros(B, H, D, D, dtype=torch.bfloat16, device='cuda')
final_state = torch.empty(B, H, D, D, dtype=torch.bfloat16, device='cuda')

# Variable-length mode
N_sequences = cu_seqlens.numel() - 1
initial_state = torch.zeros(N_sequences, H, D, D, dtype=torch.bfloat16, device='cuda')
final_state = torch.empty(N_sequences, H, D, D, dtype=torch.bfloat16, device='cuda')
```

--------------------------------

### flash_kda.fwd

Source: https://github.com/moonshotai/flashkda/blob/master/README.md

The `flash_kda.fwd` function provides a direct interface to the FlashKDA forward kernel. It allows for fine-grained control over the attention computation, including recurrent state management and variable-length batching.

```APIDOC
## flash_kda.fwd

### Description
Provides a direct interface to the FlashKDA forward kernel for advanced usage, including recurrent state management and variable-length batching.

### Method
`flash_kda.fwd(q, k, v, g, beta, scale, out, A_log, dt_bias, lower_bound, initial_state=None, final_state=None, cu_seqlens=None)`

### Parameters
#### Input Parameters
- **q** (bf16) - `[B, T, H, K]` - Query
- **k** (bf16) - `[B, T, H, K]` - Key
- **v** (bf16) - `[B, T, H, V]` - Value
- **g** (bf16) - `[B, T, H, K]` - Gate before activation
- **beta** (bf16) - `[B, T, H]` - Beta logits (pre-activation; sigmoid applied internally)
- **scale** (float) - scalar - scaling factor
- **A_log** (fp32) - `[H]` - Log-gate parameter
- **dt_bias** (fp32) - `[H, K]` - Gate bias
- **lower_bound** (float) - scalar - Gate lower bound (range from -5.0 to 0)

#### Optional Parameters
- **initial_state** (bf16/fp32/None) - `[B, H, V, K]` or `[N, H, V, K]` - (optional) Initial recurrent state. Accepts `None` (stateless), bf16, or fp32 tensors. When both `initial_state` and `final_state` are provided, their dtypes must match.
- **final_state** (bf16/fp32/None) - `[B, H, V, K]` or `[N, H, V, K]` - (optional, output) Final recurrent state. Accepts `None` (stateless), bf16, or fp32 tensors. When both `initial_state` and `final_state` are provided, their dtypes must match.
- **cu_seqlens** (int64) - `[N+1]` - (optional) Cumulative sequence lengths for variable-length batching. When provided, `B` must be 1, `T` is the total length across all sequences, and `initial_state`/`final_state` have shape `[N, H, V, K]`. When `cu_seqlens` is `None`, each batch element is treated as an independent sequence, and the state shape is `[B, H, V, K]`.

### Notes
- Currently requires `K = V = 128`.

```