Fix native static attention sliding-window masks - #22285
Conversation
Summary: The native static-attention mask kept every valid cache slot visible for every input row. For multi-token autoregressive chunks, row k therefore attended to too many cached tokens, and chunks longer than the local window also retained too many in-chunk keys. The QNN runner made the same assumption for long and sparse chunks. Derive SMART_MASK cache visibility from ring age and each row's remaining window budget, bound the causal in-chunk region, and apply equivalent oldest-visible trimming in the QNN runner. Add regression coverage for cache wrapping, mixed local/global masks, and chunk lengths both above and below the sliding window. Update the SNR evaluator to use the public decode path so it exercises the same per-row masking behavior. Differential Revision: D117411412
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22285
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 Unrelated FailureAs of commit 0181d8a with merge base d614f9d ( NEW FAILURES - The following jobs have failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
@YIWENX14 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D117411412. |
This PR needs a
|
Summary:
The native static-attention mask kept every valid cache slot visible for every input row. For multi-token autoregressive chunks, row k therefore attended to too many cached tokens, and chunks longer than the local window also retained too many in-chunk keys. The QNN runner made the same assumption for long and sparse chunks.
Derive SMART_MASK cache visibility from ring age and each row's remaining window budget, bound the causal in-chunk region, and apply equivalent oldest-visible trimming in the QNN runner. Add regression coverage for cache wrapping, mixed local/global masks, and chunk lengths both above and below the sliding window. Update the SNR evaluator to use the public decode path so it exercises the same per-row masking behavior.
Differential Revision: D117411412