Core operations for flash-linear-attention

0.4.2 · active · verified Wed Apr 15

fla-core is a Python library providing efficient, Triton-based implementations of core operations and kernels for state-of-the-art linear attention and state-space models. It serves as a minimal-dependency subset of the larger 'flash-linear-attention' project, focusing on the fundamental computational building blocks. It is currently at version 0.4.2 and follows a regular release cadence, often in conjunction with its parent project, flash-linear-attention.

Warnings

gotcha The `fla-core` package is a minimal subset of `flash-linear-attention`. It contains core kernels and modules (e.g., in `fla.ops` and `fla.modules`) but lacks higher-level layers and models (e.g., `fla.layers`, `fla.models`). Attempting to import these high-level components with only `fla-core` installed will result in an `ImportError`.
Fix: If you need layers or models, install the full `flash-linear-attention` package: `pip install flash-linear-attention`.
breaking The input tensor format for some kernels switched from 'head-first' to 'sequence-first'. This change affects how dimensions are ordered for input tensors (e.g., `(batch, heads, sequence, dim)` might become `(batch, sequence, heads, dim)`).
Fix: Review the documentation or source code for the specific kernel being used to ensure correct input tensor dimension ordering (e.g., (B, T, H, K) or (B, H, T, K)).
gotcha fla-core heavily relies on NVIDIA Triton for its optimized kernels. Specific Triton versions (>=3.0 or nightly) and correct backend installations are required, especially for AMD ROCm or Intel XPU GPUs.
Fix: Ensure Triton is installed correctly (`pip install triton`) and that the appropriate backend is configured for your GPU hardware. Check Triton FAQs for platform-specific instructions.
gotcha Requires Python 3.10 or higher. Running with older Python versions will lead to installation or runtime errors.
Fix: Upgrade your Python environment to version 3.10 or newer.
gotcha Users on H100 GPUs may encounter 'MMA Assertion Error' or 'LinearLayout Assertion Error' due to known Triton issues.
Fix: Install the latest nightly version of Triton, which often includes fixes for such hardware-specific issues. Refer to Triton's GitHub issues or FAQs for the most current solutions.
gotcha PyTorch version requirement: fla-core expects PyTorch >= 2.5. Older versions may cause compatibility issues or runtime errors.
Fix: Ensure your PyTorch installation is version 2.5 or newer: `pip install 'torch>=2.5'`.

Install

pip install fla-core Install from PyPI

Imports

FusedRMSNormGated
```
from fla.modules import FusedRMSNormGated
```
Commonly used fused normalization module from fla-core.
chunk_kda
```
from fla.ops.kda import chunk_kda
```
Example of a kernel import for Kimi Delta Attention operations.
MultiScaleRetention
wrong:
```
from fla.layers import MultiScaleRetention
```
correct:
```
from flash_linear_attention.layers import MultiScaleRetention
```
fla-core does not contain high-level layers or models; these are part of the 'flash-linear-attention' package. Importing directly from 'fla.layers' or 'fla.models' will fail if only 'fla-core' is installed.

Quickstart

This quickstart demonstrates the use of a fused normalization module from `fla-core`. It initializes `FusedRMSNormGated` and applies it to a dummy tensor on a CUDA-enabled GPU. This illustrates how to integrate low-level, optimized operations provided by `fla-core`.

import torch
from fla.modules import FusedRMSNormGated
import os

# fla-core operations require a CUDA-enabled GPU
if not torch.cuda.is_available():
    raise RuntimeError("CUDA not available. fla-core requires a CUDA-enabled GPU.")

device = torch.device("cuda")

# Define model parameters
hidden_size = 768
batch_size = 4
sequence_length = 512

# Initialize FusedRMSNormGated module from fla-core
norm_layer = FusedRMSNormGated(hidden_size).to(device)

# Create a dummy input tensor
input_tensor = torch.randn(batch_size, sequence_length, hidden_size, device=device, dtype=torch.float16)

# Perform a forward pass
output_tensor = norm_layer(input_tensor)

print(f"Input tensor shape: {input_tensor.shape}")
print(f"Output tensor shape: {output_tensor.shape}")
print("FusedRMSNormGated operation successful, demonstrating fla-core usage.")

view raw JSON →