[MLAS] Add an AVX-512 fused kernel for LinearAttention

[mirounga]

[MLAS] Add an AVX-512 fused kernel for LinearAttention
The portable linear-attention kernel walks the state matrix S once per step
of the recurrence - decay, retrieval, rank-1 update, readout - which is
roughly 4 reads and 2 writes of S per token. S is d_k*d_v floats (8-64 KB,
L1/L2 resident but far larger than the register file), so that traffic, not
the arithmetic, is what the kernel is bound by.

Add an AVX-512 kernel behind MLAS_LINEAR_ATTENTION_DISPATCH that fuses all
four steps into two passes, one read and one write, using the identity

o_g[j] = scale * ( sum_i q'_g[i]*S_old[i,j] + (q_g . k) * upd[j] )

with q'_g[i] = dec[i]*q_g[i]. The readout is therefore expressible from
S_old plus a rank-1 correction and never needs the written-back S_new, so
pass 1 can accumulate retrieved and every query head's partial readout
from a single read of S_old while pass 2 writes S_new in place. S is
processed one 32-column panel at a time, so pass 2 re-reads the panel from
L1 and the intermediates stay in ZMM registers.

Two kernels: a single-readout-head variant (MHA, and inverse GQA where one
query head is shared) and a GQA variant templated on the group size, where
one read of each S_old panel feeds retrieved plus all NOUT heads'
accumulators at 1 + NOUT FMAs per load. Templating on NOUT keeps the
per-head accumulators at compile-time indices so they stay in registers; a
runtime loop spills them. AVX-512 embedded broadcast makes the pre-weight
vectors memory operands, so they cost no registers.

Everything is intrinsics rather than autovectorized: left to the compiler,
the pass-1 accumulators spilled and q.k became a serial scalar horizontal
reduction.

Both kernels are additionally templated on whether the rule carries decay.
This is not a micro-optimization - without it the kernel lost to the
portable path on the linear rule at d_k = d_v = 128 (0.77-0.97x). That
rule is the one the portable path handles in the fewest passes, so it is
already close to memory-bound on a tuned SGEMM, and the kernel was paying
for a decay it did not have: filling dec[i] = 1, pre-multiplying both
weight vectors by 1, and doing a redundant multiply per element in pass 2.
The specialization drops all three - the kernels read q and k directly and
copy nothing, and pass 2 collapses from a multiply plus an FMA to a single
FMA - which brings linear back to parity.

The kernel is used when d_k % 16 == 0, d_k <= 256, d_v % 32 == 0 and the
group size is 1, 2, 4 or 8; anything else falls back to the portable
kernel, which has no shape restrictions. The panel loops are deliberately
unmasked and the pre-weight buffers fixed-size, which is what sets those
bounds. A group size of 16 is excluded: it would need a 16-column
single-lane variant to fit the register file, and that was benchmarked
slower than a batched GEMM.

Registration is in the AVX512F block of platform.cpp, inside
MLAS_TARGET_AMD64, matching the two cmake lists the translation unit is
added to (mlas_platform_srcs_avx512 with /arch:AVX512 for MSVC,
mlas_platform_srcs_avx512f with -mavx512f for GCC/Clang). Only AVX512F is
required - no BW/DQ/VL/VNNI - so this is not gated on Avx512Supported_.
Non-x86 targets never reference the symbol.

Measured A/B in one binary with only the registration line toggled, on a
-march=native RelWithDebInfo build, over 44 shape/rule combinations:

rule           median   range
linear          1.01x   0.92 - 9.24x
gated           1.22x   1.06 - 7.77x
delta           1.32x   1.02 - 10.21x
gated_delta     1.46x   1.09 - 7.92x

Overall median 1.32x. The largest wins are below the portable kernel's
d_k*d_v >= 4096 SGEMM threshold, where it would otherwise run scalar
loops. Three combinations remain marginally slower (0.92x, 0.98x, 0.98x,
all linear at d_k = d_v = 128); the 0.92x case is a single (batch, kv
head) pair, so it runs unthreaded against MLAS SGEMM. Routing is left
simple rather than special-casing it.

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

Extend the MLAS unit test to cover the eligibility envelope: head configs
now span group sizes 1, 2, 4, 8 and 16 (the last must fall back), and the
shape list adds {272, 32}, which satisfies d_k % 16 == 0 but exceeds the
d_k <= 256 bound. Extend the benchmark to all four update rules. Both
kernels were confirmed to actually execute rather than silently falling
back, by perturbing each call site independently and checking the unit
test fails.

All 39 ContribOpLinearAttentionTest cases pass, at 1, 4 and 8 cores.

[Copilot]

Pull request overview

This PR adds an AVX-512F optimized FP32 LinearAttention recurrence kernel to MLAS (with dispatch selection) and migrates the CPU contrib LinearAttention operator to call the new MLAS entrypoint, alongside expanded unit tests/benchmarks and updated schema/docs for inverse-GQA output packing.

Changes:

• Implement MLAS LinearAttention core + dispatch contract, plus an AVX-512F fused two-pass kernel.
• Route the CPU contrib LinearAttention operator through MlasLinearAttention (including scratch allocation and head-mapping validation).
• Extend MLAS unit tests/benchmarks and update schema/test comments for max(H_q, H_kv) * d_v output packing.

Reviewed changes

Copilot reviewed 13 out of 13 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
onnxruntime/test/mlas/unittest/test_linear_attention.cpp	Adds MLAS unit test comparing MlasLinearAttention against a scalar oracle across shapes/rules/head configs.
onnxruntime/test/mlas/bench/bench_linear_attention.cpp	Adds benchmark for scalar baseline vs MLAS LinearAttention across rules/shapes.
onnxruntime/test/contrib_ops/linear_attention_op_test.cc	Updates output-shape comment and expands contrib-op coverage to all four update rules in a config.
onnxruntime/core/mlas/lib/platform.cpp	Initializes default LinearAttention dispatch and overrides it for AVX-512F targets.
onnxruntime/core/mlas/lib/mlasi.h	Declares LinearAttention dispatch externs and adds dispatch pointer to platform struct.
onnxruntime/core/mlas/lib/linear_attention.h	Introduces work-item struct and dispatch interface for LinearAttention kernels.
onnxruntime/core/mlas/lib/linear_attention.cpp	Implements portable LinearAttention kernel + threaded driver and default dispatch object.
onnxruntime/core/mlas/lib/linear_attention_kernel_avx512.cpp	Implements AVX-512 fused recurrence kernel and AVX-512 dispatch object.
onnxruntime/core/mlas/inc/mlas.h	Adds public MLAS LinearAttention API surface and argument struct/enums.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Updates LinearAttention schema docs and output-0 shape inference to max(H_q, H_kv) * d_v.
onnxruntime/contrib_ops/cpu/bert/linear_attention.h	Stores resolved MLAS rule enum in the kernel instance.
onnxruntime/contrib_ops/cpu/bert/linear_attention.cc	Refactors CPU kernel to validate inputs and call MlasLinearAttention with allocated scratch/state/output.
cmake/onnxruntime_mlas.cmake	Adds new MLAS LinearAttention sources (including AVX-512 TU) to the build.

Suppressed comments (1)

onnxruntime/core/mlas/lib/linear_attention_kernel_avx512.cpp:215

• wkv_buf is allocated at MaxK regardless of HAS_DECAY, but it’s only used in the HAS_DECAY && needs_retrieval path. Making it conditional avoids unnecessary stack usage for the no-decay specialization and for the decay specialization when retrieval is disabled.

    float wkv_buf[MaxK];
    float wqv_buf[HAS_DECAY ? NOUT * MaxK : 1];