chore: import upstream snapshot with attribution

doc/source/llm/doc_code/serve/custom_vllm/.gitignore

@@ -0,0 +1,3 @@
+build/
+*.egg-info/
+__pycache__/

doc/source/llm/doc_code/serve/custom_vllm/custom_vllm_config.yaml

@@ -0,0 +1,20 @@
+# config.yaml
+applications:
+- args:
+    llm_configs:
+        - model_loading_config:
+            model_id: qwen3-reward
+            model_source: Qwen/Qwen3-0.6B
+          # Register the plugin architecture in the Ray Serve LLM processes.
+          server_cls: qwen3_reward_plugin.serve_hook:RewardModelServer
+          engine_kwargs:
+            runner: pooling
+            hf_overrides:
+                architectures:
+                    - Qwen3CustomRewardModel
+                num_labels: 1
+                problem_type: regression
+            max_model_len: 4096
+  import_path: ray.serve.llm:build_openai_app
+  name: custom_vllm_app
+  route_prefix: "/"

doc/source/llm/doc_code/serve/custom_vllm/custom_vllm_example.py

@@ -0,0 +1,91 @@
+"""Documentation example and CI test: custom vLLM model with Ray Serve LLM.
+
+Structure:
+1. Test-only setup: force serve.run non-blocking. The plugin is baked into the
+   cluster image, and direct streaming is enabled through the cluster
+   environment.
+2. Docs example (between __custom_vllm_example_start/end__): embedded in the
+   guide via literalinclude.
+3. Test validation (deployment status polling + reward-head assertion + cleanup).
+"""
+
+import json
+import time
+import urllib.request
+
+from ray import serve
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve.schema import ApplicationStatus
+
+_original_serve_run = serve.run
+
+def _non_blocking_serve_run(app, **kwargs):
+    """Forces blocking=False for testing."""
+    kwargs["blocking"] = False
+    return _original_serve_run(app, **kwargs)
+
+serve.run = _non_blocking_serve_run
+
+# __custom_vllm_example_start__
+from ray import serve
+from ray.serve.llm import LLMConfig, build_openai_app
+
+# Serve Qwen3-0.6B with a scalar reward head from the `qwen3_reward_plugin`
+# vLLM plugin (pip install it into your cluster image).
+llm_config = LLMConfig(
+    model_loading_config=dict(
+        model_id="qwen3-reward",
+        model_source="Qwen/Qwen3-0.6B",
+    ),
+    # vLLM's entry point registers the architecture in the engine and worker
+    # processes. Ray Serve LLM also resolves it while building the engine config.
+    server_cls="qwen3_reward_plugin.serve_hook:RewardModelServer",
+    engine_kwargs=dict(
+        runner="pooling",
+        hf_overrides=dict(
+            architectures=["Qwen3CustomRewardModel"],
+            num_labels=1,
+            problem_type="regression",
+        ),
+        max_model_len=4096,
+    ),
+)
+
+# Requires direct streaming so vLLM's native /classify route is exposed.
+app = build_openai_app({"llm_configs": [llm_config]})
+serve.run(app, blocking=True)
+# __custom_vllm_example_end__
+
+status = ApplicationStatus.NOT_STARTED
+timeout_seconds = 300
+start_time = time.time()
+
+while (
+    status != ApplicationStatus.RUNNING and time.time() - start_time < timeout_seconds
+):
+    status = serve.status().applications[SERVE_DEFAULT_APP_NAME].status
+    if status in [ApplicationStatus.DEPLOY_FAILED, ApplicationStatus.UNHEALTHY]:
+        raise AssertionError(f"Deployment failed with status: {status}")
+    time.sleep(1)
+
+if status != ApplicationStatus.RUNNING:
+    raise AssertionError(
+        f"Deployment failed to reach RUNNING status within {timeout_seconds}s. "
+        f"Current status: {status}"
+    )
+
+# Verify the reward head ran end to end: /classify returns a single scalar score
+# per input (num_labels=1), read from data[0].probs.
+body = json.dumps(
+    {"model": "qwen3-reward", "input": "The capital of France is Paris."}
+).encode()
+request = urllib.request.Request(
+    "http://localhost:8000/classify",
+    data=body,
+    headers={"Content-Type": "application/json"},
+)
+response = json.load(urllib.request.urlopen(request, timeout=60))
+reward = response["data"][0]["probs"]
+assert len(reward) == 1, f"Expected a scalar reward, got {len(reward)} values"
+
+serve.shutdown()

