[MLAS] Add a NEON fused kernel for LinearAttention

[mirounga]

Description

Fused NEON kernel for LinearAttention

Motivation and Context

3x speedup compared to generic dispatch

[Copilot]

Pull request overview

Adds an ARM64 NEON fused kernel for LinearAttention to improve recurrence performance.

Changes:

• Implements fused single- and two-pass NEON kernels.
• Adds ARM64 dispatch registration and declarations.
• Includes the kernel in Windows and non-Windows ARM64 builds.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/mlas/lib/platform.cpp	Selects the NEON dispatch on ARM64.
onnxruntime/core/mlas/lib/mlasi.h	Declares the NEON dispatch.
onnxruntime/core/mlas/lib/linear_attention_kernel_neon.cpp	Implements the fused NEON kernel and fallback logic.
cmake/onnxruntime_mlas.cmake	Adds the kernel to ARM64 builds.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Hariharan Seshadri (hariharans29)]

Review — PR #32178: [MLAS] Add a NEON fused kernel for LinearAttention

Scope

Five-file PR that adds an ARM64 NEON peer to the existing portable and AVX-512 LinearAttention kernels. New TU linear_attention_kernel_neon.cpp (~447 lines), plus the standard trio of scaffolding: declaration in mlasi.h, assignment in platform.cpp, and source additions in onnxruntime_mlas.cmake (both the Windows and non-Windows ARM64 branches). Author reports 3x speedup vs generic dispatch.

The kernel is a NEON-specific rewrite, not a port of the AVX-512 one: it takes the single-pass form where the algebra allows it, which is the right call on 128-bit vectors even though the AVX-512 kernel uses two-pass uniformly. Everything below is verifiable from the diff alone.

Two-shape design

• Single-pass (FusedTokenSinglePassNeon<HAS_DECAY>) for linear/gated. No retrieval, so upd = v is known before the panel walk starts, and each S_new[i,j] can be produced and consumed by the readout in the same iteration. 1 load + 1 store of S per element.
• Two-pass (FusedTokenTwoPassNeon<HAS_DECAY>) for delta/gated_delta. Uses the AVX-512 kernel's identity o = scale * (sum_i(dec*q * S_old) + (q.k)*upd) so pass 1 reads S_old once, and pass 2 re-reads the panel from L1 while writing S_new in place. 2 loads + 1 store of S per element.

Author's write-up in the module-header comment is accurate. Working the math for linear: two-pass would cost 2 loads + 2 FMA + 1 store per element-vector (pass 1 reads S and accumulates the readout; pass 2 reads S, updates it and stores). Single-pass drops one of those loads to 1 load + 2 FMA + 1 store — a genuine memory-op reduction, not a redistribution. On a 128-bit path where the two-pass cost was already a dead tie with the portable update+gemv, that's the decisive difference. ✓

Panel geometry and register budget

MlasLinearAttentionNeonPanelWidth = 32 floats = 8 × float32x4_t (NLANE = 8), matching AVX-512's panel. Author's L1-residency math checks out: 16 KB at d_k=128 and 32 KB at the d_k=256 bound, inside the typical 64 KB L1.

Two-pass live set at NLANE=8:

• Pass 1: 8 r[] + 8 a[] + in-flight s = ~17 architectural regs.
• Single-pass: 8 a[] + 8 vv[] + in-flight s = ~17.

Well within the 32 aarch64 SIMD registers, with 16 independent FMA chains to cover FMLA latency. This is exactly why HeadsPerGroup != 1 falls back — GQA groups would multiply the accumulator count and force a narrower panel. Correct gate to have.

vfmaq_n_f32 idiom

The scalar-broadcast FMA replaces AVX-512's embedded-broadcast set1-from-memory. Author's note "At 8 lanes each scalar is amortized over 8 FMAs, so loading weights four at a time for vfmaq_laneq_f32 would buy nothing here" is right — the laneq form would only cut scalar loads from 8 to 2 without changing the FMA count, and would eat a vector register per lane group. Skipping it is the correct call.

Shape envelope (fallback logic)

const bool shape_ok = (d_k % 4 == 0) &&
                      (d_k <= MlasLinearAttentionNeonMaxKHeadSize) &&
                      (d_v % MlasLinearAttentionNeonPanelWidth == 0);

if (!shape_ok || Work->HeadsPerGroup != 1) {
    MlasLinearAttentionProcessHead(Work);
    return;
}

The d_k % 4 vs the AVX-512 kernel's d_k % 16 is a real widening of coverage: LinearAttentionDotNeon uses a 16-wide main loop plus a 4-wide tail, so d_k=12 or d_k=20 works here where the AVX-512 kernel would decline. This isn't hypothetical — the test additions specifically cover it.

d_k bound is the 256-float staging buffer for the two-pass form (wkv_buf + wqv_buf = 2 KB on the stack). d_v panel-alignment is because the panel loops are unmasked. All three gates are motivated and documented.

q.k dot product

LinearAttentionDotNeon — four independent float32x4_t accumulators, main loop unrolled by 16, 4-wide tail, vaddvq_f32 horizontal reduce at the end. Correct FMA-pipeline-fill pattern. Author's comment "Both operands are the raw query and key rows: the decay belongs to the sum term, not to this rank-1 coefficient" is exactly the right observation — the identity o = scale * (sum_i(dec*q * S_old) + (q.k)*upd) has the decay only in the sum term, not on the q.k coefficient. Missing that would introduce a subtle correctness bug.

UnrolledLoop<N> helper

The std::index_sequence unroll is real: gcc will spill fixed-size arrays indexed by a loop variable, turning FMA-bound inner loops into load/store-bound ones. This is a documented gcc pathology. Author's rationale for repeating the helper instead of pulling in qnbitgemm_kernel_neon.h is fair — that header carries the entire QNBitGemm surface. A small shared mlasi_neon_unroll.h header would deduplicate, but not blocking.

Decay materialisation strategy

Per-key-dim decay routes through MlasComputeExp<float>(gt, decvec, d_k), which resolves to the NEON MlasComputeExpF32Kernel on ARM64. Per-head decay uses std::exp(gt[0]) splatted across decvec[i]. This unification means the panel loops see one uniform decvec[] regardless of layout — no branch on decay_per_key_dim inside the hot path. Nice.

Small nit: the per-head splat writes the same value d_k times into decvec and then reads it back. On d_k=256 that's 256 unnecessary stores per token. Would be free to just carry a scalar and use a specialized HAS_DECAY_PER_HEAD template branch, but the extra branching would explode the template matrix (currently 4 rules × 1 decay flavour = 4 instantiations, would become 6). Author correctly prioritised uniform code paths. Non-blocking.

Reassociation callout

This reassociates the floating-point sums relative to the portable kernel, so results agree to tolerance rather than bit-exactly.

Important and correct. Panel-wise vector accumulation reorders sums vs the portable per-element loop. The existing MlasLinearAttentionTest presumably uses CloseEnough/tolerance-based comparison; if not, that's a portable-vs-AVX-512 vs NEON issue predating this PR.

Test expansion in test_linear_attention.cpp

Three new shapes with a specific test-selection rationale in the comment:

• {12, 32} — d_k=12 (% 4 but not % 16): runs the 4-wide dot tail with the 16-wide main loop skipped entirely.
• {20, 32} — d_k=20 (% 4 but not % 16): runs main and tail together (16 main + 4 tail).
• {256, 32} — exact fit of both AVX-512 and NEON fixed d_k staging buffers ("an off-by-one would smash the stack").

Author's comment explicitly says "The last four entries exist for specific edges, so do not drop them without checking what they cover" — reviewer-friendly future-proofing. Coverage is thoughtful. The {272, 32} shape (already present) hits the accept side of the d_k <= 256 bound as a negative test — covers the case where a shape passes the AVX-512 envelope but fails the NEON envelope's tighter d_k rule, exercising the fallback dispatch.

Concerns and nits

