HF Attention Normalizers

Utilities for experimenting with alternative attention normalizers in HuggingFace Transformers.

Package name:

pip install hf-attention-normalizers

Python import name:

import hf_attention_normalizers

This project has two integration paths:

1. Register HuggingFace-style attention backend names such as softmax1_sdpa.
2. Patch supported model modules directly with a surgery policy. Currently the built-in policy covers Qwen3.

Supported Normalizers

Normalizer	Function	Notes
softmax1	exp(x_i) / (1 + sum_j exp(x_j))	Also called Softmax-N with n=1.
sparsemax	sparse probability transform	Can produce exact zeros in attention weights.
entmax15	1.5-entmax transform	Between softmax and sparsemax; also sparse.
vanilla / softmax	PyTorch softmax	Baseline.

Backend Support Matrix

Backend name	softmax1	sparsemax	entmax15
*_eager	supported	supported	supported
*_sdpa	supported via custom SDPA-like path	supported via custom SDPA-like path	supported via custom SDPA-like path
*_flash_attention_2	native FA2 extension if flash-attention-softmax-n is installed	supported as a Hopper-Triton compatibility route	supported as a Hopper-Triton compatibility route
*_triton	not used	fused forward and backward	fused forward and backward
*_flash_attention_3	custom Hopper online kernel	custom Hopper multi-block kernel	custom Hopper multi-block kernel
*_flex_attention	BlockMask compatibility path	BlockMask compatibility path	BlockMask compatibility path
`*_paged	*`	registered only; strict mode raises	registered only; strict mode raises

flash-attention-softmax-n implements the native FA2 Softmax-N kernel. FA2 itself cannot express sparsemax/entmax15; their *_flash_attention_2 names are HuggingFace-compatible aliases that route to the custom Hopper Triton kernels, not to official FA2 CUDA.

Install

Minimum runtime:

pip install torch transformers

Editable local install:

pip install -e .

Optional Softmax-N FlashAttention backend:

pip install "hf-attention-normalizers[flash-softmax-n]"

Optional experimental sparsemax/entmax Triton kernels:

pip install "hf-attention-normalizers[triton]"

Install every optional backend:

pip install "hf-attention-normalizers[all]"

HuggingFace-Native Usage

Register backends, then use attn_implementation just like native HuggingFace backends:

from transformers import AutoModelForCausalLM
from hf_attention_normalizers import register_softmax1_attention_backends

register_softmax1_attention_backends(mode="strict")

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-...",
    attn_implementation="softmax1_sdpa",
)

For sparsemax:

from hf_attention_normalizers import register_sparsemax_attention_backends

register_sparsemax_attention_backends(mode="strict")

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

Experimental sparsemax Triton backend:

from hf_attention_normalizers import register_sparsemax_attention_backends

register_sparsemax_attention_backends(mode="strict")

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

For entmax:

from hf_attention_normalizers import register_entmax15_attention_backends

register_entmax15_attention_backends(mode="strict")

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

Experimental entmax15 Triton backend:

from hf_attention_normalizers import register_entmax15_attention_backends

register_entmax15_attention_backends(mode="strict")

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

Softmax1 FlashAttention 2

If flash-attention-softmax-n is installed, you can use:

from hf_attention_normalizers import register_softmax1_attention_backends

register_softmax1_attention_backends(mode="strict")

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

This path calls flash_attention_softmax_n.flash_attention_n with softmax_n_param=1.

Strict vs Fallback Mode

Strict mode is safest:

register_softmax1_attention_backends(mode="strict")

Unsupported fused or paged backend names raise a clear error.

Fallback mode preserves math but may not preserve the requested backend performance:

register_softmax1_attention_backends(mode="fallback")

For unsupported fused or paged names, fallback mode routes to the configured fallback backend, defaulting to custom sdpa.

Already Loaded Models

You can switch an existing model:

from hf_attention_normalizers import set_softmax_attention_backend

set_softmax_attention_backend(
    model,
    base_backend="sdpa",
    softmax_fn="softmax1",
)

For sparsemax:

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

For entmax:

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

For the experimental Triton kernels:

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

Differentiable Paged Attention

The paged reference path supports Softmax1, sparsemax, and entmax15 with gradients accumulated back into physical K/V cache pages:

from hf_attention_normalizers import paged_attention

output = paged_attention(
    query,
    key_cache,
    value_cache,
    block_table,
    sequence_lengths,
    block_size=32,
    normalizer="sparsemax",
)

The CUDA Triton path supports Softmax1, sparsemax, and entmax15 with fused attention forward/backward, variable sequence lengths, GQA, and physical-page gradient accumulation:

from hf_attention_normalizers import paged_triton_attention

output = paged_triton_attention(
    query,
    key_cache,
    value_cache,
    block_table,
    sequence_lengths,
    block_size=32,
    normalizer="entmax15",  # or "softmax1" / "sparsemax"
)

The three normalizers also register HuggingFace continuous-batching backends (softmax1_paged_attention, sparsemax_paged_attention, and entmax15_paged_attention). They follow Transformers' paged_attention interface, update its cache, then run one packed varlen Hopper grid over the active requests. This interface is for the standard generation cache; use paged_triton_attention with DifferentiablePagedCache when gradients to physical cache pages are required during training.

