Skip to content

CPU FlashAttention silently disabled on all Linux/aarch64: Env::GetL2CacheSize() returns 0 (glibc sysconf), causing O(S²) memory fallback and OOM in MultiHeadAttention #29613

Description

@hughescr

Describe the issue

The CPU FlashAttention path for com.microsoft.MultiHeadAttention is silently disabled on every Linux/aarch64 system, because the gate requires Env::GetL2CacheSize() > 0 and the POSIX implementation of GetL2CacheSize() calls sysconf(_SC_LEVEL2_CACHE_SIZE), which glibc does not implement on aarch64 — it returns 0 ("no information available"). The kernel then falls back to the naive attention implementation, which materializes the full B·N·S·T fp32 score tensor. For long sequences this is a multi-GiB transient allocation per MHA node (exactly 4 GiB at seq 8192 with 16 heads), which OOM-kills small ARM instances, and it is also substantially slower (DRAM-bound full-tensor passes instead of L2-tiled streaming).

The same defect exists on x86 macOS/Windows? No — this is specifically an ARM-Linux hole:

  • Windows: WindowsEnv::GetL2CacheSize() uses GetLogicalProcessorInformation → works.
  • macOS/BSD: _SC_LEVEL2_CACHE_SIZE is undefined, so the code falls through to sysctl(CTL_HW, HW_L2CACHESIZE) → works (returns 4 MiB on Apple M3, verified).
  • Linux/x86-64: glibc implements the cache sysconf values via CPUID → works.
  • Linux/aarch64: glibc's aarch64 sysconf backend only implements L1 cache linesize (via CTR_EL0); all cache size queries fall through to the generic implementation, which returns 0. See glibc sysdeps/unix/sysv/linux/aarch64/sysconf.c, the long-standing RH bug 1190638, and the old libc-alpha discussion of _SC_*CACHE returning 0. getconf LEVEL2_CACHE_SIZE prints 0 on Graviton (Amazon Linux 2023) and other aarch64 distros.

Relevant code (v1.27.0; unchanged on main as of filing)

The gateonnxruntime/contrib_ops/cpu/bert/multihead_attention.cc#L199-L211, with l2_cache_size_ initialized at #L49:

l2_cache_size_ = env.GetL2CacheSize();
...
if (std::is_same_v<T, float> && !disable_flash_ && !is_unidirectional_ &&
    key_padding_mask == nullptr && attn_bias == nullptr && ... &&
    l2_cache_size_ > 0) {           // <-- always false on Linux/aarch64
  MlasFlashAttentionThreadedArgs args;   // tiled path, ~1 MB scratch per thread
  ...

The platform forkonnxruntime/core/platform/posix/env.cc#L304-L318:

int GetL2CacheSize() const override {
#ifdef _SC_LEVEL2_CACHE_SIZE
    return static_cast<int>(sysconf(_SC_LEVEL2_CACHE_SIZE));  // glibc/aarch64: returns 0
#else
    ... sysctl(CTL_HW, HW_L2CACHESIZE) ...                     // macOS: works
#endif
}

The consequence — the fallback onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h#L106-L108 materializes the whole score tensor:

size_t bytes = SafeInt<size_t>(batch_size) * num_heads_ * sequence_length * total_sequence_length * sizeof(T);
auto attention_probs = allocator->Alloc(bytes);

At batch=1, num_heads=16, S=T=8192, fp32 this is 16 · 8192 · 8192 · 4 = 4,294,967,296 bytes — exactly 4 GiB per MHA node, fully written (GEMM + in-place softmax) and re-read, every layer.

Same defect in GQAonnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h#L727 also consumes env.GetL2CacheSize(), so GroupQueryAttention tile sizing on Linux/aarch64 is fed 0 as well.

History: the CPU FlashAttention path was added in #20805 and appears to have been benchmarked on x86 Linux/Windows only, so the ARM-Linux detection hole was never visible.

Measured evidence

Workload: 28-layer encoder-style embedding model (Qwen3-0.6B architecture, 16 heads, head_size 128, bidirectional unidirectional=0), com.microsoft.MultiHeadAttention nodes with no mask/bias inputs (so the flash gate's other conditions are all satisfied), batch 1, enableCpuMemArena: false, intraOpNumThreads: 2, CPU EP, onnxruntime-node 1.27.0. Same 8192-token input on all rows.

platform flash path peak RSS outcome
macOS arm64 (M3 Max), default active (hw.l2cachesize = 4 MiB) 2.26 GB (baseline 1.63 GB) completes in 67 s
macOS arm64, ORT_DISABLE_FLASH_ATTENTION=1 forced off 5.43 GB (+3.2 GB transient) completes — reproduces the fallback cost
Linux aarch64 (AWS Graviton, AL2023, glibc), default silently off (sysconf → 0) >2 GB burst over a 1.3 GB baseline OOM-killed on a 3.7 GB instance (2.3 GB anon-rss at kill, still growing)

The macOS ORT_DISABLE_FLASH_ATTENTION=1 run is the controlled A/B: identical binary and input, only the flash path toggled — it reproduces the Linux memory blowup on macOS, confirming the fallback (not something platform-specific in MLAS or the allocator) is the mechanism.

Latency is also affected: on the Graviton box the fallback scales worse than O(n²) once the score tensor exceeds cache (640 tokens → 1.2 s, 1800 tokens → 21 s; 2.8× tokens → 17.5× time), while the macOS flash path tracks O(n²) cleanly across 640 → 8192 tokens.

Impact

Every Linux/ARM64 deployment (AWS Graviton, Ampere Altra, NVIDIA Grace, Raspberry Pi, ...) of any model whose graph uses the MHA (and, for tile sizing, GQA) CPU kernels is affected. The failure is silent — no warning is emitted; users just see multi-GiB allocation bursts and OOM kills on long sequences, or unexplained latency vs. comparable x86/macOS hardware. Given how much CPU inference is moving to arm64 servers, this quietly negates a headline CPU optimization on the whole platform class.

Suggested fix

  1. On Linux, implement GetL2CacheSize() via sysfs cacheinfo (the mechanism lscpu uses): scan /sys/devices/system/cpu/cpu0/cache/index*/, pick the entry with level == 2 and type Unified/Data, and parse size. Alternatively use the already-vendored cpuinfo library, which has cache detection for ARM.
  2. Regardless of detection method, when detection fails fall back to a conservative positive default (e.g. 1–2 MiB) instead of 0. The value only sizes flash tiles: a conservative guess costs a few percent of tile efficiency, whereas 0 silently disables the path and costs O(S²) memory and DRAM-bound latency. l2_cache_size_ > 0 as a hard feature gate turns a missing sysconf into a functional regression.

Workarounds for affected users

1. LD_PRELOAD sysconf interposition (no rebuild, immediate):

/* l2fix.c — cc -O2 -fPIC -shared -o l2fix.so l2fix.c -ldl
 * Usage: LD_PRELOAD=/path/to/l2fix.so node app.js
 * Override detection with L2FIX_BYTES=2097152 if desired. */
#define _GNU_SOURCE
#include <unistd.h>
#include <dlfcn.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

static long l2_bytes;

__attribute__((constructor)) static void l2fix_init(void) {
  const char* env = getenv("L2FIX_BYTES");
  if (env) l2_bytes = atol(env);
  if (l2_bytes <= 0) {
    /* sysfs cacheinfo: find the level-2 unified/data cache */
    for (int i = 0; i < 8 && l2_bytes <= 0; i++) {
      char path[128]; long level = 0;
      snprintf(path, sizeof path, "/sys/devices/system/cpu/cpu0/cache/index%d/level", i);
      FILE* f = fopen(path, "r");
      if (!f) continue;
      if (fscanf(f, "%ld", &level) != 1) level = 0;
      fclose(f);
      if (level != 2) continue;
      snprintf(path, sizeof path, "/sys/devices/system/cpu/cpu0/cache/index%d/size", i);
      f = fopen(path, "r");
      if (!f) continue;
      long val = 0; char unit = 0;
      if (fscanf(f, "%ld%c", &val, &unit) >= 1)
        l2_bytes = (unit == 'M') ? val * 1024 * 1024 : val * 1024;
      fclose(f);
    }
  }
  if (l2_bytes <= 0) l2_bytes = 2 * 1024 * 1024; /* conservative default */
}

long sysconf(int name) {
  static long (*real_sysconf)(int);
  if (!real_sysconf) real_sysconf = (long (*)(int))dlsym(RTLD_NEXT, "sysconf");
  if (name == _SC_LEVEL2_CACHE_SIZE) {
    long v = real_sysconf(name);
    return v > 0 ? v : l2_bytes;
  }
  return real_sysconf(name);
}

With the shim active, the same 8192-token workload runs in the flash path with an O(S) working set (a few hundred MB of Q/K/V copies instead of 4 GiB/layer of scores).

2. Application-level: chunk inputs to ≤2048 tokens (caps the fallback's score buffer at 256 MiB per node), and keep enableCpuMemArena: false — with the arena enabled, the BFC arena permanently retains the multi-GiB fallback high-water mark after the first long sequence.

To reproduce

  1. Any fp32-eligible model containing com.microsoft.MultiHeadAttention with no mask/bias inputs (e.g. a BERT/embedding-style export), CPU EP, on any Linux/aarch64 host.
  2. Confirm the detection hole: getconf LEVEL2_CACHE_SIZE0 (vs. a positive value on Linux/x86-64 or sysctl hw.l2cachesize on macOS).
  3. Run a long-sequence input (e.g. 8192 tokens) and watch RSS: peak transient ≈ num_heads · S² · 4 bytes per layer. On macOS the identical run stays flat; setting ORT_DISABLE_FLASH_ATTENTION=1 on macOS reproduces the Linux memory profile, isolating the gate as the cause.
  4. Optional: verify with the LD_PRELOAD shim above that returning a positive L2 size restores the flash path and the flat memory profile on Linux/aarch64.

Urgency

Blocks CPU inference with long sequences on small/mid ARM64 instances (OOM kill); silent quadratic-memory + latency regression on all ARM64 Linux otherwise.

Platform

Linux

OS Version

Amazon Linux 2023 (aarch64, glibc); also verified against macOS 15 arm64 for the A/B

ONNX Runtime Installation

Released Package

ONNX Runtime Version or Commit ID

1.27.0 (onnxruntime-node); code unchanged on main

ONNX Runtime API

JavaScript

Architecture

ARM64

Execution Provider

Default CPU

Metadata

Metadata

Assignees

No one assigned

    Labels

    api:Javascriptissues related to the Javascript APIplatform:webissues related to ONNX Runtime web; typically submitted using template

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions