$ show proger/accelerated-scan
Accelerated Scan is a focused, high-performance library for solving first-order linear recurrences (associative scans) on GPU. It provides a pure CUDA kernel (`warp`) that uses warp shuffle intrinsics, shared memory, and a hierarchical chunked algorithm to achieve state-of-the-art speed, and a Triton-based implementation (`scalar`) that leverages `tl.associative_scan` for portability. A reference PyTorch implementation (`ref`) implements the classic Blelloch tree scan for validation. The library is designed for integration into state space models and linear RNNs, where the recurrence `x[t] = gate[t] * x[t-1] + token[t]` is central. The codebase is small but exceptionally well-crafted, with clear separation of concerns, comprehensive tests, and a clean API. The CUDA kernel in particular is a standout example of GPU optimization, using warp-level primitives, shared memory, and chunking to handle sequences of arbitrary length (up to 65536) with minimal overhead.
The LLM identified 3 files worth highlighting for reuse:
A highly optimized CUDA kernel implementing a hierarchical associative scan using warp shuffle intrinsics (level 1), shared memory between warps (level 2), and chunked processing across thread blocks. The kernel handles both forward and backward scans, supports multiple dtypes (bfloat16, float16, float32), and uses alignment-aware tuple types for memory coalescing. The `DISPATCH_SCAN` macro selects the optimal kernel configuration based on sequence length. This component is a standalone, reusable pattern for any GPU-based prefix sum or recurrence solver.
A pure PyTorch reference implementation of the Blelloch tree-based associative scan. It correctly handles arbitrary sequence lengths by padding to the next power of two, implements the classic up-sweep/down-sweep tree algorithm, and supports both forward and reverse scans. The code is clear, well-commented, and modular (split, merge, scan1). It is highly reusable as a reference for testing GPU kernels, as a teaching tool, or as a starting point for custom scan variants.
A Triton implementation of the chunked associative scan that uses `tl.associative_scan` within each block and maintains a carry across blocks. It includes both `forward_scan` and `backward_scan` kernels, wrapped in a `torch.autograd.Function` for GPU training. The backward pass implements the reverse scan with correct inter-block carry propagation. This component is a clean, reusable pattern for implementing scan operations in Triton, and can be adapted to other recurrence relations or combined with other Triton kernels.