Serve KerasHub CausalLM models through vLLM on TPU

[anthony-etim]

Description of the change

This adds a keras_hub/src/vllm/ module (plus small, additive in-file attention routes) that serves KerasHub CausalLM models on vLLM's TPU backend (tpu-inference) through its native flax/nnx path. It reuses the existing backbone and from_preset weights, so there's no model reimplementation and no checkpoint conversion.

Usage:

import keras_hub
llm = keras_hub.vllm.KerasHubLLM("keras_hub:gemma_2b_en")
print(llm.generate(["The future of AI is"])[0].outputs[0].text)

KerasHub's own generate() uses a static, dense, per-request KV cache with no continuous batching or paged attention, so it doesn't scale well under concurrent load. This routes KerasHub models through vLLM's scheduler and paged attention on TPU. It runs as a native JAX model, and the existing generate() path is untouched.

The serving model class, KerasHubVllmModel, lives in this PR (keras_hub/src/vllm/keras_hub_vllm_model.py); the companion tpu-inference PR registers it through the standard register_model API under the KerasHubForCausalLM architecture string, so the native loader resolves it like any other model. Its
forward delegates to the backbone's own graph, it does not re-implement embeddings, layers, or norms, and only publishes a thread-local serving context carrying the paged-attention kernel, the per-layer paged KV caches, and vLLM's token positions. This PR is made up of the following:

