FlashSlice is the companion code of “Does Transolver really need a Transformer?” (Wen & Mishra, 2026). The paper finds that Transolver's accuracy comes from the slice/deslice coupling between a full-resolution point stream and a small set of physics tokens, not from the self-attention among those tokens. This repository provides the two things needed to use that result:

• Kernels that compute slice and deslice as two streaming passes over the points, without ever writing the N × H × G slice-weight tensor to memory. They reproduce eager PyTorch to within its own floating-point error, are bitwise deterministic, and train 1.35–1.7× faster in up to 25% less memory at the standard configuration and 4–6.5× faster in 85–96% less memory at large slice counts.
• The Transolver model, with every ablation of the paper available as a single flag.

It is not a training framework: there is no trainer, data pipeline or config system, so the model and the kernels drop into an existing codebase.

Contents

• Installation
• Quick start
• How it works
• Performance
• Numerical precision
• Supported shapes
• Ablations and results
• Tests and benchmarks
• Citation
• License

Installation

pip install flashslice
# or, from source
git clone https://github.com/Shizheng-Wen/flashslice.git
cd flashslice && pip install -e .

Requirements: Python ≥ 3.9, PyTorch ≥ 2.4, Triton ≥ 3.0 and an NVIDIA GPU. Measurements were taken on two GPUs: an NVIDIA GH200 with PyTorch 2.5 and Triton 3.0 (the paper's numbers), and an NVIDIA RTX 4090 with PyTorch 2.8 and Triton 3.4.

Supported GPUs. Tile configurations are tuned per GPU class and picked automatically from the device's shared memory per block: a Hopper table (H100/GH200, swept on a GH200) and an Ada table (swept on an RTX 4090), which also serves other GPUs with less shared memory, such as consumer Ampere and the A100. flashslice.kernels.tile_table() reports the class in use; set_tile_table(...) or FLASHSLICE_TILE_TABLE forces one. The default ieee dot mode needs no tensor cores; tf32/bf16 need Ampere or newer. For a new GPU, re-tune with bench/bench_kernels.py and bench/pick_tiles.py.

Quick start

The model

import torch
from flashslice import Transolver

model = Transolver(
    space_dim=3, fun_dim=1, out_dim=4,   # coordinates, per-point inputs, outputs
    n_hidden=256, n_heads=8, n_layers=8,
    slice_num=32,                        # G, slices per head
    use_fused_slice=True,                # route slice/deslice through the kernels
).cuda()

x  = torch.randn(1, 200_000, 3, device="cuda")   # coordinates
fx = torch.randn(1, 200_000, 1, device="cuda")   # per-point inputs
y  = model(x, fx)                                # (1, 200000, 4)

use_fused_slice is an implementation switch, not a model change: outputs match the eager path, and checkpoints are interchangeable in both directions.

The kernels on their own

x_mid (B, N, H, D) is what the slice weights are computed from, fx_mid (B, N, H, DV) is what is pooled, W is the slot projection, b a per-slot bias (or None) and tau (H,) the temperature.

from flashslice.kernels import fused_slice, fused_deslice

# slice, then deslice (Transolver's order); the deslice reuses the slice's statistics
z_num, s, stats = fused_slice(x_mid, fx_mid, W, b, tau, return_stats=True)
tokens = mix(z_num / (s + 1e-5)[..., None])                   # (B, H, G, DV)
out = fused_deslice(x_mid, W, b, tau, tokens, stats=stats)    # (B, N, H, DV)

# deslice, then slice (a persistent point stream)
out, stats = fused_deslice(x_mid, W, b, tau, tokens, return_stats=True)
z_num, s = fused_slice(x_mid, fx_mid, W, b, tau, stats=stats)

W may be shared (G, D), per head (H, G, D) or per sample and head (B, H, G, D) (e.g. W = k(tokens)); its gradient is returned in the same shape, and b follows likewise. DV may differ from D.

How it works

The eager layer materializes the slice weights w of shape (B, H, N, G) and keeps them for the backward pass; for N in the millions this is the largest activation of the layer and the only one that grows with G. FlashSlice applies FlashAttention's trade to this pooling: each pass streams over the points in tiles, forms the tile of w in registers, uses it and discards it, and the backward pass recomputes it. The token state between the passes is only H × G × D floats (32 KB at the standard configuration).

The kernels use no atomics, so two runs are bitwise identical, and they are registered as PyTorch custom operators, so the model compiles into a single torch.compile graph.

Performance

One GH200, eager PyTorch against FlashSlice, no compilation or activation checkpointing. Full model (L=8, C=256, H=8, G=32):

One layer (H=8, D=32) at large slice counts, bf16 inputs and dots:

What changes most is how the layer scales:

• Memory is flat in the slice count. From G=16 to G=128 a bf16 training step stays at 15.4 GB with FlashSlice, while eager grows from 17.9 to 36.9 GB. One layer holds 2 GB whether G is 48 or 1024.
• Depth reaches further. Within 95 GB, FlashSlice trains 32 layers where eager fits 16 (fp32) and 48 where eager fits 32 (bf16).
• It composes with the usual tools. Against a torch.compiled eager baseline the kernels still train 1.24–1.26× faster; with activation checkpointing they are 1.31–1.42× faster; sharded over four GPUs they train 16.8M points at 14.6M points/s, 2.6× the best achievable eager pipeline.

Sizing rule: fused memory scales with B·N and the number of layers, not with G — about 8 GB per million points per layer for a bf16 training step and 2 GB per million for inference (H=8, D=32).

On an RTX 4090 (24 GB, torch 2.8, Triton 3.4, the Ada tile table), one layer (H=8, D=32), eager against FlashSlice in the same process:

bf16 rows use bf16 dots. The Hopper table does not transfer: several of its tiles exceed the 4090's 99 KB of shared memory, and among those that fit, the tuned table is up to 4.1× faster (G=64, fp32 training step: 18.1 against 74.2 ms). Every case of bench/parity_test.py passes on this GPU. Two checks of bench/parity_ops.py do not: the temperature gradient of the blocked kernels with per-sample weights, slice-first order, lands at 2.8–3.4× eager's error against the 2× gate (relative error ~2×10⁻⁵), with either tile table, so it is a numerics difference of this GPU and Triton version rather than a tiling one.

Numerical precision

Inputs are fp32, or bf16 under torch.autocast; parameters stay fp32 and accumulation is fp32 in every mode. The precision of the dots is chosen with set_dot_mode(...) or the FLASHSLICE_DOT_MODE environment variable:

The logits dot stays in fp32 below the bf16 level because the softmax Jacobian amplifies noise in the slice weights. Every mode is held to a parity gate against an fp64 reference: the fused error on every output and parameter gradient must stay within 1.25× of eager fp32's error (2× of eager bf16's for the bf16 modes), and two fused runs must agree bitwise. The learned temperature, whose gradient cancels row by row in exact arithmetic, is the most sensitive quantity and the one the gate watches most closely.

