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)