PaperlayerView exact revision on arXiv

Source-linked AI summary

Mamba-Family State-Space Model Kernels on a Programmable CGLA

Takuto Ando, Yasuhiko Nakashima

arXiv:2609.27437v1cs.AR

TL;DR

Edge inference is constrained by power and data movement, while Mamba-family models combine dense projections, short-reduction SSD kernels, and recurrent-state updates. This paper maps and measures those kernel groups on programmable IMAX from kernel execution to token generation. Projections fit the long-reduction pipeline, while SSD Step-1 and decode-time projection expose boundary and GEMV bottlenecks.

  • Problem

    Edge deployment of sequence models is constrained by inference power and data movement, and Mamba-family inference combines distinct projection, SSD, and recurrent-state kernel demands.

    Source: §I · p. 1

  • Method

    The paper implements Mamba-family kernel groups on programmable IMAX and measures them from kernel execution through token-level integration.

    Source: §I · p. 1

  • Results

    Projections fit the long-reduction IMAX pipeline, unfused SSD Step-1 is limited by short reductions and kernel boundaries, and Mamba-130M decoding reaches 0.345 tokens/s with projection GEMV dominant.

    Source: B. IMAX Kernel Measurements · p. 5, §VI · p. 6

  • Takeaways & Limitations

    The results support programmable CGLAs for long-reduction projections, while identifying boundary reduction and persistent-weight execution as needs for SSD and decode-time projection support.

    Source: §VI · p. 6

    The token run has no matched end-to-end throughput baseline, so the study does not demonstrate system-level acceleration.

    Source: §VI · p. 6

Abstract

from arXiv · show

Edge and embedded inference is constrained by power and data movement. Mamba-family state-space models replace attention with sequence-linear recurrence, but their inference path combines dense projections, short-reduction SSD kernels, and recurrent-state updates. This paper maps these kernel groups onto IMAX, a programmable CPU-Grounded Linear Array (CGLA), and measures them from kernel execution to token-level integration. Projection kernels match the long-reduction IMAX pipeline, whereas SSD Step-1 is limited by short reductions and kernel-boundary overheads. Mamba-130M token-level integration identifies projection GEMV as the decode bottleneck. These results show that programmable CGLAs fit long-reduction projection kernels, while SSD and decode-time projection support require boundary reduction and persistent-weight execution.

I. INTRODUCTION

The paper treats Mamba-2 as a mix of kernel types with distinct hardware demands, then measures their fit on programmable IMAX from individual kernels through token generation. Its results identify both suitable paths and bottlenecks in the current CGLA stack.

I. INTRODUCTION

  • Mamba-2 layers combine long dense projections, short-reduction SSD kernels, and sequential recurrent-state updates, with different reduction, pipeline, data-movement, and recurrence demands.The study characterizes which programmable-CGLA paths match these kernel types without requiring a fixed-function Mamba datapath.Source: §I · p. 1
  • The paper implements projection, SSD Step-1, Hadamard, and recurrent-state kernels on IMAX under a common profiling path.It measures IMAX execution time, DMA-inclusive time, ARM Cortex-A72 software reference, and token-integration behavior.Source: §I · p. 1
  • The study integrates IMAX offload into Mamba-130M autoregressive decoding and finds projection GEMV, not recurrent-state update, is the dominant decode bottleneck.The same programmable CGLA stack is used across projection, SSD Step-1, recurrent-state update, and token generation.Source: §I · p. 1

II. RELATED WORK

Prior Mamba hardware work includes specialized FPGA and ASIC designs, while this paper uses Mamba as a diagnostic workload for an existing programmable CGLA. Its focus is identifying which phases fit that pipeline and which encounter bottlenecks.

II. RELATED WORK

  • IMAX belongs to the programmable spatial-fabric line, and prior IMAX studies evaluate GEMM, LLM, and Whisper workloads.Table I summarizes how Mamba/SSM accelerators differ from this work's programmable-CGLA diagnostic scope.Source: §II · p. 1, Table I · p. 2
  • Prior Mamba hardware studies target quantized FPGA dataflows, operator specialization or fusion, speculative decoding, and recurrent integration alongside systolic GEMM.These designs treat Mamba inference as more than a GEMM problem, with long reductions, recurrences, element-wise operations, and memory traffic posing different datapath needs.Source: §II · p. 1
  • Unlike Mamba-specific FPGA and ASIC accelerators, this work evaluates an existing programmable CGLA without model changes or fixed-function SSM units.Table I positions the paper as a programmable-CGLA diagnostic study rather than a dedicated Mamba accelerator design.Source: §II · pp. 1–2, Table I · p. 2

A. Mamba-2 Layer with Three Kernel Shapes

