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

Metadata

Release files for kl-div-attention 0.1.6

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for kl-div-attention 0.1.6
File	Size	Uploaded
kl_div_attention-0.1.6.tar.gz	36.6 MB	Feb 15, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for kl-div-attention 0.1.6
File	Interpreter	ABI	Platform	Reset
kl_div_attention-0.1.6-py3-none-any.whl 11.5 kB Feb 15, 2026	Python 3	none	any	Details

Total release size: 36.6 MB

Release files / kl_div_attention-0.1.6.tar.gz

Download URL	kl_div_attention-0.1.6.tar.gz
Size	36.6 MB
Tags	Source
SHA-256 checksum How to use checksums	`96332fc96df992df02f36ae50e1f04339713298bbad32071742a6b3022228533`
BLAKE2b-256 checksum How to use checksums	`986642ab2aac0bbd2a41ffa7ecb458980c6eba1fd4f774f101593bfa12cf4932`
Upload date	Feb 15, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.10.19`

Release files / kl_div_attention-0.1.6-py3-none-any.whl

Download URL	kl_div_attention-0.1.6-py3-none-any.whl
Size	11.5 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`49a21e6198d996e2352426509a04dfe9547766109acc97b36ac448ab02cbd254`
BLAKE2b-256 checksum How to use checksums	`ce27783d898a1d6b7ac20c0366b88ff4871572227abcb3932cb178cd6e98bfba`
Upload date	Feb 15, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.2.0 CPython/3.10.19`