attorch: A Readable Triton Rewrite of a PyTorch nn Subset

The gap attorch fills between pure PyTorch and hand-written CUDA

The README frames the target user precisely: someone who wants to develop custom deep learning operations, is not satisfied with the speed of a pure PyTorch implementation, and does not have the technical expertise or resources to write CUDA kernels. attorch sits in that space by expressing neural network modules in Triton, a Python-embedded language for GPU kernels. The project describes itself as an easily hackable, self-contained, and readable collection of modules, and the stated intent is to be forkable rather than to be a framework you build a product on.

The audience matters because it shapes what counts as success. A team shipping a model does not need readable kernels; it needs stable APIs and predictable memory behavior. A researcher or engineer prototyping a fused operation does need to see how a layer normalization kernel loads rows, transforms features, and stores results without reading a thousand lines of CUDA. attorch is written for the second group. The README also positions it against existing Triton-based projects such as kernl, xFormers, Unsloth, and fla, noting that most of those concentrate on Transformers and NLP, while attorch includes layers aimed at other areas such as computer vision. Conv1d, Conv2d, AvgPool1d, and AvgPool2d are the concrete evidence for that claim.

What is actually implemented, and what the layer list implies

The README enumerates the implemented layers. Convolutions and pooling cover the vision side. MultiheadAttention covers the attention case. A long list of activations (ELU, GELU, SiLU, Mish, Hardswish, and roughly a dozen more) each optionally fuse dropout. Normalization includes BatchNorm1d, BatchNorm2d, LayerNorm, and RMSNorm. Linear supports an optional bias and an optional fused activation. Losses include L1Loss, MSELoss, CrossEntropyLoss, NLLLoss, HuberLoss, and SmoothL1Loss, with class reweighting available on the cross entropy and negative log likelihood variants.

The README states that unless a docstring says otherwise, these layers behave identically to their PyTorch equivalents. That sentence carries more weight than it appears to. Behavioural equivalence is a claim about numerics, broadcasting, and edge cases, and it is the kind of claim best verified per layer rather than assumed. The fusion options are the more interesting part of the design: BatchNorm1d and BatchNorm2d can fuse an activation and add a residual to the pre-activation result, and the activation layers can absorb dropout. Those fusions are where a Triton rewrite can plausibly pay off, because they remove intermediate memory traffic that a sequence of separate PyTorch ops would incur.

attorch.math and the load-math-store kernel structure

The README describes Triton kernels as composed of two parts: one that handles loading and storing tensors, and one that transforms data with mathematical functions. A layer normalization kernel reads one or several rows (load), standardizes the features (math), and writes results into a container (store). attorch.math exposes a selection of those pure math functions so that custom kernels and operation fusion are easier to write.

The design constraint here is explicit. Only forward passes are available in attorch.math, but because the functions are pure and perform no I/O, their gradients can be derived automatically through the triton-autodiff library. The README also notes that large portions of attorch's kernels could be refactored to call attorch.math instead of inlining the math, but that doing so would sacrifice the single-file, self-contained design. That is a real trade-off, not a marketing line: the project chooses duplication over a shared dependency to keep each kernel readable in isolation. If you fork attorch and start sharing math helpers across kernels, you are moving away from the property the project was built to preserve, and you should decide that consciously.

Getting it running: two pinned dependencies and a clone

Installation is deliberately minimal. The README states the only dependencies are torch==2.4.0 and triton==3.0.0, and instructs you to install those specified versions and clone the repository. There is no package on an index described in the material, no console entry point, and no configuration file with keys to set. Usage is by import: attorch.Conv2d, attorch.LayerNorm, attorch.MultiheadAttention, and so on.

The pinned versions are the practical constraint. torch==2.4.0 and triton==3.0.0 are exact pins, not minimums, and the README does not describe a compatibility range. In an environment where another package forces a different torch build, attorch is the thing that has to give. The absence of retrieved releases compounds this: there is no versioned artifact to pin against, so the commit you clone is the version you get, and upgrades are whatever lands on the main branch. The repository was last pushed in August 2026 according to the supplied metadata, but no release tags were retrieved, so treat the git history as the changelog.

The PyTorch fallback and what it tells you about coverage

