[JacobSzwejbka]

Summary

• skip fully masked KV blocks and trim masked boundaries in CPU custom SDPA
• split wrapped ring windows around large masked gaps while preserving additive-mask semantics
• size local KV caches as sliding window plus the maximum in-flight prefill chunk, capped by full context
• support regular, custom-op, and quantized ring caches

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22358

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 115 Pending

As of commit 84502cd with merge base 5b77187 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[meta-codesync]

@JacobSzwejbka has imported this pull request. If you are a Meta employee, you can view this in D118840114.

[digantdesai]

can this blow up memory accidentally?

[JacobSzwejbka]

the previous sizing was 2x window size which would produce incorrect outputs on really long sequences. It can meaningfully regress memory on enormous context lengths if we arent chunking the input but for anything in mobile world it should be fine.

[digantdesai]

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 5m 11s —— View job

---

PR Review: Optimize CPU SDPA for ring attention

• Gather PR diff and context
• Review op_sdpa_impl.h (core mask-skipping logic)
• Review cache sizing changes (custom_kv_cache.py, attention_sink.py, attention.py)
• Review tests
• Post review feedback

Overall this is a well-constructed change. The mask-range precompute is shared across batch/heads (correct, since the mask is 2D), the bit-state logic is sound, and splitting a wrapped ring window into two flash-attention sub-blocks is mathematically valid because online softmax composes over disjoint column sets. Below are the findings, most important first.

---

1. Memory: answering @digantdesai's question on attention.py:273

For the ring path this is safe — _get_ring_cache_size returns min(max_context_length, window_size + max_seq_len). When max_seq_len is None it defaults to max_context_length, so the worst case is exactly max_context_length, i.e. the same size as a plain full-context cache. It can never blow up; the only downside is silently losing the SWA memory savings, which is already logged in replace_kv_cache_with_ring_kv_cache.

However, the attention-sink sizing is not symmetrically capped and is where memory can grow unexpectedly. _get_attention_sink_cache_size (attention_sink.py:27-48) returns:

return sink_size + window_size + max_seq_len   # no min(..., max_context_length)

with the same max_seq_len = max_context_length default (line 45-46). The only guard is assert sink_size + window_size <= max_context_length. So with max_seq_len unspecified, the cache becomes sink_size + window_size + max_context_length, which is larger than a full non-windowed cache — the opposite of the intended optimization, and it also enlarges the precomputed RoPE tables (RopeWithAttentionSink precomputes freqs up to cache_size, attention_sink.py:91-98).

Recommendation: either cap the sink cache the same way the ring cache is capped, or document why the asymmetry is intentional. If the default is meant to be "the largest in-flight chunk," defaulting to the full context is surprising. Fix this →

---

2. Precompute overhead when no skipping is possible (op_sdpa_impl.h:1011-1176)

The mask-range precompute runs whenever has_attn_mask is true, allocating mask_block_ranges and column_states and scanning every (row, col) of the mask — even for a dense bidirectional mask where use_mask_ranges ends up false and none of the work is used. For non-ring models that pass a full mask this is pure overhead (one extra O(qSize·kvSize) pass plus two temp allocations).

It's likely negligible in practice since the scan is done once and amortized across all batch·head iterations, but worth confirming with a benchmark on a plain masked-SDPA path (not just ring), and/or a comment noting the tradeoff. The +178 deletions suggest this replaced the old loop, so a quick before/after on a dense-mask model would settle it.

---

3. Correctness spot-checks (all passed, noting for the record)

• Bit semantics (:1071-1082): bit0 = "some row attends" (state |= mask_value != neg_inf), bit1 = "all rows have exactly-zero additive mask" (state &= 1 on any nonzero). A column with mixed -inf/0 across rows correctly keeps bit0 set and clears bit1, so the mask add is not skipped for it. Sound.
• range_mask_is_zero fast path (:1139-1153, :1264): skipping the additive-mask add when the whole retained sub-block is exactly zero is correct, and the per-row causal -inf fill still runs independently, so causality is preserved.
• causal_end in precompute (:1052) matches num_keys in the main loop (:1229-1230) — both resolve to min(start_pos + m + qBlockSize, kvSize), so the redundant std::min(kvBlockEnd, num_keys) at :1259 is just belt-and-suspenders.
• Scratch sizing: column_states_bytes = kvSplitSize * num_thread + qSlice (:1029) matches the useful_mask_ranges offset (:1042-1043); per-thread stride thread_index * kvSplitSize (:1046-1047) holds up to block_size <= kvSplitSize. No overflow.

One thing to double check: if a query row is fully masked across all kv blocks, qk_sum_data[row] stays 0 and the final 1 / qk_sum_data[row] (:1490) is inf/nan. This matches the pre-existing behavior (the old path also leaves max=-inf → sum=0), and for causal/ring masks a row always attends to at least its own position, so it shouldn't trigger — but if custom_sdpa can receive an arbitrary additive mask that fully masks a row, that row's output is nan. Consider whether an explicit guard is warranted for the "arbitrary additive mask" use case the comments advertise.

---

4. Minor

• custom_kv_cache.py:757 assert message references self.k_cache.size(2) but the check is on size(1); same stale index appears in a couple of the ring/sink cache asserts. Cosmetic, but the message will print the wrong dimension.
• RingKVCache.__init__ places a large docstring-style string literal (attention.py:303-352) after executable statements — it's a no-op string expression, not a docstring. Fine functionally, but consider moving explanatory prose into a real comment so it isn't mistaken for a docstring.

The 2 red CI jobs (gemma3-4b multimodal, samsung-models) look like infra/exit-137 (OOM) failures rather than assertion failures — but given finding #1, it's worth confirming the samsung/multimodal OOM isn't the sink-cache sizing actually blowing up memory in a configuration that omits max_seq_len.
· branch perf/cpu-ring-attention-mask-skipping

[digantdesai]

easy to introduce bugs here, I guess make sure we have enough test cases