doc/source/llm/doc_code/serve/custom_vllm/custom_vllm_yaml_config_example.py

@@ -0,0 +1,63 @@
+"""Documentation example and CI test: custom vLLM model via a YAML config.
+
+Structure:
+1. Test-only setup: force serve.run non-blocking and strip accelerator
+   requirements for CI. The plugin is baked into the cluster image, and direct
+   streaming is enabled through the cluster environment.
+2. Load the YAML config and deploy it with build_openai_app.
+3. Test validation (deployment status polling + reward-head assertion + cleanup).
+"""
+
+import json
+import os
+import time
+import urllib.request
+
+import yaml
+
+from ray import serve
+from ray.serve import llm
+from ray.serve._private.constants import SERVE_DEFAULT_APP_NAME
+from ray.serve.schema import ApplicationStatus
+
+config_path = os.path.join(os.path.dirname(__file__), "custom_vllm_config.yaml")
+with open(config_path, "r") as f:
+    config_dict = yaml.safe_load(f)
+
+llm_configs = config_dict["applications"][0]["args"]["llm_configs"]
+app = llm.build_openai_app({"llm_configs": llm_configs})
+serve.run(app, blocking=False)
+
+status = ApplicationStatus.NOT_STARTED
+timeout_seconds = 300
+start_time = time.time()
+
+while (
+    status != ApplicationStatus.RUNNING and time.time() - start_time < timeout_seconds
+):
+    status = serve.status().applications[SERVE_DEFAULT_APP_NAME].status
+    if status in [ApplicationStatus.DEPLOY_FAILED, ApplicationStatus.UNHEALTHY]:
+        raise AssertionError(f"Deployment failed with status: {status}")
+    time.sleep(1)
+
+if status != ApplicationStatus.RUNNING:
+    raise AssertionError(
+        f"Deployment failed to reach RUNNING status within {timeout_seconds}s. "
+        f"Current status: {status}"
+    )
+
+# Verify the reward head ran end to end: /classify returns a single scalar score
+# per input (num_labels=1), read from data[0].probs.
+body = json.dumps(
+    {"model": "qwen3-reward", "input": "The capital of France is Paris."}
+).encode()
+request = urllib.request.Request(
+    "http://localhost:8000/classify",
+    data=body,
+    headers={"Content-Type": "application/json"},
+)
+response = json.load(urllib.request.urlopen(request, timeout=60))
+reward = response["data"][0]["probs"]
+assert len(reward) == 1, f"Expected a scalar reward, got {len(reward)} values"
+
+serve.shutdown()

doc/source/llm/doc_code/serve/custom_vllm/qwen3_reward_plugin/__init__.py

@@ -0,0 +1,11 @@
+# __register_start__
+"""vLLM plugin entry point: register the custom architecture."""
+def register() -> None:
+    from vllm import ModelRegistry
+
+    if "Qwen3CustomRewardModel" not in ModelRegistry.get_supported_archs():
+        ModelRegistry.register_model(
+            "Qwen3CustomRewardModel",
+            "qwen3_reward_plugin.qwen3_rm:Qwen3CustomRewardModel",
+        )
+# __register_end__

doc/source/llm/doc_code/serve/custom_vllm/qwen3_reward_plugin/qwen3_rm.py

