ScanWeaver: Compiler-Driven Parallelization of Affine Recurrences via Associative Scan Lowering
For ML systems researchers, ScanWeaver provides a general compiler-driven approach to parallelize input-dependent recurrences, which are a bottleneck in modern sequence models like Mamba.
ScanWeaver is a compiler framework that transforms affine recurrences (e.g., Mamba-style selective scans) into associative scan representations and lowers them to GPU programs, achieving up to 2x speedup over sequential baselines on forward selective-scan workloads.
Selective state-space models such as Mamba highlight the practical importance of input-dependent scan recurrences, which preserve linear-time sequence modeling while improving language modeling capabilities. However, these recurrences introduce stricter sequential dependencies than classical structured SSMs, limiting parallel execution on modern accelerators. We present \textbf{ScanWeaver}, a compiler framework that transforms recurrence-based computations into associative scan representations and lowers them end-to-end to executable GPU programs. We use Mamba-style selective scan as a motivating example of a broader class of affine recurrences that arise in modern ML workloads. Rather than targeting a single model family, ScanWeaver elevates this recurrence structure to a first-class compiler abstraction, enabling systematic MLIR-based lowering to compiler-generated Blelloch scan execution on GPUs. Across forward selective-scan workloads with matched local recurrence semantics, we validate affine recurrence decomposition, Blelloch lowering, MLIR GPU lowering, executable artifact generation, and actual GPU execution from generated MLIR. We benchmark the resulting ScanWeaver GPU execution against PyTorch and CUDA sequential baselines, and include the Mamba kernel as a fused production baseline for systems context.