The Triton implementation reads physical K/V pages directly through the block_table; it does not materialize gathered K/V tensors or the quadratic attention matrix. Backward atomically accumulates K/V gradients directly into their physical cache pages.

For training batches whose block_table has no shared physical pages, paged_triton_attention(..., max_key_length=..., assume_unique_pages=True) uses a two-stage KV-tile reduction without atomics and supports BF16 torch.compile(fullgraph=True). The flag is an explicit correctness contract: leave it at its default (False) whenever pages can be shared.

For incremental HuggingFace training, DifferentiablePagedCache uses functional page writes so gradients remain connected across cache updates:

from hf_attention_normalizers import DifferentiablePagedCache

cache = DifferentiablePagedCache(
    num_hidden_layers=model.config.num_hidden_layers,
    block_table=block_table,
    block_size=32,
    num_key_value_heads=model.config.num_key_value_heads,
    head_dim=model.config.head_dim,
    dtype=torch.float16,
    device=torch.device("cuda"),
)

outputs = model(
    input_ids,
    past_key_values=cache,
    use_cache=True,
    cache_position=cache_position,
)

Unlike HuggingFace's inference-oriented paged cache, writes are out-of-place index_copy operations. This preserves gradients to both earlier cached K/V states and newly written K/V states.

Packed Variable-Length Attention

Softmax1, sparsemax, and entmax15 support FlashAttention-style packed Q/K/V layouts with cumulative sequence lengths:

from hf_attention_normalizers import varlen_hopper_attention

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

The packed path supports GQA, different Q/K lengths, per-request bottom-right causal alignment, and full backward. Each normalizer launches one combined Hopper grid across the packed batch; inactive programs/tiles are masked from the cumulative sequence-length arrays, without padded Q/K/V staging.

Qwen3 Surgery Path

For models that do not use HuggingFace AttentionInterface, use the surgery path:

from transformers import AutoModelForCausalLM
from hf_attention_normalizers import apply_softmax_attention

model = AutoModelForCausalLM.from_pretrained(model_id)
model = apply_softmax_attention(
    model,
    softmax_fn="softmax1",
    attn_implementation="sdpa",
)

Current built-in surgery policies:

from hf_attention_normalizers import supported_attention_policies

print(supported_attention_policies())
# ("qwen3",)

Additional model families can be added by registering an AttentionReplacementPolicy.

The old Qwen_attention.py module is kept as a compatibility shim, but new code should import from hf_attention_normalizers.

Important Limitations

• Native PyTorch SDPA and native FlashAttention kernels do not expose a softmax_fn argument.
• softmax1_flash_attention_3 uses an online tiled Hopper kernel. It treats the Softmax1 denominator term as a virtual logit-0/value-0 sink, saves row max/denominator statistics, and recomputes probabilities during fused backward.
• sparsemax_flash_attention_3 and entmax15_flash_attention_3 use this project's Hopper-oriented multi-block Triton kernels. They stream K/V tiles, solve the global normalization threshold by tiled bisection, save one threshold per query row, and recompute tiles during fused backward. They are FA3-style custom-normalizer kernels, not wrappers around the official fixed-softmax FA3 CUDA kernel.
• Native FlexAttention exposes score and mask modifiers but keeps softmax fixed. Pure causal, causal sliding-window, and standard causal + right-padding HuggingFace BlockMask objects route to Hopper forward/backward without dense-mask expansion. Explicit full-block CSR masks (BlockMask.from_kv_blocks with noop_mask) also avoid dense expansion and preserve backward through direct block gathers. The right-padding path currently dispatches one fused kernel per batch row because each row has a different valid K prefix. Left padding, sliding-window + padding, and arbitrary element-level mask_mod functions use the differentiable dense compatibility path.
• *_sdpa in this project is a custom SDPA-like implementation that materializes attention weights so it can apply another normalizer.
• softmax1_flash_attention_2 uses flash-attention-softmax-n, not HuggingFace's native FlashAttention 2 kernel.
• sparsemax_triton and entmax15_triton fuse QK, normalization, and PV in the forward pass without materializing the attention matrix.
• Their fused backward kernels recompute probabilities row-by-row instead of storing the quadratic attention matrix, then directly accumulate Q/K/V gradients.
• The single-block experimental Triton kernels require CUDA tensors, no dropout, no external padding mask on the fast path, matching Q/K/V head dimensions, and key_length <= max_block_n where the default is 4096. The Hopper FA3-style kernels support training dropout with a saved Philox-style seed and recompute the mask in backward without storing a dense attention mask.
• If a mask/dropout/CPU path is encountered through the HuggingFace wrapper, the Triton wrapper falls back to the custom SDPA implementation to preserve math.
• The Hopper multi-block kernels remove the single-block 4096-key limit. Further work remains for native BlockMask traversal and deeper Hopper-specific scheduling/autotuning.

On Hopper (H100), the FA3-style kernels are covered in BF16 with full forward/backward and torch.compile(fullgraph=True) regression tests for all three normalizers. Varlen and direct physical-page Paged paths are also validated in BF16. Validation on non-Hopper GPU architectures remains pending.

• Paged Softmax1, sparsemax, and entmax15 kernels read K/V directly through block tables and scatter K/V gradients directly into physical cache pages.