kl-div-attention

Just another attention variant, where the query key distance uses KL divergence, for testing Gitlab workflow

Install

pip install kl-div-attention

Usage

import torch
from kl_div_attention import KLDivAttention

attn = KLDivAttention(
    dim = 512,
    heads = 8,
    dim_head = 64,
    causal = True,
    prenorm = True, 
    fused_mode = 'flash'    # use fused triton flash attention
).cuda()

tokens = torch.randn(1, 1024, 512).cuda()

out = attn(tokens) + tokens

assert out.shape == tokens.shape

out.sum().backward()

Training

You can train a small transformer on Enwik8 with KL-divergence attention using the provided training script.

Eg. to use the fused Triton kernel

uv run train_enwik8.py --fused-mode flash

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

kl_div_attention-0.1.6.tar.gz (36.6 MB view details)

Uploaded Feb 15, 2026 Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

The dropdown lists show the available interpreters, ABIs, and platforms. Enable javascript to be able to filter the list of wheel files.

kl_div_attention-0.1.6-py3-none-any.whl (11.5 kB view details)

Uploaded Feb 15, 2026 Python 3

File details

Details for the file kl_div_attention-0.1.6.tar.gz.

File metadata

Download URL: kl_div_attention-0.1.6.tar.gz
Upload date: Feb 15, 2026
Size: 36.6 MB
Tags: Source
Uploaded using Trusted Publishing? No
Uploaded via: twine/6.2.0 CPython/3.10.19

File hashes

Hashes for kl_div_attention-0.1.6.tar.gz
Algorithm	Hash digest
SHA256	`96332fc96df992df02f36ae50e1f04339713298bbad32071742a6b3022228533`
MD5	`8415469e917dc20fcb0cef3f563f1b03`
BLAKE2b-256	`986642ab2aac0bbd2a41ffa7ecb458980c6eba1fd4f774f101593bfa12cf4932`

See more details on using hashes here.

File details

Details for the file kl_div_attention-0.1.6-py3-none-any.whl.

File metadata

Download URL: kl_div_attention-0.1.6-py3-none-any.whl
Upload date: Feb 15, 2026
Size: 11.5 kB
Tags: Python 3
Uploaded using Trusted Publishing? No
Uploaded via: twine/6.2.0 CPython/3.10.19

File hashes

Hashes for kl_div_attention-0.1.6-py3-none-any.whl
Algorithm	Hash digest
SHA256	`49a21e6198d996e2352426509a04dfe9547766109acc97b36ac448ab02cbd254`
MD5	`749d0ba45cde4fd1c0f070ebffc1c6e3`
BLAKE2b-256	`ce27783d898a1d6b7ac20c0366b88ff4871572227abcb3932cb178cd6e98bfba`