Skip to content

Force contiguous K/V on sliding-window SDPA to avoid the strided-K/V NaN (#311) - #344

Closed
stikves wants to merge 1 commit into
apple:mainfrom
stikves:sukru/fix-issue-311
Closed

stikves wants to merge 1 commit into
apple:mainfrom
stikves:sukru/fix-issue-311

Conversation

@stikves

@stikves stikves commented Oct 11, 2026

Copy link
Copy Markdown
Contributor

Fixes #311.

The fused fp16 sliding-window SDPA kernel (is_causal + window_size) returns NaN on strided (non-contiguous) K/V views — e.g. narrow views of a growing packed cache — for a short query chunk once the cache holds more than ~3584 keys in deep (28+ layer) dynamic graphs, which collapses the stream to <pad>. The SDPA primitive now materializes contiguous K/V on sliding-window layers (window_size > 0) — a numerical no-op that copies only when the input is actually strided; global-attention layers are untouched. This is one of the reporter's verified-clean variants; the runtime kernel defect is tracked via Feedback Assistant.

Distinct from (but same kernel family as) #339/#314, which fixes Muse's ring write-order. This one covers the Gemma-family / gpt_oss sliding layers that drive the fused is_causal+window kernel directly with strided K/V.

The Core AI/MPSGraph fused fp16 sliding-window SDPA kernel returns NaN on strided
(non-contiguous) K/V views past ~3584 keys in deep dynamic graphs. Materialize
contiguous K/V on sliding-window layers (window_size > 0) to sidestep the faulty
lowering -- a numerical no-op that copies only when the input is actually strided.
Global-attention layers are untouched.

Workaround for apple#311 (framework kernel bug, tracked via Feedback Assistant).
@tjia1818 tjia1818 closed this Oct 12, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Sliding-window SDPA (is_causal + window_size) on strided K/V views returns NaN for a short chunk over more than 3,584 keys in deep dynamic graphs

2 participants