Pattern-Aware Softmax Scheduling

• The paper introduces a dataflow method that fuses softmax into the attention stage using pattern-aware scheduling, cutting exponentiation costs by 61.4%.
• It employs a QO-stationary dataflow that retains query and output tiles in situ while streaming key/value tiles, significantly reducing on-chip data movement.
• By leveraging predictable attention score patterns, the approach achieves up to 3.86× energy savings and nearly 2× speedup in large language model inference.

Pattern-aware Softmax scheduling is a dataflow and nonlinearity optimization technique for LLM attention accelerators, designed to fuse the Softmax computation directly into tile-wise data paths by exploiting typical structure (patterns) in attention score distributions. This approach enables deeply fused compute-in-memory (CIM) pipelines, supporting significant improvements in inference energy and latency by minimizing repeated memory access, on-chip buffering, and redundant nonlinear operations (Xuan et al., 28 Apr 2026).

1. Foundations and Motivation

The computational bottleneck in LLM transformers is dominated by the scaled dot-product attention block, comprising three stages:

1. $S = QK^\top$: Inner-product of Queries and Keys, producing attention scores.
2. $P = \mathrm{SoftMax}(S)$: Nonlinear normalization.
3. $O = PV$: Weighted aggregation via output projection.

Classical accelerator dataflows—such as weight-stationary, output-stationary, or row-stationary—struggle to efficiently fuse these stages, particularly incurring significant cost in softmax computation and data movement, especially as model and sequence size scale.

Pattern-aware Softmax scheduling is formulated within a "QO-stationary" dataflow: Query tiles remain resident in IP-CIM (inner-product) SRAM, while Output tiles are stationary in OP-CIM (outer-product) memory. The softmax and associated nonlinearities are fused into the pipeline, leveraging typical max-patterns in attention scores to reduce redundant exponential and normalization steps (Xuan et al., 28 Apr 2026).

2. QO-Stationary Dataflow and Softmax Fusion

Formally, the QO-stationary schedule partitions the matrices Q, K, and V into tiles. For each query tile $Q[t]$ and output tile $O[t]$, every key/value tile $K[s], V[s]$ is processed as follows:

$$\begin{aligned} \text{for } t &= 0 \ldots F-1 \ &\quad Q[t] \ \text{stationary in IP-CIM} \ &\quad O[t] \ \text{stationary in OP-CIM} \ &\quad \text{for } s = 0 \ldots G-1 \ &\qquad S[t,s] = Q[t]K[s]^\top \qquad\text{(IP-CIM)} \ &\qquad P[t,s] = \mathrm{SoftMax}(S[t,s]) \qquad\text{(in-pipeline)} \ &\qquad O[t] = O[t] + P[t,s]V[s] \qquad\text{(OP-CIM)} \ &\quad \text{end for} \ \text{end for} \end{aligned}$$

The Query and Output tiles remain in situ while K and V tiles stream through inner and outer product CIM arrays, respectively. Intermediate attention probability tensors $P[t,s]$ are produced, normalized online by the Softmax, and consumed without external buffering. This arrangement allows the nonlinear normalization to be applied in-place, minimizing memory traffic and eliminating the need for explicit on-chip transposes or off-chip storage of intermediate results.

3. Pattern-Aware Scheduling and Nonlinearity Minimization

Pattern-awareness in Softmax scheduling refers to exploiting the typical structure of attention-score matrices during fusion:

• Attention scores $S[t,s]$ often produce "early" large-max patterns; the maximum logit is frequently encountered at a predictable or patterned location within the key tile sequence.
• Rather than traversing K/V tiles in a strict order ($s = 0...G-1$), a permutation $P = \mathrm{SoftMax}(S)$0 of processing order is selected—diagonal-first or reverse—for increased likelihood of encountering the max early.

As a result, the online Softmax normalization can:

• Trigger exponential rescaling for fewer rows, as soon as an early maximum is seen.
• Reduce the total number of exponentiations and normalization steps required, thus lowering nonlinear unit workload and energy consumption.

Empirically, this strategy yields a 61.4% reduction in Softmax-related exponentiation cost (Xuan et al., 28 Apr 2026). The underlying mechanism is to match score loading and normalization to the observed regularity in attention matrices—this is the essence of pattern-aware softmax: a data-driven scheduling of nonlinear operations aligned with score distributions.

4. Performance, Bandwidth, and Energy Implications

Pattern-aware softmax scheduling, when combined with QO-stationary dataflow, produces substantial resource and performance benefits:

• On-chip bandwidth reduction: By ensuring Q and O tiles are stationary, and K/V/probability tiles are streamed exactly as needed, on-chip memory accesses are reduced by 64.3–71.0% across sequence lengths 256–4096.
• DRAM-access minimization: Total DRAM fetches for Q, K, V drop from $P = \mathrm{SoftMax}(S)$1 loads (naive) to only $P = \mathrm{SoftMax}(S)$2 loads.
• Energy saving: Combined, QO-stationary and pattern-aware softmax yield a 3.86× energy reduction over static-KV baselines, with Softmax pattern-awareness alone cutting the nonlinear operation energy by 61% (Xuan et al., 28 Apr 2026).
• Speedup: By overlapping compute and nonlinear stages in the pipeline, and avoiding repeated buffer movement, total latency is reduced by a factor of 1.98× at sequence length 4096.

The hybrid CIM architecture, with an IP-CIM for QKᵀ and OP-CIM for PV, integrates these pattern-aware scheduling strategies at the circuit and pipeline levels.

5. Comparison with Classical Scheduling Strategies

Pattern-aware softmax scheduling within the QO-stationary dataflow demonstrates theoretical and empirical advantages over classical alternatives:

Dataflow	Q Reuse	K/V Reuse	Softmax Fusion	Intermediate Buffering/Transpose
Weight-stationary(KV)	Low	High	No	Yes
Output-stationary	Low	Low	No	Yes
Row-stationary	Moderate	Moderate	No	Yes
QO-stationary + Pattern-aware Softmax	High	Efficient	Yes	No

Classical stationary dataflows require redundant movement of Q/K/V tiles and intermediate Softmax results, as well as explicit buffer transposes and costly normalization steps. Pattern-aware scheduling eliminates these, permitting the direct and minimal-dataflow fusion of all nonlinear and linear stages.

6. Integration into CIM Accelerator Architecture

Pattern-aware softmax scheduling is implemented as an integral part of the FusionCIM architecture. The pipeline proceeds as follows:

• IP-CIM loads a Query tile (Q stationary); streams K tiles for QKᵀ computation
• Softmax normalization is performed online, with scheduling tailored to pattern structures in incoming scores
• OP-CIM accumulates the result by combining normalized probabilities with V tiles (O stationary), with all accumulation and normalization completed before write-back.

The scheduling enables simultaneous vector-level and bit-serial-level pipelining, with each sub-operation mapped to a hardware macro. The coordination with pattern-aware Softmax minimization ensures no idle cycles and no excess nonlinear activations.

7. Broader Implications and Applicability

Pattern-aware softmax scheduling represents a convergence between algorithmic learning of data distribution characteristics and hardware pipeline scheduling in matrix-heavy AI workloads. Its principles apply wherever data-dependent nonlinearities are a bottleneck, and matrix product chains can be fused under reuse-optimal tiling. The success of this strategy in reducing resource usage in LLM inference suggests its applicability in future CIM-based neural architectures and potentially in quantum-classical co-optimization pipelines, if appropriate stationary dataflow representations and pattern-aware scheduling semantics are adopted (Xuan et al., 28 Apr 2026).