Skip to content

[Native] Support sliding-window attention with the off-graph KV cache - #24944

Draft
metascroy wants to merge 1 commit into
offgraph-kv-ptn-metadatafrom
offgraph-kv-swa
Draft

metascroy wants to merge 1 commit into
offgraph-kv-ptn-metadatafrom
offgraph-kv-swa

Conversation

@metascroy

@metascroy metascroy commented Nov 12, 2026 •

Copy link
Copy Markdown
Contributor

Stack

Off-graph KV cache for the native ggml engine, through serving Muse Glimmer on it. Review bottom-up; each PR builds and passes its tests on its own.

Summary

--offgraph_cache now exports models with sliding-window layers instead of rejecting them. The window belongs to the cache, not the graph, so it travels in the package's metadata.

  • Export: writes the cache geometry (get_n_caches, get_kv_heads, get_head_dims, get_windows). A layer's window follows HF: layer_types when the config has them, else sliding_window on every layer (Mistral). Golden samples come from a reference cache whose sliding layers are rings.
  • Runner: reads the windows, checks the shapes against the graph, and builds ring layers sized for the method's token input (max_write). A package without get_windows is refused rather than run as full attention. --warmup_prefill clears the cache after its warmup prefill, since a ring layer can't rewind to the start of a prompt longer than its window.
  • ggml cache: GgmlSequenceCache holds a ring layer at window + max_write - 1 rows from the start; flat layers still grow on demand.
  • Executor: fills one set of row and mask uploads per window from the cache's plan. Each row is masked by the logical position it holds, so a ring layer attends its rows in any physical order, and every query sees only its own window, including within a prefill wider than the window.

Test plan

  • Engine tests: ring attention against a naive windowed reference (wraparound, windows 2 and 5, intra-step windowing, fp32/fp16), flat and ring layers sharing one sequence, rewind limited to the rows a ring holds. Also passes on Metal.
  • Export tests compare off-graph golden samples with HF's own sliding mask for Mistral and Gemma3; they fail with the windows dropped.
  • End to end: tiny Mistral and Gemma3 (window 8) golden samples replay on CPU and Metal. --warmup_prefill with a chunked 37-token prompt on Mistral works; without the clear it fails with "cannot rewind to position 0". A 121-token chunked prefill matches HF eager to 2.4e-7 on Mistral (0.175 without the windows).

Authored with Claude Code.

@pytorch-bot

pytorch-bot Bot commented Nov 12, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 137 New Failures, 7 Unrelated Failures, 10 Unclassified Failures

As of commit 3f908ee with merge base c7f7931 (image):

NEW FAILURES - The following jobs have failed:

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Nov 12, 2026
This was referenced Nov 12, 2026
--offgraph_cache now exports models with sliding-window layers instead of
rejecting them. The window is a property of the cache, not the graph, so the
export writes the cache geometry (get_n_caches, get_kv_heads, get_head_dims,
get_windows) to the package's metadata, and the golden samples come from a
reference cache whose sliding layers are rings. A layer's window follows HF:
layer_types when the config has them, else sliding_window on every layer.

The runner reads the windows from the metadata, checks the shapes against the
graph, and builds ring layers sized for the method's token input.
GgmlSequenceCache holds a ring layer at window + max_write - 1 rows from the
start. The executor fills one set of row and mask uploads per window from the
cache's plan: each row is masked by the logical position it holds, so a ring
layer attends all its rows in any physical order, and a query sees only its
own window, including within a prefill wider than the window.

This branch had an error being deployed

1 failed deployment
cadence — 3f908ee7 Deployed Nov 14, 2026 by metascroy via hifi5s-build / hifi5s #38324
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant