Dispatch via keras.ops.set_attention_override instead of per-layer wrapping

Author: anthony-etim

keras_hub/src/vllm/attention.py

@@ -117,6 +117,11 @@ def vllm_paged_attention(
         cache_index = ctx.layer_index
         ctx.layer_index = cache_index + 1
         cache = ctx.kv_caches[cache_index]
+        # Sliding windows can't ride through the attention op's arguments,
+        # so the adapter probes them per layer at load time and they're
+        # applied here by layer index.
+        if sliding_window is None and ctx.per_layer_sliding is not None:
+            sliding_window = ctx.per_layer_sliding[cache_index]
 
     # vLLM convention: one flattened token dim, channels = heads*head_dim.
     query = ops.reshape(query, (-1, num_heads * head_dim))

keras_hub/src/vllm/context.py

@@ -26,6 +26,11 @@ class VLLMContext(threading.local):
             by `set_vllm_context` at the start of every forward step.
         num_kv_heads: The model's true KV head count (from the model config),
             used by the bridge to undo GQA head expansion generically.
+        per_layer_sliding: Optional per-layer sliding-window sizes, aligned
+            with `kv_caches` order. Sliding windows can't be expressed
+            through the attention op's arguments (its mask only covers the
+            new tokens, not the paged history), so the adapter probes them
+            once at load time and the bridge applies them by layer index.
         active: Boolean indicating if the context is currently active.
     """
 
@@ -42,6 +47,7 @@ def __init__(self) -> None:
         self.updated_kv_caches: Optional[list] = None
         self.layer_index: int = 0
         self.num_kv_heads: Optional[int] = None
+        self.per_layer_sliding: Optional[list] = None
         self.active: bool = False
 
 
@@ -57,6 +63,7 @@ def set_vllm_context(
     positions: Optional[Any] = None,
     kv_caches: Optional[list] = None,
     num_kv_heads: Optional[int] = None,
+    per_layer_sliding: Optional[list] = None,
 ) -> None:
     """Sets the thread-local vLLM context parameters.
 
@@ -72,6 +79,7 @@ def set_vllm_context(
         kv_caches: Per-layer paged KV caches, in transformer-layer order.
         num_kv_heads: The model's true KV head count, for generic GQA
             un-expansion in the bridge.
+        per_layer_sliding: Per-layer sliding-window sizes, in layer order.
     """
     _vllm_context.block_tables = block_tables
     _vllm_context.slot_mapping = slot_mapping
@@ -85,6 +93,7 @@ def set_vllm_context(
     )
     _vllm_context.layer_index = 0
     _vllm_context.num_kv_heads = num_kv_heads
+    _vllm_context.per_layer_sliding = per_layer_sliding
     _vllm_context.active = True
 
 
@@ -100,6 +109,7 @@ def clear_vllm_context() -> None:
     _vllm_context.updated_kv_caches = None
     _vllm_context.layer_index = 0
     _vllm_context.num_kv_heads = None
+    _vllm_context.per_layer_sliding = None
     _vllm_context.active = False