# TPA is nanoGPT with an attention swap and an eight-GPU bill > The official code for a NeurIPS 2025 Spotlight paper on Tensor Product Attention, built on nanoGPT, with data preparation, pretraining and evaluation scripts. FlashTPA decoding came in two pieces months after the code did. The entry requirement is eight 80GB GPUs. **tensorgi/TPA** — [NeurIPS 2025 Spotlight] TPA: Tensor ProducT ATTenTion Transformer (https://arxiv.org/abs/2501.06425) - Repository: https://github.com/tensorgi/TPA - Website: https://tensorgi.github.io/TPA - Stars: 463 · Forks: 41 - Language: Python - License: MIT - Published: 2026-09-10 · Updated: 2026-09-10 · Language: en - Canonical page: https://hysenlabs.com/projects/tensorgi-tpa ## The codebase is nanoGPT with a different attention The acknowledgment section gives the foundation away: Karpathy's nanoGPT is described as the foundational codebase this repository is built upon. That explains the shape of everything else. The training entry points are `train_adam.py` and `train_adam_fw.py`, there is a `pretrain.sh` to drive them, and a `configurator.py` at the root, all of which follow nanoGPT's conventions rather than a framework's. The contribution is the attention mechanism. TPA stands for Tensor Product Attention, and the repository is the official code for the paper Tensor Product Attention Is All You Need, whose stated aim is to improve performance while reducing KV cache size. The model built on it is called T6, for Tensor ProducT ATTenTion Transformer, and the acronym is spelled out in the title. The scope is narrow and stated plainly: tools for data preparation, model pretraining and evaluation. There is no inference entry point, no quantisation story and no checkpoint download, which is consistent with a paper repository rather than a released model. Seven authors are listed, two of them marked as equal contributors, and the paper appears in the NeurIPS 2025 proceedings as a spotlight. ## Eight 80GB cards is the entry fee The hardware section is two sentences and both of them matter. A100 and H100 are the recommended cards, and at least eight 80GB of VRAM is needed. That requirement is visible elsewhere too, because it is baked into the example configuration. The manual training command points at a file named `config/train_T6_medium_adam_80g8.py`, where 80g reads as the GPU memory and 8 as the GPU count, with medium naming the model size and adam the optimiser. Pretraining is run one of two ways. The provided script is a single command: ```bash bash pretrain.sh ``` For more control you invoke `torchrun` yourself, replacing the configuration file with your own: ```bash torchrun --standalone --nproc_per_node=8 \ train_adam_fw.py \ config/train_T6_medium_adam_80g8.py ``` The `--nproc_per_node=8` flag is annotated in the README as the number of processes, typically matching the number of GPUs. So the parallelism model is data parallel across devices, one process each, which is the nanoGPT arrangement. The dependency set is short: torch pinned to 2.4.0, plus numpy, transformers, datasets, tiktoken, wandb and tqdm. Python 3.10 or newer is required, and a virtual environment is recommended. ## FlashTPA arrived in two pieces, months apart The news section is the best record of how this project developed, and it is dated. The code was open-sourced and the paper posted on 11 January 2025. The paper was then updated on arXiv on 29 May 2025 specifically for FlashTPA decoding, which is the point where the extra algorithm enters. Decoding code for FlashTPA followed on 13 June 2025, and prefilling code on 4 July 2025. So the two halves of the fast path were not released together. Anyone who cloned in February had the base architecture and neither FlashTPA half, and anyone who cloned in June had decoding without prefilling. The implementations live in a `decode` directory at the repository root, in both Python and Triton, and the algorithms themselves are referenced by number: Algorithms 2 and 3 in the paper PDF. The PDF is committed to the repository rather than only linked, which is helpful because the code in that directory is not self-explanatory without it. Triton is what makes the second implementation worth having, and it is worth noting where it appears. The requirements file lists triton as a dependency, while the pip install line in the README does not mention it. The paper's acceptance as a NeurIPS 2025 spotlight was announced on 18 September 2025. ## Two dataset paths with different scripts Data preparation is where the scale of this project becomes concrete, and the two supported datasets are prepared by separate scripts in different directories. For FineWeb-Edu-100B, the large educational dataset hosted on Hugging Face, you change into the data directory, run the preparation script and then move the result up a level: ```bash cd data/fineweb-edu python fineweb-edu.py mv fineweb-edu100B .. cd ../.. ``` The move step is not cosmetic: the output directory name is versioned with the dataset size in it, and the training scripts expect to find it one level up. For OpenWebText, described as an open reproduction of OpenAI's WebText dataset, it is a single command from the repository root: ```bash python data/openwebtext/prepare.py ``` That path has a warning attached, and the warning is the point: ensure you have sufficient storage and computational resources, as OpenWebText is sizable. No sampling, mixing or tokenisation configuration is documented here. The scripts are run as they are, which is the nanoGPT assumption: the preparation step produces what the trainer wants and nothing else is exposed. ## Evaluation means cd into a directory and read it Evaluation is the thinnest part of the documentation, and the reason is that the harness is in the repository. There is a top-level `lm-evaluation-harness` directory, credited to EleutherAI in the acknowledgments, and the instruction is to navigate into it and follow the instructions within that directory. There is no command given, no configuration file named and no list of tasks. The only caveat offered is that your model should be compatible with the evaluation harness requirements, which is the kind of sentence that means there is an adapter to write. Vendoring the harness rather than depending on it is a deliberate choice and has a cost. It means the results are comparable against whatever version of the harness is committed here rather than the current release, and it means adapting a new model to the harness is manual work inside a directory you now maintain. The rest of the acknowledgments cover the data side: Hugging Face for the FineWeb-Edu-100B dataset, the OpenWebText team for the WebText replication, and nanoGPT as the base codebase. The citation block points at the NeurIPS proceedings entry, volume 38, with a specific page range and a 2026 year. ## No releases, and a paper PDF in the repository root Two facts about the distribution are worth stating before anyone plans around this project. There are no GitHub releases. Nothing is tagged, so there is no version to pin and no artefact to download. The last push to the default branch was on 2026-09-04, which is after the spotlight announcement, so the branch has moved since the paper was published without any marker telling you what changed. The model weights are not here either. The README links to a Hugging Face paper page, not a model repository, which is consistent with a paper release rather than a model release. What is in the repository root: a licence, the README, the paper as a PDF, then the directories that matter, being `config/` for the training configurations, `data/` for the preparation scripts, `decode/` for the FlashTPA implementations, `lm-evaluation-harness/` for evaluation and `model/` for the architecture, alongside the two training scripts, the shell wrapper and the configurator. That layout is short enough to read in a minute, which is a fair summary of how much this repository asks of you beyond the hardware. ## Conclusion Adopt TPA if you are studying attention architectures and have the hardware to match, since the repository is a research codebase with a training entry point rather than something you can load as a model. Do not expect released checkpoints or packaged weights, because there are no GitHub releases and the README documents training and evaluation rather than inference. Verify two things first. Count your GPUs against the stated requirement of at least eight 80GB cards, and note that the example configuration is named for exactly that setup. Then read the paper PDF committed at the repository root, since the FlashTPA algorithms are referenced by number in it rather than explained in the README. ## FAQ ### What is TPA and what does it change? TPA is the official code for the paper Tensor Product Attention Is All You Need, a NeurIPS 2025 spotlight. It implements Tensor Product Attention to improve model performance while reducing KV cache size, and the model built on it is called T6, the Tensor ProducT ATTenTion Transformer. ### What hardware does TPA need? A100 and H100 are recommended, and at least eight 80GB of VRAM is required. The example configuration is named for that setup, train_T6_medium_adam_80g8.py, and the manual training command uses torchrun with --nproc_per_node=8. ### What is FlashTPA in TPA? It is the extra decoding and prefilling algorithm described as Algorithms 2 and 3 in the paper. The decoding code was released on 2025-06-13 and the prefilling code on 2025-07-04, and both Python and Triton implementations live in the decode directory. ### How do I pretrain the T6 model? Either run the provided pretrain.sh script, or invoke torchrun yourself with train_adam_fw.py and a configuration file, replacing config/train_T6_medium_adam_80g8.py with your own. Preparation for FineWeb-Edu-100B and OpenWebText is done with separate scripts under data/. ### Does TPA publish model checkpoints? The repository has no GitHub releases and no tagged versions, and it links to a Hugging Face paper page rather than a model repository. It documents data preparation, pretraining and evaluation rather than inference, and the paper's PDF is committed at the repository root. ## Sources - [Issues](https://github.com/tensorgi/TPA/issues) - [License: MIT](https://github.com/tensorgi/TPA/blob/main/LICENSE) - [Project website](https://tensorgi.github.io/TPA) - [README](https://github.com/tensorgi/TPA/blob/main/README.md) - [tensorgi/TPA on GitHub](https://github.com/tensorgi/TPA) --- Hysen Labs editorial analysis, written from the project's own repository and release notes. Cite the canonical page: https://hysenlabs.com/projects/tensorgi-tpa