@@ -0,0 +1,91 @@
+"""Qwen3 with the LM head replaced by a scalar reward head, as a vLLM plugin."""
+
+import os
+from collections.abc import Iterable
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.pooler import DispatchPooler
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
+from vllm.model_executor.models.interfaces_base import default_pooling_type
+from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM
+from vllm.model_executor.models.utils import AutoWeightsLoader, maybe_prefix
+
+logger = init_logger(__name__)
+
+REWARD_HEAD_FILE = "reward_head.pt"
+REWARD_HEAD_KEYS = ("linear.weight", "weight", "score.weight")
+
+
+@default_pooling_type(seq_pooling_type="LAST")
+class Qwen3CustomRewardModel(Qwen3ForCausalLM):
+    """Qwen3 backbone + scalar reward head, served as a pooling model."""
+
+    is_pooling_model = True
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        # A single scalar output (set problem_type="regression" in the config
+        # for an identity activation, so /classify returns the raw reward).
+        vllm_config.model_config.hf_config.num_labels = 1
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+
+        model_config = vllm_config.model_config
+
+        # Scalar reward head. Replicated because the output dim is 1.
+        self.score = ReplicatedLinear(
+            model_config.hf_config.hidden_size,
+            1,
+            bias=False,
+            params_dtype=model_config.head_dtype,
+            return_bias=False,
+            prefix=maybe_prefix(prefix, "score"),
+        )
+
+        # LAST-token pooling -> reward head -> raw score. Advertising the
+        # "classify" task mounts vLLM's /classify route.
+        pooler_config = model_config.pooler_config
+        assert pooler_config is not None
+        self.pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        skip_prefixes = ["score."]  # reward head loads separately, below
+        if self.config.tie_word_embeddings:
+            skip_prefixes.append("lm_head.")  # tied: no lm_head weight to load
+        loader = AutoWeightsLoader(self, skip_prefixes=skip_prefixes)
+        loaded = loader.load_weights(weights)
+        if self._load_reward_head():
+            loaded.add("score.weight")
+        return loaded
+
+    def _load_reward_head(self) -> bool:
+        """Load ``reward_head.pt``; return False (keep init values) if absent."""
+        env_path = os.environ.get("RM_REWARD_HEAD_PATH")
+        model_dir = self.vllm_config.model_config.model
+        head_path = (
+            env_path
+            if env_path and os.path.isfile(env_path)
+            else os.path.join(model_dir, REWARD_HEAD_FILE)
+        )
+        if not os.path.isfile(head_path):
+            logger.warning(
+                "Reward head %r not found; scores are not meaningful.", head_path
+            )
+            return False
+
+        state = torch.load(head_path, map_location="cpu", weights_only=True)
+        weight = (
+            state
+            if isinstance(state, torch.Tensor)
+            else next((state[k] for k in REWARD_HEAD_KEYS if k in state), None)
+        )
+        if weight is None:
+            raise KeyError(f"{REWARD_HEAD_FILE} has none of {REWARD_HEAD_KEYS}")
+
+        param = self.score.weight
+        weight_loader = getattr(param, "weight_loader", default_weight_loader)
+        weight_loader(param, weight.reshape(1, -1).to(param.dtype))
+        logger.info("Loaded reward head from %s.", head_path)
+        return True

doc/source/llm/doc_code/serve/custom_vllm/qwen3_reward_plugin/serve_hook.py

@@ -0,0 +1,13 @@
+# __serve_hook_start__
+"""LLMServer subclass for ``LLMConfig.server_cls``: importing it registers the plugin."""
+
+from ray.llm._internal.serve.core.server.llm_server import LLMServer
+
+from qwen3_reward_plugin import register
+
+register()
+
+
+class RewardModelServer(LLMServer):
+    pass
+# __serve_hook_end__

doc/source/llm/doc_code/serve/custom_vllm/setup.py

@@ -0,0 +1,16 @@
+# __setup_start__
+from setuptools import setup
+
+setup(
+    name="qwen3-reward-plugin",
+    version="0.1.0",
+    packages=["qwen3_reward_plugin"],
+    # vLLM discovers and runs this entry point in every process (the engine core
+    # and each rank worker) via load_general_plugins().
+    entry_points={
+        "vllm.general_plugins": [
+            "qwen3_reward = qwen3_reward_plugin:register",
+        ],
+    },
+)
+# __setup_end__