Supported shapes

Two kernel families sit behind use_fused_slice, and the layer picks one by shape:

D > 256 falls back to the eager path visibly: the layer logs a warning, sets use_fused_slice = False on the model and records the reason in fused_slice_fallback, so an inert flag can never silently train the baseline. set_kernel_mode("blocked") forces the blocked family (for timing and parity runs), and set_block_g overrides its block size.

Ablations and results

Every variant of the paper is one flag on Transolver; at most one may be set at a time.

use_fused_slice combines with every flag except share_slice_across_layers and slice_once, which need the slice weights the kernels never materialize; those combinations raise at construction.

Best validation relative L¹ on nine benchmarks in fluid dynamics and industrial aerodynamics, with up to 1.3×10⁸ mesh points per sample. Baseline and NoTokenAttn are mean ± std over three seeds, the others one seed; bold marks the best per row. Three significant figures, standard deviations to one.

• Token attention is removable. NoTokenAttn is within −2.6% to +1.8% of the baseline on seven benchmarks and at +3.8% and +7.9% on the other two, and it is the best variant on four (tied with the baseline on SHIFT-Wing volume).
• The coupling is not. Removing it (MlpOnly) costs 1.04–5.6×, and widening the pointwise model to 112% of the baseline's parameters closes almost none of that gap.
• Nor is the point stream. SliceOnce keeps the attention and has more parameters than the baseline, but collapses the points into token space after one slice; it is the worst variant everywhere, at 2.2–6.8×.

The controls behind these readings: (a) whether the projections are tied or untied does not interact with removing the attention; (b) NoTokenAttn stays within 1.5% of the baseline for every slice count from 8 to 256; (c) a wider pointwise model does not recover the coupling; (d) on every aerodynamic benchmark, removing the coupling hurts the pressure, the globally determined field, more than the velocity or wall shear beside it.

Tests and benchmarks

pytest tests/                                            # dims contract, routing, visible fallback, eager equivalence
python bench/parity_test.py                              # the layer: fused vs eager, fp64-referenced, both families
python bench/parity_ops.py                               # the ops: weight layouts, value widths, both call orders
python bench/bench_layer.py --slices 256 --dtype bf16    # one layer, eager vs fused, any shape
python bench/bench_kernels.py --defaults-only            # per-kernel time, bandwidth, TFLOP/s
python bench/bench_kernels.py --family blocked --slices 256   # tile sweep for the blocked kernels
python bench/pick_tiles.py bench/results/sweep_*.json    # sweep results -> tile-table entry

bench/README.md lists what each script measures and the pitfalls that matter when timing these kernels. Tag v0.1.0 is the tree the paper's numbers were measured on; CHANGELOG.md lists what main adds, and ROADMAP.md what comes next.

Citation

If you use FlashSlice or build on its findings, please cite:

@article{wen2026transolver,
  title   = {Does Transolver really need a Transformer?},
  author  = {Wen, Shizheng and Mishra, Siddhartha},
  journal = {arXiv preprint arXiv:2609.32525},
  year    = {2026},
  url     = {https://arxiv.org/abs/2609.32525}
}

License

Apache License 2.0; see LICENSE.

The physics-attention layer, the MLP and the block structure derive from Transolver (Copyright © 2024 THUML @ Tsinghua University), used under the MIT License; those files carry a header saying so, and NOTICE reproduces the MIT notice in full. The Triton kernels, the ablation flags and the instrumentation are original work. LinearNO (Hu et al., AAAI 2026), which the paper compares against, is not included; see the authors' release at https://github.com/HiPRL/LinearNO.