The README has a dedicated PyTorch Fallback section in its table of contents, which is a signal worth reading carefully in the repository itself. A fallback path exists because Triton kernels are not available for every input shape, dtype, or device configuration, and a subset implementation will inevitably encounter cases its kernels do not handle. The supplied material does not describe the fallback's selection logic, its performance characteristics, or which layers use it, so that is something to inspect in the source before you rely on attorch in a training loop.

The practical consequence is that a benchmark of attorch's fast path does not tell you what your workload will do. If your batch shapes or dtypes route to the fallback, you are running PyTorch with extra indirection. The README also states that attorch fully supports both forward and backward passes, so it can be used during training, but that inference performance is generally not on par with dedicated inference engines. Combined with the fallback, this suggests attorch is best understood as a training-time and experimentation tool, not a serving runtime.

Where attorch is the wrong tool

Three limitations follow directly from the material. First, scope: this is a subset of torch.nn, so any model depending on a layer outside the enumerated list has no attorch equivalent, and the project's own framing as a subset means that gap is by design rather than a backlog item. Second, inference: the README concedes that performance for inference is generally not on par with dedicated inference engines, so choosing attorch for a latency-sensitive serving path is choosing against the project's stated strengths. Third, dependency rigidity: exact pins on torch and triton make attorch a poor fit for environments that must track upstream PyTorch releases.

There is also a maintenance question the material cannot answer. With no retrieved releases, no homepage, and a README that presents the project as forkable, the realistic adoption model is that you take ownership of the code you depend on. That is fine for a research group that wants a starting point for custom kernels. It is a poor fit for a team that needs a vendor-like dependency with a release cadence and a deprecation policy. If you need the latter, attorch is not the thing you are looking for, and no amount of reading the layer list will change that.

PyTorch itself, and where the two approaches diverge

The most direct alternative is plain PyTorch nn, which attorch is explicitly modelled on and claims behavioural equivalence to. The difference in approach is not accuracy but extensibility. PyTorch's kernels are compiled from C++ and CUDA and are not editable in a Python file; when you need an operation PyTorch does not provide, you write a custom extension in C++/CUDA or reach for torch.compile. attorch's kernels are Python, so the loop from reading a kernel to modifying it is much shorter. That is the whole proposition.

The cost is coverage and maturity. PyTorch supports far more modules, more dtypes, and more backends, and its behaviour is documented as a compatibility contract. attorch supports the layers it lists, with a fallback for the rest, and its equivalence claim is scoped by per-layer docstrings. If your work is standard model training on supported layers, PyTorch is the lower-risk choice and attorch offers nothing you cannot get elsewhere. If your work is writing or modifying the kernel itself, attorch gives you a readable reference implementation in the same language you would write your own kernels in, which is a genuinely different starting point from a CUDA codebase.

Licence, upgrade cost, and what to check before adopting

attorch is MIT licensed, which permits commercial use, modification, and redistribution provided the copyright notice and permission notice are preserved. That is permissive and unsurprising for a project of this kind. This is a description of the licence text, not legal advice; if you are folding attorch into a product, have your own counsel review the notice requirements and any third-party code the repository vendors.

The upgrade cost is dominated by the pinned dependencies. Moving to a newer torch or triton means validating that the kernels still compile and produce equivalent results, and because there are no retrieved releases, there is no changelog to scan for breaking changes. The README's Tests section is the mechanism the project provides for that validation, and the Examples section is where you would look for usage patterns. Before adopting, verify three things in the repository: which layers have working backward passes, how the PyTorch fallback decides to engage, and whether the tests cover your target layer, dtype, and device combination. If the tests do not cover your case, you are the test.

Editorial conclusion

Adopt attorch if you want to read, fork, or extend Triton kernels for standard layers and you can pin torch==2.4.0 and triton==3.0.0 in your environment. Do not adopt it as a drop-in replacement for a production training stack: the layer set is a subset, the README states inference performance is generally not on par with dedicated inference engines, and there are no retrieved releases to pin against. Before committing, check three things in the repository itself: which layers have working backward passes, how the PyTorch fallback is selected at runtime, and whether the test suite covers the specific layer and dtype combination you depend on.

Official sources

1. BobMcDear/attorch on GitHub
2. Issues
3. License: MIT
4. README