Address review feedback on the serving context, config, samplers, and guards

Author: anthony-etim

keras_hub/src/layers/modeling/cached_multi_head_attention.py

@@ -93,7 +93,16 @@ def call(
                     "`return_attention_scores=True` is not supported while "
                     "serving through vLLM paged attention."
                 )
-            return self._compute_vllm_attention(query, value, key, cache)
+            # Self-attention passes the same tensor for query and value;
+            # cross-attention (e.g. TransformerDecoder's second attention
+            # block) must not push encoder K/V into the paged cache.
+            if value is not query:
+                raise ValueError(
+                    "Cross-attention is not supported under vLLM serving; "
+                    "only decoder self-attention can use the paged KV "
+                    "cache."
+                )
+            return self._compute_vllm_attention(query, key, value, cache)
 
         query = self._query_dense(query)
 
@@ -146,7 +155,7 @@ def call(
         else:
             return attention_output
 
-    def _compute_vllm_attention(self, query, value, key, cache):
+    def _compute_vllm_attention(self, query, key, value, cache):
         """vLLM paged-attention route (serving on TPU only).
 
         Projects Q/K/V and hands them to the shared paged-attention bridge.

keras_hub/src/layers/modeling/position_embedding.py

@@ -147,6 +147,16 @@ def _vllm_serving_positions():
     layout the serving model feeds in.
     """
     ctx = get_vllm_context()
-    if ctx is None or ctx.positions is None:
+    if ctx is None:
+        return None
+    if ctx.positions is None:
+        # A kernel with no positions would silently misplace every learned
+        # embedding; fail loudly instead of falling back to 0..seq_len.
+        if ctx.paged_attention_func is not None:
+            raise ValueError(
+                "vLLM serving context is active but carries no per-token "
+                "positions; learned position embeddings cannot be "
+                "computed."
+            )
         return None
     return ops.reshape(ctx.positions, (-1, 1))

keras_hub/src/vllm/attention_test.py

@@ -56,11 +56,7 @@ def __call__(
 def activate_serving(kernel, kv_caches=None, positions=None):
     """Publishes a serving context wired to `kernel` for a test."""
     vllm_context.set_vllm_context(
-        block_tables=None,
-        slot_mapping=None,
-        attention_metadata="META",
         paged_attention_func=kernel,
-        mesh="MESH",
         positions=positions,
         kv_caches=kv_caches,
     )
@@ -71,6 +67,23 @@ def tearDown(self):
         vllm_context.clear_vllm_context()
         super().tearDown()
 
+    def test_returns_none_when_no_kernel_injected(self):
+        # Context active, but no paged-attention function published: the
+        # second half of the entry guard.
+        activate_serving(None, kv_caches=["C0"])
+        z = ops.zeros((1, 2, 4, 8))
+        self.assertIsNone(vllm_paged_attention(z, z, z, 0.125))
+
+    def test_raises_when_layers_exceed_caches(self):
+        # More attention layers than allocated caches is the layer-count
+        # mismatch a wrong config produces; it must fail loudly.
+        kernel = RecordingKernel()
+        activate_serving(kernel, kv_caches=["C0"])
+        z = ops.zeros((1, 2, 4, 8))
+        vllm_paged_attention(z, z, z, 0.125)
+        with self.assertRaisesRegex(ValueError, "caches"):
+            vllm_paged_attention(z, z, z, 0.125)
+
     def test_returns_none_when_context_inactive(self):
         vllm_context.clear_vllm_context()
         z = ops.zeros((1, 2, 4, 8))

keras_hub/src/vllm/context.py

@@ -2,11 +2,11 @@
 
 The serving model (``KerasHubForCausalLM`` in this package, registered
 with vLLM by tpu-inference's plugin hook) publishes this
-context around each forward step with ``vllm_context_scope``, which clears
-it on exit even when the forward raises, all on one thread. KerasHub
-attention layers, the paged-attention bridge, and ``PositionEmbedding`` read
-it in-place, so serving metadata never has to thread through layer
-signatures.
+context around each forward step with ``vllm_context_scope``, which restores
+the prior state on exit even when the forward raises, all on one thread.
+KerasHub attention layers, the paged-attention bridge, and
+``PositionEmbedding`` read it in-place, so serving metadata never has to
+thread through layer signatures.
 
 It is thread-local so concurrent requests on separate worker threads never see
 each other's caches, positions, or kernel: each thread gets its own inactive
@@ -16,16 +16,31 @@
 import contextlib
 import threading
 
+# The single enumeration of the context's fields and their inactive values.
+# `VLLMContext.__init__`, `clear_vllm_context`, and the scope's save/restore
+# all derive from it, so adding a field means touching exactly one place.
+_INACTIVE_STATE = {
+    "paged_attention_func": None,
+    "positions": None,
+    "kv_caches": None,
+    "updated_kv_caches": None,
+    "layer_index": 0,
+    "active": False,
+}
+
+
+def _reset(context):
+    """Puts ``context`` into the inactive state."""
+    for name, value in _INACTIVE_STATE.items():
+        setattr(context, name, value)
+
 
 class VLLMContext(threading.local):
     """Thread-local context for passing serving metadata to Keras layers.
 
     Attributes:
-        block_tables: The block tables tensor for paged attention.
-        slot_mapping: The slot mapping tensor for paged attention.
-        attention_metadata: The full attention metadata object from vLLM.
-        paged_attention_func: The compiled paged-attention kernel.
-        mesh: The JAX device mesh required by the paged attention kernel.
+        paged_attention_func: The paged-attention closure published by the
+            serving model; engine details are already bound inside it.
         positions: vLLM's per-token absolute position ids (for RoPE models).
         kv_caches: Per-layer paged KV caches for the current forward step.
             Attention layers consume these in call order via `layer_index`,
@@ -40,69 +55,40 @@ class VLLMContext(threading.local):
     def __init__(self):
         """Initializes an empty inactive context."""
         super().__init__()
-        self.block_tables = None
-        self.slot_mapping = None
-        self.attention_metadata = None
-        self.paged_attention_func = None
-        self.mesh = None
-        self.positions = None
-        self.kv_caches = None
-        self.updated_kv_caches = None
-        self.layer_index = 0
-        self.active = False
+        _reset(self)
 
 
 _vllm_context = VLLMContext()
 
 
 def set_vllm_context(
-    block_tables,
-    slot_mapping,
-    attention_metadata=None,
     paged_attention_func=None,
-    mesh=None,
     positions=None,
     kv_caches=None,
 ):
     """Sets the thread-local vLLM context parameters.
 
     Args:
-        block_tables: Array representing memory blocks for key/value caching.
-        slot_mapping: Array mapping sequence tokens to cache slots.
-        attention_metadata: Additional hardware/framework specific metadata.
-        paged_attention_func: The function to use for paged attention.
-        mesh: The JAX device mesh the paged attention kernel shards across.
+        paged_attention_func: The paged-attention closure to dispatch to.
         positions: vLLM's per-token absolute position ids (used by RoPE models
             to apply rotary embeddings at the correct positions under paged /
             continuous-batched decode).
         kv_caches: Per-layer paged KV caches, in transformer-layer order.
     """
-    _vllm_context.block_tables = block_tables
-    _vllm_context.slot_mapping = slot_mapping
-    _vllm_context.attention_metadata = attention_metadata
     _vllm_context.paged_attention_func = paged_attention_func
-    _vllm_context.mesh = mesh
     _vllm_context.positions = positions
     _vllm_context.kv_caches = list(kv_caches) if kv_caches is not None else None
     _vllm_context.updated_kv_caches = (
         list(kv_caches) if kv_caches is not None else None
     )
     _vllm_context.layer_index = 0
+    # Set last, so a failure partway through can never look half-active.
     _vllm_context.active = True
 
 
 def clear_vllm_context():
-    """Clears the thread-local vLLM context."""
-    _vllm_context.block_tables = None
-    _vllm_context.slot_mapping = None
-    _vllm_context.attention_metadata = None
-    _vllm_context.paged_attention_func = None
-    _vllm_context.mesh = None
-    _vllm_context.positions = None
-    _vllm_context.kv_caches = None
-    _vllm_context.updated_kv_caches = None
-    _vllm_context.layer_index = 0
-    _vllm_context.active = False
+    """Resets the thread-local vLLM context to its inactive state."""
+    _reset(_vllm_context)
 
 
 def get_vllm_context():
@@ -116,30 +102,25 @@ def get_vllm_context():
 
 @contextlib.contextmanager
 def vllm_context_scope(
-    block_tables,
-    slot_mapping,
-    attention_metadata=None,
     paged_attention_func=None,
-    mesh=None,
     positions=None,
     kv_caches=None,
 ):
     """Publishes the serving context for the enclosed forward step.
 
-    Sets the thread-local context on entry and always clears it on exit,
-    including when the forward raises, so a reused worker thread can never
-    observe a stale context. Arguments match ``set_vllm_context``.
+    Saves the prior state on entry and always restores it on exit,
+    including when the forward raises — so a reused worker thread can never
+    observe a stale context, and nested scopes hand back the outer scope's
+    state instead of wiping it. Arguments match ``set_vllm_context``.
     """
+    previous = {name: getattr(_vllm_context, name) for name in _INACTIVE_STATE}
     set_vllm_context(
-        block_tables=block_tables,
-        slot_mapping=slot_mapping,
-        attention_metadata=attention_metadata,
         paged_attention_func=paged_attention_func,
-        mesh=mesh,
         positions=positions,
         kv_caches=kv_caches,
     )
     try:
         yield
     finally:
-        clear_vllm_context()
+        for name, value in previous.items():
+            setattr(_vllm_context, name, value)

keras_hub/src/vllm/context_test.py

@@ -17,36 +17,33 @@ def test_inactive_by_default(self):
 
     def test_set_get_roundtrip(self):
         vllm_context.set_vllm_context(
-            block_tables="BT",
-            slot_mapping="SM",
-            attention_metadata="META",
             paged_attention_func="FUNC",
-            mesh="MESH",
+            positions="POS",
+            kv_caches=["C0"],
         )
         ctx = vllm_context.get_vllm_context()
         self.assertIsNotNone(ctx)
-        self.assertEqual(ctx.block_tables, "BT")
-        self.assertEqual(ctx.slot_mapping, "SM")
-        self.assertEqual(ctx.attention_metadata, "META")
         self.assertEqual(ctx.paged_attention_func, "FUNC")
-        self.assertEqual(ctx.mesh, "MESH")
+        self.assertEqual(ctx.positions, "POS")
+        self.assertEqual(ctx.kv_caches, ["C0"])
         self.assertTrue(ctx.active)
 
     def test_clear_resets_everything(self):
-        vllm_context.set_vllm_context("BT", "SM", "META", "FUNC", "MESH")
+        vllm_context.set_vllm_context(
+            paged_attention_func="FUNC", positions="POS", kv_caches=["C0"]
+        )
         vllm_context.clear_vllm_context()
         self.assertIsNone(vllm_context.get_vllm_context())
-        # The singleton's fields are reset too.
-        self.assertIsNone(vllm_context._vllm_context.mesh)
-        self.assertIsNone(vllm_context._vllm_context.paged_attention_func)
-        self.assertIsNone(vllm_context._vllm_context.kv_caches)
-        self.assertEqual(vllm_context._vllm_context.layer_index, 0)
+        # Every field named in _INACTIVE_STATE is back at its inactive
+        # value, so nothing can survive a "clear".
+        for name, value in vllm_context._INACTIVE_STATE.items():
+            self.assertEqual(
+                getattr(vllm_context._vllm_context, name), value, name
+            )
 
     def test_kv_caches_copied_and_lifecycle_fields_initialized(self):
         caches = ["C0", "C1"]
-        vllm_context.set_vllm_context(
-            None, None, kv_caches=caches, positions="POS"
-        )
+        vllm_context.set_vllm_context(kv_caches=caches, positions="POS")
         ctx = vllm_context.get_vllm_context()
         # Layer counter starts fresh; updated caches seed from the input.
         self.assertEqual(ctx.layer_index, 0)
@@ -62,24 +59,44 @@ def test_kv_caches_copied_and_lifecycle_fields_initialized(self):
     def test_reset_on_new_forward_step(self):
         # A second forward step must reset the per-step layer counter, even if
         # the previous step was never explicitly cleared.
-        vllm_context.set_vllm_context(None, None, kv_caches=["C0"])
+        vllm_context.set_vllm_context(kv_caches=["C0"])
         vllm_context.get_vllm_context().layer_index = 5
-        vllm_context.set_vllm_context(None, None, kv_caches=["C0"])
+        vllm_context.set_vllm_context(kv_caches=["C0"])
         self.assertEqual(vllm_context.get_vllm_context().layer_index, 0)
 
+    def test_scope_clears_on_normal_exit(self):
+        with vllm_context.vllm_context_scope(paged_attention_func="KERNEL"):
+            self.assertIsNotNone(vllm_context.get_vllm_context())
+        self.assertIsNone(vllm_context.get_vllm_context())
+
     def test_scope_clears_on_exception(self):
         with self.assertRaisesRegex(RuntimeError, "boom"):
-            with vllm_context.vllm_context_scope(
-                None, None, paged_attention_func="KERNEL"
-            ):
+            with vllm_context.vllm_context_scope(paged_attention_func="KERNEL"):
                 self.assertIsNotNone(vllm_context.get_vllm_context())
                 raise RuntimeError("boom")
         self.assertIsNone(vllm_context.get_vllm_context())
 
+    def test_nested_scopes_restore_outer(self):
+        # The inner scope's exit must hand back the outer scope's state, not
+        # wipe it — otherwise the rest of the outer forward would silently
+        # fall back to the dense path.
+        with vllm_context.vllm_context_scope(
+            paged_attention_func="OUTER", kv_caches=["O0"]
+        ):
+            with vllm_context.vllm_context_scope(paged_attention_func="INNER"):
+                ctx = vllm_context.get_vllm_context()
+                self.assertEqual(ctx.paged_attention_func, "INNER")
+                self.assertIsNone(ctx.kv_caches)
+            ctx = vllm_context.get_vllm_context()
+            self.assertIsNotNone(ctx)
+            self.assertEqual(ctx.paged_attention_func, "OUTER")
+            self.assertEqual(ctx.kv_caches, ["O0"])
+        self.assertIsNone(vllm_context.get_vllm_context())
+
     def test_context_is_thread_local(self):
         # A context set on the main thread must be invisible to another thread,
         # so concurrent requests can never see each other's serving state.
-        vllm_context.set_vllm_context("BT", "SM", "META", "FUNC", "MESH")
+        vllm_context.set_vllm_context(paged_attention_func="FUNC")
         seen = {}
 
         def worker():