flash-xlstm-exp 0.1.0

Fork-and-upgrade experiments on top of NX-AI TFLA mLSTM Triton kernels.

flash_xlstm_exp (fsp)

Experimental fork-and-upgrade of NX-AI's TFLA (Tiled Flash Linear Attention) mLSTM kernels, targeting further GPU/Triton-level optimization on top of the official implementation.

Built with technology from NXAI. See NOTICE and LICENSE — this repo is derived from NX-AI/mlstm_kernels, licensed under the NXAI Community License Agreement, not a permissive OSS license. Redistributing fsp (or derivatives of it) requires keeping the LICENSE + NOTICE files and the "Built with technology from NXAI" attribution intact. See the NXAI Community License for the (currently inapplicable to solo/small-scale use) large-enterprise commercial terms.

Upgrade roadmap

#	Upgrade	Status
1	Dynamic tile-size autotuner (triton.autotune in place of the static get_head_dim_block_size / hardcoded siz_b_LQ=32,siz_b_LKV=32 heuristic)	Implemented — fsp/kernels/autotune.py
2	Fused recurrent+parallel forward kernel (avoid materializing intermediate chunk states to HBM between the two kernel launches)	Planned
3	Quantized (fp8/int8) memory state C for long-context inference	Planned

Why GPU-only

TFLA's kernels are written in Triton and target NVIDIA/AMD GPUs via CUDA/ROCm. There is currently no Triton backend for mobile/ARM64-without-GPU environments (e.g. Termux on Android), so fsp's kernels can be written and code-reviewed in such an environment, but must be benchmarked and numerically validated on an actual CUDA (or ROCm) GPU — a local NVIDIA GPU, or a rented cloud GPU instance.

Layout

fsp/
├── vendor/                  # Baseline TFLA, vendored from NX-AI/mlstm_kernels
│   ├── triton_kernels/      #   - the @triton.jit kernels themselves
│   ├── torch_fw/            #   - PyTorch launchers/wrappers (forward pass only, so far)
│   └── utils/               #   - small shared dtype/validation helpers
├── kernels/
│   └── autotune.py          # Upgrade #1: autotuned forward-parallel kernel
└── benchmark/
    └── bench_autotune_vs_baseline.py
tests/
└── test_autotune_equivalence.py

Scope note: only the forward pass (recurrent + parallel kernels) is vendored so far. Backward-pass kernels (bw_kernel_parallel_dK/dQ/dV, bw_kernel_recurrent) will be vendored in a later batch once the forward-pass upgrades are validated on real hardware.

Install

pip install torch triton
pip install -e .

Usage

from fsp.kernels.autotune import mlstm_chunkwise__parallel_fw_Hintra_autotuned

matH_out, vecN_out, vecM_out = mlstm_chunkwise__parallel_fw_Hintra_autotuned(
    matQ, matK, matV, vecI, vecF,
    matC_states, vecN_states, scaMinter_states,
    chunk_size=128,
)

Drop-in replacement for fsp.vendor.torch_fw.fw_parallel.mlstm_chunkwise__parallel_fw_Hintra — same inputs/outputs, same math, but tile sizes are chosen empirically per shape instead of by static heuristic.

Running tests / benchmarks

Both require a CUDA-capable GPU:

pytest tests/test_autotune_equivalence.py
python -m fsp.benchmark.bench_autotune_vs_baseline