### Install FlashQLA

Source: https://github.com/qwenlm/flashqla/blob/main/README.md

Clone the repository and install the library using pip. Ensure you are in the root directory of the cloned repository.

```bash
git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA
pip install -v .
```

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

### Install FlashQLA

Source: https://context7.com/qwenlm/flashqla/llms.txt

Clone the repository, navigate to the directory, and install using pip. Ensure CUDA 12.8+ and PyTorch 2.8+ are installed.

```bash
# Requirements: SM90+, CUDA 12.8+, PyTorch 2.8+
git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA
pip install -v .
```

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

### Run Benchmark Script

Source: https://github.com/qwenlm/flashqla/blob/main/README.md

Execute the benchmark script to compare FlashQLA performance against baseline implementations. Ensure flash linear attention and flashinfer are installed.

```bash
# require flash linear attention and flashinfer for comparison
pip install flash_linear_attention==0.5.0 flashinfer-python==0.6.9

cd benchmark
python bench_gated_delta_rule.py
```

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

### Run Tests for FlashQLA

Source: https://github.com/qwenlm/flashqla/blob/main/README.md

Execute tests for FlashQLA, including development, variable-length sequence, profiling, and product comparisons. Ensure flash linear attention is installed.

```bash
# require flash linear attention for comparison
pip install flash_linear_attention==0.5.0

cd tests
python test_gdr.py --set develop
python test_gdr.py --set varlen --num-heads 32
python test_gdr.py --set profile --num-heads 32
python test_gdr.py --set product --ref-dtype float32 --num-heads 32
```

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

### Prepare Chunk Offsets

Source: https://context7.com/qwenlm/flashqla/llms.txt

Prepares chunk offsets based on cumulative sequence lengths and chunk size. Helps in managing chunk boundaries within sequences.

```python
chunk_offsets  = prepare_chunk_offsets(cu_seqlens, chunk_size=64)
# tensor([0, 4, 8, 16])
```

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

### Run Gated Delta Rule Benchmarks

Source: https://context7.com/qwenlm/flashqla/llms.txt

Runs benchmarks for the gated delta rule, comparing against FlashInfer. Requires flash_linear_attention==0.5.0 and flashinfer-python==0.6.9.

```bash
# Requires both fla and flashinfer for full comparison
pip install flash_linear_attention==0.5.0 flashinfer-python==0.6.9

cd benchmark
python bench_gated_delta_rule.py               # all modes
python bench_gated_delta_rule.py --mode fwd    # forward only
python bench_gated_delta_rule.py --mode bwd    # backward only
python bench_gated_delta_rule.py --skip-fi     # skip FlashInfer
```

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

### Prepare Chunk Indices and Offsets Utilities

Source: https://context7.com/qwenlm/flashqla/llms.txt

Helper functions `prepare_chunk_indices` and `prepare_chunk_offsets` compute per-sequence chunk indices and cumulative chunk counts for variable-length dispatch in TileLang kernels. These results are LRU-cached.

```python
from flash_qla.utils import prepare_chunk_indices, prepare_chunk_offsets
import torch

cu_seqlens = torch.tensor([0, 256, 512, 1024], dtype=torch.int32, device="cuda")


```

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

### prepare_chunk_indices / prepare_chunk_offsets

Source: https://context7.com/qwenlm/flashqla/llms.txt

Cached utility functions for preparing chunk indices and offsets, used by TileLang kernels for variable-length dispatch.

```APIDOC
## prepare_chunk_indices / prepare_chunk_offsets

### Description
Provides cached utilities for computing per-sequence chunk indices and cumulative chunk counts for variable-length sequence processing.

### Parameters
- **cu_seqlens** (Tensor) - Cumulative sequence lengths tensor `[B + 1]`.

### Returns
- **prepare_chunk_indices** returns chunk indices for each sequence.
- **prepare_chunk_offsets** returns cumulative chunk offsets.
```

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

### Run GDR Tests (Develop Mode)

Source: https://context7.com/qwenlm/flashqla/llms.txt

Executes correctness and speedup tests for gated delta rule on small development shapes. Requires flash_linear_attention==0.5.0.

```bash
# Requires flash_linear_attention==0.5.0 for reference comparison
pip install flash_linear_attention==0.5.0

cd tests
# Correctness + speedup on small development shapes
python test_gdr.py --set develop
```

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

### Run GDR Tests (Custom Options)

Source: https://context7.com/qwenlm/flashqla/llms.txt

