Serve KerasHub CausalLM models through vLLM on TPU

Author: anthony-etim

keras_hub/api/__init__.py

@@ -10,6 +10,7 @@
 from keras_hub import samplers as samplers
 from keras_hub import tokenizers as tokenizers
 from keras_hub import utils as utils
+from keras_hub import vllm as vllm
 from keras_hub.src.utils.preset_utils import upload_preset as upload_preset
 from keras_hub.src.version import __version__ as __version__
 from keras_hub.src.version import version as version

keras_hub/api/vllm/__init__.py

@@ -0,0 +1,7 @@
+"""DO NOT EDIT.
+
+This file was autogenerated. Do not edit it by hand,
+since your modifications would be overwritten.
+"""
+
+from keras_hub.src.vllm.registry import KerasHubLLM as KerasHubLLM

keras_hub/src/layers/modeling/cached_multi_head_attention.py

@@ -2,6 +2,8 @@
 from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.vllm.attention import vllm_paged_attention
+from keras_hub.src.vllm.context import get_vllm_context
 
 
 @keras_hub_export("keras_hub.layers.CachedMultiHeadAttention")
@@ -79,6 +81,20 @@ def call(
         if key is None:
             key = value
 
+        # Route to vLLM paged attention while serving on TPU; otherwise the
+        # original dense path below runs unchanged.
+        vllm_context = get_vllm_context()
+        if (
+            vllm_context is not None
+            and vllm_context.paged_attention_func is not None
+        ):
+            if return_attention_scores:
+                raise ValueError(
+                    "`return_attention_scores=True` is not supported while "
+                    "serving through vLLM paged attention."
+                )
+            return self._compute_vllm_attention(query, value, key, cache)
+
         query = self._query_dense(query)
 
         # If cache is not `None`, we will use the cache to compute the final key
@@ -129,3 +145,24 @@ def call(
             return attention_output, cache
         else:
             return attention_output
+
+    def _compute_vllm_attention(self, query, value, key, cache):
+        """vLLM paged-attention route (serving on TPU only).
+
+        Projects Q/K/V and hands them to the shared paged-attention bridge.
+        Unlike the RoPE model routes, no positional encoding is applied here:
+        models built on this layer (GPT-2, OPT) add learned position
+        embeddings to the token embeddings before the transformer layers run,
+        so the inputs to this layer already carry position information. The
+        softmax scale matches Keras MHA's ``1/sqrt(key_dim)``.
+        """
+        attention_output = vllm_paged_attention(
+            self._query_dense(query),
+            self._key_dense(key),
+            self._value_dense(value),
+            float(self._key_dim) ** -0.5,
+        )
+        attention_output = self._output_dense(attention_output)
+        if cache is not None:
+            return attention_output, cache
+        return attention_output

keras_hub/src/layers/modeling/position_embedding.py

@@ -2,6 +2,7 @@
 from keras import ops
 
 from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.vllm.context import get_vllm_context
 
 
 @keras_hub_export("keras_hub.layers.PositionEmbedding")
@@ -103,6 +104,13 @@ def call(self, inputs, start_index=0, positions=None):
         # trim to match the length of the input sequence, which might be less
         # than the sequence_length of the layer.
         position_embeddings = ops.convert_to_tensor(self.position_embeddings)
+        if positions is None:
+            # When serving through vLLM, the engine assigns each token an
+            # absolute position (driving the paged KV cache) that a plain
+            # 0..seq_len slice would not match. If a serving context is
+            # active, take those positions from it; otherwise this is a
+            # no-op and the standard slice path below runs.
+            positions = _vllm_serving_positions()
         if positions is None:
             position_embeddings = ops.slice(
                 position_embeddings,
@@ -122,3 +130,23 @@ def call(self, inputs, start_index=0, positions=None):
 
     def compute_output_shape(self, input_shape):
         return input_shape
+
+
+def _vllm_serving_positions():
+    """Absolute position ids assigned by vLLM, or ``None`` outside serving.
+
+    Under vLLM's paged / continuous-batched decode, tokens from many requests
+    are packed into one flat sequence and each carries its own absolute
+    position, so learned position embeddings cannot use a contiguous
+    ``0..seq_len`` slice — they must index by these per-token positions to stay
+    aligned with the paged KV cache.
+
+    Returns ``None`` unless a serving context is active, so normal training and
+    inference are unaffected. When active, the positions are reshaped to
+    ``(num_tokens, 1)`` to match the flattened ``(num_tokens, 1, hidden)`` token
+    layout the serving model feeds in.
+    """
+    ctx = get_vllm_context()
+    if ctx is None or ctx.positions is None:
+        return None
+    return ops.reshape(ctx.positions, (-1, 1))

keras_hub/src/models/gemma/gemma_attention.py

@@ -10,6 +10,8 @@
 from keras_hub.src.utils.keras_utils import gpu_supports_fused_attention_op
 from keras_hub.src.utils.keras_utils import running_on_gpu
 from keras_hub.src.utils.keras_utils import running_on_tpu
+from keras_hub.src.vllm.attention import vllm_paged_attention
+from keras_hub.src.vllm.context import get_vllm_context
 
 
 class CachedGemmaAttention(keras.layers.Layer):
@@ -97,9 +99,16 @@ def build(self, inputs_shape):
 
         self.built = True
 
-    def _apply_rope(self, x, start_index):
-        """Rope rotate q or k."""
-        x = self.rope_layer(x, start_index=start_index)
+    def _apply_rope(self, x, start_index, positions=None):
+        """Rope rotate q or k.
+
+        When ``positions`` is given (vLLM serving), rotate at those absolute
+        per-token positions instead of a contiguous ``start_index`` range.
+        """
+        if positions is not None:
+            x = self.rope_layer(x, positions=positions)
+        else:
+            x = self.rope_layer(x, start_index=start_index)
         # Gemma uses a different layout for positional embeddings.
         # The transformation below ensures the embeddings are numerically
         # equivalent to the original gemma implementation.
@@ -226,6 +235,43 @@ def _mask_sliding_window(
         sliding_mask = ops.expand_dims(sliding_mask, 0)
         return ops.logical_and(attention_mask, ops.cast(sliding_mask, "bool"))
 
+    def _compute_vllm_attention(self, x, cache, vllm_context):
+        """vLLM paged-attention route (serving on TPU only).
+
+        Projects Q/K/V, applies Gemma's RoPE (with its split/reshape layout) at
+        the engine's per-token positions, and hands the unexpanded K/V to the
+        shared paged-attention bridge. Gemma's softmax scale, sliding window,
+        and optional logit soft cap (Gemma 2) are all supplied here.
+        """
+        positions = ops.reshape(vllm_context.positions, (-1, 1))
+        query = self._apply_rope(self.query_dense(x), 0, positions=positions)
+        key = self._apply_rope(self.key_dense(x), 0, positions=positions)
+        value = self.value_dense(x)
+
+        if self.query_head_dim_normalize:
+            scale = 1.0 / np.sqrt(self.head_dim)
+        else:
+            scale = 1.0 / np.sqrt(self.hidden_dim // self.num_query_heads)
+        sliding_window = (
+            self.sliding_window_size
+            if self.use_sliding_window_attention
+            else None
+        )
+
+        attention_vec = vllm_paged_attention(
+            query,
+            key,
+            value,
+            scale,
+            num_kv_heads=self.num_key_value_heads,
+            soft_cap=self.logit_soft_cap,
+            sliding_window=sliding_window,
+        )
+        attention_output = self.output_dense(attention_vec)
+        if cache is not None:
+            return attention_output, cache
+        return attention_output
+
     def call(
         self,
         x,
@@ -234,6 +280,15 @@ def call(
         cache_update_index=0,
         training=False,
     ):
+        # Route to vLLM paged attention while serving on TPU; otherwise the
+        # original dense path below runs unchanged.
+        vllm_context = get_vllm_context()
+        if (
+            vllm_context is not None
+            and vllm_context.paged_attention_func is not None
+        ):
+            return self._compute_vllm_attention(x, cache, vllm_context)
+
         query = self.query_dense(x)
         query = self._apply_rope(query, cache_update_index)
 

keras_hub/src/models/gemma3/gemma3_attention.py

@@ -11,6 +11,8 @@
 from keras_hub.src.utils.keras_utils import gpu_supports_fused_attention_op
 from keras_hub.src.utils.keras_utils import running_on_gpu
 from keras_hub.src.utils.keras_utils import running_on_tpu
+from keras_hub.src.vllm.attention import vllm_paged_attention
+from keras_hub.src.vllm.context import get_vllm_context
 
 
 class CachedGemma3Attention(keras.layers.Layer):
@@ -150,6 +152,59 @@ def _apply_rope(self, x, start_index):
         x = self.rope_layer(x, start_index=start_index)
         return x
 
+    def _compute_vllm_attention(self, x, cache, vllm_context):
+        """vLLM paged-attention route for this layer.
+
+        Used only while serving through vLLM on TPU (when a serving context
+        is active). It mirrors `call`'s projection + RoPE + output steps but
+        applies RoPE at the engine's per-token positions and hands Q/K/V to
+        the shared paged-attention bridge instead of the dense path. All
+        Gemma3-specific choices (Q/K norm, the softmax scale, the sliding
+        window, and the logit soft cap) are made here, in the model.
+        """
+        # RoPE at vLLM's absolute positions: one flattened batch mixes many
+        # requests, each at its own position, so the usual contiguous
+        # start_index does not apply.
+        positions = ops.reshape(vllm_context.positions, (-1, 1))
+
+        query = self.query_dense(x)
+        if self.use_query_key_norm:
+            query = self.query_norm(query)
+        query = self.rope_layer(query, positions=positions)
+
+        key = self.key_dense(x)
+        if self.use_query_key_norm:
+            key = self.key_norm(key)
+        key = self.rope_layer(key, positions=positions)
+
+        value = self.value_dense(x)
+
+        if self.query_head_dim_normalize:
+            scale = 1.0 / np.sqrt(self.head_dim)
+        else:
+            scale = 1.0 / np.sqrt(self.hidden_dim // self.num_query_heads)
+
+        sliding_window = (
+            self.sliding_window_size
+            if self.use_sliding_window_attention
+            else None
+        )
+
+        attention_vec = vllm_paged_attention(
+            query,
+            key,
+            value,
+            scale,
+            num_kv_heads=self.num_key_value_heads,
+            soft_cap=self.logit_soft_cap,
+            sliding_window=sliding_window,
+        )
+
+        attention_output = self.output_dense(attention_vec)
+        if cache is not None:
+            return attention_output, cache
+        return attention_output
+
     def _use_fused_attention_op(self):
         if not fused_attention_op_available():
             return False
@@ -323,6 +378,15 @@ def call(
         cache_update_mask=None,
         training=False,
     ):
+        # Route to vLLM paged attention while serving on TPU; otherwise the
+        # original dense path below runs unchanged.
+        vllm_context = get_vllm_context()
+        if (
+            vllm_context is not None
+            and vllm_context.paged_attention_func is not None
+        ):
+            return self._compute_vllm_attention(x, cache, vllm_context)
+
         query = self.query_dense(x)
 
         if self.use_query_key_norm:

keras_hub/src/models/llama/llama_attention.py

@@ -8,6 +8,8 @@
 )
 from keras_hub.src.utils.keras_utils import clone_initializer
 from keras_hub.src.utils.keras_utils import fused_attention_op_available
+from keras_hub.src.vllm.attention import vllm_paged_attention
+from keras_hub.src.vllm.context import get_vllm_context
 
 
 class LlamaAttention(keras.layers.Layer):
@@ -135,6 +137,17 @@ def call(
         cache_update_index=None,
         training=None,
     ):
+        # Route to vLLM paged attention while serving on TPU; otherwise the
+        # original dense path below runs unchanged.
+        vllm_context = get_vllm_context()
+        if (
+            vllm_context is not None
+            and vllm_context.paged_attention_func is not None
+        ):
+            return self._compute_vllm_attention(
+                hidden_states, cache, vllm_context
+            )
+
         start_index = (
             cache_update_index if cache_update_index is not None else 0
         )
@@ -197,6 +210,34 @@ def _masked_softmax(self, attention_scores, attention_mask=None):
             )
         return self._softmax(attention_scores)
 
+    def _compute_vllm_attention(self, hidden_states, cache, vllm_context):
+        """vLLM paged-attention route (serving on TPU only).
+
+        Projects Q/K/V, applies RoPE at the engine's per-token positions,
+        and hands the (unexpanded) K/V to the shared paged-attention bridge.
+        The kernel handles grouped-query attention via ``num_kv_heads``.
+        """
+        positions = ops.reshape(vllm_context.positions, (-1, 1))
+        query = self.rotary_embedding_layer(
+            self._query_dense(hidden_states), positions=positions
+        )
+        key = self.rotary_embedding_layer(
+            self._key_dense(hidden_states), positions=positions
+        )
+        value = self._value_dense(hidden_states)
+
+        attention_output = vllm_paged_attention(
+            query,
+            key,
+            value,
+            self._inv_norm_factor,
+            num_kv_heads=self.num_key_value_heads,
+        )
+        attention_output = self._output_dense(attention_output)
+        if cache is not None:
+            return attention_output, cache
+        return attention_output
+
     def _compute_attention(self, query, key, value, attention_mask=None):
         if fused_attention_op_available():
             # Use `dot_product_attention` with Flash Attention support if