ArmScan

Inspiration

Transformers have reshaped AI, but their memory requirements grow with context because they retain an expanding key-value cache. State-space models such as Mamba offer a compelling alternative: they process sequences in linear time while maintaining a fixed-size recurrent state.

That architecture should be a natural fit for CPUs. However, the selective-scan operation at the heart of Mamba remains heavily GPU-oriented. Existing PyTorch implementations of Mamba-1 fall back to a slow sequential recurrence on CPUs, while the official Mamba-3 implementations rely on Triton, TileLang, or CuTe kernels and provide no native CPU path.

We built ArmScan around a simple question: what could Mamba do if Arm64 cloud CPUs were treated as a first-class AI target rather than a fallback?

Arm processors already power an expanding share of cloud and data-center infrastructure, including AWS Graviton instances. Yet optimized AI inference still often assumes access to GPUs, CUDA, and a GPU-specific software stack. For many models and workloads, that adds cost, operational complexity, and a significant barrier to experimentation.

ArmScan is our submission to the Cloud AI track because it focuses directly on this gap: accelerating AI inference on scalable Arm64 cloud infrastructure while remaining inside the PyTorch ecosystem developers already use.

Our goal was not to create another standalone inference engine or require developers to convert their models into a new format. We wanted existing PyTorch applications and checkpoints to benefit from an optimized Arm implementation with minimal code changes.

What it does

ArmScan is a PyTorch-native selective-scan engine optimized for Arm CPUs. It combines hand-written Rust kernels, Arm NEON vectorization, and multicore execution to accelerate the recurrent operation at the center of Mamba models.

For Hugging Face Mamba-1 models, enabling the optimized path takes one call:

import arm_scan

arm_scan.patch()

Developers can continue using their existing PyTorch models and weights without converting them to another runtime or model format.

ArmScan implements separate native kernels for Mamba-1 and Mamba-3, whose recurrences and state representations differ. The library supports:

Mamba-1 selective scan
Mamba-3 SISO
Rank-r Mamba-3 MIMO
Forward and reverse traversal
Bidirectional and non-causal scans
Four-direction 2D cross-scans
PyTorch and NumPy interfaces

The same underlying scan primitive can therefore support language sequences, bidirectional models, and two-dimensional token grids.

To prove that ArmScan works beyond synthetic kernel tests, we ran three existing pretrained checkpoints through it:

Hugging Face mamba-130m-hf
state-spaces/mamba3-siso-187m
state-spaces/mamba3-mimo-187m

We did not create, train, fine-tune, or modify the weights of these models. The 130M and 187M labels simply describe the approximate number of parameters in the published checkpoints. We used their original pretrained weights as real-world integration and correctness tests.

For Mamba-1, ArmScan replaces the slow selective-scan prefill path while leaving the surrounding Hugging Face model unchanged. For Mamba-3, we reproduced the published model’s forward-pass plumbing on CPU and routed its selective scan through our kernel, again leaving the pretrained weights unchanged.

Cloud performance

We collected our headline measurements on a dedicated AWS Graviton4 c8g.16xlarge instance with 64 Arm Neoverse-V2 CPU cores. This is the environment the project is designed for: multicore Arm64 cloud infrastructure running AI inference through a standard framework.

ArmScan achieved:

38.69× multicore scaling from one to 64 cores
A 26× optimization ladder from sequential scalar execution to NEON plus multicore execution
Up to 18.89× faster Mamba-3 SISO scanning than the PyTorch recurrence
3,820 tokens per second on the real Mamba-3 SISO 187M model at a sequence length of 1,024
A 10.14× faster prefill and 3.77× end-to-end speedup on Hugging Face mamba-130m-hf with a 2,048-token prompt
Up to 92.5× faster 2D Mamba-3 scanning than the reference recurrence

The existing Mamba-3 SISO and MIMO checkpoints run end to end through our CPU implementation with 98.05% and 96.48% logit-argmax agreement, respectively, against the official GPU-backed implementation.

The official GPU implementation is itself not perfectly reproducible across independent processes because its Triton kernels are autotuned. Our results fall within the reproducibility range we measured from the official implementation itself.

Why ArmScan fits the Cloud AI track

ArmScan aligns directly with the Cloud AI track in four ways.

Arm64 cloud infrastructure

The kernel is designed for server-class Arm processors and measured on AWS Graviton4. Rayon distributes independent channels, batches, and heads across the available CPU cores, allowing the implementation to scale from smaller instances to many-core cloud systems.

Our 38.69× scaling result on 64 cores demonstrates that ArmScan is not merely code that happens to compile on Arm. It was designed around the parallelism and vector capabilities of Arm cloud CPUs.

AI inference performance

ArmScan accelerates the inference-time recurrence at the center of Mamba. It targets prefill and long-sequence processing, where the unoptimized PyTorch recurrence becomes increasingly expensive.

The project includes operator-level benchmarks, complete-model benchmarks, long-context measurements, and comparisons against both eager PyTorch and torch.compile. Every performance measurement is gated by correctness before timing is accepted.

Framework integration

