Fix CPU FlashAttention disabled on Linux/aarch64 by detecting L2 cache via cpuinfo - #29621
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.
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.
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.
tianleiwu
left a comment
There was a problem hiding this comment.
The fix is correct and well-scoped. GetL2CacheSize() now calls sysconf(_SC_LEVEL2_CACHE_SIZE) first and returns early when the result is positive, so Linux/x86 and other platforms where sysconf works keep their existing values. cpuinfo is consulted only when sysconf reports <= 0 (the glibc aarch64 case that returns 0 and silently disabled CPU FlashAttention / degenerated GQA tiling). The ORT_USE_CPUINFO guard, cpuinfo_available_ check, null/positive-size checks, and narrow<int> narrowing all match the existing conventions in this file.
My earlier concern about Linux/x86 being changed to the cpuinfo path, and the checked-narrowing suggestion, are both addressed on this head. The final return narrow<int>(l2_cache_size) only runs when the value is <= 0, so it preserves the original sysconf semantics without throwing.
Minor note (non-blocking): on heterogeneous aarch64 (big.LITTLE), cpuinfo_get_l2_cache(0) picks cache instance 0, which may correspond to the smaller cluster's L2. That is a conservative choice for tiling (no OOM risk) and consistent with the rest of the file, so no change needed.
LGTM.
|
@Sammy-Dabbas, could you merge latest main branch to pass CI? |
|
Done, merged latest main. Thanks for the review. |
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.