Apply local sliding-window mask in StaticAttentionIOManager single-chunk prefill (#21384)

[YIWENX14]

Summary:

StaticAttentionIOManager.prefill applied a full-causal input mask to every
layer. For local (sliding-window) attention layers, the window is normally
realized via KV-cache eviction across prefill chunks. But the numeric-comparison
harness prefills the whole prompt in a single chunk, so the cache-eviction
window is never exercised and local layers attend to the entire causal history
instead of their cache_len-sized window. This diverged from the rlformers
reference (which applies the correct local window via BlockDiagonalCausalMask)
and made local-global (LGA) checkpoints look badly degraded in numeric SNR
comparisons, while global-only checkpoints were unaffected.

Fix: for local layers (0 < cache_len < input_len), also mask keys older than
the window by adding tril(diagonal=-cache_len). No-op for global layers
(cache_len >= input_len) and on-device chunked prefill (input_len <= cache_len); skip layers (cache_len == 0) are excluded.

Input mask, input_len=5, local window (cache_len)=2:

Before (full causal, applied to ALL layers):

       j=0   1    2    3    4
i=0     0  -inf -inf -inf -inf
i=1     0    0  -inf -inf -inf
i=2     0    0    0  -inf -inf
i=3     0    0    0    0  -inf
i=4     0    0    0    0    0     <- attends to all of j=0..4

After (local layer, window=2):

       j=0   1    2    3    4
i=0     0  -inf -inf -inf -inf
i=1     0    0  -inf -inf -inf
i=2   -inf   0    0  -inf -inf
i=3   -inf -inf   0    0  -inf
i=4   -inf -inf -inf   0    0     <- attends to j=3,4 only (window=2)

Reviewed By: billmguo

Differential Revision: D113574455

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 1fb1945 with merge base fa4508e ():

NEW FAILURE - The following job has failed:

• pull / android / run-emulator (gh)
  The process '/usr/bin/sh' failed with exit code 1

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

[meta-codesync]

@YIWENX14 has exported this pull request. If you are a Meta employee, you can view the originating Diff in D113574455.

[YIWENX14]

@pytorchbot label "release notes: none"