Serve KerasHub models on vLLM's GPU engine - #2943
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a GPU serving wrapper (KerasHubTorchModel) in keras_hub/src/vllm/torch_wrapper.py along with unit tests to enable serving KerasHub CausalLM models on vLLM's PyTorch engine. It also updates the vLLM plugin registration and registry to support this new backend path. The review feedback is highly constructive and identifies several critical issues: 1) registration logic in plugin.py can bypass GPU model registration if tpu_inference is present; 2) backend and load format checks in registry.py are incorrectly skipped in mixed environments; 3) eager imports of Keras/CausalLM at the module level in torch_wrapper.py can cause premature Keras initialization during vLLM startup; 4) attention layer attribute resolution should support self_attention as a fallback; and 5) exact float comparison for scale should be replaced with math.isclose to avoid precision issues.
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request adds support for serving KerasHub models on vLLM's GPU (PyTorch) engine, complementing the existing TPU path. It introduces KerasHubTorchModel as a PyTorch-based wrapper, registers a custom keras_hub load format, and updates the plugin registration logic to dynamically branch based on the hardware platform. Feedback focuses on improving the robustness of the TPU platform detection fallback and handling variable sequence lengths more safely in the model's forward pass.
bfloat16 needs compute capability 8.0, so defaulting to it made KerasHubLLM fail to start on a T4 with a message pointing at a vLLM CLI flag this API does not take. The TPU arm keeps bfloat16; the GPU arm asks the device and takes float16 below Ampere. An explicit dtype still wins.
Capture forbids host-to-device copies, and KerasHub layers make them: the causal mask builds a tensor from a Python scalar on every call, which stopped gpt2 during capture. Every family does something of the kind, so capture stays off for all of them rather than one line at a time. vLLM already declines to torch.compile these models, so only capture is lost.
The published function needs to know what its layer was built with, and was reading it back off attributes set on vLLM's Attention. Carry a record through the serving context instead, so nothing is stored on an object this integration does not own.
Gemma and Qwen alternate windowed and full layers, so they gate a width on a boolean. Mistral windows every layer and only holds the width, so the wrapper read no window, built its Attention modules without one, and the route then disagreed with them at the first call.
1 participant
Serves KerasHub
CausalLMmodels on vLLM's GPU engine. The backbone runs on Keras's torch backend, so its tensors are already the engine's. Same approach as the TPU path in #2794. The bridge, serving context and attention routes are unchanged.KerasHubTorchModelbuilds the backbone withCausalLM.from_preset, then onevllm.Attentionper layer for the engine to bind KV cache to. Those modules go through the serving context as the per-layer caches, so each route calls the module for its own layer.Sliding window and soft cap come from the built attention layer, not from config. The published function checks what the route passes against what the layer holds, and raises if they differ.
KerasHubPresetLoaderadds akeras_hubload format.plugin.pyregisters the TPU model on TPU and the torch one otherwise.KerasHubLLMerrors if Keras is not on the torch backend.7 presets match
CausalLM.generate()on an H100: GPT-2 base and large, Qwen 2.5, Llama 3.2, Gemma 1, 2 and 3. TPU still matches on v6.