Executes tests with custom options, including disabling intra-card CP, forward only mode, and hiding accuracy output. Requires flash_linear_attention==0.5.0.

```bash
python test_gdr.py --set develop --no-cp --skip-bwd --hide-acc
```

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

### Profile Kernel Latency

Source: https://context7.com/qwenlm/flashqla/llms.txt

Runs a callable with PyTorch Profiler to collect per-kernel CUDA times. Returns average device time in milliseconds for specified kernels.

```python
from flash_qla.utils import profile
from flash_qla import chunk_gated_delta_rule_fwd
import torch

# ... prepare q, k, v, g, beta, h0 as above ...

results = profile(
    chunk_gated_delta_rule_fwd,
    inputs=[q, k, v, g, beta, D**-0.5, h0, None, True, False, True],
    wait=50, warmup=50, rep=100,
)
print(results["tilelang_fused_chunk_gdr_fwd_kernel_kernel"])  # ms for main fwd kernel
print(results["total"])                                        # end-to-end ms
```

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

### chunk_gated_delta_rule

Source: https://context7.com/qwenlm/flashqla/llms.txt

The primary end-user entry point for FlashQLA. It wraps the full forward and backward passes as a custom `torch.autograd.Function`, enabling seamless integration with PyTorch's autograd system. Inputs must be in `bfloat16` or `float16`.

```APIDOC
## chunk_gated_delta_rule

### Description
This function serves as the main entry point for users, encapsulating the complete forward and backward passes of the gated linear attention mechanism as a `torch.autograd.Function`. This allows for automatic gradient computation when `.backward()` is called.

### Method
`torch.autograd.Function` wrapper

### Parameters
- **q** (Tensor) - Query tensor.
- **k** (Tensor) - Key tensor.
- **v** (Tensor) - Value tensor.
- **g** (Tensor) - Gating tensor, typically derived from `F.logsigmoid`.
- **beta** (Tensor) - Beta parameter for the gated delta rule.
- **scale** (float, optional) - Scaling factor for attention. Defaults to `D**-0.5`.
- **initial_state** (Tensor, optional) - Initial hidden state for the attention mechanism.
- **output_final_state** (bool, optional) - If True, returns the final hidden state.
- **use_qk_l2norm_in_kernel** (bool, optional) - If True, applies L2 normalization within the kernel. Defaults to False.
- **cu_seqlens** (Tensor, optional) - Packed sequence lengths for variable-length sequences.

### Returns
- **o** (Tensor) - Output tensor of shape `[B, T, Hv, D]` in `bfloat16`.
- **final_state** (Tensor, optional) - The final hidden state of shape `[B, Hv, D, D]` in `float32`, returned if `output_final_state` is True.

### Request Example
```python
import torch
import torch.nn.functional as F
from flash_qla import chunk_gated_delta_rule

device = "cuda"
B, T, Hk, Hv, D = 1, 4096, 16, 64, 128

q = torch.nn.functional.normalize(torch.randn(B, T, Hk, D, device=device, dtype=torch.bfloat16), dim=-1)
k = torch.nn.functional.normalize(torch.randn(B, T, Hk, D, device=device, dtype=torch.bfloat16), dim=-1)
v = torch.randn(B, T, Hv, D, device=device, dtype=torch.bfloat16)
g = F.logsigmoid(torch.randn(B, T, Hv, device=device, dtype=torch.float32)) / 16
beta = torch.randn(B, T, Hv, device=device, dtype=torch.float32).sigmoid()
h0 = torch.zeros(B, Hv, D, D, device=device, dtype=torch.float32)

o, final_state = chunk_gated_delta_rule(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=D ** -0.5,
    initial_state=h0,
    output_final_state=True,
    use_qk_l2norm_in_kernel=False,
    cu_seqlens=None
)
```

### Response Example
```json
{
  "o": "[B, T, Hv, D] bfloat16",
  "final_state": "[B, Hv, D, D] float32"
}
```
```

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

### Run GDR Tests (Profiling Shapes)

Source: https://context7.com/qwenlm/flashqla/llms.txt

Executes tests using profiling shapes and a specified number of heads. Requires flash_linear_attention==0.5.0.

```bash
python test_gdr.py --set profile --num-heads 32
```

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

### Run GDR Tests (Production Shapes)

Source: https://context7.com/qwenlm/flashqla/llms.txt