Cloud AI infrastructure is only useful when developers can access it through the frameworks they already use. ArmScan registers its native kernel as a PyTorch custom operation and provides a drop-in patch for Hugging Face Mamba-1 models.

Developers do not need to convert their weights, move to a standalone runtime, or rewrite their applications around a proprietary interface. ArmScan works within the existing PyTorch ecosystem and composes with torch.compile.

Production-oriented developer workflows

ArmScan includes more than a kernel. It provides:

A versioned C ABI
Safe Rust implementations
PyTorch and NumPy interfaces
Platform-specific wheel packaging
Deterministic golden-vector validation
Continuous integration on Linux Arm64, Apple Silicon, and x86
Reproducible benchmark scripts
Fresh-machine setup and validation instructions
End-to-end checks using real pretrained models

This makes ArmScan a reusable Cloud AI software component rather than a one-machine optimization demonstration.

The project may eventually enable local and edge use cases, but those are not the basis of this submission. We did not optimize for phone battery life, mobile user experiences, robotic control, sensors, autonomy, or functional safety. Our implementation, measurements, integrations, and developer workflow are centered on Arm64 cloud inference, making Cloud AI the clearest and most accurate track.

Why it matters

ArmScan is not itself a language, healthcare, or financial model. It is enabling infrastructure: an optimized execution layer that can help future Mamba-based applications run on widely available Arm CPUs instead of depending exclusively on specialized GPU hardware.

Access to AI is often limited not only by the availability of models, but also by the cost and complexity of the hardware required to run them. Arm-based cloud instances provide another deployment option, particularly for inference workloads that do not require a GPU.

Potential applications include:

Healthcare and scientific computing: Mamba-based systems could process long medical records, biosignals, image sequences, or reconstruction workloads on private or institution-controlled Arm cloud infrastructure. This could be useful where GPU access is limited or sensitive data must remain within a controlled environment. Our MRI pipeline demonstrates the 2D execution topology, not a clinically validated medical product.
Finance and enterprise analytics: State-space models could analyze long transaction histories, time series, reports, or event streams on private cloud infrastructure where data governance and deployment control matter.
Cloud services: Mamba-based inference services could run on Arm cloud instances without requiring every workload to provision a GPU. ArmScan’s multicore design allows the same implementation to use different instance sizes.
Software development: Developers could test Mamba applications, run continuous integration, and reproduce model behavior on CPU-based cloud runners without maintaining a separate CUDA environment.
Education and research: The committed Mamba-3 golden vectors and CPU reference allow researchers and students to study and validate the recurrence without owning a GPU.
High-performance and long-sequence computing: Mamba’s linear-time recurrence and fixed-size state make it interesting for long sequences, time series, scientific data, and other workloads where a growing transformer KV cache can become costly.

ArmScan does not make every AI model small enough for every computer, and it does not replace GPUs for workloads that genuinely require them. Its contribution is more focused: it turns Mamba’s selective scan from a weak or missing CPU path into a native, validated Arm implementation.

Because ArmScan stays inside PyTorch and works with existing checkpoints, developers do not need to adopt a new runtime or convert their model weights. It lowers one infrastructure barrier between emerging Mamba architectures and production Arm64 cloud deployment.

How we built it

ArmScan has four main layers:

A safe Rust core implements the Mamba-1 and Mamba-3 recurrences.
Arm NEON kernels vectorize the computationally expensive operations.
Rayon distributes independent work across CPU cores.
A small C interface and Python integration expose the kernels as PyTorch custom operations.

A call flows from PyTorch through a registered custom operation, across a narrow C ABI, and into the Rust kernel. Raw pointers are confined to the interface boundary, while the numerical implementation remains in safe Rust.

Mamba-1 and Mamba-3 use different mathematical state representations, so we implemented them as separate kernels. Mamba-1 maintains a vector state for each channel, while Mamba-3 maintains a matrix state for each head. The two kernels share packaging, threading infrastructure, Python integration, and the correctness framework.

For Mamba-1, the optimized kernel uses a chunked two-pass design. It first performs the expensive direction-independent discretization and exponential work, then executes the recurrent scan over compact intermediate data. NEON vectorizes the numerical hot paths, while Rayon distributes independent channels and batches across CPU cores.

We designed bidirectional and 2D execution as traversal patterns over the core scan primitive. A 2D cross-scan flattens a token grid in four directions—left to right, right to left, top to bottom, and bottom to top—and combines the results.

This separation keeps the recurrence independent from the traversal topology. Once forward and reverse scans worked correctly, bidirectional, non-causal, and four-direction 2D execution became composition and layout problems rather than entirely new mathematical kernels.

Correctness was treated as part of the architecture rather than as a final testing step. Our validation system includes:

Golden outputs captured from the official Mamba-3 GPU kernels
An independent, high-precision CPU reference
Scalar, cache-blocked, and NEON parity checks
Golden vectors replayed through the real C interface
Property tests for forward and reverse traversal
Bit-identical output across different thread counts
Model-level checks using the original pretrained weights
Continuous integration on Linux Arm64, Apple Silicon, and x86

The committed golden vectors allow contributors to validate Mamba-3 without owning a GPU or downloading a model.

Challenges we ran into

The largest challenge was that Mamba-3 had no authoritative CPU reference implementation. Its upstream model code imports GPU kernels and cannot simply fall back to a plain PyTorch version. We therefore had to capture ground truth directly from the official kernels and independently reconstruct the recurrence from the paper and implementation.

Even generating that ground truth was difficult. A released upstream package predated a fix for silent forward-pass corruption on newer Blackwell GPUs. If we had recorded those corrupted results, we could have validated every downstream implementation against incorrect ground truth. After identifying the issue, we rebuilt the environment using a corrected upstream revision before capturing our final test data.

Reproducibility created another unexpected problem. Random inputs generated through framework-specific generators did not remain identical across versions and architectures. Floating-point transcendental functions also produced slightly different values on different machines. We replaced them with deterministic, version-stable input generation so the same golden cases could be replayed everywhere.

NEON optimization introduced expected numerical trade-offs. Polynomial exponential approximations and fused multiply-add operations do not reproduce a scalar reference bit for bit. Instead of hiding that difference, every golden test records its numerical error floor, and changes are gated against measured floating-point tolerances.

The official Mamba-3 GPU implementation introduced another complication. Its Triton kernels select configurations through timing-based autotuning, so two independent runs can produce slightly different results. We measured that reproducibility floor and evaluated ArmScan against it rather than claiming that an unreachable 100% agreement was required.

Real hardware also overturned some of our early conclusions. A traversal-pair optimization that appeared faster on a four-core x86 machine slightly regressed on a 64-core Graviton4 system because it exposed less parallel work. That result reinforced why headline performance numbers must come from the target hardware.

Finally, no official 2D Mamba-3 weights or authoritative public 2D implementation are currently available. We validated our implementation by comparing two independent algorithms to approximately 2.99e-16, but we do not claim downstream vision accuracy without published weights and an external oracle.

Accomplishments that we're proud of

We are especially proud that ArmScan is more than an isolated kernel benchmark. Existing Mamba-1 and Mamba-3 checkpoints run through the same PyTorch-accessible CPU library without changing their pretrained weights.

To the best of our knowledge, ArmScan provides:

The first Arm NEON selective-scan kernel exposed as a PyTorch custom operation
The first PyTorch-callable, NEON-optimized Mamba-3 scan
The first fast CPU implementation of the SS2D cross-scan
The first CPU implementation of 2D Mamba-3 in causal and non-causal forms

We are also proud of the strength of the correctness work. The optimized Rust implementation is checked against independent CPU references, captured official GPU outputs, real model weights, the public Python operation, and the actual C ABI used by applications.

Our Mamba-3 MIMO path deserves an honest distinction: it is a correctness result, not yet a performance result. Both published Mamba-3 model families run end to end on an Arm CPU, but the current MIMO kernel is scalar and remains slower than the optimized SISO path.

We also report an unexplained benchmark anomaly, document an optimization that reversed on real Arm hardware, and distinguish validated 2D operator correctness from unavailable 2D model accuracy.

For us, trustworthy Cloud AI infrastructure means publishing surprising or negative results alongside the headline numbers.

What we learned

We learned that implementing a mathematical recurrence is only a small part of building a dependable cloud inference component. Reproducible test inputs, independent references, ABI validation, tensor-shape checks, thread determinism, realistic model integration, packaging, and measurements on the target hardware were just as important as the NEON instructions themselves.

We also learned how reusable a well-designed primitive can be. Once forward and reverse scans were exposed cleanly, bidirectional, non-causal, and four-direction 2D execution became composition and layout problems rather than entirely new kernels.

Another major lesson was that optimization results do not always transfer between machines. A change that improves performance on a small x86 system may reduce available parallelism on a many-core Arm server. Cloud AI optimization must be evaluated on the server architecture and instance type where the software will actually run.

Most importantly, we learned that CPUs still have an important role in modern AI inference. State-space models have characteristics—linear-time execution, fixed-size state, and long-context efficiency—that align naturally with CPUs. Giving them a serious Arm implementation creates another option for scalable cloud inference, particularly when a GPU is unavailable or unnecessary.

What's next for ArmScan

Our next priority is a fully fused 2D selective-scan kernel. Graviton4 measurements show that layout and dispatch overhead can reach 46.1% at real application shapes, making fusion the clearest remaining performance opportunity.

We also plan to:

Add cache blocking and NEON vectorization to the rank-r MIMO kernel
Investigate BF16 state storage and SVE2 acceleration for exponential evaluation
Expand prebuilt Arm64 wheels and platform coverage
Integrate ArmScan with more Mamba-based PyTorch architectures
Add containerized and automated deployment paths for Arm cloud instances
Continue benchmarking across Graviton, Ampere, and other Arm cloud CPUs
Evaluate 2D Mamba-3 with real vision checkpoints when authoritative weights become available

ArmScan began as an attempt to accelerate one missing operation. It has grown into a reusable Cloud AI foundation for running the broader PyTorch Mamba ecosystem natively on Arm64 servers—using existing models, existing weights, scalable cloud infrastructure, and developer workflows designed for reproducibility.