The Mamba-2 layer is described as three kernel shapes: dense projections, intra-chunk SSD Step-1, and sequential inter-chunk state updates. Their reduction lengths and recurrence dependencies distinguish their mapping demands.

A. Mamba-2 Layer with Three Kernel Shapes

  • The dense in-proj maps input u to concatenated features [z, x, B, C, ∆], and out-proj is a second dense GEMM.For Mamba-370M, the in-proj inner dimension is 1,024 and the out-proj inner dimension is 2,048; both projection kernels map to the IMAX GEMM path.Source: A. Mamba-2 Layer with Three Kernel… · p. 2
  • SSD Step-1 computes intra-chunk operations after dividing the sequence into chunks of length Q, using GEMM1 and GEMM2 reductions of N and Q.Both reductions are short relative to the 64-stage IMAX FMA chain; the representative measured configuration uses N=Q=64.Source: A. Mamba-2 Layer with Three Kernel… · p. 2
  • Step-3 updates recurrent state h sequentially across chunks because the chunk dependency prevents parallel execution across them.A_c is an input-dependent diagonal matrix.Source: A. Mamba-2 Layer with Three Kernel… · p. 2

B. IMAX Architecture and Compilation Model

The evaluated IMAX3 prototype runs one 64-stage CGLA lane, with a host CPU coordinating transfers and kernel launches. Its compiler assigns static pipeline stages, and latency reporting separates execution from transfer-related phases.

B. IMAX Architecture and Compilation Model

C. IMAX GEMM Mapping

The IMAX GEMM template maps output rows, reduction accumulation, and output features onto a linear pipeline. Long reductions can keep stages active, while short reductions increase the relative weight of pipeline and data-movement overheads.

C. IMAX GEMM Mapping

  • IMAX maps GEMM output rows across parallel work, accumulates along the reduction dimension, then generates output features.Decode-time projection GEMV is the M2 = 1 special case of the published GEMM mapping.Source: C. IMAX GEMM Mapping · p. 3
  • For long reductions, operand loading, FMA partial-sum updates, and result drain can keep the 64-stage pipeline active for many cycles.Published IMAX measurements report high occupancy when M1, L, and M2 are each at least 480.Source: C. IMAX GEMM Mapping · p. 3
  • At L = 64, few partial-sum updates make pipeline fill, drain, and data movement relatively costly.The projection mapping distributes operand load, FMA updates, and result drain across the pipeline.Source: C. IMAX GEMM Mapping · p. 3

IV. MAPPING MAMBA-2 KERNELS ONTO IMAX

The mapping reuses the standard IMAX GEMM template for projections and implements SSD Step-1 as GEMM1, Hadamard, and GEMM2. Short reductions motivate batching and make kernel-boundary overhead important to the Step-1 path.