Executes tests with production shapes, float32 reference dtype, and a specified number of heads. Requires flash_linear_attention==0.5.0.

```bash
python test_gdr.py --set product --ref-dtype float32 --num-heads 32
```

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

### FlashQLA Citation

Source: https://github.com/qwenlm/flashqla/blob/main/README.md

Use this BibTeX entry for citing the FlashQLA project in academic work.

```bibtex
@misc{flashqla2025,
    title={FlashQLA: Flash Qwen Linear Attention},
    author={Zhang, Chengruidong and Lin, Xi and Jiang, Huiqiang and Wang, Zekun and Li, Xiao and Cao, Yizhong and Zhuang, Bohan and Men, Rui and Zhang, Jianwei and Zheng, Bo and Lin, Junyang and Liu, Dayiheng and Zhou, Jingren},
    year={2026},
    publisher={GitHub},
    howpublished={\url{https://github.com/QwenLM/FlashQLA}},
}
```

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

### Run GDR Tests (Variable Length)

Source: https://context7.com/qwenlm/flashqla/llms.txt

Executes tests for variable-length sequences with a specified number of heads. Requires flash_linear_attention==0.5.0.

```bash
python test_gdr.py --set varlen --num-heads 32
```

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

### High-level Differentiable Forward Pass with chunk_gated_delta_rule

Source: https://context7.com/qwenlm/flashqla/llms.txt

Use this function as the primary entry point for end-users. It wraps the forward and backward passes as a torch.autograd.Function for seamless integration with PyTorch's autograd system. Inputs must be bfloat16 or float16; head_first layout is not supported. Compiler tracing is disabled.

```python
import torch
import torch.nn.functional as F
from flash_qla import chunk_gated_delta_rule

device = "cuda"
B, T, Hk, Hv, D = 1, 4096, 16, 64, 128   # batch, tokens, QK-heads, V-heads, head-dim

q = torch.nn.functional.normalize(
    torch.randn(B, T, Hk, D, device=device, dtype=torch.bfloat16), dim=-1
)
k = torch.nn.functional.normalize(
    torch.randn(B, T, Hk, D, device=device, dtype=torch.bfloat16), dim=-1
)
v     = torch.randn(B, T, Hv, D,  device=device, dtype=torch.bfloat16)
g     = F.logsigmoid(torch.randn(B, T, Hv, device=device, dtype=torch.float32)) / 16
beta  = torch.randn(B, T, Hv, device=device, dtype=torch.float32).sigmoid()
h0    = torch.zeros(B, Hv, D, D,  device=device, dtype=torch.float32)   # optional initial state

# scale defaults to D**-0.5 if None
o, final_state = chunk_gated_delta_rule(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=D ** -0.5,
    initial_state=h0,
    output_final_state=True,
    use_qk_l2norm_in_kernel=False,   # set True to apply L2-norm inside the kernel
    cu_seqlens=None,                 # set for variable-length packed sequences
)
# o:           [B, T, Hv, D]  bfloat16
# final_state: [B, Hv, D, D]  float32

# Backward works automatically
loss = o.sum()
loss.backward()
```

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

### Chunk-wise Local Cumulative Sum with `chunk_local_cumsum`

Source: https://context7.com/qwenlm/flashqla/llms.txt

The `chunk_local_cumsum` kernel computes prefix cumulative sums independently within each 64-token chunk. It supports both fixed-length and variable-length inputs and includes a `reverse` mode for backward pass computations.

```python
from flash_qla.ops.utils import chunk_local_cumsum
import torch

B, T, H = 2, 4096, 64
g = torch.randn(B, T, H, device="cuda", dtype=torch.bfloat16)
g_cumsum = chunk_local_cumsum(g, chunk_size=64)
# g_cumsum[b, t, h] = sum of g[b, t', h] for t' in the same 64-token chunk, t' <= t

# Reverse mode (used in backward for gradient propagation)
g_cumsum_rev = chunk_local_cumsum(g, chunk_size=64, reverse=True)
```

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

### Low-level API for Forward and Backward Passes

Source: https://github.com/qwenlm/flashqla/blob/main/README.md

For separate control over forward and backward computations, use chunk_gated_delta_rule_fwd and chunk_gated_delta_rule_bwd. The forward pass returns intermediate values like attention matrix A and hidden states.

