Fix CPU FlashAttention disabled on Linux/aarch64 by detecting L2 cache via cpuinfo#29621
Fix CPU FlashAttention disabled on Linux/aarch64 by detecting L2 cache via cpuinfo#29621Sammy-Dabbas wants to merge 2 commits into
Conversation
…e via cpuinfo glibc's aarch64 sysconf backend does not implement the cache-size queries, so PosixEnv::GetL2CacheSize() returns 0 and the l2_cache_size_ > 0 gate in MultiHeadAttention silently disables CPU FlashAttention. Query cpuinfo for the L2 size first, mirroring the existing cpuinfo_available_ usage in this file for physical-core count and thread affinity, and fall back to sysconf/sysctl. Guarded by ORT_USE_CPUINFO; x86-64 and Windows are unaffected.
There was a problem hiding this comment.
Pull request overview
This PR fixes a Linux/aarch64 performance/functional regression where PosixEnv::GetL2CacheSize() returns 0 (due to glibc sysconf(_SC_LEVEL2_CACHE_SIZE) not being implemented on aarch64), which in turn disables L2-tiled CPU kernels like the CPU FlashAttention path in com.microsoft.MultiHeadAttention and leads to large O(S²) memory fallbacks.
Changes:
- Prefer cpuinfo-based L2 cache size detection (when available) in
PosixEnv::GetL2CacheSize(). - Keep existing
sysconf/sysctlfallbacks when cpuinfo is unavailable or returns 0.
| if (cpuinfo_available_ && cpuinfo_get_l2_caches_count() > 0) { | ||
| const auto* l2_cache = cpuinfo_get_l2_cache(0); | ||
| if (l2_cache != nullptr && l2_cache->size > 0) { | ||
| return static_cast<int>(l2_cache->size); | ||
| } | ||
| } |
tianleiwu
left a comment
tianleiwu
left a comment
There was a problem hiding this comment.
The cpuinfo fallback addresses the Linux/aarch64 detection gap cleanly. I found one compatibility concern: the new preference is also active on Linux x86, where the existing platform query already works; details are inline. The separate integer-narrowing concern is already covered by an existing thread.
| } | ||
|
|
||
| int GetL2CacheSize() const override { | ||
| #ifdef ORT_USE_CPUINFO |
There was a problem hiding this comment.
ORT_USE_CPUINFO is also enabled on supported Linux/x86 builds, so this changes those systems from their existing working sysconf result to cpuinfo cache instance 0, despite the PR description saying x86-64 is unaffected. On heterogeneous CPUs, instance 0 may not represent every core used by the thread pool. Could we preserve a positive sysconf result and use cpuinfo only when that query returns <= 0 (or guard this path to Linux/aarch64)? That keeps the fix scoped to the detection failure while retaining current behavior elsewhere.
There was a problem hiding this comment.
Thanks, both are addressed in the latest commit.
@tianleiwu you're right that ORT_USE_CPUINFO is also on for Linux x86, so preferring cpuinfo was wrong. GetL2CacheSize now runs sysconf first and returns its value whenever it is greater than zero, so x86 keeps its existing result and cpuinfo is reached only when sysconf returns zero or less, the glibc aarch64 case. The Apple/BSD fallback is unchanged, which also makes the "x86 unaffected" note accurate.
sysconf(_SC_LEVEL2_CACHE_SIZE) is now queried first, so platforms where it works, such as Linux/x86, keep their existing L2 cache values unchanged. cpuinfo is used only as a fallback when sysconf returns a value of zero or less, which is the glibc aarch64 case where the cache-size queries are unimplemented and CPU FlashAttention was being disabled. Both return paths now use narrow<int> instead of static_cast<int> so that narrowing the cache size to int is checked rather than silently truncated.
On Linux/aarch64,
PosixEnv::GetL2CacheSize()returnssysconf(_SC_LEVEL2_CACHE_SIZE), but glibc's aarch64 sysconf backend does not implement the cache-size queries and returns 0. That makes thel2_cache_size_ > 0gate in MultiHeadAttention always false, so the CPU FlashAttention path is silently disabled and the kernel falls back to materializing the full attention score tensor (multi-GiB at long sequence lengths, and OOM on small ARM instances). GroupQueryAttention divides by the same 0 and ends up with a degenerate tile size of 1.This queries cpuinfo for the L2 cache size before falling back to sysconf/sysctl, mirroring the existing
cpuinfo_available_usage in the same file for physical-core count and thread affinity. cpuinfo reads sizes from sysfs cacheinfo and works on aarch64. The change is guarded byORT_USE_CPUINFOand only returns early on a positive result, so x86-64 (working sysconf) and Windows (separateWindowsEnvimplementation) are unaffected.Fixes #29613.