Add guide: Serving KerasHub models with vLLM using Kinetic

[Dhiraj099]

What this adds

• examples/vllm_serving.py + examples/export_worker.py — a self-contained example that exports a KerasHub causal LM with export_to_transformers() and serves it with vLLM in a single Kinetic job, on TPU or GPU, with any Keras backend (jax / torch / tensorflow).
• docs/guides/vllm_serving.md — the guide.

How it works

One Kinetic job: download the KerasHub preset → export with export_to_transformers() in a child process → the child exits and the OS reclaims its framework state and device memory → vLLM starts on the exported directory in the now-clean process → continuous-batched completions are returned to the local machine. DEVICE = "tpu" | "gpu" selects the accelerator, requirements, and serve config; BACKEND selects the Keras export backend (defaults to the one natural to the device, which needs no extra deps).

Design notes

Why a subprocess for the export (instead of doing it in-process). Exporting and serving in the same process only works cleanly for one combination — torch on GPU (no libtpu, and torch.cuda.empty_cache() returns VRAM). Everywhere else, in-process breaks or is fragile:

Device	Backend	In-process result
TPU	torch / tensorflow	Hard crash — keras-hub loads TensorFlow, which can't coexist with libtpu: protobuf double-registration (timestamp.proto via tpu_metric_service.pb.cc) + a fatal SIGSEGV in GOMP_parallel at weight load. Reproduced with the working serve path plus a single import tensorflow.
TPU	jax	Runs, but XLA never frees the export's device memory without process exit, so the resident export competes with vLLM's KV cache and OOMs for non-trivial models .
GPU	torch	Works in-process with del + empty_cache.
GPU	jax / tensorflow	Needs that backend's CUDA wheel (else Unable to initialize backend 'rocm'), and has the same no-free-without-exit issue.

Question for Reviewers: Single job vs. two jobs (GCS)
Both isolate the export. I chose the single-job subprocess for the example because it's one command, keeps the checkpoint local, and needs no GCS round-trip. The two-job split (export job → KINETIC_OUTPUT_DIR on GCS → serve job) is the better fit when we want to export once and serve the same checkpoint many times, and it puts vllm and vllm-tpu in separate images naturally; the guide documents it as that variant. Please let me know if we should proceed with the two-job split.

Running the export in a child process that exits first reclaims its entire memory/state and is the single pattern correct for every device × backend, with no hardcoded backend and no manual cleanup. (Kinetic unpacks job files onto sys.path; the parent resolves the worker with find_spec without importing it, so Keras/TF never load in the serve process.)

architectures key — handled upstream, not patched here. vLLM routes checkpoints via the architectures key in config.json. Gemma 3's exporter writes it, so this example (published for Gemma 3) needs no patch. The GPT-2 / Qwen / Gemma exporters don't write it yet — that should be fixed in those keras-hub exporters (follow-up) rather than worked around in user-facing example code.

_setup_gpu_runtime() (GPU only) — exposes the GKE-mounted NVIDIA driver on LD_LIBRARY_PATH (inherited by the export child, so the export loads onto the GPU, not host RAM) and disables FlashInfer's nvcc JIT. Runs before the export. No-op once the GPU image ships the driver on the path (#269).

Validated runs

• TPU (subprocess) — gemma3_4b on tpu-v5litepod-1
• GPU — gemma3_4b on L4.

Contributor Agreement

Please check all boxes below before submitting your PR for review:

• I am a human, and not a bot.
• I will be responsible for responding to review comments in a timely manner.
• I will work with the maintainers to push this PR forward until submission.
• I will test the changes on my cloud setup and provide proof of successful validation.

Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed.

[gemini-code-assist]

Code Review

This pull request introduces a new guide and example script for exporting KerasHub models to Hugging Face Transformers format and serving them with vLLM on Cloud GPUs or TPUs using Kinetic. The feedback focuses on improving robustness and resource utilization: first, by uploading the exported model directory directly to GCS using transfer_manager instead of creating a large .tar archive in /tmp which could cause GKE pod eviction; second, by persisting the exported model to GCS before initializing the high-risk vLLM engine; and finally, updating the documentation to reflect these direct directory transfer workflows.

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.