```python
from flash_qla import chunk_gated_delta_rule_fwd, chunk_gated_delta_rule_bwd

# Forward
g, A, o, h, final_state = chunk_gated_delta_rule_fwd(
    q, k, v, g, beta, scale=scale, initial_state=h0, cu_seqlens=cu_seqlens
)

# Backward
dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
    q, k, v, g, beta, A, do, dht=dht, scale=scale, initial_state=h0, cu_seqlens=cu_seqlens
)
```

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

### Pack and Unpack Variable-Length Sequences with cu_seqlens

Source: https://context7.com/qwenlm/flashqla/llms.txt

Use `pack` and `unpack` helpers to convert between padded tensors and the packed format required for variable-length sequences specified by `cu_seqlens`. The batch dimension must be 1 for packed tensors.

```python
from flash_qla import chunk_gated_delta_rule
from flash_qla.utils import pack, unpack
import torch

device = "cuda"
# Two sequences of length 3000 and 1096 (total 4096 tokens)
seqlens = [3000, 1096]
cu_seqlens = torch.tensor([0, 3000, 4096], dtype=torch.int32, device=device)

Hk, Hv, D = 16, 64, 128
B_real = len(seqlens)

# Create padded tensors then pack them
q_pad  = torch.nn.functional.normalize(torch.randn(B_real, max(seqlens), Hk, D, device=device, dtype=torch.bfloat16), dim=-1)
k_pad  = torch.nn.functional.normalize(torch.randn(B_real, max(seqlens), Hk, D, device=device, dtype=torch.bfloat16), dim=-1)
v_pad  = torch.randn(B_real, max(seqlens), Hv, D, device=device, dtype=torch.bfloat16)
g_pad  = torch.nn.functional.logsigmoid(torch.randn(B_real, max(seqlens), Hv, device=device, dtype=torch.float32)) / 16
b_pad  = torch.randn(B_real, max(seqlens), Hv, device=device, dtype=torch.float32).sigmoid()

q = pack(q_pad, cu_seqlens)   # [1, 4096, Hk, D]
k = pack(k_pad, cu_seqlens)
v = pack(v_pad, cu_seqlens)
g = pack(g_pad, cu_seqlens)
beta = pack(b_pad, cu_seqlens)

h0 = torch.zeros(B_real, Hv, D, D, device=device, dtype=torch.float32)

o_packed, final_states = chunk_gated_delta_rule(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=D ** -0.5,
    initial_state=h0,
    output_final_state=True,
    cu_seqlens=cu_seqlens,
)
# o_packed:     [1, 4096, Hv, D]
# final_states: [2, Hv, D, D]    one per sequence

o_unpadded = unpack(o_packed, cu_seqlens)  # [2, 3000, Hv, D]
```

```python
from flash_qla.utils import pack, unpack
import torch

cu_seqlens = torch.tensor([0, 512, 768, 1024], dtype=torch.int32, device="cuda")
x_padded = torch.randn(3, 512, 8, device="cuda", dtype=torch.bfloat16)  # [B, T_max, H]

x_packed   = pack(x_padded, cu_seqlens)   # [1, 1024, 8]
x_restored = unpack(x_packed, cu_seqlens) # [3, 512, 8]  (padded back)
```

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

### Prepare Chunk Indices

Source: https://context7.com/qwenlm/flashqla/llms.txt

Prepares chunk indices for sequences based on cumulative sequence lengths and chunk size. Used for processing long sequences in chunks.

```python
chunk_indices  = prepare_chunk_indices(cu_seqlens, chunk_size=64)
# tensor([[0, 0], [0, 1], [0, 2], [0, 3],   # seq-0: 256/64=4 chunks
#         [1, 0], [1, 1], [1, 2], [1, 3],   # seq-1: 4 chunks
#         [2, 0], ..., [2, 7]])              # seq-2: 1024/64=8 chunks
```

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

### pack / unpack

Source: https://context7.com/qwenlm/flashqla/llms.txt

Utilities for converting between padded tensors and packed tensors using cumulative sequence lengths. Essential for handling variable-length sequences efficiently.

```APIDOC
## pack / unpack

### Description
Converts tensors between a padded format `[B, T_max, ...]` and a packed format `[1, T_total, ...]` using `cu_seqlens`.

### Parameters
- **pack**:
  - **tensor** (Tensor) - Padded input tensor `[B, T_max, ...]`.
  - **cu_seqlens** (Tensor) - Cumulative sequence lengths `[B + 1]`.
- **unpack**:
  - **packed_tensor** (Tensor) - Packed input tensor `[1, T_total, ...]`.
  - **cu_seqlens** (Tensor) - Cumulative sequence lengths `[B + 1]`.

### Returns
- **pack** returns a packed tensor `[1, T_total, ...]`.
- **unpack** returns an un-packed (padded) tensor `[B, T_max, ...]`.
```

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

### Low-level API: Forward and Backward Passes

Source: https://github.com/qwenlm/flashqla/blob/main/README.md

Provides separate functions for the forward and backward passes of the chunked gated delta rule operation, allowing for more granular control.

```APIDOC
## Low-level API: Forward and Backward Passes

### Description
Provides separate functions for the forward and backward passes of the chunked gated delta rule operation, allowing for more granular control.

### Forward Pass

#### Method
`chunk_gated_delta_rule_fwd`

#### Parameters
- **q** (torch.Tensor) - Query tensor.
- **k** (torch.Tensor) - Key tensor.
- **v** (torch.Tensor) - Value tensor.
- **g** (torch.Tensor) - Gate tensor.
- **beta** (torch.Tensor) - Beta parameter.
- **scale** (float, optional) - Scaling factor.
- **initial_state** (torch.Tensor, optional) - Initial state for the recurrence.
- **cu_seqlens** (torch.Tensor, optional) - Cumulative sequence lengths for variable-length sequences.

#### Response
- **g** (torch.Tensor) - Gate output.
- **A** (torch.Tensor) - Attention matrix.
- **o** (torch.Tensor) - Output tensor.
- **h** (torch.Tensor) - Hidden state.
- **final_state** (torch.Tensor) - The final state of the recurrence.

### Backward Pass

#### Method
`chunk_gated_delta_rule_bwd`

#### Parameters
- **q** (torch.Tensor) - Query tensor.
- **k** (torch.Tensor) - Key tensor.
- **v** (torch.Tensor) - Value tensor.
- **g** (torch.Tensor) - Gate tensor.
- **beta** (torch.Tensor) - Beta parameter.
- **A** (torch.Tensor) - Attention matrix from forward pass.
- **do** (torch.Tensor) - Gradient of the output.
- **dht** (torch.Tensor, optional) - Gradient of the final hidden state.
- **scale** (float, optional) - Scaling factor.
- **initial_state** (torch.Tensor, optional) - Initial state from forward pass.
- **cu_seqlens** (torch.Tensor, optional) - Cumulative sequence lengths for variable-length sequences.

#### Response
- **dq** (torch.Tensor) - Gradient w.r.t. query.
- **dk** (torch.Tensor) - Gradient w.r.t. key.
- **dv** (torch.Tensor) - Gradient w.r.t. value.
- **db** (torch.Tensor) - Gradient w.r.t. beta.
- **dg** (torch.Tensor) - Gradient w.r.t. gate.
- **dh0** (torch.Tensor) - Gradient w.r.t. initial state.
```

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

### Low-level Forward Pass with chunk_gated_delta_rule_fwd

Source: https://context7.com/qwenlm/flashqla/llms.txt

This function executes the forward pass and returns intermediate tensors required for the backward pass when manually managing the autograd graph. It supports optional intra-card context parallelism (CP) via `auto_cp=True`.

```python
from flash_qla import chunk_gated_delta_rule_fwd

g_cumsum, A, o, h, final_state = chunk_gated_delta_rule_fwd(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=D ** -0.5,
    initial_state=h0,
    cu_seqlens=None,
    output_final_state=True,
    output_h=True,     # True to also return per-chunk hidden states
    auto_cp=True,      # automatically enables intra-card CP when beneficial
)
# g_cumsum: [B, T, Hv]         float32  — chunk-local prefix sums of g
# A:        [B, num_chunks, Hk, chunk_size, chunk_size]  — KKT solve matrix
# o:        [B, T, Hv, D]      bfloat16 — output
# h:        per-chunk hidden states (None if output_h=False)
# final_state: [B, Hv, D, D]   float32
```

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

### High-level API for Chunked Gated Delta Rule

Source: https://github.com/qwenlm/flashqla/blob/main/README.md

Use the high-level API for a simplified interface to the chunk_gated_delta_rule function. It handles forward and backward passes internally. Optional parameters include initial_state and cu_seqlens for variable-length sequences.

```python
import torch
from flash_qla import chunk_gated_delta_rule

o, final_state = chunk_gated_delta_rule(
    q=q,          # [B, T, H_q, K]
    k=k,          # [B, T, H_q, K]
    v=v,          # [B, T, H_v, V]
    g=g,          # [B, T, H_v]
    beta=beta,    # [B, T, H_v]
    scale=scale,
    initial_state=initial_state,   # optional, [B, H_v, K, V]
    output_final_state=True,
    cu_seqlens=cu_seqlens,         # optional, for variable-length sequences
)
```

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

### chunk_local_cumsum

Source: https://context7.com/qwenlm/flashqla/llms.txt

Computes prefix cumulative sums within fixed-size chunks. Supports both forward and reverse modes, useful for attention mechanisms and their backward passes.

```APIDOC
## chunk_local_cumsum

### Description
Calculates cumulative sums within each chunk independently. Supports a reverse mode for backward computations.

### Parameters
- **g** (Tensor) - Input tensor.
- **chunk_size** (int) - The size of each chunk for cumulative sum calculation.
- **reverse** (bool, optional) - If True, computes the cumulative sum in reverse order. Defaults to False.

### Returns
- **g_cumsum** (Tensor) - Tensor with chunk-wise local cumulative sums.
```

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

### High-level API: chunk_gated_delta_rule

Source: https://github.com/qwenlm/flashqla/blob/main/README.md

This function provides a high-level interface for the chunked gated delta rule operation, combining forward and backward passes implicitly.

```APIDOC
## High-level API: chunk_gated_delta_rule

### Description
This function provides a high-level interface for the chunked gated delta rule operation, combining forward and backward passes implicitly.

### Method
`chunk_gated_delta_rule`

### Parameters
- **q** (torch.Tensor) - Query tensor.
- **k** (torch.Tensor) - Key tensor.
- **v** (torch.Tensor) - Value tensor.
- **g** (torch.Tensor) - Gate tensor.
- **beta** (torch.Tensor) - Beta parameter.
- **scale** (float, optional) - Scaling factor.
- **initial_state** (torch.Tensor, optional) - Initial state for the recurrence.
- **output_final_state** (bool) - Whether to output the final state.
- **cu_seqlens** (torch.Tensor, optional) - Cumulative sequence lengths for variable-length sequences.

### Response
- **o** (torch.Tensor) - Output tensor.
- **final_state** (torch.Tensor) - The final state of the recurrence, if `output_final_state` is True.
```

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

### Low-level Backward Pass with chunk_gated_delta_rule_bwd

Source: https://context7.com/qwenlm/flashqla/llms.txt

This function computes all input gradients using the intermediates from the forward pass. It recomputes hidden states on-the-fly using `fused_gdr_h` and then calls the fused backward kernel. Gradients are returned in the same dtype as their corresponding inputs.

```python
from flash_qla import chunk_gated_delta_rule_bwd

do  = torch.randn_like(o)   # gradient w.r.t. output
dht = torch.randn(B, Hv, D, D, device=device, dtype=torch.float32) / 8  # grad w.r.t. final_state

dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
    q=q, k=k, v=v,
    g=g_cumsum,   # pass the cumsum g returned by fwd, NOT the raw g
    beta=beta,
    A=A,
    do=do,
    dht=dht,
    scale=D ** -0.5,
    initial_state=h0,
    cu_seqlens=None,
)
# dq:  [B, T, Hk, D]   — same dtype as q
# dk:  [B, T, Hk, D]   — same dtype as k
```

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

### chunk_gated_delta_rule_fwd

Source: https://context7.com/qwenlm/flashqla/llms.txt

The low-level forward pass primitive. It computes and returns intermediate tensors required for the backward pass, such as cumulative sums of `g`, the KKT solve matrix `A`, the output `o`, and hidden states `h` and `final_state`. Supports optional intra-card context parallelism (CP).

```APIDOC
## chunk_gated_delta_rule_fwd

### Description
This function executes the forward pass of the gated linear attention mechanism and returns all intermediate tensors necessary for manual gradient management via the backward pass. It also supports enabling intra-card context parallelism (CP) through the `auto_cp` flag.

### Method
`chunk_gated_delta_rule_fwd`

### Parameters
- **q** (Tensor) - Query tensor.
- **k** (Tensor) - Key tensor.
- **v** (Tensor) - Value tensor.
- **g** (Tensor) - Gating tensor.
- **beta** (Tensor) - Beta parameter.
- **scale** (float, optional) - Scaling factor.
- **initial_state** (Tensor, optional) - Initial hidden state.
- **cu_seqlens** (Tensor, optional) - Packed sequence lengths.
- **output_final_state** (bool, optional) - If True, returns the final hidden state.
- **output_h** (bool, optional) - If True, returns per-chunk hidden states.
- **auto_cp** (bool, optional) - Automatically enables intra-card CP if beneficial. Defaults to True.

### Returns
- **g_cumsum** (Tensor) - Chunk-local prefix sums of `g` in `float32`.
- **A** (Tensor) - KKT solve matrix.
- **o** (Tensor) - Output tensor in `bfloat16`.
- **h** (Tensor, optional) - Per-chunk hidden states, returned if `output_h` is True.
- **final_state** (Tensor, optional) - The final hidden state in `float32`, returned if `output_final_state` is True.

### Request Example
```python
from flash_qla import chunk_gated_delta_rule_fwd

# Assuming q, k, v, g, beta, h0 are defined as in the high-level example
g_cumsum, A, o, h, final_state = chunk_gated_delta_rule_fwd(
    q=q, k=k, v=v, g=g, beta=beta,
    scale=D ** -0.5,
    initial_state=h0,
    cu_seqlens=None,
    output_final_state=True,
    output_h=True,
    auto_cp=True
)
```

### Response Example
```json
{
  "g_cumsum": "[B, T, Hv] float32",
  "A": "[B, num_chunks, Hk, chunk_size, chunk_size]",
  "o": "[B, T, Hv, D] bfloat16",
  "h": "per-chunk hidden states (None if output_h=False)",
  "final_state": "[B, Hv, D, D] float32"
}
```
```

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

### L2 Normalize Last Dimension with `l2norm`

Source: https://context7.com/qwenlm/flashqla/llms.txt

The `l2norm` utility performs compiled L2 normalization on the last axis of a tensor. It is commonly used for creating unit-norm query and key vectors in gated linear attention configurations.

```python
from flash_qla.utils import l2norm
import torch

x = torch.randn(2, 512, 16, 128, device="cuda", dtype=torch.bfloat16)
x_normed = l2norm(x)   # [2, 512, 16, 128], each 128-dim vector has unit L2 norm
# Equivalent to: x / (x.norm(dim=-1, keepdim=True) + 1e-6)
```

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

### chunk_gated_delta_rule_bwd

Source: https://context7.com/qwenlm/flashqla/llms.txt

The low-level backward pass primitive. It computes gradients for all inputs (`dq`, `dk`, `dv`, `db`, `dg`, `dh0`) using the intermediate tensors generated by `chunk_gated_delta_rule_fwd`. It recomputes hidden states on-the-fly using `fused_gdr_h`.

```APIDOC
## chunk_gated_delta_rule_bwd

### Description
This function computes the gradients with respect to all input parameters of the gated linear attention mechanism. It requires the intermediate tensors produced by `chunk_gated_delta_rule_fwd` and uses a recomputation strategy for hidden states before executing the fused backward kernel. Gradients are returned in the same data type as their corresponding inputs.

### Method
`chunk_gated_delta_rule_bwd`

### Parameters
- **q** (Tensor) - Query tensor.
- **k** (Tensor) - Key tensor.
- **v** (Tensor) - Value tensor.
- **g** (Tensor) - Cumulative sum of gating tensor (`g_cumsum` from fwd pass).
- **beta** (Tensor) - Beta parameter.
- **A** (Tensor) - KKT solve matrix from fwd pass.
- **do** (Tensor) - Gradient with respect to the output `o`.
- **dht** (Tensor) - Gradient with respect to the final hidden state `final_state`.
- **scale** (float, optional) - Scaling factor.
- **initial_state** (Tensor, optional) - Initial hidden state used in fwd pass.
- **cu_seqlens** (Tensor, optional) - Packed sequence lengths.

### Returns
- **dq** (Tensor) - Gradient w.r.t. query tensor `q`.
- **dk** (Tensor) - Gradient w.r.t. key tensor `k`.
- **dv** (Tensor) - Gradient w.r.t. value tensor `v`.
- **db** (Tensor) - Gradient w.r.t. beta parameter.
- **dg** (Tensor) - Gradient w.r.t. gating tensor `g`.
- **dh0** (Tensor) - Gradient w.r.t. initial hidden state `h0`.

### Request Example
```python
from flash_qla import chunk_gated_delta_rule_bwd