1. default in the rule switch. linear_attention_kernel_neon.cpp uses:

   switch (Work->Rule) {
       case MlasLinearAttentionRuleLinear:      ProcessHeadNeon<false, false>(Work); break;
       case MlasLinearAttentionRuleGated:       ProcessHeadNeon<true,  false>(Work); break;
       case MlasLinearAttentionRuleDelta:       ProcessHeadNeon<false, true >(Work); break;
       default: /* MlasLinearAttentionRuleGatedDelta */
                                                ProcessHeadNeon<true,  true >(Work); break;
   }

   The default will silently route any future new rule (say RuleFooBar) to the gated-delta path. A fourth explicit case MlasLinearAttentionRuleGatedDelta: plus a default: MLAS_UNREACHABLE() (or assert(false)) would fail loudly instead. Same trap exists in the AVX-512 kernel if it uses the same pattern, but worth fixing here before it copies further. Non-blocking.

2. PR description very terse. Motivation ("3x speedup compared to generic dispatch") doesn't say which device, which rule, which (d_k, d_v), or which sequence length. Given the two-pass vs single-pass distinction, the delta/gated-delta speedup will be different from the linear/gated one. A one-liner per rule (or a small table) in the PR description would help downstream release-notes writers and future perf-regression bisectors.

3. Aarch32 exclusion is implicit. The CMake diff adds linear_attention_kernel_neon.cpp to source lists that are already ARM64-only (they carry other ARM64-only sources like sqnbitgemm_kernel_neon_int8_i8mm.cpp), so aarch32 doesn't compile this. Module header explicitly says "cmake compiles it only for ARM64, hence no MLAS_TARGET_ARM64 guard here." Verified via the surrounding source-list membership. ✓

4. Per-head decay could avoid the splat write. As above — a HAS_DECAY_PER_HEAD template axis would let the inner loops carry a scalar decay and skip the 256-write splat. Doubles the template instantiation matrix, so likely not worth it, but flagging for future consideration if profiling shows the splat is visible.

5. extern const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchNeon is declared in mlasi.h between the AVX-512 dispatch and the depthwise conv dispatch. Consistent with existing pattern. ✓

6. Init syntax. const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchNeon = { MlasLinearAttentionProcessHeadNeon }; uses positional init. If the dispatch struct grows a second field later, this silently either fails to compile (if the new field is non-optional) or leaves it zero-initialized (which may be the intent). Designated initializers would document intent: = { .ProcessHead = MlasLinearAttentionProcessHeadNeon }. Same nit as would apply to the AVX-512 dispatch. Non-blocking.

Non-issues I verified

• The buffer sizing constexpr size_t StageK = HAS_DECAY ? 256 : 1; in the two-pass form correctly sizes the two staging buffers away to 1-float stubs when there's no decay. No wasted stack. ✓
• The HAS_DECAY ? decvec : nullptr local dec is only dereferenced inside if constexpr (HAS_DECAY) branches in both pass helpers, so the nullptr case never dereferences. ✓
• The panel-major traversal of S[i * d_v + j0] means Si advances by d_v * 4 bytes per outer step, so the panel does stay hot in L1 as claimed (the working set for one panel is d_k * PW * 4 bytes). ✓
• LinearAttentionDotNeon is called before pass 1 in the two-pass form, so its 4-way parallel reduce hits raw q0/kt which are still hot from the caller (Query + t * QueryTokenStride, Key + ...). ✓
• d_k % 4 == 0 guarantees the 4-wide tail cleanly exits — no per-lane predication needed. ✓

Recommendation

Approve. High-quality NEON kernel with the right trade-offs called out and defended in code comments:

• Single-pass vs two-pass split by algorithmic necessity, correctly biased to single-pass where possible on 128-bit NEON.
• 32-column panel + UnrolledLoop<8> gives 16 independent FMA chains and keeps the panel L1-resident.
• Shape envelope widens coverage over the AVX-512 kernel (d_k % 4 vs % 16) with the 4-wide dot tail.
• Fallback for HeadsPerGroup > 1 correctly prevents register-file blowup.
• Test additions cover the newly accepted d_k-eligibility class plus the d_k=256 exact-fit boundary.

Suggested pre-merge tightening (all non-blocking):

• Replace the switch default with an explicit MlasLinearAttentionRuleGatedDelta case plus a default: MLAS_UNREACHABLE() / assert.
• Expand the PR description with per-rule perf numbers and the target device.
• Consider consolidating the UnrolledLoop helper into a shared mlasi_neon_unroll.h in a follow-up.

[Hariharan Seshadri (hariharans29)]

Re-review — PR #32178: [MLAS] Add a NEON fused kernel for LinearAttention

One new commit since the previous review: 9a605f7 — "NEON LinearAttention: exhaustive rule switch and named dispatch". Addresses two of the concrete nits from the prior review.

Nit 1 (rule switch default) — addressed, and improved beyond the suggestion

Previous form used default: for MlasLinearAttentionRuleGatedDelta. New form:

switch (Work->Rule) {
    case MlasLinearAttentionRuleLinear:      ProcessHeadNeon<false, false>(Work); return;
    case MlasLinearAttentionRuleGated:       ProcessHeadNeon<true,  false>(Work); return;
    case MlasLinearAttentionRuleDelta:       ProcessHeadNeon<false, true >(Work); return;
    case MlasLinearAttentionRuleGatedDelta:  ProcessHeadNeon<true,  true >(Work); return;
}

//
// Deliberately no default label above: -Wswitch turns a newly added rule
// into a compile error here rather than silently routing it to one of the
// existing specializations. A value outside the enum can still arrive at
// runtime, so defer to the portable kernel rather than guess at its
// semantics.
//
MlasLinearAttentionProcessHead(Work);

This is strictly safer than the MLAS_UNREACHABLE() I suggested:

• Compile-time: no default + -Wswitch (or MSVC C4062/C4061 if enabled) turns any newly added enum member into a build error at exactly this call-site. Extension caught mechanically.
• Runtime: an out-of-enum Work->Rule value (bit-flip, uninitialised, cross-version ABI drift) falls through to the portable kernel — defined behaviour rather than UB. MLAS_UNREACHABLE() would have made that scenario a UB hazard.
• Every case ends in return, so no fall-through traps.

The docstring comment explains why the missing default is intentional, so future contributors won't "helpfully" add one back. Good defensive design.

One follow-up worth verifying at merge time: ORT's MSVC build config should have C4062 (enum not handled in switch) elevated to warning-as-error, otherwise this contract is Clang/GCC-only. If ORT currently only enforces -Wswitch-equivalent behaviour on GCC/Clang, the runtime fallback still keeps the build safe — but the compile-time signal would be one-sided. Non-blocking; worth a follow-up if MSVC drift becomes a concern.

Nit 6 (positional initializer) — addressed

const MLAS_LINEAR_ATTENTION_DISPATCH MlasLinearAttentionDispatchNeon = {
    .ProcessHead = MlasLinearAttentionProcessHeadNeon
};

Designated initializer as suggested. If the dispatch struct grows a second field, this stays a valid initialization with the new field default-initialized, and reviewers can see at a glance which slot is being wired. Same treatment should be applied to MlasLinearAttentionDispatchAvx512F and MlasLinearAttentionDispatchDefault in a follow-up (both currently positional). Non-blocking.

Nits still open

• Perf numbers. PR description still says "3x speedup compared to generic dispatch" without device / rule / (d_k, d_v) / sequence length. Not blocking, but the two-pass form (delta/gated_delta) and the single-pass form (linear/gated) will exhibit different ratios by construction — the single-pass path drops one full S load — so a two-line breakdown would be more informative for release notes and future perf-regression bisection.
• Per-head decay splat. Still writes the same value 256 times into decvec when d_k=256. Not addressed. As noted before, a HAS_DECAY_PER_HEAD template axis would fix it but double the instantiation matrix. Fine to defer.

Non-issues re-verified against the new diff

Kernel body is identical to the prior review (single-pass / two-pass split, LinearAttentionDotNeon structure, UnrolledLoop<8> unroll trick, HAS_DECAY ? 256 : 1 staging sizing, HAS_DECAY ? decvec : nullptr local dereference discipline, HeadsPerGroup != 1 fallback). No regressions from the switch/init cleanup. ✓

Recommendation

Approve. Both concrete nits from the prior review are addressed, and the switch fix is strictly better than what I proposed. Remaining suggestions (perf breakdown, decay-per-head optimization, applying the designated-initializer style to the AVX-512 / Default dispatches) are all non-blocking follow-ups.