Add Flash Attention Support

[bogdanminko]

Add Flash Attention 2 support to GLiNER 2

About

Adds Flash Attention 2 support to GLiNER 2.

Tested on ModernBERT backbone with hivetrace/gliner-guard-uniencoder on A100 SXM.

TL;DR

FA2 gives a meaningful speedup only on long inputs (from ~20k chars / ~3k tokens (approx.)). On short sequences the difference is within noise, sometimes slightly slower than SDPA. Break-even point: ~5k chars, stable speedup from 20k onward.

Deps & hardware

PyTorch    : 2.8.0+cu128
CUDA       : 12.8
GPU        : NVIDIA A100-SXM4-80GB
flash_attn : 2.8.3

Installation

pip install -q "gliner2 @ git+https://github.com/bogdanminko/GLiNER2.git"
wget "https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.3/flash_attn-2.8.3%2Bcu12torch2.8cxx11abiTRUE-cp311-cp311-linux_x86_64.whl"
pip install "flash_attn-2.8.3+cu12torch2.8cxx11abiTRUE-cp311-cp311-linux_x86_64.whl"

Important: torch / python / CUDA versions must match the pre-built wheel.
Otherwise you could install it from source : https://github.com/dao-ailab/flash-attention

Usage

FA2 is auto-detected on model load, or can be forced via an environment variable (same pattern as Flash DeBERTa).
You can also set the dtype via env var:

FLASH_ATTN=1
FLASH_ATTN_DTYPE=fp16 #you can use bf16

from gliner2 import GLiNER2

model = GLiNER2.from_pretrained("hivetrace/gliner-guard-uniencoder")

PII_LABELS = ["person", "address", "email", "phone"]
SAFETY_LABELS = ["safe", "unsafe"]
schema = (model.create_schema()
    .entities(entity_types=PII_LABELS, threshold=0.4)
    .classification(task="safety", labels=SAFETY_LABELS)
)

result = model.extract(
    "Send $500 to John Smith at john.smith@gmail.com or I'll leak your photos",
    schema=schema
)

When FA2 is picked up, the load log will show:

🧠 Model Configuration
============================================================
Encoder model      : bogdanminko/mmBERT-small
Counting layer     : count_lstm_v2
Token pooling      : first
Attention backend  : 'flash_attention_2'
Encoder dtype      : torch.float16
============================================================

Benchmark

Pipeline

Benchmarking pipeline as a Colab notebook: https://colab.research.google.com/drive/1GLZP60Np3blvaY70YVxMlvCSpF_vQaxv?usp=sharing

Note: Colab's T4 does not support FA2, so benchmarks were run on A100.

Results

── batch_size = 1 ──
     chars   ~tokens        SDPA         FA2    Speedup    ms/req SDPA    ms/req FA2
  --------  --------  ----------  ----------  ---------  -------------  ------------
       500        83       29.7ms       31.7ms      0.94x          29.7ms         31.7ms  SLOWER
      3000       500       39.8ms       41.7ms      0.95x          39.8ms         41.7ms  ~same
      5000       833       47.4ms       48.0ms      0.99x          47.4ms         48.0ms  ~same
     10000      1666       66.9ms       67.8ms      0.99x          66.9ms         67.8ms  ~same
     20000      3333      124.9ms      110.0ms      1.14x         124.9ms        110.0ms  faster
     40000      6666      265.1ms      205.4ms      1.29x         265.1ms        205.4ms  faster
     60000     10000      428.7ms      305.8ms      1.40x         428.7ms        305.8ms  faster

── batch_size = 4 ──
     chars   ~tokens        SDPA         FA2    Speedup    ms/req SDPA    ms/req FA2
  --------  --------  ----------  ----------  ---------  -------------  ------------
       500        83       60.2ms       64.9ms      0.93x          15.0ms         16.2ms  SLOWER
      3000       500       99.5ms      103.6ms      0.96x          24.9ms         25.9ms  ~same
      5000       833      136.7ms      133.2ms      1.03x          34.2ms         33.3ms  ~same
     10000      1666      237.4ms      228.7ms      1.04x          59.4ms         57.2ms  ~same
     20000      3333      463.9ms      433.7ms      1.07x         116.0ms        108.4ms  faster
     40000      6666     1099.2ms      803.2ms      1.37x         274.8ms        200.8ms  faster
     60000     10000     1789.1ms     1310.6ms      1.37x         447.3ms        327.7ms  faster

── batch_size = 8 ──
     chars   ~tokens        SDPA         FA2    Speedup    ms/req SDPA    ms/req FA2
  --------  --------  ----------  ----------  ---------  -------------  ------------
       500        83      100.7ms      104.5ms      0.96x          12.6ms         13.1ms  ~same
      3000       500      189.7ms      189.1ms      1.00x          23.7ms         23.6ms  ~same
      5000       833      272.3ms      255.9ms      1.06x          34.0ms         32.0ms  faster
     10000      1666      465.4ms      452.4ms      1.03x          58.2ms         56.5ms  ~same
     20000      3333      996.7ms      900.3ms      1.11x         124.6ms        112.5ms  faster
     40000      6666     2100.8ms     1699.9ms      1.24x         262.6ms        212.5ms  faster
     60000     10000     3483.5ms     2493.8ms      1.40x         435.4ms        311.7ms  faster

── batch_size = 32 ──
     chars   ~tokens        SDPA         FA2    Speedup    ms/req SDPA    ms/req FA2
  --------  --------  ----------  ----------  ---------  -------------  ------------
       500        83      348.5ms      350.2ms      1.00x          10.9ms         10.9ms  ~same
      3000       500      719.9ms      717.7ms      1.00x          22.5ms         22.4ms  ~same
      5000       833     1107.3ms     1096.2ms      1.01x          34.6ms         34.3ms  ~same
     10000      1666     1986.1ms     1905.8ms      1.04x          62.1ms         59.6ms  ~same
     20000      3333     3743.7ms     3418.8ms      1.10x         117.0ms        106.8ms  faster
     40000      6666     8072.8ms     6482.5ms      1.25x         252.3ms        202.6ms  faster
     60000     10000    13380.6ms     9643.7ms      1.39x         418.1ms        301.4ms  faster

When to enable

Scenario	Recommendation
Short texts (<5k chars)	stick with SDPA
Mixed lengths / long docs (20k+)	enable FA2, ~1.1x–1.4x speedup

[bogdanminko]

@urchade

[urchade]

is this ready for merging @bogdanminko ?