• keras_hub_vllm_model.py — the serving model (KerasHubVllmModel, an nnx adapter) that tpu-inference's native loader drives
• attention.py — the generic bridge into the paged-attention kernel
• context.py — thread-local serving state read by the attention layers
• hf_config.py — the neutral KerasHubConfig (model_type = "keras_hub") and the KerasHubForCausalLM architecture name, so no real model name is borrowed
• registry.py — KerasHubLLM and setup_vllm_model (writes the config). generate() without explicit SamplingParams defaults to the model's own compiled sampler (e.g. GPT-2's top_k), translated into vLLM settings
• tokenizer.py — KerasHubTokenizer: the preset's own KerasHub tokenizer on vLLM's tokenizer protocol, registered as tokenizer_mode="keras_hub" (the same registry mechanism vLLM uses for Mistral's tokenizer), no HF conversion or export, so serving tokenizes identically to generate()
• CPU unit tests, flat in the package (run on jax, tensorflow, and torch backends)
• keras_hub/api/vllm/ — the generated public API (keras_hub.vllm.KerasHubLLM)

Serving is layered into the models through explicit in-file routes: each supported attention layer's call() starts with a guarded check that, while serving, delegates to a small _compute_vllm_attention reproducing that
model's Q/K/V + positional logic and calling the shared bridge; otherwise the original dense path runs unchanged (byte-identical off-path).

Routes are in CachedMultiHeadAttention (GPT-2/OPT), LlamaAttention, QwenAttention, CachedGemmaAttention (Gemma 1 and 2), and CachedGemma3Attention; PositionEmbedding also reads vLLM's positions so learned-position models line up with the paged cache.

All routed families are validated end-to-end on TPU v6e:
GPT-2 (base and large, byte-level BPE), Llama 3.2 (instruct 1b), Gemma 1, Gemma 2 (attention soft cap and sliding window), Gemma 3, and Qwen 2.5 (coder 0.5b). Single chip TPU (TP=1).

Families that use cross-layer KV-sharing (Gemma 4, Gemma 3n) and ALiBi families (Bloom, Falcon) are not yet supported: they need bridge-side design work, not just a route.

Reference

Companion PR in vllm-project/tpu-inference registers this PR's KerasHubVllmModel (with the keras_hub model_type and tokenizer modes) through the standard mechanisms, importing keras_hub as an optional dependency, and threads the attention soft cap to its kernel; its shared loader is unchanged. The two PRs are used together: vllm-project/tpu-inference#3104.

Colab Notebook

colab_jax_native_benchmark.ipynb runs the model end to end and benchmarks it, with three configs: A = pure KerasHub generate(), B = vLLM-native (vLLM's own model for the same weights), C = this integration.

On correctness: C's next token matches native KerasHub exactly (greedy, bf16).

Benchmarks below are TPU v6e, bf16, output tok/s. prompt is the prompt size in words; conc is the number of simultaneous requests.

gpt2_large_en:

prompt	conc	KerasHub	Pure vLLM	KerasHub on vLLM	KerasHub on vLLM/KerasHub	KerasHub on vLLM/Pure vLLM
32	1	527	355	391	0.74×	1.10×
32	32	2077	5720	5508	2.65×	0.96×
32	64	2291	6312	6162	2.69×	0.98×
512	1	469	354	373	0.80×	1.06×
512	32	637	5372	5384	8.46×	1.00×
512	64	OOM	5926	6033	(KerasHub OOM)	1.02×

For GPT-2, C matches vLLM-native to within a few percent (0.96–1.10×) and is 2.6–8.5× faster than pure KerasHub under concurrency, and keeps serving long-context loads where pure KerasHub runs out of memory.

gemma3_instruct_1b:

prompt	conc	KerasHub	Pure vLLM	KerasHub on vLLM	KerasHub on vLLM/KerasHub	KerasHub on vLLM/Pure vLLM
32	1	514	147	331	0.64×	2.25×
32	32	11060	3863	7355	0.67×	1.90×
32	64	15507	6656	11205	0.72×	1.68×
512	1	508	150	329	0.65×	2.20×
512	32	5334	3933	6904	1.29×	1.76×
512	64	6258	6773	10363	1.66×	1.53×

For Gemma 3, C is consistently faster than vLLM-native (1.5–2.25×). Against pure KerasHub, C wins on the 512-word prompts under load (up to 1.66×); at short prompts KerasHub's static batching is very fast for this small model, so C's
paged-attention overhead only pays off as prompt length and concurrency grow.

llama3.2_instruct_1b:

prompt	conc	KerasHub	Pure vLLM	KerasHub on vLLM	KerasHub on vLLM/KerasHub	KerasHub on vLLM/Pure vLLM
32	1	431	399	380	0.88×	0.95×
32	32	5663	9598	8328	1.47×	0.87×
32	64	6898	15090	12725	1.84×	0.84×
512	1	402	374	369	0.92×	0.99×
512	32	1885	9250	7820	4.15×	0.85×
512	64	1919	14138	11864	6.18×	0.84×

For Llama 3.2, C tracks vLLM-native closely (0.84–0.99×) and is 1.5–6.2× faster than pure KerasHub under concurrency, with the gain over KerasHub growing at longer prompts (up to 6.2× at 512 words).

Checklist

• I have added all the necessary unit tests for my change.
• I have verified that my change does not break existing code and works with all backends. (The serving module is JAX/TPU-only; the in-file attention routes are additive and no-op off the serving path, byte-identical, and CPU-tested on jax/tf/torch.)
• My PR is based on the latest changes of the main branch.
• I have followed the Keras Hub Model contribution guidelines.
• I have followed the Keras Hub API design guidelines.
• I have signed the Contributor License Agreement.

[gemini-code-assist]

Warning

Gemini encountered an error creating the review. You can try again by commenting /gemini review.

[divyashreepathihalli]

The PR is hyper focused on a few models. There is model specific code in generic functions and calls. This way of coding is not scalable.

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces an integration module for serving Keras Hub models through vLLM on TPU, adding thread-local context management, a shared dispatch for paged attention, layer adapters, a JAX-native flax/nnx model adapter, and a tokenizer adapter, along with comprehensive unit tests. The review feedback focuses on enhancing the robustness of the implementation, including suggestions to clean up temporary directories using an atexit handler, add proactive bounds checking for the KV cache index, provide a fallback identity function when no final norm layer is found, defensively check for None samplers, and use robust UTF-8 decoding with errors='replace' to handle partial characters during streaming.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces an integration module for serving Keras Hub models through vLLM on TPU. It adds routing logic to vLLM's paged-attention kernel within several attention layers (GPT-2, Gemma 3, LLaMA, Qwen) and the position embedding layer when a thread-local serving context is active. It also implements a JAX-native flax/nnx model adapter, a neutral Hugging Face configuration, and a tokenizer adapter to bridge Keras Hub models to vLLM. Feedback on the pull request highlights a few issues: a namespace error when importing Tokenizer in registry.py, a potential type error in tokenizer.py where bytes keys are not decoded to strings, and several violations of the repository's style guide regarding the use of type hints in function signatures.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request implements JAX-native (flax/nnx) integration for serving Keras Hub models through vLLM on TPU. It introduces a thread-local serving context, a paged-attention dispatch bridge, and a tokenizer adapter, alongside routing hooks in attention and position embedding layers for models like GPT-2, Gemma, Llama, and Qwen. The review feedback highlights critical improvements: utilizing preset_utils.load_weights in the NNX adapter to prevent memory duplication and potential OOM errors, and ensuring proper decoding of bytes to str for special tokens and vocabularies in the registry and tokenizer modules to avoid JSON serialization crashes downstream.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces JAX-native integration for serving Keras Hub models through vLLM on TPU. It adds a thread-local context to pass serving metadata, routes various attention layers (GPT-2, Gemma, Gemma 3, LLaMA, Qwen) to a shared paged-attention bridge, and implements a flax/nnx adapter alongside tokenizer export utilities. The review feedback highlights several critical bugs and robustness issues: the Unigram constructor and decoders.Sequence in the tokenizer export use non-existent arguments and classes (byte_fallback and Strip), an unhandled IndexError could occur when querying CPU devices, and zipping weights during loading lacks length validation. Additionally, suggestions were made to cache configuration parsing and to use a more robust check for decoding byte-like tokens.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces a comprehensive integration module for serving Keras Hub models through vLLM on TPU. It adds routing hooks in several attention layers (CachedMultiHeadAttention, Gemma, Gemma3, Llama, and Qwen) and PositionEmbedding to delegate to vLLM's paged-attention kernel and handle absolute positions via a thread-local context. It also implements a JAX-native flax/nnx model adapter, a neutral HF config carrier, and tokenizer export utilities. The review feedback is highly constructive and should be fully addressed: it highlights several style guide violations regarding the use of type hints in public signatures, the missing @keras_hub_export decorator on the public KerasHubLLM class, and robustness issues such as unhandled exceptions in file deletion, fragile fallback tokenizers, and potential disk space leaks from relying solely on atexit for directory cleanup.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces support for serving KerasHub models through vLLM on TPU by adding a new keras_hub/src/vllm/ integration module. The changes include a thread-local serving context, a paged-attention bridge, and updates to various attention layers and the position embedding layer to support vLLM's paged KV cache and per-token absolute positions. The reviewer noted that the docstrings for compute_logits and load_weights in keras_hub/src/vllm/nnx_adapter.py are missing required Args and Returns sections and should be updated to follow the repository's Google-style docstring requirements.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[divyashreepathihalli]

/Gemini review

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request integrates KerasHub with vLLM to enable TPU serving for causal language models. It introduces a thread-local serving context, a paged-attention bridge, a custom tokenizer adapter, and a neutral Hugging Face configuration, while routing attention layers in GPT-2, Gemma, LLaMA, and Qwen to the paged-attention bridge when serving is active. The review feedback highlights critical robustness improvements: mapping max_sequence_length to max_position_embeddings to prevent out-of-bounds crashes, using tempfile.TemporaryDirectory instead of mkdtemp and atexit to avoid memory leaks, and updating the tokenizer to robustly handle NumPy integer types and 0-D arrays to prevent type errors.

[gemini-code-assist]

Code Review

This pull request introduces integration between KerasHub and vLLM to enable serving KerasHub models on TPUs. It adds a thread-local serving context, a paged-attention bridge, a custom vLLM-compatible tokenizer wrapper, and configuration setup utilities, while routing attention layers (Gemma, Gemma3, LLaMA, Qwen, and CachedMHA) and position embeddings to use vLLM's metadata when active. The review feedback highlights a potential issue in multi-node TPU environments where local temporary directories are inaccessible, suggesting the addition of a parent_dir parameter. It also notes a style guide violation regarding missing type information in a docstring, and points out an inconsistency in bos_token_id fallback logic between the model setup and the tokenizer.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request integrates KerasHub with vLLM to enable TPU serving of CausalLM models via a thread-local context, a paged-attention bridge, and a custom tokenizer wrapper. Several model attention layers and position embeddings are updated to route to the vLLM serving path when active. The review feedback focuses on enhancing robustness and API consistency, specifically by raising an error when unsupported attention scores are requested in CachedMultiHeadAttention, matching the generate signature of KerasHubLLM with the base class, and casting token IDs to standard Python integers in KerasHubTokenizer to prevent backend compatibility issues.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request integrates KerasHub with vLLM to enable serving KerasHub causal language models on TPUs. It introduces a thread-local serving context, a paged-attention bridge, and custom routing logic within attention layers (MHA, Gemma, Gemma 3, Llama, Qwen) and position embeddings to delegate computations to vLLM's paged-attention kernel. Additionally, it registers a custom KerasHubTokenizer and a neutral Hugging Face configuration to support loading KerasHub presets directly. The review feedback highlights a potential serialization issue in setup_vllm_model where passing a non-string torch.dtype object could cause json.dump to raise a TypeError, and suggests normalizing the dtype parameter to a string to ensure robustness.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request integrates KerasHub with vLLM to enable serving CausalLM models on TPU. It introduces a thread-local serving context, a paged-attention bridge (vllm_paged_attention), a neutral Hugging Face configuration class for KV-cache profiling, and a custom KerasHubTokenizer that wraps KerasHub tokenizers for vLLM compatibility. Existing attention layers (MHA, Gemma, Gemma 3, Llama, Qwen) and PositionEmbedding are updated to route to the vLLM paged-attention path when a serving context is active. The review feedback correctly identifies that KerasHubTokenizer.convert_ids_to_tokens should return a single string (or None) instead of a list when given a scalar integer input to comply with the Hugging Face tokenizer API expected by vLLM, and suggests updating the corresponding unit test.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces integration for serving KerasHub models through vLLM on TPU. It implements a thread-local serving context, a paged-attention bridge, a custom tokenizer wrapper, and a neutral Hugging Face configuration carrier. Additionally, it routes attention layers (including Gemma, Gemma3, LLaMA, and Qwen) and position embeddings to delegate to the vLLM paged-attention kernel when active. The feedback suggests handling dtype=None robustly in setup_vllm_model to prevent string conversion to "None", which would cause vLLM loading failures.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces integration with vLLM to serve KerasHub models on TPU. It implements a thread-local context to pass serving metadata, a paged-attention bridge that routes attention computations to vLLM's Pallas kernel, and a custom KerasHubTokenizer to serve tokenizers directly through vLLM's registry. The attention routing is integrated into several model attention layers (MHA, Gemma, Gemma 3, LLaMA, Qwen) and PositionEmbedding. Feedback on the changes highlights a few robustness issues: potential crashes from duplicate registrations of "keras_hub" in interactive environments, invalid torch_dtype values when using Keras DTypePolicy or "auto", and a potential KeyError when passing bytes tokens to convert_tokens_to_ids.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces integration between KerasHub and vLLM to serve KerasHub models on TPU. It implements a thread-local serving context, a paged-attention bridge that routes KerasHub attention layers (MHA, Gemma, Gemma3, Llama, Qwen) and position embeddings to vLLM's TPU kernels, and a custom KerasHubTokenizer to serve tokenizers directly through vLLM. The feedback highlights two key improvements: first, handling task presets in registry.py where the backbone configuration is nested under the "backbone" key; second, robustly converting tensor inputs to standard Python lists in tokenizer.py using ops.convert_to_numpy to prevent tracing failures and performance issues.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request integrates KerasHub with vLLM to enable serving CausalLM models on TPUs. It introduces a thread-local serving context, a paged-attention bridge, a custom tokenizer adapter, and routes various model attention layers (Gemma, Gemma 3, LLaMA, Qwen, and CachedMultiHeadAttention) to the vLLM paged-attention path when serving. The feedback recommends using a deterministic directory path instead of a random temporary directory to prevent failures in multi-node TPU pods, and converting framework tensors to numpy in _as_id_list to ensure backend-agnostic behavior and avoid slow device-to-host syncs.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request implements KerasHub integration with vLLM for TPU serving, introducing a thread-local VLLMContext, a paged-attention bridge (vllm_paged_attention) that routes attention layers (such as Gemma, Llama, and Qwen) to vLLM's Pallas kernel, and a KerasHubTokenizer adapter. The review feedback identifies two robustness issues: first, the dtype parsing logic in setup_vllm_model is fragile and should use keras.backend.standardize_dtype; second, KerasHubTokenizer.all_special_tokens should use convert_ids_to_tokens to leverage the existing fallback mechanism and avoid crashes on tokenizers that lack id_to_token support.

[anthony-etim]

/gemini review

[gemini-code-assist]

Code Review

This pull request introduces integration between KerasHub and vLLM to serve KerasHub models on TPU. It adds a thread-local serving context, a paged-attention bridge, and custom tokenizer support, routing attention computations for models like Gemma, Llama, and Qwen to vLLM's injected Pallas kernel. The review feedback focuses on improving robustness and preventing side effects: specifically, wrapping tokenizer and renderer registrations in try-except blocks to handle duplicate registration, converting token IDs to tensors before detokenization to ensure backend-agnostic compatibility, and avoiding global modifications to os.environ by configuring implementation types locally.

[divyashreepathihalli]

Thank you for the updates @anthony-etim !
Left some comments.