From 380638afbb2c5658f22124f6d1d38fcd464ea7d9 Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Tue, 13 Oct 2026 13:46:37 +0800 Subject: [PATCH] [Model] Add Clef-Flash support (#61308) Signed-off-by: Tianyu Guo Signed-off-by: Isotr0py Co-authored-by: Codex Co-authored-by: Isotr0py --- .../online_serving/structured_decisions.md | 26 +- tests/config/test_model_arch_config.py | 34 ++ .../structured_decisions/test_strategies.py | 65 ++-- .../multimodal/pooling/test_clef_omni.py | 7 +- tests/models/registry.py | 4 + .../structured_decisions/strategies.py | 56 +++ vllm/model_executor/models/clef_flash.py | 44 +++ .../models/clef_joint_schema.py | 343 ++++++++++++++++++ vllm/model_executor/models/clef_omni.py | 339 +---------------- vllm/model_executor/models/registry.py | 1 + .../model_arch_config_convertor.py | 22 ++ 11 files changed, 580 insertions(+), 361 deletions(-) create mode 100644 vllm/model_executor/models/clef_flash.py create mode 100644 vllm/model_executor/models/clef_joint_schema.py diff --git a/docs/serving/online_serving/structured_decisions.md b/docs/serving/online_serving/structured_decisions.md index ce09d6f81f..4a7ec00dd2 100644 --- a/docs/serving/online_serving/structured_decisions.md +++ b/docs/serving/online_serving/structured_decisions.md @@ -130,6 +130,30 @@ the request lists them, a `noul` labels its options `yes` and `no`, and a | options per `choice` | 26 | | levels per `score` | 10 | +## Clef-Flash + +[Clef-Flash](https://huggingface.co/Cloudflare/clef-flash) uses a Qwen3.5-9B +backbone and a joint schema head to score all questions in one forward pass. +vLLM detects the head and serves the checkpoint as a pooling model. + +```bash +vllm serve Cloudflare/clef-flash --max-model-len 16384 +``` + +Send requests to `/v1/systemone` using `state` and `questions`, as in the +examples above. `choice`, `noul` and `score` questions are supported. + +- Text and JSON states use the checkpoint's fixed prompt. JSON is compact + with sorted keys. `instructions` returns 400; `chat_template_kwargs` + has no effect. +- The state is truncated to fit 16,384 tokens, or `max_model_len` if smaller. + A schema that exceeds this bound returns 400. +- `images` and `videos` are placed before the state. Videos are sampled at + 2 fps and processed by the Qwen3.5 vision processor. Audio inputs return + 400, and video soundtracks are ignored. +- Each request performs one pooling read. `usage.input_tokens` counts the + prompt once, `output_tokens` is 0, and `label_mass` is 1. + ## Clef-Omni [Clef-Omni](https://huggingface.co/Cloudflare/clef-omni) is a decision model: @@ -188,7 +212,7 @@ vllm serve LiquidAI/d1-3B | `instructions` | optional context placed ahead of the questions | | `chat_template_kwargs` | passed to the chat template, for example `{"enable_thinking": false}` | | `cache_salt` | optional [prefix-cache salt](../../usage/security.md#cache-salting) applied to every question | -| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Omni](#clef-omni), [Liquid d1](#liquid-d1), images only) | +| `images`, `audio`, `videos` | optional media URLs or data URLs, for models that take them ([Clef-Flash](#clef-flash), images/video; [Clef-Omni](#clef-omni); [Liquid d1](#liquid-d1), images only) | A question with any field other than `type`, `instructions` and `criteria` is rejected with a 400. diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index 8ec16ee7f9..2ff992a205 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -70,6 +70,40 @@ def _load_groundtruth(filename: str) -> dict: return json.load(f) +@pytest.mark.parametrize("has_joint_head", [False, True]) +@pytest.mark.parametrize( + "architecture", ["Qwen3_5ForConditionalGeneration", "ColQwen3_5"] +) +def test_qwen35_joint_head_selects_clef_flash(tmp_path, has_joint_head, architecture): + """Only a Qwen3.5 generation checkpoint with a joint head selects Clef.""" + from transformers import Qwen3_5Config + + head_config = { + "hidden_size": 4096, + "width": 1024, + "routing_layers": 2, + "layers": 4, + "heads": 16, + "feedforward": 4096, + } + if has_joint_head: + (tmp_path / "joint_head_config.json").write_text(json.dumps(head_config)) + config = Qwen3_5Config(architectures=[architecture]) + config.name_or_path = str(tmp_path) + model_config = SimpleNamespace( + hf_config=config, + hf_text_config=config.text_config, + _supports_multimodal_for_mm_prefix=lambda: True, + ) + converted = ModelConfig.get_model_arch_config(model_config) + if has_joint_head and architecture == "Qwen3_5ForConditionalGeneration": + assert converted.architectures == ["ClefFlashForJointSchema"] + assert config.joint_head_config == head_config + else: + assert converted.architectures == [architecture] + assert not hasattr(config, "joint_head_config") + + def _assert_model_arch_config( model_config, expected: dict, check_head_size: bool = True ): diff --git a/tests/entrypoints/generate/structured_decisions/test_strategies.py b/tests/entrypoints/generate/structured_decisions/test_strategies.py index dae545d1f0..505f0e35ee 100644 --- a/tests/entrypoints/generate/structured_decisions/test_strategies.py +++ b/tests/entrypoints/generate/structured_decisions/test_strategies.py @@ -33,6 +33,7 @@ from vllm.entrypoints.generate.structured_decisions.serving import ( parse_questions, ) from vllm.entrypoints.generate.structured_decisions.strategies import ( + ClefFlashStrategy, DiffusionGemmaCanvasStrategy, JointSchemaStrategy, LiquidStrategy, @@ -94,6 +95,8 @@ def test_strategy_selection(): # Clef-Omni scores options with its head, so it reads no logprobs. clef = model("ClefOmniForJointSchema", "raw_logits") assert select_read_strategy(clef, ("token_classify",)) is JointSchemaStrategy + clef_flash = model("ClefFlashForJointSchema", "raw_logits") + assert select_read_strategy(clef_flash, ("token_classify",)) is ClefFlashStrategy def test_route_is_registered_by_default(): @@ -366,8 +369,8 @@ JOINT_SCHEMA_REQUEST: dict[str, Any] = dict( ) -@pytest.fixture -def joint_schema(qwen, engine_and_renderer): +@pytest.fixture(params=[JointSchemaStrategy, ClefFlashStrategy]) +def joint_schema(request, qwen, engine_and_renderer): tokenizer, _ = qwen engine, online_renderer = engine_and_renderer # The probability the head gives each option, by option id. @@ -382,9 +385,7 @@ def joint_schema(qwen, engine_and_renderer): yield PoolingRequestOutput(request_id, PoolingOutput(data), ids, 0, True) engine.encode.side_effect = encode - strategy = JointSchemaStrategy( - ReadContext(engine, online_renderer, None, "auto", {}) - ) + strategy = request.param(ReadContext(engine, online_renderer, None, "auto", {})) return serving(engine, strategy), engine, tokenizer @@ -468,7 +469,10 @@ async def test_media_needs_a_model_that_takes_it(serving_and_engine): @pytest.mark.asyncio -async def test_joint_schema_reads_media_before_the_state(joint_schema, monkeypatch): +@pytest.mark.parametrize("num_images", [1, 2]) +async def test_joint_schema_reads_media_before_the_state( + joint_schema, monkeypatch, num_images +): serving, engine, tokenizer = joint_schema strategy = serving.strategy strategy.__dict__["_connector"] = MediaConnector() @@ -491,23 +495,30 @@ async def test_joint_schema_reads_media_before_the_state(joint_schema, monkeypat Image.new("RGB", (8, 8), "red").save(buffer, format="PNG") image = "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() response = await serving.create_decision( - StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image]) + StructuredDecisionRequest(**JOINT_SCHEMA_REQUEST, images=[image] * num_images) ) assert isinstance(response, StructuredDecisionResponse) - placeholder = tokenizer.encode( - "<|vision_start|><|image_pad|><|vision_end|>\n", add_special_tokens=False - ) + if isinstance(strategy, ClefFlashStrategy): + media_text = "<|vision_start|><|image_pad|><|vision_end|>" * num_images + "\n" + else: + media_text = "<|vision_start|><|image_pad|><|vision_end|>\n" * num_images + placeholder = tokenizer.encode(media_text, add_special_tokens=False) assert rendered["prompt_token_ids"] == strategy.prefix + placeholder assert rendered["multi_modal_data"]["image"][0].size == (8, 8) - assert rendered["mm_processor_kwargs"] == { - "images_kwargs": {"size": strategy.IMAGE_SIZE} - } + if isinstance(strategy, ClefFlashStrategy): + assert rendered["mm_processor_kwargs"] == { + "videos_kwargs": {"do_sample_frames": False} + } + else: + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE} + } prompt = engine.encode.call_args.args[0] assert prompt["type"] == "multimodal" ids = prompt["prompt_token_ids"] - expanded = len(strategy.prefix) + len(placeholder) + 2 - assert ids[:expanded].count(image_pad) == 3 + expanded = len(strategy.prefix) + len(placeholder) + 2 * num_images + assert ids[:expanded].count(image_pad) == 3 * num_images state_and_questions = JOINT_SCHEMA_PROMPT.split("STATE:\n", 1)[1] assert tokenizer.decode(ids[expanded:]) == state_and_questions assert response.answers["team"]["choice"] == "billing" @@ -515,7 +526,9 @@ async def test_joint_schema_reads_media_before_the_state(joint_schema, monkeypat @pytest.mark.asyncio -async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): +async def test_joint_schema_video_audio_follows_model_modalities( + joint_schema, monkeypatch +): serving, engine, tokenizer = joint_schema strategy = serving.strategy video = MediaWithBytes( @@ -532,6 +545,7 @@ async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): return video async def fetch_audio_async(self, url): + assert not isinstance(strategy, ClefFlashStrategy) if url == "https://example.org/silent.mp4": raise ValueError("No audio stream found.") return track @@ -552,11 +566,17 @@ async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): ) ) assert isinstance(response, StructuredDecisionResponse) - assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} - assert rendered["mm_processor_kwargs"] == { - "images_kwargs": {"size": strategy.IMAGE_SIZE}, - "use_audio_in_video": True, - } + if isinstance(strategy, ClefFlashStrategy): + assert rendered["multi_modal_data"] == {"video": [video]} + assert rendered["mm_processor_kwargs"] == { + "videos_kwargs": {"do_sample_frames": False} + } + else: + assert rendered["multi_modal_data"] == {"video": [video], "audio": [track]} + assert rendered["mm_processor_kwargs"] == { + "images_kwargs": {"size": strategy.IMAGE_SIZE}, + "use_audio_in_video": True, + } # A video without a decodable soundtrack is watched silently. response = await serving.create_decision( @@ -566,7 +586,8 @@ async def test_joint_schema_hears_video_soundtracks(joint_schema, monkeypatch): ) assert isinstance(response, StructuredDecisionResponse) assert rendered["multi_modal_data"] == {"video": [video]} - assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False + if not isinstance(strategy, ClefFlashStrategy): + assert rendered["mm_processor_kwargs"]["use_audio_in_video"] is False # Audio clips cannot be read with videos that have soundtracks. response = await serving.create_decision( diff --git a/tests/models/multimodal/pooling/test_clef_omni.py b/tests/models/multimodal/pooling/test_clef_omni.py index ccdc913b31..381ac23d53 100644 --- a/tests/models/multimodal/pooling/test_clef_omni.py +++ b/tests/models/multimodal/pooling/test_clef_omni.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""The Clef-Omni pooler against a port of the checkpoint's joint schema head, +"""The shared Clef pooler against a port of the checkpoint's joint schema head, which loops over questions where the pooler batches them.""" import math @@ -11,7 +11,10 @@ import pytest import torch import torch.nn.functional as F -from vllm.model_executor.models.clef_omni import JointSchemaHead, JointSchemaPooler +from vllm.model_executor.models.clef_joint_schema import ( + JointSchemaHead, + JointSchemaPooler, +) from vllm.pooling_params import PoolingParams from vllm.v1.pool.metadata import PoolingMetadata, PoolingStates diff --git a/tests/models/registry.py b/tests/models/registry.py index 0a7027461e..b08efd9fb0 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -802,6 +802,10 @@ _MULTIMODAL_EXAMPLE_MODELS = { "Salesforce/blip2-opt-2.7b", extras={"6b": "Salesforce/blip2-opt-6.7b"}, ), + "ClefFlashForJointSchema": _HfExamplesInfo( + "Cloudflare/clef-flash", + max_model_len=4096, + ), "Cohere2VisionForConditionalGeneration": _HfExamplesInfo( "CohereLabs/command-a-vision-07-2025", extras={"command-a-plus": "CohereLabs/command-a-plus-05-2026-bf16"}, diff --git a/vllm/entrypoints/generate/structured_decisions/strategies.py b/vllm/entrypoints/generate/structured_decisions/strategies.py index 5567cb808a..66238cf7f1 100644 --- a/vllm/entrypoints/generate/structured_decisions/strategies.py +++ b/vllm/entrypoints/generate/structured_decisions/strategies.py @@ -809,7 +809,63 @@ class JointSchemaStrategy(ReadStrategy): return reads +class ClefFlashStrategy(JointSchemaStrategy): + """Clef-Flash's joint schema prompt with Qwen3.5 image/video processing.""" + + modalities = frozenset({"image", "video"}) + MAX_LENGTH = 16384 + PLACEHOLDERS = { + "image": "<|vision_start|><|image_pad|><|vision_end|>", + "video": "<|vision_start|><|video_pad|><|vision_end|>", + } + + async def _render_media( + self, media: Mapping[str, Sequence[str]], cache_salt: str | None + ) -> MultiModalInput: + connector = self._connector + video_io = VideoMediaIO(ImageMediaIO(), num_frames=-1, fps=self.VIDEO_FPS) + try: + images, videos = await asyncio.gather( + asyncio.gather( + *map(connector.fetch_image_async, media.get("image", ())) + ), + self._fetch( + media.get("video", ()), video_io, envs.VLLM_VIDEO_FETCH_TIMEOUT + ), + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + text = ( + self.PLACEHOLDERS["image"] * len(images) + + self.PLACEHOLDERS["video"] * len(videos) + + "\n" + ) + mm_data: dict[str, Any] = {} + if images: + mm_data["image"] = images + if videos: + mm_data["video"] = videos + prompt = TokensPrompt( + prompt_token_ids=self.prefix + self._tokens(text), + multi_modal_data=mm_data, + mm_processor_kwargs={"videos_kwargs": {"do_sample_frames": False}}, + ) + if cache_salt is not None: + prompt["cache_salt"] = cache_salt + try: + engine_input = ( + await self.context.online_renderer.renderer.process_for_engine_async( + prompt, time.time() + ) + ) + except ValueError as e: + raise StructuredDecisionError(f"media: {e}") from e + assert engine_input["type"] == "multimodal" + return engine_input + + READ_STRATEGIES: dict[str, type[ReadStrategy]] = { + "ClefFlashForJointSchema": ClefFlashStrategy, "ClefOmniForJointSchema": JointSchemaStrategy, "Lfm2VlForConditionalGeneration": LiquidStrategy, "Qwen3ForCausalLM": NextTokenStrategy, diff --git a/vllm/model_executor/models/clef_flash.py b/vllm/model_executor/models/clef_flash.py new file mode 100644 index 0000000000..32a9e0a590 --- /dev/null +++ b/vllm/model_executor/models/clef_flash.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Clef-Flash: the Qwen3.5 backbone with a joint schema decision head.""" + +import torch + +from vllm.config import VllmConfig +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.default_loader import DefaultModelLoader +from vllm.multimodal import MULTIMODAL_REGISTRY + +from .clef_joint_schema import JOINT_HEAD_WEIGHTS, JointSchemaHead, JointSchemaPooler +from .interfaces_base import default_pooling_type +from .qwen3_5 import Qwen3_5ForConditionalGeneration, Qwen3_5ProcessingInfo +from .qwen3_vl import Qwen3VLDummyInputsBuilder, Qwen3VLMultiModalProcessor + + +@default_pooling_type(tok_pooling_type="ALL") +@MULTIMODAL_REGISTRY.register_processor( + Qwen3VLMultiModalProcessor, + info=Qwen3_5ProcessingInfo, + dummy_inputs=Qwen3VLDummyInputsBuilder, +) +class ClefFlashForJointSchema(Qwen3_5ForConditionalGeneration): + is_pooling_model = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + model_config = vllm_config.model_config + self.pooler = JointSchemaPooler( + JointSchemaHead(**model_config.hf_config.joint_head_config), + self._embed_lexical, + ) + self.secondary_weights = [ + DefaultModelLoader.Source( + model_or_path=model_config.model, + revision=model_config.revision, + prefix="pooler.head.", + allow_patterns_overrides=[JOINT_HEAD_WEIGHTS], + ) + ] + + def _embed_lexical(self, token_ids: torch.Tensor) -> torch.Tensor: + return VocabParallelEmbedding.forward(self.language_model.lm_head, token_ids) diff --git a/vllm/model_executor/models/clef_joint_schema.py b/vllm/model_executor/models/clef_joint_schema.py new file mode 100644 index 0000000000..fad2beb098 --- /dev/null +++ b/vllm/model_executor/models/clef_joint_schema.py @@ -0,0 +1,343 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Shared joint schema head for Clef decision models. + +The head reads the final hidden states of the whole prompt, so it runs as a +``token_classify`` pooler. + +Each request carries its schema in ``PoolingParams.extra_kwargs`` under +``"joint_schema"``: + +- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. +- ``question_spans``: the ``[start, end)`` of each question's instructions. +- ``option_spans``: the ``[start, end)`` of each option, question by question. +- ``option_counts``: the number of options of each question. + +Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), +so they hold after multimodal placeholders are expanded. The output has one +value per option, in ``option_spans`` order: its probability within its +question, or its logit with ``use_activation=False``. +""" + +import math +from collections.abc import Callable, Set +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate +from vllm.model_executor.layers.pooler.tokwise.methods import AllPool +from vllm.tasks import PoolingTask +from vllm.utils.torch_utils import async_tensor_h2d +from vllm.v1.outputs import PoolerOutput +from vllm.v1.pool.metadata import PoolingMetadata + +JOINT_HEAD_WEIGHTS = "joint_head.safetensors" +NUM_QUESTION_TYPES = 3 +MAX_LOGIT_SCALE = math.log(100.0) + + +def segment_mean( + values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor +) -> torch.Tensor: + sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) + sums.index_add_(0, segment_ids, values.float()) + return (sums / lengths[:, None]).to(values.dtype) + + +@dataclass +class SchemaIndices: + """Where one prompt's questions and options are, on the model's device.""" + + positions: torch.Tensor + """Every token of every span, questions first.""" + span_ids: torch.Tensor + """The span of each of ``positions``.""" + span_lengths: torch.Tensor + """``[Q + O]``.""" + question_types: torch.Tensor + """``[Q]``.""" + option_question: torch.Tensor + """The question of each option, ``[O]``.""" + option_slots: torch.Tensor + """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" + option_mask: torch.Tensor + """``[Q, max_options]``, true at the slots of options.""" + + +class EvidenceRoutingLayer(nn.Module): + def __init__(self, width: int, heads: int, feedforward: int) -> None: + super().__init__() + self.query_norm = nn.LayerNorm(width) + self.memory_norm = nn.LayerNorm(width) + self.attention = nn.MultiheadAttention(width, heads, batch_first=True) + self.feedforward_norm = nn.LayerNorm(width) + # The Identity takes the place of the checkpoint's dropout. + self.feedforward = nn.Sequential( + nn.Linear(width, feedforward), + nn.GELU(), + nn.Identity(), + nn.Linear(feedforward, width), + ) + + def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: + # Separate key and value tensors, like the reference, keep attention's + # unpacked input projections. + routed, _ = self.attention( + self.query_norm(queries), + self.memory_norm(memory), + self.memory_norm(memory), + need_weights=False, + ) + queries = queries + routed + return queries + self.feedforward(self.feedforward_norm(queries)) + + +class JointSchemaHead(nn.Module): + def __init__( + self, + hidden_size: int, + width: int, + routing_layers: int, + layers: int, + heads: int, + feedforward: int, + ) -> None: + super().__init__() + self.hidden_norm = nn.LayerNorm(hidden_size) + self.memory_projection = nn.Linear(hidden_size, width, bias=False) + self.question_projection = nn.Linear(hidden_size, width, bias=False) + self.option_question_projection = nn.Linear(hidden_size, width, bias=False) + self.global_projection = nn.Linear(hidden_size, width, bias=False) + self.option_context_projection = nn.Linear(hidden_size, width, bias=False) + self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) + self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) + self.evidence_layers = nn.ModuleList( + EvidenceRoutingLayer(width, heads, feedforward) + for _ in range(routing_layers) + ) + self.option_summary_norm = nn.LayerNorm(width) + self.layers = nn.ModuleList( + nn.TransformerDecoderLayer( + d_model=width, + nhead=heads, + dim_feedforward=feedforward, + dropout=0.0, + activation="gelu", + batch_first=True, + norm_first=True, + ) + for _ in range(layers) + ) + self.field_norm = nn.LayerNorm(width) + self.option_norm = nn.LayerNorm(width) + self.residual_scorer = nn.Sequential( + nn.Linear(width * 4, width), + nn.GELU(), + nn.Identity(), + nn.Linear(width, 1), + ) + self.prior_logit_scale = nn.Parameter(torch.zeros(())) + self.joint_logit_scale = nn.Parameter(torch.zeros(())) + self.residual_gate = nn.Parameter(torch.zeros(())) + + def forward( + self, + hidden_states: torch.Tensor, + lexical: torch.Tensor, + schema: SchemaIndices, + ) -> torch.Tensor: + """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` + final hidden states and each option's ``[O, hidden]`` mean output + embedding.""" + num_questions = schema.question_types.shape[0] + option_question = schema.option_question + hidden_states = self.hidden_norm(hidden_states) + memory = self.memory_projection(hidden_states).unsqueeze(0) + global_vector = hidden_states[-1] + means = segment_mean( + hidden_states[schema.positions], schema.span_ids, schema.span_lengths + ) + question_vectors, contexts = means[:num_questions], means[num_questions:] + + options = ( + self.option_context_projection(contexts) + + self.option_lexical_projection(lexical) + + self.option_question_projection(question_vectors)[option_question] + ).unsqueeze(0) + for layer in self.evidence_layers: + options = layer(options, memory) + options = options[0] + + # Each field summarizes its options, weighted by their match to it. + base_fields = self.question_projection(question_vectors) + width = options.shape[-1] + grid = options.new_zeros(schema.option_mask.numel(), width) + grid[schema.option_slots] = options + grid = grid.view(*schema.option_mask.shape, width) + routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) + routing = routing.masked_fill(~schema.option_mask, float("-inf")) + summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) + fields = ( + base_fields + + self.option_summary_norm(summaries) + + self.global_projection(global_vector).unsqueeze(0) + + self.type_embedding(schema.question_types) + ).unsqueeze(0) + for layer in self.layers: + fields = layer(fields, memory) + fields = self.field_norm(fields[0])[option_question] + + anchors = F.normalize(question_vectors + global_vector, dim=-1) + prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + prior = prior_scale * ( + F.normalize(lexical, dim=-1) * anchors[option_question] + ).sum(dim=-1) + options = self.option_norm(options) + cosine = F.cosine_similarity(fields, options, dim=-1) + features = torch.cat( + [fields, options, fields * options, torch.abs(fields - options)], dim=-1 + ) + residual = self.residual_scorer(features).squeeze(-1) + joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() + joint = joint_scale * cosine + residual + return prior + torch.sigmoid(self.residual_gate) * joint + + +def parse_joint_schema( + extra_kwargs: dict[str, Any] | None, prompt_len: int +) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: + """The question types, the absolute spans of the questions then of the + options, and the option counts. None if the schema is missing or invalid.""" + schema = (extra_kwargs or {}).get("joint_schema") + if not isinstance(schema, dict): + return None + try: + types = [int(t) for t in schema["question_types"]] + counts = [int(c) for c in schema["option_counts"]] + question_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["question_spans"] + ] + option_spans = [ + (int(start) + prompt_len, int(end) + prompt_len) + for start, end in schema["option_spans"] + ] + except (KeyError, TypeError, ValueError): + return None + spans = question_spans + option_spans + if ( + not types + or len(counts) != len(types) + or len(question_spans) != len(types) + or len(option_spans) != sum(counts) + or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) + or any(c < 1 for c in counts) + or any(not 0 <= start < end <= prompt_len for start, end in spans) + ): + return None + return types, spans, counts + + +class JointSchemaPooler(Pooler): + """Runs the joint schema head on each finished prompt. + + It never raises, because it runs in the engine core: a prompt without a + valid schema, such as a profiling dummy, gets an empty output. + """ + + def __init__( + self, + head: JointSchemaHead, + embed_lexical: Callable[[torch.Tensor], torch.Tensor], + ) -> None: + super().__init__() + self.pooling = AllPool() + self.head = head + self.embed_lexical = embed_lexical + + def get_supported_tasks(self) -> Set[PoolingTask]: + return {"token_classify"} + + def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: + return PoolingParamsUpdate(requires_token_ids=True) + + def forward( + self, + hidden_states: torch.Tensor, + pooling_metadata: PoolingMetadata, + ) -> PoolerOutput: + pooled = self.pooling(hidden_states, pooling_metadata) + prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() + outputs: list[torch.Tensor | None] = [] + for hidden, token_ids, params in zip( + pooled, prompt_token_ids, pooling_metadata.pooling_params + ): + if hidden is None: + outputs.append(None) + continue + parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) + if parsed is None: + outputs.append(hidden.new_empty(0, dtype=torch.float32)) + continue + schema, lexical_ids, lexical_span_ids = self._indices( + *parsed, hidden.device + ) + lexical = segment_mean( + self.embed_lexical( + async_tensor_h2d(token_ids[lexical_ids], hidden.device) + ), + lexical_span_ids, + schema.span_lengths[len(parsed[0]) :], + ) + logits = self.head(hidden, lexical, schema).float() + if params.use_activation is not False: + mask = schema.option_mask + grid = logits.new_full((mask.numel(),), float("-inf")) + grid[schema.option_slots] = logits + probs = grid.view(mask.shape).softmax(dim=-1) + logits = probs.view(-1)[schema.option_slots] + outputs.append(logits) + return outputs + + @staticmethod + def _indices( + types: list[int], + spans: list[tuple[int, int]], + counts: list[int], + device: torch.device, + ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: + """The schema's indices, plus the option tokens' positions (on the CPU) + and their option spans, for the lexical embeddings.""" + num_questions = len(types) + lengths = torch.tensor([end - start for start, end in spans]) + span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) + positions = torch.cat([torch.arange(start, end) for start, end in spans]) + num_question_tokens = int(lengths[:num_questions].sum()) + + counts_cpu = torch.tensor(counts) + max_options = max(counts) + option_question = torch.repeat_interleave( + torch.arange(num_questions), counts_cpu + ) + option_slots = torch.cat( + [torch.arange(count) + q * max_options for q, count in enumerate(counts)] + ) + schema = SchemaIndices( + positions=async_tensor_h2d(positions, device), + span_ids=async_tensor_h2d(span_ids, device), + span_lengths=async_tensor_h2d(lengths, device), + question_types=async_tensor_h2d(types, device, dtype=torch.long), + option_question=async_tensor_h2d(option_question, device), + option_slots=async_tensor_h2d(option_slots, device), + option_mask=async_tensor_h2d( + torch.arange(max_options) < counts_cpu[:, None], device + ), + ) + lexical_span_ids = async_tensor_h2d( + span_ids[num_question_tokens:] - num_questions, device + ) + return schema, positions[num_question_tokens:], lexical_span_ids diff --git a/vllm/model_executor/models/clef_omni.py b/vllm/model_executor/models/clef_omni.py index 7cd0d67687..1533dde273 100644 --- a/vllm/model_executor/models/clef_omni.py +++ b/vllm/model_executor/models/clef_omni.py @@ -1,47 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Clef-Omni (https://huggingface.co/Cloudflare/clef-omni). - -The Qwen3-Omni thinker with a joint schema head that scores every option of -every question in one forward pass. The head reads the final hidden states of -the whole prompt, so it runs as a ``token_classify`` pooler. - -Each request carries its schema in ``PoolingParams.extra_kwargs`` under -``"joint_schema"``: - -- ``question_types``: 0 (noul), 1 (choice) or 2 (score) per question. -- ``question_spans``: the ``[start, end)`` of each question's instructions. -- ``option_spans``: the ``[start, end)`` of each option, question by question. -- ``option_counts``: the number of options of each question. - -Spans count from the end of the prompt (``-len(prompt) <= start < end <= 0``), -so they hold after multimodal placeholders are expanded. The output has one -value per option, in ``option_spans`` order: its probability within its -question, or its logit with ``use_activation=False``. -""" +"""The Qwen3-Omni thinker with a Clef joint schema head.""" import math -from collections.abc import Callable, Mapping, Set -from dataclasses import dataclass +from collections.abc import Mapping from typing import Any import numpy as np import torch import torch.nn.functional as F -from torch import nn from transformers.feature_extraction_utils import BatchFeature from vllm.config import VllmConfig -from vllm.model_executor.layers.pooler import Pooler, PoolingParamsUpdate -from vllm.model_executor.layers.pooler.tokwise.methods import AllPool from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding from vllm.model_executor.model_loader.default_loader import DefaultModelLoader from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.tasks import PoolingTask -from vllm.utils.torch_utils import async_tensor_h2d -from vllm.v1.outputs import PoolerOutput -from vllm.v1.pool.metadata import PoolingMetadata +from .clef_joint_schema import JOINT_HEAD_WEIGHTS, JointSchemaHead, JointSchemaPooler from .interfaces_base import default_pooling_type from .qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerDummyInputsBuilder, @@ -50,314 +25,6 @@ from .qwen3_omni_moe_thinker import ( Qwen3OmniMoeThinkerProcessingInfo, ) -JOINT_HEAD_WEIGHTS = "joint_head.safetensors" -NUM_QUESTION_TYPES = 3 -MAX_LOGIT_SCALE = math.log(100.0) - - -def segment_mean( - values: torch.Tensor, segment_ids: torch.Tensor, lengths: torch.Tensor -) -> torch.Tensor: - sums = values.new_zeros(lengths.shape[0], values.shape[-1], dtype=torch.float32) - sums.index_add_(0, segment_ids, values.float()) - return (sums / lengths[:, None]).to(values.dtype) - - -@dataclass -class SchemaIndices: - """Where one prompt's questions and options are, on the model's device.""" - - positions: torch.Tensor - """Every token of every span, questions first.""" - span_ids: torch.Tensor - """The span of each of ``positions``.""" - span_lengths: torch.Tensor - """``[Q + O]``.""" - question_types: torch.Tensor - """``[Q]``.""" - option_question: torch.Tensor - """The question of each option, ``[O]``.""" - option_slots: torch.Tensor - """Each option's flat index in a ``[Q, max_options]`` grid, ``[O]``.""" - option_mask: torch.Tensor - """``[Q, max_options]``, true at the slots of options.""" - - -class EvidenceRoutingLayer(nn.Module): - def __init__(self, width: int, heads: int, feedforward: int) -> None: - super().__init__() - self.query_norm = nn.LayerNorm(width) - self.memory_norm = nn.LayerNorm(width) - self.attention = nn.MultiheadAttention(width, heads, batch_first=True) - self.feedforward_norm = nn.LayerNorm(width) - # The Identity takes the place of the checkpoint's dropout. - self.feedforward = nn.Sequential( - nn.Linear(width, feedforward), - nn.GELU(), - nn.Identity(), - nn.Linear(feedforward, width), - ) - - def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor: - # Separate key and value tensors, like the reference, keep attention's - # unpacked input projections. - routed, _ = self.attention( - self.query_norm(queries), - self.memory_norm(memory), - self.memory_norm(memory), - need_weights=False, - ) - queries = queries + routed - return queries + self.feedforward(self.feedforward_norm(queries)) - - -class JointSchemaHead(nn.Module): - def __init__( - self, - hidden_size: int, - width: int, - routing_layers: int, - layers: int, - heads: int, - feedforward: int, - ) -> None: - super().__init__() - self.hidden_norm = nn.LayerNorm(hidden_size) - self.memory_projection = nn.Linear(hidden_size, width, bias=False) - self.question_projection = nn.Linear(hidden_size, width, bias=False) - self.option_question_projection = nn.Linear(hidden_size, width, bias=False) - self.global_projection = nn.Linear(hidden_size, width, bias=False) - self.option_context_projection = nn.Linear(hidden_size, width, bias=False) - self.option_lexical_projection = nn.Linear(hidden_size, width, bias=False) - self.type_embedding = nn.Embedding(NUM_QUESTION_TYPES, width) - self.evidence_layers = nn.ModuleList( - EvidenceRoutingLayer(width, heads, feedforward) - for _ in range(routing_layers) - ) - self.option_summary_norm = nn.LayerNorm(width) - self.layers = nn.ModuleList( - nn.TransformerDecoderLayer( - d_model=width, - nhead=heads, - dim_feedforward=feedforward, - dropout=0.0, - activation="gelu", - batch_first=True, - norm_first=True, - ) - for _ in range(layers) - ) - self.field_norm = nn.LayerNorm(width) - self.option_norm = nn.LayerNorm(width) - self.residual_scorer = nn.Sequential( - nn.Linear(width * 4, width), - nn.GELU(), - nn.Identity(), - nn.Linear(width, 1), - ) - self.prior_logit_scale = nn.Parameter(torch.zeros(())) - self.joint_logit_scale = nn.Parameter(torch.zeros(())) - self.residual_gate = nn.Parameter(torch.zeros(())) - - def forward( - self, - hidden_states: torch.Tensor, - lexical: torch.Tensor, - schema: SchemaIndices, - ) -> torch.Tensor: - """The ``[O]`` option logits of one prompt, from its ``[L, hidden]`` - final hidden states and each option's ``[O, hidden]`` mean output - embedding.""" - num_questions = schema.question_types.shape[0] - option_question = schema.option_question - hidden_states = self.hidden_norm(hidden_states) - memory = self.memory_projection(hidden_states).unsqueeze(0) - global_vector = hidden_states[-1] - means = segment_mean( - hidden_states[schema.positions], schema.span_ids, schema.span_lengths - ) - question_vectors, contexts = means[:num_questions], means[num_questions:] - - options = ( - self.option_context_projection(contexts) - + self.option_lexical_projection(lexical) - + self.option_question_projection(question_vectors)[option_question] - ).unsqueeze(0) - for layer in self.evidence_layers: - options = layer(options, memory) - options = options[0] - - # Each field summarizes its options, weighted by their match to it. - base_fields = self.question_projection(question_vectors) - width = options.shape[-1] - grid = options.new_zeros(schema.option_mask.numel(), width) - grid[schema.option_slots] = options - grid = grid.view(*schema.option_mask.shape, width) - routing = torch.einsum("qow,qw->qo", grid, base_fields) / math.sqrt(width) - routing = routing.masked_fill(~schema.option_mask, float("-inf")) - summaries = torch.einsum("qo,qow->qw", routing.softmax(dim=-1), grid) - fields = ( - base_fields - + self.option_summary_norm(summaries) - + self.global_projection(global_vector).unsqueeze(0) - + self.type_embedding(schema.question_types) - ).unsqueeze(0) - for layer in self.layers: - fields = layer(fields, memory) - fields = self.field_norm(fields[0])[option_question] - - anchors = F.normalize(question_vectors + global_vector, dim=-1) - prior_scale = self.prior_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() - prior = prior_scale * ( - F.normalize(lexical, dim=-1) * anchors[option_question] - ).sum(dim=-1) - options = self.option_norm(options) - cosine = F.cosine_similarity(fields, options, dim=-1) - features = torch.cat( - [fields, options, fields * options, torch.abs(fields - options)], dim=-1 - ) - residual = self.residual_scorer(features).squeeze(-1) - joint_scale = self.joint_logit_scale.clamp(max=MAX_LOGIT_SCALE).exp() - joint = joint_scale * cosine + residual - return prior + torch.sigmoid(self.residual_gate) * joint - - -def parse_joint_schema( - extra_kwargs: dict[str, Any] | None, prompt_len: int -) -> tuple[list[int], list[tuple[int, int]], list[int]] | None: - """The question types, the absolute spans of the questions then of the - options, and the option counts. None if the schema is missing or invalid.""" - schema = (extra_kwargs or {}).get("joint_schema") - if not isinstance(schema, dict): - return None - try: - types = [int(t) for t in schema["question_types"]] - counts = [int(c) for c in schema["option_counts"]] - question_spans = [ - (int(start) + prompt_len, int(end) + prompt_len) - for start, end in schema["question_spans"] - ] - option_spans = [ - (int(start) + prompt_len, int(end) + prompt_len) - for start, end in schema["option_spans"] - ] - except (KeyError, TypeError, ValueError): - return None - spans = question_spans + option_spans - if ( - not types - or len(counts) != len(types) - or len(question_spans) != len(types) - or len(option_spans) != sum(counts) - or any(not 0 <= t < NUM_QUESTION_TYPES for t in types) - or any(c < 1 for c in counts) - or any(not 0 <= start < end <= prompt_len for start, end in spans) - ): - return None - return types, spans, counts - - -class JointSchemaPooler(Pooler): - """Runs the joint schema head on each finished prompt. - - It never raises, because it runs in the engine core: a prompt without a - valid schema, such as a profiling dummy, gets an empty output. - """ - - def __init__( - self, - head: JointSchemaHead, - embed_lexical: Callable[[torch.Tensor], torch.Tensor], - ) -> None: - super().__init__() - self.pooling = AllPool() - self.head = head - self.embed_lexical = embed_lexical - - def get_supported_tasks(self) -> Set[PoolingTask]: - return {"token_classify"} - - def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate: - return PoolingParamsUpdate(requires_token_ids=True) - - def forward( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput: - pooled = self.pooling(hidden_states, pooling_metadata) - prompt_token_ids = pooling_metadata.get_prompt_token_ids_cpu() - outputs: list[torch.Tensor | None] = [] - for hidden, token_ids, params in zip( - pooled, prompt_token_ids, pooling_metadata.pooling_params - ): - if hidden is None: - outputs.append(None) - continue - parsed = parse_joint_schema(params.extra_kwargs, hidden.shape[0]) - if parsed is None: - outputs.append(hidden.new_empty(0, dtype=torch.float32)) - continue - schema, lexical_ids, lexical_span_ids = self._indices( - *parsed, hidden.device - ) - lexical = segment_mean( - self.embed_lexical( - async_tensor_h2d(token_ids[lexical_ids], hidden.device) - ), - lexical_span_ids, - schema.span_lengths[len(parsed[0]) :], - ) - logits = self.head(hidden, lexical, schema).float() - if params.use_activation is not False: - mask = schema.option_mask - grid = logits.new_full((mask.numel(),), float("-inf")) - grid[schema.option_slots] = logits - probs = grid.view(mask.shape).softmax(dim=-1) - logits = probs.view(-1)[schema.option_slots] - outputs.append(logits) - return outputs - - @staticmethod - def _indices( - types: list[int], - spans: list[tuple[int, int]], - counts: list[int], - device: torch.device, - ) -> tuple[SchemaIndices, torch.Tensor, torch.Tensor]: - """The schema's indices, plus the option tokens' positions (on the CPU) - and their option spans, for the lexical embeddings.""" - num_questions = len(types) - lengths = torch.tensor([end - start for start, end in spans]) - span_ids = torch.repeat_interleave(torch.arange(len(spans)), lengths) - positions = torch.cat([torch.arange(start, end) for start, end in spans]) - num_question_tokens = int(lengths[:num_questions].sum()) - - counts_cpu = torch.tensor(counts) - max_options = max(counts) - option_question = torch.repeat_interleave( - torch.arange(num_questions), counts_cpu - ) - option_slots = torch.cat( - [torch.arange(count) + q * max_options for q, count in enumerate(counts)] - ) - schema = SchemaIndices( - positions=async_tensor_h2d(positions, device), - span_ids=async_tensor_h2d(span_ids, device), - span_lengths=async_tensor_h2d(lengths, device), - question_types=async_tensor_h2d(types, device, dtype=torch.long), - option_question=async_tensor_h2d(option_question, device), - option_slots=async_tensor_h2d(option_slots, device), - option_mask=async_tensor_h2d( - torch.arange(max_options) < counts_cpu[:, None], device - ), - ) - lexical_span_ids = async_tensor_h2d( - span_ids[num_question_tokens:] - num_questions, device - ) - return schema, positions[num_question_tokens:], lexical_span_ids - - MAX_FRAME_PIXELS = 256 * 32 * 32 diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 36f9bd6e6f..cd0557ac45 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -287,6 +287,7 @@ _REWARD_MODELS = { _TOKEN_CLASSIFICATION_MODELS = { "BertForTokenClassification": ("bert", "BertForTokenClassification"), + "ClefFlashForJointSchema": ("clef_flash", "ClefFlashForJointSchema"), "ClefOmniForJointSchema": ("clef_omni", "ClefOmniForJointSchema"), "ModernBertForTokenClassification": ( "modernbert", diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 147b65759d..06a9376971 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -647,6 +647,27 @@ class MimoV2ModelArchConfigConvertor(ModelArchConfigConvertorBase): _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config) +class Qwen3_5ModelArchConfigConvertor(ModelArchConfigConvertorBase): + def __init__( + self, + hf_config: PreTrainedConfig, + hf_text_config: PreTrainedConfig, + revision: str | None = None, + ): + if getattr(hf_config, "architectures", None) == [ + "Qwen3_5ForConditionalGeneration" + ]: + model_path = hf_config.name_or_path + if model_path and file_or_path_exists( + model_path, "joint_head_config.json", revision + ): + hf_config.architectures = ["ClefFlashForJointSchema"] + hf_config.joint_head_config = get_hf_file_to_dict( + "joint_head_config.json", model_path, revision + ) + super().__init__(hf_config, hf_text_config, revision) + + class Qwen3OmniMoeModelArchConfigConvertor(ModelArchConfigConvertorBase): def __init__( self, @@ -850,6 +871,7 @@ MODEL_ARCH_CONFIG_CONVERTORS = { "nemotron-nas": NemotronNasModelArchConfigConvertor, "bailing_hybrid_v3_mtp": BailingHybridV3MTPModelArchConfigConvertor, "pangu_ultra_moe_mtp": PanguUltraMoeMTPModelArchConfigConvertor, + "qwen3_5": Qwen3_5ModelArchConfigConvertor, "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor, "qwen3_omni_moe": Qwen3OmniMoeModelArchConfigConvertor, "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor,