Attention Performance Roadmap#

Date:

2026-08

Updated:

2026-09-05

complete

onnx-light-cpu provides materialized and bounded-memory streaming Attention for the roadmap types and cache modes. The optimization sequence adds a single-key value-copy path, tiled execution starting at query length 16, shape-adaptive query/KV blocks, and a tiled FP16 prefill path which converts half inputs while packing GEMM panels instead of widening complete Q and K tensors. These follow-ups were delivered by #559, #569, and #578. #599 closes the remaining AVX-512 shape and scheduling gaps measured by the published dashboard. #605 and #608 establish the AVX2 scheduling and dedicated decode path; continued AVX2 tuning is tracked separately by the AVX2 performance follow-up.

Objective#

The objective is a correct, bounded-memory CPU implementation of tensor-based ONNX Attention with performance parity against the ONNX Runtime CPU execution provider. Parity means a corpus median speed-up of at least 1.0x and no priority case below 0.9x on dedicated hardware.

This roadmap starts from the shared matrix engine delivered through Roadmap PR10.5. Attention reuses its packing, dot-product micro-kernels, type conversion, and scheduler. The materialized correctness path may compose those primitives, but the optimized path is not implemented as two ordinary materialized MatMul calls.

Scope#

The adapter implements ai.onnx::Attention v23 and v24. It handles Q/K/V head geometry, scaling, masks, causal behavior, grouped-query and multi-query head mapping, and optional tensor cache state according to the selected opset.

The implementation supports:

  • rank-4 [B, H, L, D] and rank-3 packed [B, L, H * D] layouts;

  • MHA, GQA, and MQA through one internal descriptor, including a value head dimension different from the Q/K head dimension;

  • broadcastable rank-2 through rank-4 boolean and additive attn_mask tensors, causal masking, their intersection, and zero output for a fully masked query row;

  • v24 nonpad_kv_seqlen external-cache masking and v23/v24 internal past_key/past_value concatenation with optional present outputs;

  • scale, softcap, qk_matmul_output_mode 0 through 3, and the optional observable qk_matmul_output;

  • prefill, short-query lengths 2-16, and single-token decode;

  • equal-type FP32, FP16, and BF16 Q/K/V with default or explicit FP32 softmax accumulation;

  • zero-copy GQA/MQA head mapping, without physically repeating K or V;

  • exact ONNX layouts, output shapes, empty sequences, and validation errors.

The ONNX schema also permits DOUBLE, a V type different from Q/K, and softmax precisions other than FP32. Those combinations are outside this roadmap and are not currently supported by the portable kernel either. CPU registration must not advertise unsupported Q/K types; combinations that can reach the CPU adapter through a supported Q type fail with an explicit unsupported error rather than silently changing precision or returning an approximation. A requested qk_matmul_output necessarily materializes that observable tensor; the dispatcher therefore selects the materialized path for it. Sliding-window or sparse behavior is in scope only when represented by a valid attn_mask. It does not introduce a separate non-standard mask contract.

Benchmark contract#

Optimization begins with reproducible end-to-end measurements against ONNX Runtime and isolated measurements for the score, softmax, and value-update stages.

  • Use identical tensors, masks, layouts, CPU affinity, and correctness tolerances.

  • The primary comparison sets both runtimes to one thread. --threads N applies the same explicit count to both runtimes for scaling diagnostics.

  • Warm up every candidate equally, alternate candidate order, and report median, p90 tail latency, and dispersion rather than the best observation.

  • Run on an idle, pinned machine with a fixed power policy and record CPU, cache, ISA, compiler, and build metadata.

  • Cover query lengths 1, 2, 8, 16, 128, and 512; KV lengths 1, 128, 1024, 4096, and 8192 when machine memory permits; MHA 12/12, GQA 16/4, and MQA 16/1 query/KV-head ratios; head dimensions 64 and 128; rank-3 and rank-4 layouts; and FP32, FP16, and BF16.

  • Include full-attention layers from Qwen3.6-27B (24 query heads, 4 KV heads) and Qwen3.6-35B-A3B (16 query heads, 2 KV heads), both with head dimension 256, using rank-3 FP16 prefill and tensor-cache decode cases. These geometries come from the published model configurations; the 128/1024/4096 sequence points are serving-oriented benchmark choices rather than model configuration values.

  • Pair the full attribute and mask matrix with small correctness cases. The timed priority corpus covers no mask, causal, boolean, additive, v24 nonpad_kv_seqlen, and internal tensor-cache cases without taking their impractical Cartesian product.

  • Report time to first token, per-token decode latency, tokens/second, peak temporary memory, bytes copied for tensor past/present, and effective KV-cache bandwidth.

  • Measure materialized and streaming paths separately. End-to-end results must include every conversion, allocation, mask, and tensor past/present copy visible to the caller.

  • The priority platforms are x86-64 AVX2, x86-64 AVX-512 when dedicated hardware is available, and Arm64 NEON. SVE measurements are reported separately until a dedicated SVE machine is part of the regular gate.

Backend test corpus#

Roadmap PR11 adds onnx-light-cpu backend test cases for Attention in both correctness and TestMode::BENCHMARK modes. The benchmark cases:

  • are registered through RegisterCpuAttentionCases and the CPU backend collector, so the standard onnx-light backend API, the onnx-light-cpu benchmark runner, and the dashboard consume the same models;

  • produce inputs and expected outputs lazily so large context cases are not allocated while the registry is collected;

  • use globally unique test_cpu_attention_*_benchmark names encoding opset, rank-3/rank-4 layout, MHA/GQA/MQA geometry, query and KV lengths, cache and mask mode, and element type;

  • include the complete timed priority corpus defined above, including KV length 8192 only in the opt-in large corpus when machine memory permits.

  • include Qwen3-8B and both published Qwen3.6 full-attention geometries, with one 128-token prefill case and decode cache lengths 128, 1024, and 4096.

Unit tests execute bounded representative benchmark cases through the registered CPU kernel. Metadata tests cover the complete corpus without materializing every workload and enforce unique names, lazy construction, declared input/output element counts, every priority type and geometry, and the opt-in 8192 cases. The benchmark runner records the backend case name in raw output so published dashboard rows remain traceable to the registered model.

Final parity gate#

The equal-thread corpus is run on a pinned dedicated machine with:

python tools/benchmark_attention_parity.py --cpus 0-15 \
    --output attention_parity_results.json --enforce

The default excludes the memory-intensive KV-length 8192 cases; add --large when machine memory permits. --threads N produces a separate equal-thread scaling run.

Every JSON row retains the globally unique backend case name, raw alternating samples and candidate order, median, p90 tail latency, interquartile dispersion, effective worker count, conservative peak streaming scratch and score-tile bytes from the kernel’s allocation model, tensor-cache bytes copied, and effective KV bandwidth. Report metadata includes affinity, CPU and cache topology, ISA flags, compiler and flags, package versions, platform, and the exact git revision. The summary applies the 1.0x median and 0.9x per-case median and p90 thresholds independently to FP32, FP16, and BF16 and also requires the streaming-memory bound. Shared CI checks the corpus structure, globally unique names, gate arithmetic, dispatch correctness, and bounded streaming implementation; it does not run --enforce.

Phase 1: plan and materialized correctness path#

An immutable AttentionDescriptor is built when the node is initialized. It contains only information known without runtime inputs:

  • opset, attributes, optional input/output positions, and observable outputs;

  • statically known rank, type, head counts, and head dimensions when present;

  • candidate dot-product, conversion, and packing functions allowed by the CPU.

Every invocation validates the actual Q/K/V, mask, and cache tensors and builds a lightweight AttentionPlan from their concrete shapes and strides. It records:

  • batch size, query-head and KV-head counts, head dimensions, and GQA ratio;

  • input/output layouts and strides, query and KV lengths, and cache mode;

  • resolved scale, causal bounds, mask representation, and output obligations;

  • materialized, prefill-streaming, short-query, or decode algorithm;

  • query-row and KV-column block sizes and the available outer task count.

The plan does not prepack Q, K, or V at model load because those tensors are unknown then. It emits work items to the onnx-light executor; it does not store or impose a fixed thread count. The runtime policy decides how many workers are admitted. Shape-specialized decisions may be cached only by a complete key and must not retain invocation tensors.

The first implementation is a materialized correctness path:

S = scale * Q @ transpose(K)
S = apply_mask_and_causality(S)
P = softmax(S)
O = P @ V

It validates all shape, mask, head-mapping, precision, optional-output, and tensor cache semantics against ONNX Runtime. It remains the deterministic CPU fallback for supported types. DOUBLE, mixed Q/K and V types, and unsupported softmax_precision combinations remain outside the registered support boundary and produce an explicit error if encountered.

Phase 2: streaming and online softmax#

The optimized path fuses Q @ K^T, masking, softmax, and P @ V by blocks. It never materializes the complete [batch, heads, query_length, kv_length] score or probability tensors.

For each query block, KV blocks are processed left to right while maintaining the running maximum m, denominator l, and unnormalized output o:

m_new = max(m, row_max(S))
correction = exp(m - m_new)
p = exp(S - m_new)
l = correction * l + row_sum(p)
o = correction * o + p @ V_block
m = m_new

output = o / l

The CPU engine requires:

  • a SIMD Q x K score kernel that fuses scale, mask, causal bounds, and row maximum;

  • vector exponential and reductions with FP32 accumulation;

  • a probability-by-V kernel that updates the output accumulator directly;

  • cache-aware Br and Bc block sizes;

  • causal, sliding-window, and sparse tile skipping;

  • batch/head/query-block scheduling without nested thread pools;

  • dedicated MHA/GQA/MQA algorithms for query lengths 1 and 2-16;

  • FP16/BF16 score and V-update kernels with the materialized path as fallback.

Here sliding-window and sparse tile skipping are optimizations inferred from a valid boolean or additive attn_mask. Arbitrary masks remain correct even when no skip structure can be inferred. Internal past tensors and external v24 caches may be consumed block by block, but requested present or qk_matmul_output tensors still incur the allocations required by ONNX.

Temporary score storage falls from O(B * Hq * Lq * Lkv) to O(Br * Bc) per worker. The expected gain is 0-20% when ONNX Runtime already uses a fused path and 1.2-2x when it materializes scores or probabilities; these are targets to measure, not guarantees.

Acceptance criteria#

Area

Exit criterion

Correctness

Differential tests against ONNX Runtime and the portable kernel cover v23/v24, stateless and both cache modes, MHA/GQA/MQA, every attribute and optional output, mask/causal composition, fully masked rows, rank-3 and rank-4 layouts, empty sequences, and every dispatched type.

Fallback

The materialized path handles every optimized-type combination not handled by streaming. Types and precision modes outside CPU scope are not advertised and fail explicitly if dispatch still reaches the CPU adapter.

Memory

Streaming Attention never materializes the complete score or probability tensor; temporary memory is bounded by worker count and Br x Bc blocks.

Scaling

Prefill exposes enough independent work to scale through the physical cores when useful, while short-query and decode avoid harmful parallel overhead. The runtime, not the kernel, controls admitted workers.

Performance

On each priority platform/type corpus, median speed-up is at least 1.0x ONNX Runtime and no priority case is below 0.9x.

Reproducibility

Dedicated-machine results retain raw samples and environment metadata; shared CI enforces correctness but does not decide narrow performance regressions.

Pull-request sequence#

PR

Deliverable

Merge criterion

Dependency

Status

