FlashQLA Documentation

To verify the installation and run performance comparisons:

Run Unit Tests:

python -m pytest tests/test_gdr_unit.py -v

Run Profiling: Note: Requires flash_linear_attention==0.5.0 for comparison.

pip install flash_linear_attention==0.5.0
python profile/profile_gdr.py --set develop
python profile/profile_gdr.py --set develop --skip-bwd

Run Benchmarks: Note: Requires flash_linear_attention==0.5.0 and flashinfer-python==0.6.13 for comparison.

pip install flash_linear_attention==0.5.0 flashinfer-python==0.6.13
python benchmark/bench_gated_delta_rule.py

FlashQLA can be installed via pip or built from source.

Requirements

• GPU Architecture: SM90, SM100, SM103, or SM120
• CUDA: 12.8 or above
• PyTorch: 2.8 or above

# Install via pip
pip install flash-qla

# Or build from source
git clone https://github.com/QwenLM/FlashQLA.git
cd FlashQLA
pip install -v .

If you need to manually manage the forward and backward steps (e.g., for custom training loops), use chunk_gated_delta_rule_fwd and chunk_gated_delta_rule_bwd.

Forward Pass (chunk_gated_delta_rule_fwd): Returns g, A, o, h, final_state.

Backward Pass (chunk_gated_delta_rule_bwd): Requires the outputs from the forward pass (like A and do) to compute gradients. Returns dq, dk, dv, db, dg, dh0.

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
)

The prepare_chunk_indices function generates indices used for chunked kernel execution. It takes cu_seqlens and a chunk_size and returns a 2D tensor where each row contains [start_index_of_chunk, chunk_index]. This is used to map tokens to their respective chunks.

import torch
from flash_qla.utils import prepare_chunk_indices

cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.long)
chunk_size = 2
indices = prepare_chunk_indices(cu_seqlens, chunk_size)
# Returns a tensor mapping tokens to chunk starts and IDs

The prepare_lens function calculates the individual sequence lengths from a tensor of cumulative sequence lengths (cu_seqlens). It uses the @tensor_cache decorator to optimize repeated calls with the same input.

import torch
from flash_qla.utils import prepare_lens

# Example: cu_seqlens representing sequences of length 3 and 5
cu_seqlens = torch.tensor([0, 3, 8], dtype=torch.long)
lens = prepare_lens(cu_seqlens)
# lens will be tensor([3, 5])

The prepare_chunk_offsets function calculates the starting positions of chunks in a flattened buffer. It returns a tuple containing:

1. chunk_offsets: A torch.LongTensor of the same shape as cu_seqlens indicating the offset for each sequence.
2. total_chunks: An integer representing the total number of chunks across all sequences.

This is essential for indexing into global chunk buffers during kernel execution.

import torch
from flash_qla.utils import prepare_chunk_offsets

cu_seqlens = torch.tensor([0, 4, 8], dtype=torch.long)
chunk_size = 2
offsets, total_chunks = prepare_chunk_offsets(cu_seqlens, chunk_size)

The @tensor_cache decorator provides single-entry caching for functions that take torch.Tensor inputs. It stores the most recent result for a specific set of input tensors. If the function is called again with the exact same tensor objects (checked via identity is), it returns the cached result instead of recomputing. This is useful for avoiding redundant kernel preparation steps in iterative loops.

import torch
from flash_qla.utils import tensor_cache

@tensor_cache
def my_tensor_func(x: torch.Tensor) -> torch.Tensor:
    return x * 2

a = torch.randn(10)
b = my_tensor_func(a)  # Computes
c = my_tensor_func(a)  # Returns cached result