Rename the serving model to KerasHubVllmModel

Author: anthony-etim

keras_hub/src/vllm/__init__.py

@@ -4,7 +4,7 @@
 tpu-inference's native flax/nnx path. This package provides the serving
 context, the paged-attention bridge, the vLLM tokenizer (the preset's own
 KerasHub tokenizer, served via `tokenizer_mode="keras_hub"`), and the
-`KerasHubLLM` entry point; the serving model class (`KerasHubForCausalLM`)
+`KerasHubLLM` entry point; the serving model class (`KerasHubVllmModel`)
 lives in tpu-inference.
 """
 

keras_hub/src/vllm/attention.py

@@ -97,7 +97,7 @@ def vllm_paged_attention(
     value = ops.reshape(value, (-1, kv_heads * head_dim))
 
     # `paged_attention_func` comes from the serving model
-    # (keras_hub_for_causal_lm.py).
+    # (keras_hub_vllm_model.py).
     # It takes only what an attention layer knows; engine details (the
     # attention metadata, the mesh) are already bound inside it. It returns
     # `(new_kv_cache, output)`.

keras_hub/src/vllm/context.py

@@ -1,6 +1,6 @@
 """Thread-local context for passing vLLM metadata to Keras layers.
 
-The serving model (``KerasHubForCausalLM`` in this package, registered
+The serving model (``KerasHubVllmModel`` in this package, registered
 with vLLM by tpu-inference's plugin hook) publishes this
 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.

keras_hub/src/vllm/hf_config.py

@@ -7,7 +7,7 @@
 ``"KerasHubForCausalLM"``.
 
 This config is never used to build a model. tpu-inference's native loader routes
-any config carrying ``keras_hub_preset`` to its ``KerasHubForCausalLM`` model
+any config carrying ``keras_hub_preset`` to its ``KerasHubVllmModel`` model
 class (the real KerasHub backbone on the flax/nnx path); this object only
 carries the dims vLLM reads for KV-cache profiling, which
 ``setup_vllm_model`` fills from each preset's own config.

…_hub/src/vllm/keras_hub_for_causal_lm.py …ras_hub/src/vllm/keras_hub_vllm_model.pykeras_hub/src/vllm/keras_hub_for_causal_lm.py renamed to keras_hub/src/vllm/keras_hub_vllm_model.py keras_hub/src/vllm/keras_hub_for_causal_lm.py renamed to keras_hub/src/vllm/keras_hub_vllm_model.py

@@ -1,6 +1,6 @@
 """The serving model that runs a KerasHub `CausalLM` on vLLM's TPU backend.
 
-`KerasHubForCausalLM` lives here in keras-hub and is registered with
+`KerasHubVllmModel` lives here in keras-hub and is registered with
 tpu-inference's model registry (and vLLM's) by tpu-inference's plugin hook,
 so the native `flax_nnx` loader resolves it by architecture name like any
 other model. It is only ever instantiated on the serving path, where flax
@@ -24,12 +24,16 @@
 from keras_hub.src.vllm.context import vllm_context_scope
 
 
-class KerasHubForCausalLM(_NnxModule):
+class KerasHubVllmModel(_NnxModule):
     """Serves a KerasHub `CausalLM` on tpu-inference's native JAX path.
 
+    Registered under the ``"KerasHubForCausalLM"`` architecture string
+    (``KERAS_HUB_ARCHITECTURE``), which follows the HF naming convention
+    configs use on the wire; the class itself carries the keras-hub-style
+    name.
+
     An adapter, not a conversion: it implements the model interface the
-    native `flax_nnx` runner drives (resolved by the `KerasHubForCausalLM`
-    architecture name through the standard model registry), reusing the
+    native `flax_nnx` runner drives, reusing the
     preset's existing backbone and weights. Keras's NNX mode
     (`KERAS_NNX_ENABLED=true`) makes the backbone's variables nnx state, so
     the runner's `nnx.split`/`nnx.merge` machinery carries the weights with

keras_hub/src/vllm/registry.py

@@ -2,7 +2,7 @@
 Registry module for Keras Hub to vLLM Integration.
 
 Materializes a model directory for `LLM(model="keras_hub:preset_name")` so that
-tpu-inference's native flax/nnx path serves it via `KerasHubForCausalLM`.
+tpu-inference's native flax/nnx path serves it via `KerasHubVllmModel`.
 """
 
 import inspect
@@ -48,7 +48,7 @@
 # teaches transformers the `model_type`, and tpu-inference registers the
 # architecture with vLLM. Both only carry KV dims for cache profiling; the real
 # model is selected by `keras_hub_preset` and served by tpu-inference's
-# `KerasHubForCausalLM` model class.
+# `KerasHubVllmModel` model class.
 
 
 def _normalize_dtype(dtype):