From b5446d4b6dcd35ee4bd9d7481f7da8b02f415bb7 Mon Sep 17 00:00:00 2001 From: Yiwen Xie Date: Mon, 27 Jul 2026 11:06:15 -0700 Subject: [PATCH] Apply local sliding-window mask in StaticAttentionIOManager single-chunk prefill (#21384) 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 --- examples/models/llama/static_attention.py | 29 +++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/examples/models/llama/static_attention.py b/examples/models/llama/static_attention.py index 8e985239651..e3b0001f78c 100644 --- a/examples/models/llama/static_attention.py +++ b/examples/models/llama/static_attention.py @@ -429,13 +429,34 @@ def prefill( if self.cache_full: raise RuntimeError("KV cache is full.") + # Global (full-attention) layers use the largest cache; local + # (sliding-window) layers use a strictly smaller cache. In a global-only + # model every layer shares the same cache, so this stays equal to it. + global_cache_len = max( + (mask.cache_len for mask in self._masks.values()), default=0 + ) for mask in self._masks.values(): - mask.set_input_mask( - torch.triu( + input_mask = torch.triu( + torch.full((1, self.input_len, self.input_len), self.mask_val), + diagonal=1, + ) + # Local (sliding-window) layers use a per-layer cache_len smaller than + # the prompt. When the whole prompt is prefilled in a single chunk, the + # cache-eviction window is never exercised, so also mask keys older than + # the window here to reproduce the local attention instead of full causal. + # Only genuine local layers (cache smaller than the global cache) are + # windowed; a global layer is left full-causal even when its (nominal) + # cache is smaller than the prompt. No-op for chunked prefill + # (input_len <= cache_len). + if ( + 0 < mask.cache_len < self.input_len + and mask.cache_len < global_cache_len + ): + input_mask = input_mask + torch.tril( torch.full((1, self.input_len, self.input_len), self.mask_val), - diagonal=1, + diagonal=-mask.cache_len, ) - ) + mask.set_input_mask(input_mask) if isinstance(tokens, list): tokens = torch.tensor([tokens], dtype=torch.int32)