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

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)