onnxruntime

64618ee7 - [Performance] Improve CPU masked attention performance (#29719)

[Performance] Improve CPU masked attention performance (#29719) ## Description This change improves the CPU execution path shared by `com.microsoft.Attention` and `MultiHeadAttention` when masks or attention bias are present. It removes redundant per-head mask copies, exposes score GEMMs to full thread-pool parallelism, and avoids materializing query-independent padding masks across the query dimension. Fixes #29655. ## Summary of Changes ### Attention computation | File | Change | |------|--------| | `onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h` | Replace per-head score GEMMs with `MlasGemmBatch`, fuse mask/bias application and optional `output_qk` emission into row-wise softmax processing, and use checked temporary-buffer arithmetic. | | `onnxruntime/contrib_ops/cpu/bert/attention_helper.h` | Add direct `[B,T]` preparation for query-independent key-padding masks. | | `onnxruntime/contrib_ops/cpu/bert/attention_utils.cc` | Handle empty transformed tensors without dispatching transpose work. | | `onnxruntime/contrib_ops/cpu/bert/multihead_attention.cc` | Exclude empty matrices from the CPU Flash Attention path. | ### Correctness and robustness - Preserve the existing full `[B,S,T]` path for causal and 3D masks. - Preserve key/value cache concatenation before batched score GEMM execution. - Handle empty batch, query sequence, and key/value sequence dimensions without zero-dimension MLAS or thread-pool dispatches. - Initialize MLAS GEMM descriptors in framework-managed temporary storage and use `SafeInt` for allocation sizes. ### Tests and benchmarks | File | Change | |------|--------| | `onnxruntime/test/contrib_ops/multihead_attention_op_test.cc` | Cover combined padding mask, attention bias, and `output_qk`, plus empty batch/query/key-value dimensions. | | `docs/contrib_ops/cpu/mha_experiments.md` | Document the benchmark setup, intermediate experiments, realistic BERT/DistilBERT/GPT-2 shapes, and comparisons against `origin/main`. | Representative results on an AMD EPYC 7763, pinned to 16 physical cores: - Causal attention: up to 19% faster for the issue shape and up to 25% faster for popular model shapes at 8–16 threads. - Non-causal padding-mask attention: 6–17% faster across measured BERT and DistilBERT shapes. - Single-thread and low-thread results remain neutral or within measurement noise. ## Testing - `lintrunner -a` - Built `onnxruntime_provider_test` from `build/cpu/Release` after rebasing onto the latest `origin/main`. - Focused regressions: 4/4 passed. - Broad filter `*Attention*:*BeamSearch*`: 342 tests ran, 299 passed, 43 skipped because unavailable providers/features were not built, and 0 failed. ## Motivation and Context The previous CPU masked-attention path expanded masks to `[B,S,T]`, copied each `[S,T]` tile once per head, and accumulated score GEMMs with `beta=1`. Its parallelism was capped at `batch_size * num_heads`, which caused poor scaling and could make fused attention slower than an equivalent unfused graph. The new path uses clean `beta=0` batched GEMMs and applies mask/bias data while score rows are consumed by softmax. ## Checklist - [x] Tests added/updated - [x] Documentation updated - [x] No breaking changes - [ ] CI passes