Roadmap PR11 (#387)

Adapter, planning contract, backend corpus, and FP32 materialized baseline.

AttentionDescriptor parses v23/v24 attributes and optional IO. Per-invocation planning validates concrete tensors without prepacking unknown inputs or fixing threads. Rank-3/rank-4 stateless FP32 MHA/GQA/MQA, scale, causal and boolean/additive masks pass differential tests. Lazy correctness and TestMode::BENCHMARK backend cases are available through the shared collector and enforce the priority corpus. Features not implemented yet delegate to the portable kernel only when that kernel already supports them; otherwise they fail explicitly.

Shared GEMM/MatMul engine

Complete

Roadmap PR12 (#389)

Complete materialized semantics and type gate.

FP16/BF16, softcap, mask/causal intersection, fully masked rows, all qk_matmul_output_mode values, optional outputs, internal tensor past/present, and v24 external-cache nonpad_kv_seqlen match ONNX Runtime. Explicit FP32 softmax_precision matches the default accumulation. DOUBLE, mixed-type, and other softmax-precision cases prove the declared unsupported boundary.

PR11 / #387

Complete

Roadmap PR13 (#388)

FP32 online Attention compute engine.

The online recurrence matches the materialized path for stateless prefill, short-query, and decode. SIMD score kernels fuse scale, masks, causal bounds, softcap, and row maximum; vector exponential and reductions are accurate; probability-by-V updates output directly.

PR12 / #389

Complete

Roadmap PR14 (#391)

Streaming scheduling, cache modes, and low-precision types.

Runtime-owned scheduling covers MHA/GQA/MQA without nested pools or a fixed worker count. Internal/external cache blocks, inferable mask tile skipping, and FP16/BF16 score/V-update kernels match the materialized fallback. Observable full-tensor outputs select materialized execution.

PR13 / #388

Complete

Roadmap PR15 (#390)

Final parity and memory gate.

Every priority platform/type corpus has bounded temporary memory, reaches at least 1.0x median speed-up over ONNX Runtime, and has no priority case below 0.9x. Raw equal-thread samples and environment metadata are published; controlled-thread runs remain diagnostic.

PR14 / #391

Complete through #599; dedicated-machine reruns remain optional cross-machine validation

Roadmap PR15 (#390) and the #578 corrective pass delivered the implementation, benchmarking, and memory work. The later #599 pass corrected the AVX-512 shape and participant-selection regressions exposed by the expanded dashboard. The roadmap is complete; operator-specific AVX2 gaps now belong to the AVX2 performance follow-up.

AVX2 gap-closing pass (#635)#

This pass focuses on the AVX2+FMA Q == 1 decode fast path added by #605 and #608 (ComputeAttentionDecodeRowAVX2FMA in attention_plan.cc). Profiling with a standalone microbenchmark that drives ComputeAttentionFloat32Streaming directly (bypassing ONNX Runtime and onnx-light dependencies, which are not available in every build environment) isolated the score dot product (AttentionDotFloat32_AVX2_FMA) as latency-bound: its single-accumulator FMA chain must wait roughly 4-5 cycles for each vfmadd231ps to retire before issuing the next one, even though the core can issue two FMAs per cycle.

A first attempt used four independent accumulators (matching the technique already used by the online-softmax block reduction). Measured in isolation it was faster for every head_dim, but measured through the full decode recurrence it regressed the head_dim == 128 case by 3-8% – consistent, reproducible, and attributable purely to the dot product change (confirmed by swapping only that one function back to the single-accumulator baseline in an otherwise-identical build). This is a case the roadmap’s benchmark contract exists to catch: an isolated micro-improvement does not guarantee an end-to-end one, and the fix was reverted rather than shipped. A second, smaller two-accumulator version was measured safe: it improves head_dim == 64 without regressing head_dim == 128, and was kept.

A similar two-accumulator restructuring was applied to AttentionSoftmaxBlockFloat32_AVX2_FMA’s running-maximum and exponential-sum reductions (also single-accumulator chains, but over the KV block size – up to 256 – rather than head_dim); it showed the same direction of effect in isolation, and was folded into the same kernel file alongside the dot product change.

Alternating raw samples (AVX2-ceiling Release build, ONNX_LIGHT_CPU_HAVE_AVX2_FMA, single participant, this sandbox’s 4-core x86-64 host – not the dedicated machine the final parity gate requires) for ComputeAttentionFloat32Streaming with batch=1, heads=32, q_length=1, no mask, internal past_key/past_value cache, before vs. after both the dot-product and softmax-block two-accumulator changes:

head_dim / kv_len

before (median of 8, us/call)

after (median of 8, us/call)

delta

regression gate

64 / 1024

409.42

396.30

-3.2%

improved

64 / 4096

2290.31

2279.32

-0.5%

within dispersion

128 / 1024

934.60

945.57

+1.2%

within dispersion

128 / 4096

4329.93

4325.45

-0.1%

within dispersion

The combined change was only kept once this full-recurrence benchmark confirmed no case regressed beyond the issue’s 2% dispersion tolerance.

Both changes are scoped to the existing AVX2 decode kernel file (impl/attention/avx2/attention_kernel_avx2_fma.cc); no new scheduler, descriptor, or dispatch path was added, and the existing differential tests (ComputeAttentionFloat32Streaming against ComputeAttentionFloat32Materialized, including the vector-tail and mask regression cases already covering GQA/MQA, causal bounds, and fully masked rows) continue to pass unchanged.

Remaining gaps: this pass did not have access to ONNX Runtime, a dedicated pinned machine, or hardware performance counters, so it could not produce ONNX Runtime ratios or attribute the head_dim == 4096-KV cases (which showed no measurable change) to a specific bottleneck; they are most likely memory-bandwidth-bound rather than compute-bound at that working-set size. Every case below 0.9x ONNX Runtime on the dedicated-machine gate remains to be listed once that gate is run against this change.