A. Projection GEMM (in-proj and out-proj)

  • Projection reuses the standard GEMM template across stages 0–63, with a 60-stage FMA chain and compile-time parameter changes.Stages 1 and 62 load operands and perform the final add/store; stages 0 and 63 provide iteration control.Source: A. Projection GEMM (in-proj and… · p. 3, Figure 3 · p. 4

B. SSD Step-1 Kernels

  • SSD Step-1 maps to GEMM1, an element-wise Hadamard mask, and GEMM2, using seven-stage GEMM pipelines and a three-stage Hadamard pipeline.Stage 63 in the projection template is iteration control, not an SSM recurrence loop.Source: B. SSD Step-1 Kernels · pp. 3–4, Figure 3 · p. 4
  • For the single-head Step-1 sweep, Q and the reduction length are both 64, so the outer batch axis is increased to amortize launch cost.The evaluation distinguishes the nominal all-head point from the largest sweep point.Source: B. SSD Step-1 Kernels · p. 4
  • The IMAX Hadamard path masks R in place with two parallel OP_FML lanes, removing one DMA round trip per chunk at 0.25 FLOP/byte.Measured boundary overhead still motivates a fused Hadamard–GEMM2 path.Source: B. SSD Step-1 Kernels · p. 4
  • The Step-1 evaluation reports EXEC and wall-clock time for all three IMAX kernels, alongside a per-layer FLOP-distribution table for Mamba-370M.The table caption specifies dMODEL = 1,024, dINNER = 2,048, H = 32, dHEAD = 64, N = 128, and Q = 64.Source: B. SSD Step-1 Kernels · p. 4, Table III · p. 4

C. Step-3 Recurrent State Update

The inter-chunk recurrence is a small, sequentially dependent matrix-vector product, so transfer scheduling matters more than peak compute. In the layer composition, Step-3 remains on the host CPU; during decode, projection GEMV is still affected by transfer overhead.

C. Step-3 Recurrent State Update

  • The inter-chunk recurrence is a small N × d matrix-vector product per chunk whose sequential dependency makes transfer scheduling, rather than peak compute, determine cost.For Mamba-370M, the per-head shape is 128 × 64, too small to saturate the 64-stage FMA pipeline.Source: C. Step-3 Recurrent State Update · p. 4
  • Step-3 stays on the host CPU in the layer-level composition; double-buffered LMM use and ping-pong scheduling are identified as next directions for moving more of this phase to IMAX.The evaluation includes a table of platforms and ASIC model context.Source: C. Step-3 Recurrent State Update · p. 4, Table IV · p. 5
  • At decode T = 1, projection GEMV becomes less compute dense, but transfer overhead still dominates the current FPGA implementation.Source: C. Step-3 Recurrent State Update · p. 4

V. EXPERIMENTS

The evaluation characterizes Mamba-370M kernels on an IMAX FPGA prototype against an ARM Cortex-A72 software reference, with a separate Mamba-130M token-generation run.

A. Experimental Setup

  • IMAX latency and correctness are measured on the FPGA prototype, while the ARM Cortex-A72 path provides the correctness and CPU-timing reference.The ASIC context comes from core synthesis, not silicon or board-power measurements, and is not used for cross-platform system-energy claims.Source: A. Experimental Setup · p. 5
  • All kernels use FP32, and outputs are checked against the ARM Cortex-A72 software reference.The evaluation does not establish behavior for FP16, BF16, or quantized paths.Source: A. Experimental Setup · p. 5
  • Mamba-370M is used for layer-level characterization, whereas Mamba-130M is used for token generation.Projection values average three archived runs; Step-1 and token-generation values are representative runs without claimed confidence intervals.Source: A. Experimental Setup · p. 5

B. IMAX Kernel Measurements

The measurements distinguish pipeline-suited long-reduction projections from short-reduction SSD Step-1 and transfer-scheduled recurrent-state updates. At the pipeline-aligned setting, projections dominate IMAX execution, while SSD Step-1 remains constrained by kernel boundaries.

B. IMAX Kernel Measurements

  • 37.8 ms EXEC and 1.17 s DMA-inclusive total time are recorded for in-proj, so DMA dominates its measured total time.Projection is the long-reduction CGLA fit in this workload; the DMA-inclusive spread across repeated runs is below 0.1% for both projection kernels.Source: B. IMAX Kernel Measurements · p. 5
  • 48,922 µs aggregate wall-clock time for unfused SSD Step-1 exceeds the matched CPU reference's 19,407 µs despite GEMM1, Hadamard, and GEMM2 EXEC times of 7,068, 1,178, and 7,052 µs.Boundary transfers raise the total; a separate single-GEMM sweep reaches 1.43× ARM speedup for EXEC at M1=4,096, but does not overcome boundary and transfer overheads.Source: B. IMAX Kernel Measurements · p. 5
  • Block autotuning changes recurrent-state wall-clock time from 29.3 to 28.1 ms for Mamba-370M and from 124.4 to 107.0 ms for Mamba-1.4B, a 1.04–1.16× speedup range.The 1.04× Mamba-370M change is treated as no measurable gain without run-to-run variance support; the Mamba-1.4B trend identifies Step-3 as a transfer-scheduling target.Source: B. IMAX Kernel Measurements · p. 5

C. Phase Bottlenecks and Scope

  • At T=120, Q=60, projection remains the dominant FLOP component and dominates IMAX EXEC time, while SSD Step-1 is smaller and limited by short reductions and kernel boundaries.The reported comparison uses the colocated ARM reference rather than a GPU system with a different power boundary; only FP32 is evaluated.Source: C. Phase Bottlenecks and Scope · p. 5

D. Token-Level Integration with IMAX Offload

In Mamba-130M decoding on the IMAX offload path, projection GEMV dominates decode time, while the run preserves the ARM reference token sequence. The conclusion frames this as a kernel-suitability and bottleneck study, not demonstrated system-level acceleration.

D. Token-Level Integration with IMAX Offload

VI. CONCLUSION

  • The paper characterizes long-reduction projections as pipeline-suited, unfused SSD Step-1 as limited by short reductions and kernel boundaries, and decoding as projection-GEMV-dominated.The conclusion describes the unfused Step-1 path as slower than its ARM reference after boundary transfers.Source: §VI · p. 6
  • Persistent projection weights, wider row grouping, Step-1 fusion, and reduced-precision execution are identified as future directions.The conclusion calls for reduced-precision evaluation with matched system-level baselines.Source: §VI · p. 6, D. Token-Level Integration with IMAX… · p. 6

Related papers

Loading 2609.27437v1…