FiX: Introducing Fine-grained Forget Gate into Softmax Attention

ICML 2026

This repository contains the official implementation of the paper "FiX: Introducing Fine-grained Forget Gate into Softmax Attention", accepted at ICML 2026.

Overview

FiX introduces a fine-grained forget gate mechanism into standard softmax attention, enabling the model to dynamically control the retention and forgetting of information at a granular level. The core implementation can be found in fla/ops/fix_attn.

Code Structure

Module	Description
parallel_fix_fwd.py	Forward pass of the FiX attention kernel
parallel_fix_bwd_intra.py	Backward pass for intra-chunk computation
parallel_fix_bwd_inter_dkv.py	Backward pass for inter-chunk dKV computation
parallel_fix_bwd_inter_dq.py	Backward pass for inter-chunk dQ computation
parallel_fix_bwd_preprocess.py	Backward pass preprocessing
inter_chunk_scale_do.py	Inter-chunk scaling for dO
inter_chunk_scale_v.py	Inter-chunk scaling for V
intra_chunk_preprocess_fwd.py	Intra-chunk forward preprocessing
parallel.py	Parallel attention wrapper

Quick Start

Install fla following the original FLA instructions, then use the parallel_fix_attn function as a drop-in attention replacement:

import torch
from fla.ops.fix_attn import parallel_fix_attn

# Input shapes:
#   q: [batch, seqlen, num_q_heads, head_dim]
#   k: [batch, seqlen, num_kv_heads, head_dim]
#   v: [batch, seqlen, num_kv_heads, head_dim_v]
#   g: [batch, seqlen, num_kv_heads, head_dim_v]  (forget gate, same shape as v)
batch, seqlen, n_q_heads, n_kv_heads, dim_k, dim_v = 2, 2048, 8, 4, 128, 128
device, dtype = 'cuda', torch.bfloat16

q = torch.randn(batch, seqlen, n_q_heads, dim_k, device=device, dtype=dtype)
k = torch.randn(batch, seqlen, n_kv_heads, dim_k, device=device, dtype=dtype)
v = torch.randn(batch, seqlen, n_kv_heads, dim_v, device=device, dtype=dtype)
g = torch.randn(batch, seqlen, n_kv_heads, dim_v, device=device, dtype=torch.float32)

o = parallel_fix_attn(q, k, v, g)
# o: [batch, seqlen, num_q_heads, dim_v]

Variable-length sequences (e.g., for packing) are supported via cu_seqlens, consistent with FlashAttention's API:

cu_seqlens = torch.tensor([0, 1024, 2048], dtype=torch.long, device=device)
o = parallel_fix_attn(q, k, v, g, cu_seqlens=cu_seqlens)

Output normalization and gating can be optionally enabled:

o = parallel_fix_attn(
    q, k, v, g,
    o_norm=True,                              # apply RMSNorm on output
    o_norm_weight=torch.randn(dim_v, device=device, dtype=dtype),
    o_gate=torch.randn(batch, seqlen, n_q_heads, dim_v, device=device, dtype=dtype),
)

Note: The forget gate g must be float32 to preserve numerical precision, and only head dimensions in [16, 32, 64, 128] are currently supported.

Installation and Setup

This repo uses the same setup as FLA. See the original installation guide for environment requirements and install options.

Acknowledgements

This repository is forked from Flash Linear Attention (FLA), a comprehensive library for efficient Triton-based implementations of linear attention models. We are grateful to the FLA team for their excellent work and infrastructure, which made this project possible.

Citation

If you find our paper or this repository useful, please cite our work:

@inproceedings{
li2026fix,
title={{F}i{X}: Introducing Fine-grained Forget Gate into Softmax Attention},
author={Runzhong Li and Renjie Liu and Qing Li and Bo Tang},
booktitle={Forty-third International Conference on Machine Learning},
year={2026},
url={https://openreview.net/forum?id=WsNpCXq6SG}
}