# Assuming q, k, v, g_cumsum, beta, A, h0, o are defined from fwd pass
do = torch.randn_like(o)   # gradient w.r.t. output
dht = torch.randn(B, Hv, D, D, device=device, dtype=torch.float32) / 8  # grad w.r.t. final_state

dq, dk, dv, db, dg, dh0 = chunk_gated_delta_rule_bwd(
    q=q, k=k, v=v,
    g=g_cumsum,   # pass the cumsum g returned by fwd, NOT the raw g
    beta=beta,
    A=A,
    do=do,
    dht=dht,
    scale=D ** -0.5,
    initial_state=h0,
    cu_seqlens=None,
)
```

### Response Example
```json
{
  "dq": "[B, T, Hk, D] same dtype as q",
  "dk": "[B, T, Hk, D] same dtype as k",
  "dv": "[B, T, Hv, D] same dtype as v",
  "db": "[B, T, Hv] same dtype as beta",
  "dg": "[B, T, Hv] same dtype as g",
  "dh0": "[B, Hv, D, D] same dtype as initial_state"
}
```
```

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

### Grouped Head Reduction with `group_reduce_vector`

Source: https://context7.com/qwenlm/flashqla/llms.txt

The `group_reduce_vector` function reduces a tensor with `H` value-heads to `Hg` query/key-heads by summing within groups. This is used in the backward pass for grouped-query attention when `num_v_heads > num_qk_heads`.

```python
from flash_qla.ops.utils import group_reduce_vector
import torch

# 4 V-heads mapped to 2 QK-heads (group_size=2)
B, T, H_full, D = 1, 2048, 4, 128
buffer = torch.randn(B, T, H_full, D, device="cuda", dtype=torch.bfloat16)
result = group_reduce_vector(buffer, Hg=2)
# result: [B, T, 2, D]  —  result[:, :, i, :] = buffer[:, :, 2*i, :] + buffer[:, :, 2*i+1, :]
```

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

### chunk_gated_delta_rule

Source: https://context7.com/qwenlm/flashqla/llms.txt

Applies the chunked gated delta rule for attention, supporting variable-length sequences via `cu_seqlens`. This function computes attention outputs and optionally the final state.

```APIDOC
## chunk_gated_delta_rule

### Description
Computes attention outputs and final states using the chunked gated delta rule, optimized for variable-length sequences.

### Parameters
- **q** (Tensor) - Query tensor, packed format [1, T_total, Hk, D]
- **k** (Tensor) - Key tensor, packed format [1, T_total, Hk, D]
- **v** (Tensor) - Value tensor, packed format [1, T_total, Hv, D]
- **g** (Tensor) - Gate tensor, packed format [1, T_total, Hv]
- **beta** (Tensor) - Beta tensor, packed format [1, T_total, Hv]
- **scale** (float) - Scaling factor for attention scores (e.g., D ** -0.5)
- **initial_state** (Tensor, optional) - Initial state for the attention computation [B_real, Hv, D, D]. Defaults to None.
- **output_final_state** (bool) - If True, returns the final state of the attention computation. Defaults to False.
- **cu_seqlens** (Tensor) - Cumulative sequence lengths tensor [B_real + 1]. Required for packed inputs.

### Returns
- **o_packed** (Tensor) - Packed attention output [1, T_total, Hv, D]
- **final_states** (Tensor, optional) - Final states for each sequence [B_real, Hv, D, D], returned if `output_final_state` is True.
```

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

### group_reduce_vector

Source: https://context7.com/qwenlm/flashqla/llms.txt

Reduces a tensor with multiple value heads down to a specified number of query/key heads by summing within groups. Used in grouped-query attention backward passes.

```APIDOC
## group_reduce_vector

### Description
Reduces value heads by summing across groups to match the number of query/key heads.

### Parameters
- **buffer** (Tensor) - Input tensor with multiple value heads.
- **Hg** (int) - The target number of query/key heads after reduction.

### Returns
- **result** (Tensor) - Tensor reduced to `Hg` heads.
```

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

### l2norm

Source: https://context7.com/qwenlm/flashqla/llms.txt

L2-normalizes the last dimension of a tensor. This is commonly used for query and key vectors in attention mechanisms to ensure unit norm.

```APIDOC
## l2norm

### Description
Performs L2 normalization on the last dimension of the input tensor.

### Parameters
- **x** (Tensor) - Input tensor.

### Returns
- **x_normed** (Tensor) - Tensor with the last dimension L2-normalized.
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.