This repository is the accompanying implementation of our KATA paper 🔗.

• Kernels: kata/parallel_kata_attn.py (flash forward/backward), kata/spd/, kata/kernels/
• Model: kata/layer.py, kata/modeling.py registers with HF Auto* as model_type="kata"
• Benchmarks: bench/mqar_overwrite/ reproduces the paper's MQAR and overwrite tables

The environment is bootstrapped by a single script (uses uv:

• Python 3.11
• Torch 2.10 (cu128)
• fla==0.5.0
• prebuilt causal_conv1d / mamba_ssm wheels (for the Mamba2 baseline).

git clone https://github.com/ayghri/KATA && cd KATA
bash scripts/setup.sh            # ~5 min; re-running is safe
source .venv/bin/activate

Set your own credentials before training (the tokenizer and datasets are pulled from the Hub; metrics log to Weights & Biases):

export HF_TOKEN=<your_hf_token>          # HF Hub access (tokenizer + datasets)
export WANDB_API_KEY=<your_wandb_key>    # or run with wandb.mode=disabled

The trainer is a manual-DDP loop (train_llm.py), Hydra-driven and launched via torchrun (scripts/train.sh wraps it). Model/data/train configs compose from configs/ (the default train recipe is configs/train/manual.yaml). Two steps: pre-tokenize the corpus once, then launch a run.

# SlimPajama 15B tokens (the 340M paper corpus). Cache lands under data/.
bash scripts/prepare.sh data=slimpajama_15bt data.num_proc=16

data=fineweb_edu_10bt is also available (10B-token FineWeb-Edu). Tokenizer is the LLaMA SentencePiece vocab (32k); context length 2048.

# KATA-cat-M2, 340M, 4 GPUs. Runs a preflight check, then torchrun train_llm.py.
NUM_GPUS=4 bash scripts/train.sh kata_cat2_340m data=slimpajama_15bt

# equivalently, directly:
#   torchrun --nproc_per_node=4 train_llm.py model=kata_cat2_340m data=slimpajama_15bt

# any Hydra field can be overridden on the CLI:
NUM_GPUS=4 bash scripts/train.sh kata_delta_m1_340m train.lr=3e-4 wandb.mode=disabled

All 340M configs are iso-architecture (24 layers, hidden_size=1024, 16 heads, SwiGLU MLP, ctx 2048, tied embeddings) so only the token-mixer differs:

See configs/model/ for the complete set (short-conv kataconv_* variants, hybrids, etc.).

Downstream lm-eval-harness tasks (PIQA, HellaSwag, WinoGrande, ARC-e/c, BoolQ, LAMBADA, WikiText) run inline at fractions of training (configured in configs/main.yaml eval:). For standalone recall evaluation of a checkpoint:

python eval_recall.py checkpoint=/path/to/run/checkpoint-XXXX

Multi-node convention: launch one arch per node (each node's GPUs do DDP); runs are independent and log to the same W&B project.

These are controlled, state-matched comparisons on synthetic recall, run inside an unmodified official Zoology checkout plus a small patch that lives entirely in bench/mqar_overwrite/. They fill paper/tables/{mqar,overwrite}.tex and paper/figures/state_vs_score.pdf.

• MQAR — multi-query associative recall; measures recall capacity and length extrapolation (train $L\le256$, test $L$ up to 4096 = 16× training).
• Overwrite — each key is written twice; the target is the most recent value. Measures content-based forget/erase (train $L\le512$, test $L$ up to 3072).

See bench/mqar_overwrite/README.md for the full protocol and the train/test distributions (also documented in paper/appendix_benchmarks.tex).

KATA=$(pwd)                                      # this repo
bash bench/mqar_overwrite/setup.sh /path/to/work   # clones + patches official Zoology
cd /path/to/work/zoology
export PYTHONPATH=.:$KATA                         # so `import kata` resolves the kernels
export ZOOLOGY_RESULTS_DIR=/path/to/results

# run all methods (or a subset via MATCHED_METHODS=kata_cat_M2,delta_net_h64,...)
python resilient_launch.py zoology/experiments/matched_mqar.py
python resilient_launch.py zoology/experiments/matched_overwrite.py

# regenerate the paper tables + figure (point the last arg at your paper/ dir)
python $KATA/bench/mqar_overwrite/gen_tables.py $ZOOLOGY_RESULTS_DIR  <paper>/tables
python $KATA/bench/mqar_overwrite/plot_state.py $ZOOLOGY_RESULTS_DIR  <paper>/figures

Each method trains best-of-three over LR ${1,2,8}{\times}10^{-3}$; resilient_launch.py is idempotent (skips completed runs) and crash-isolated. The KATA kernels run on Ampere (sm_80, Triton 3.6+); the Hopper TMA path is auto-gated. Tables use \usepackage{xcolor} (best = bold+blue, second = underline).

from transformers import AutoModelForCausalLM, AutoConfig
import kata                      # registers model_type="kata" with HF Auto*

cfg = AutoConfig.for_model("kata", hidden_size=1024, num_hidden_layers=24,
                           num_heads=16, feature_map="kata_quadratic", spd_num_groups=2)
model = AutoModelForCausalLM.from_config(cfg)

Or call the flash SPD attention kernel directly:

from kata.parallel_kata_attn import parallel_kata_attn
o = parallel_kata_attn(q, k, v, num_groups=2)     # concat-SPD flash attention, (B,T,H,D) -> (B,T,H,D)

kata/                 production KATA package (kernels + HF model)
  parallel_kata_attn.py   flash SPD forward/backward (Triton)
  layer.py, modeling.py   HF-registered KataAttention / KataForCausalLM
  spd/, kernels/          SPD + delta kernels
configs/              Hydra configs: model/ (340M + 200M), data/, train/
scripts/              setup.sh, prepare.sh, train.sh, preflight.sh
train_llm.py          Hydra + manual-DDP trainer entry point (torchrun)
bench/mqar_overwrite/ synthetic recall benchmarks (self-contained, patches Zoology)

@misc{kata26,
  title={Kernelized Linear Attention: Breaking the Capacity Wall with Symmetric Cones}, 
  author={Ayoub Ghriss and Sourav Chakraborty},
  year={2026},
  eprint={2607.17419},
  archivePrefix={arXiv},
  primaryClass={cs.LG},
  url={https://arxiv.org/abs/2607.17419}, 
}