Skip to content

Fix CPU FlashAttention disabled on Linux/aarch64 by detecting L2 cache via cpuinfo#29621

Open
Sammy-Dabbas wants to merge 1 commit into
microsoft:mainfrom
Sammy-Dabbas:fix-l2-cache-detection-aarch64
Open

Fix CPU FlashAttention disabled on Linux/aarch64 by detecting L2 cache via cpuinfo#29621
Sammy-Dabbas wants to merge 1 commit into
microsoft:mainfrom
Sammy-Dabbas:fix-l2-cache-detection-aarch64

Conversation

@Sammy-Dabbas

Copy link
Copy Markdown
Contributor

On Linux/aarch64, PosixEnv::GetL2CacheSize() returns sysconf(_SC_LEVEL2_CACHE_SIZE), but glibc's aarch64 sysconf backend does not implement the cache-size queries and returns 0. That makes the l2_cache_size_ > 0 gate 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 by ORT_USE_CPUINFO and only returns early on a positive result, so x86-64 (working sysconf) and Windows (separate WindowsEnv implementation) are unaffected.

Fixes #29613.

…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.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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/sysctl fallbacks when cpuinfo is unavailable or returns 0.

Comment on lines +311 to +316
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);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants