Open-source toolkitApache-2.0

Alternative attention normalizers for Hugging Face Transformers

Softmax1, Sparsemax, and Entmax15 behind one attention switch.

View on GitHub PyPI

pip install hf-attention-normalizers

Overview

Change how attention normalizes its scores without rewriting the model.

Hugging Face models choose their attention kernel through attn_implementation. The package registers new names in that registry, such as softmax1_sdpa or sparsemax_flash_attention_3, so a model can load with a different normalizer in place of softmax. For models that do not route attention through that interface, a second path patches the attention modules directly. Its built-in policy covers Qwen3.

The implementations range from a reference path that materializes the attention matrix to fused kernels for Hopper GPUs that do not.

Supported normalizers
Name	Probability for score \(z_i\)	Behavior
`softmax1`	\(\dfrac{\exp z_i}{1+\sum_j \exp z_j}\)	Softmax-N with \(n=1\). Weights can sum to less than one, which leaves room for a near-zero update.
`sparsemax`	\([\,z_i-\tau\,]_+\)	Euclidean projection onto the simplex. Can assign exact zeros.
`entmax15`	\([\,z_i/2-\tau\,]_+^{2}\)	1.5-entmax, between softmax and sparsemax. Also sparse.
`softmax` / `vanilla`	\(\dfrac{\exp z_i}{\sum_j \exp z_j}\)	PyTorch softmax, as a baseline.

For sparsemax and entmax15, \(\tau\) is the threshold that makes the output sum to one. Definitions follow the package README and vutils/ reference implementations.

Where these normalizers appear in our work

The package collects the attention variants used across a line of papers on attention outliers, so the same model code can be run with each of them.

OutEffHopDerives Softmax1 as the retrieval dynamics of an outlier-efficient modern Hopfield model. ICML 2024.
GERMReplaces the attention in a genomic foundation model with the outlier-free Hopfield layer, for cheaper low-rank adaptation and quantization. ICML 2025.
FROSTFine-tunes reasoning models with Softmax1 to suppress low-attention reasoning steps. ICLR 2026.
OASISUses Softmax1 null routes over tokens and depth, and compares Sparsemax and Entmax15 as alternative normalizers in its Table 3. NeurIPS 2026.

Backends

Each normalizer is registered under the standard backend names, prefixed with its own.

Native PyTorch SDPA, FlashAttention, and FlexAttention all fix softmax. Every custom name therefore maps to this project's own implementation: an SDPA-like path that materializes attention weights, an external Softmax-N FlashAttention 2 kernel, or Triton kernels written for Hopper.

Backend support · name pattern `<normalizer>_<backend>`
Backend	`softmax1`	`sparsemax`	`entmax15`
`*_eager`	Reference path	Reference path	Reference path
`*_sdpa`	Custom SDPA-like path	Custom SDPA-like path	Custom SDPA-like path
`*_flash_attention_2`	Softmax-N FA2 kernel, with `flash-attention-softmax-n` installed	Compatibility name, routed to the Hopper Triton kernel	Compatibility name, routed to the Hopper Triton kernel
`*_flash_attention_3`	Online tiled Hopper kernel	Hopper multi-block kernel	Hopper multi-block kernel
`*_triton`	Not provided	Fused forward and backward, single-block	Fused forward and backward, single-block
`*_flex_attention`	BlockMask path	BlockMask path	BlockMask path
`_paged_attention`, `_paged\|…`	Packed varlen Hopper grid	Packed varlen Hopper grid	Packed varlen Hopper grid

Transcribed from HF_ATTENTION_BACKENDS and _make_softmax_attention_forward in backends.py. The paged row follows the code and the README's continuous-batching section; the README's summary matrix still lists paged names as registration-only. Source ↗

How the fused kernels avoid the attention matrix

The Softmax1 FlashAttention 3 kernel treats the extra 1 in the denominator as a virtual key with logit 0 and value 0. It saves the row maximum and denominator, then recomputes probabilities in the fused backward pass. The sparsemax and entmax15 Hopper kernels stream key and value tiles, find each row's threshold \(\tau\) by tiled bisection, save one threshold per query row, and recompute tiles in the backward pass. The single-block Triton kernels require key_length ≤ 4096 by default; the Hopper kernels remove that limit.

Validated: Hopper (H100), BF16: full forward and backward and torch.compile(fullgraph=True) tests for all three normalizers; varlen and paged paths also in BF16
Pending: Validation on non-Hopper GPUs
FlexAttention: Causal, causal sliding window, and causal with right padding run without expanding a dense mask; other mask_mod functions use a dense compatibility path
Paged training: paged_triton_attention and DifferentiablePagedCache keep gradients flowing into physical K/V cache pages

From the README's “Important Limitations” section. Read the full list ↗

Usage

Loading with a registered name works like any native Hugging Face backend. An already loaded model can switch in place, and models outside the registry can be patched through a surgery policy.

Load a model with a custom normalizer

from transformers import AutoModelForCausalLM
from hf_attention_normalizers import register_softmax1_attention_backends

register_softmax1_attention_backends(mode="strict")

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    attn_implementation="softmax1_sdpa",
)

Switch a model that is already loaded

from hf_attention_normalizers import set_softmax_attention_backend

set_softmax_attention_backend(model, base_backend="sdpa", softmax_fn="entmax15")
set_softmax_attention_backend(model, base_backend="triton", softmax_fn="sparsemax")

Strict or fallback

Strict mode raises when a registered name has no custom kernel. Fallback mode keeps the math by routing that name to a fallback backend, sdpa by default, which may be slower than the backend requested.

register_softmax1_attention_backends(mode="strict")
register_softmax1_attention_backends(mode="fallback", fallback_backend="sdpa")

Patch attention modules directly (Qwen3)

from hf_attention_normalizers import apply_softmax_attention, supported_attention_policies

model = apply_softmax_attention(model, softmax_fn="softmax1", attn_implementation="sdpa")
print(supported_attention_policies())  # ("qwen3",)

Other model families can be added by registering an AttentionReplacementPolicy.

Packed variable-length attention

from hf_attention_normalizers import varlen_hopper_attention

output = varlen_hopper_attention(
    query, key, value,          # [total_tokens, heads, head_dim]
    cu_seqlens_q, cu_seqlens_k,
    normalizer="softmax1",
    max_seqlen_q=max_q, max_seqlen_k=max_k,
    is_causal=True,
)

Examples adapted from the package README. The package also exports paged_attention, paged_triton_attention, and DifferentiablePagedCache for paged K/V caches. Paged attention usage ↗

Install

The base package needs only PyTorch and Transformers. Optional extras add the external Softmax-N FlashAttention kernel and Triton.

Base: pip install hf-attention-normalizers
Softmax-N FA2: pip install "hf-attention-normalizers[flash-softmax-n]"
Triton kernels: pip install "hf-attention-normalizers[triton]"
Everything: pip install "hf-attention-normalizers[all]"
From source: pip install -e . in a clone of the repository

Requires Python 3.10 or later. The current release on PyPI is 0.1.0. The import name is hf_attention_normalizers.

Maintained by @robinzixuan on GitHub.

Cite

BibTeX · software release

Download

@software{luo_hf_attention_normalizers_2026,
  author  = {Luo, Robin},
  title   = {hf-attention-normalizers},
  version = {0.1.0},
  year    = {2026},
  url     = {https://github.com/robinzixuan/hf-attention-normalizers}
}