iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Model] Add Clef-Flash support (#61308)

Browse Source 
Signed-off-by: Tianyu Guo <guoty@inferact.ai>
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Co-authored-by: Codex <noreply@openai.com>
Co-authored-by: Isotr0py <Isotr0py@outlook.com>
This commit is contained in:
Tianyu Guo
2026-10-13 05:46:37 +00:00
committed by GitHub
co-authored by Codex Isotr0py
parent 0d2689e0df
commit 380638afbb
11 changed files with 580 additions and 361 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/serving/online_serving/structured_decisions.md
+25 -1
View File
@@ -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.
tests/config/test_model_arch_config.py
+34
View File
@@ -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
):
tests/entrypoints/generate/structured_decisions/test_strategies.py
+43 -22
View File
@@ -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(
tests/models/multimodal/pooling/test_clef_omni.py
+5 -2
View File
@@ -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
tests/models/registry.py
+4
View File
@@ -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"},
vllm/entrypoints/generate/structured_decisions/strategies.py
+56
View File
@@ -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,
vllm/model_executor/models/clef_flash.py
+44
View File
@@ -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)
vllm/model_executor/models/clef_joint_schema.py
+343
View File
@@ -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
vllm/model_executor/models/clef_omni.py
+3 -336
View File
@@ -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
vllm/model_executor/models/registry.py
+1
View File
@@ -287,6 +287,7 @@ _REWARD_MODELS = {
_TOKEN_CLASSIFICATION_MODELS = {
"BertForTokenClassification": ("bert", "BertForTokenClassification"),
"ClefFlashForJointSchema": ("clef_flash", "ClefFlashForJointSchema"),
"ClefOmniForJointSchema": ("clef_omni", "ClefOmniForJointSchema"),
"ModernBertForTokenClassification": (
"modernbert",
vllm/transformers_utils/model_arch_config_convertor.py
+22
View File
@@ -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,