[Model] Add support for openai/privacy-filter (#41026)
Signed-off-by: Fabian Joswig <fjosw@users.noreply.github.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>
Fabian Joswig committed
Jun 29, 2026 at 01:18 UTC
4dfbf1503b4bae722743c483a0079ce2f0633f4c
8 files changed
+225
-5
docs/models/pooling_models/token_classify.md
+1
@@ -45,6 +45,7 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m
45
| ------------ | ------ | ----------------- | --------------------------- | --------------------------------------- |
46
| `BertForTokenClassification` | bert-based | `boltuix/NeuroBERT-NER` (see note), etc. | | |
47
| `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | |
48
+| `OpenAIPrivacyFilterForTokenClassification` | gpt-oss-based encoder | `openai/privacy-filter` | | |
49
| `Qwen3ForTokenClassification`<sup>C</sup> | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | |
50
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
51
tests/models/language/pooling/test_token_classification.py
+47
@@ -116,6 +116,53 @@ def test_modernbert_models(
116
torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
117
118
119
+PRIVACY_FILTER_PROMPTS = [
120
+ "My name is Harry Potter.",
121
+ "Email me at harry.potter@hogwarts.edu.",
122
+ "Call me on +44 20 7946 0958 tomorrow.",
123
+ "My account number is 12345678 and the API key is sk-live-abc123def456.",
124
+ "I live at 4 Privet Drive, Little Whinging.",
125
+ "Visit https://example.com/profile/harry for more info.",
126
+ "We met on 12 January 2024.",
127
+]
128
+
129
+
130
+@pytest.mark.parametrize("model", ["openai/privacy-filter"])
131
+@pytest.mark.parametrize("dtype", ["bfloat16"])
132
+@torch.inference_mode
133
+def test_openai_privacy_filter(
134
+ hf_runner,
135
+ vllm_runner,
136
+ model: str,
137
+ dtype: str,
138
+) -> None:
139
+ with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
140
+ vllm_outputs = vllm_model.token_classify(PRIVACY_FILTER_PROMPTS)
141
+
142
+ hf_model_kwargs = {}
143
+ if current_platform.is_rocm():
144
+ hf_model_kwargs["attn_implementation"] = "eager"
145
+
146
+ with hf_runner(
147
+ model,
148
+ dtype=dtype,
149
+ auto_cls=AutoModelForTokenClassification,
150
+ model_kwargs=hf_model_kwargs,
151
+ ) as hf_model:
152
+ tokenizer = hf_model.tokenizer
153
+ hf_outputs = []
154
+ for prompt in PRIVACY_FILTER_PROMPTS:
155
+ inputs = tokenizer([prompt], return_tensors="pt")
156
+ inputs = hf_model.wrap_device(inputs)
157
+ output = hf_model.model(**inputs)
158
+ hf_outputs.append(softmax(output.logits[0]))
159
+
160
+ for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
161
+ hf_output = hf_output.detach().clone().cpu().float()
162
+ vllm_output = vllm_output.detach().clone().cpu().float()
163
+ torch.testing.assert_close(hf_output, vllm_output, atol=0.1, rtol=1e-2)
164
+
165
+
166
@pytest.mark.parametrize("model", ["bd2lcco/Qwen3-0.6B-finetuned"])
167
@pytest.mark.parametrize("dtype", ["float"])
168
@torch.inference_mode
tests/models/registry.py
+4
@@ -707,6 +707,10 @@ _TOKEN_CLASSIFICATION_EXAMPLE_MODELS = {
707
"ModernBertForTokenClassification": _HfExamplesInfo(
708
"disham993/electrical-ner-ModernBERT-base"
709
),
710
+ "OpenAIPrivacyFilterForTokenClassification": _HfExamplesInfo(
711
+ "openai/privacy-filter",
712
+ min_transformers_version="5.6.0.dev0",
713
+ ),
714
}
715
716
_SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+9
@@ -11,6 +11,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11
from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.config.kernel import MoEBackend
13
from vllm.logger import init_logger
14
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
15
from vllm.model_executor.layers.fused_moe.all2all_utils import (
16
maybe_make_prepare_finalize,
17
)
@@ -83,6 +84,14 @@ def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBac
84
if moe_config.moe_parallel_config.dp_size > 1:
85
_move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)
86
87
+ # HACK: unquantized FlashInfer aliases SWIGLUOAI to plain Swiglu
88
+ # (swiglu_alpha/limit only set on the MXFP4 branch). Route to
89
+ # Triton's swigluoai_and_mul until that's plumbed through. Same
90
+ # demotion pattern as the Qwen3.5/dp_size hack above.
91
+ if moe_config.activation == MoEActivation.SWIGLUOAI:
92
+ _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_TRTLLM)
93
+ _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)
94
+
95
elif current_platform.is_xpu():
96
_AVAILABLE_BACKENDS = [UnquantizedMoeBackend.XPU]
97
elif current_platform.is_cpu():
vllm/model_executor/models/gpt_oss.py
+32
-5
@@ -70,6 +70,10 @@ from .utils import (
70
71
72
class OAIAttention(nn.Module):
73
+ # Override to switch RoPE convention. gpt-oss uses NeoX (chunk halves);
74
+ # privacy-filter and similar derivatives use GPT-J (interleaved pairs).
75
+ rope_is_neox_style: bool = True
76
+
77
def __init__(
78
self,
79
config: GptOssConfig,
@@ -99,7 +103,7 @@ class OAIAttention(nn.Module):
103
"beta_slow": config.rope_parameters["beta_slow"],
104
"truncate": config.rope_parameters.get("truncate", True),
105
},
102
- is_neox_style=True,
106
+ is_neox_style=self.rope_is_neox_style,
107
)
108
109
tp_size = get_tensor_model_parallel_world_size()
@@ -133,9 +137,25 @@ class OAIAttention(nn.Module):
137
self.num_local_attention_heads = config.num_attention_heads // tp_size
138
self.num_local_key_value_heads = config.num_key_value_heads // tp_size
139
140
+ self.attn = self._build_attention(
141
+ config=config,
142
+ cache_config=cache_config,
143
+ quant_config=quant_config,
144
+ prefix=prefix,
145
+ )
146
+
147
+ def _build_attention(
148
+ self,
149
+ config: GptOssConfig,
150
+ cache_config: CacheConfig | None,
151
+ quant_config: QuantizationConfig | None,
152
+ prefix: str,
153
+ ) -> Attention:
154
+ # Override to swap in an encoder-only attention or alter the
155
+ # per-layer sliding-window policy.
156
# Only apply sliding window to every other layer
157
sliding_window = config.sliding_window if self.layer_idx % 2 == 0 else None
138
- self.attn = Attention(
158
+ return Attention(
159
self.num_local_attention_heads,
160
self.head_dim,
161
self.scaling,
@@ -222,6 +242,10 @@ class MLPBlock(torch.nn.Module):
242
243
244
class TransformerBlock(torch.nn.Module):
245
+ # Override to swap attention/MLP without re-implementing the block.
246
+ attention_cls: type[nn.Module] = OAIAttention
247
+ mlp_cls: type[nn.Module] = MLPBlock
248
+
249
def __init__(
250
self,
251
vllm_config: VllmConfig,
@@ -234,13 +258,13 @@ class TransformerBlock(torch.nn.Module):
258
cache_config = vllm_config.cache_config
259
260
self.layer_idx = extract_layer_index(prefix)
237
- self.attn = OAIAttention(
261
+ self.attn = self.attention_cls(
262
config,
263
prefix=f"{prefix}.attn",
264
quant_config=quant_config,
265
cache_config=cache_config,
266
)
243
- self.mlp = MLPBlock(vllm_config, self.layer_idx, prefix=f"{prefix}.mlp")
267
+ self.mlp = self.mlp_cls(vllm_config, self.layer_idx, prefix=f"{prefix}.mlp")
268
self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
269
self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
270
@@ -266,6 +290,9 @@ class TransformerBlock(torch.nn.Module):
290
291
@support_torch_compile
292
class GptOssModel(nn.Module, EagleModelMixin):
293
+ # Override to swap in an alternative TransformerBlock subclass.
294
+ block_cls: type[nn.Module] = TransformerBlock
295
+
296
def __init__(
297
self,
298
*,
@@ -282,7 +309,7 @@ class GptOssModel(nn.Module, EagleModelMixin):
309
)
310
self.start_layer, self.end_layer, self.layers = make_layers(
311
self.config.num_hidden_layers,
285
- lambda prefix: TransformerBlock(
312
+ lambda prefix: self.block_cls(
313
vllm_config,
314
prefix=prefix,
315
quant_config=self.quant_config,
vllm/model_executor/models/openai_privacy_filter.py
new
+127
@@ -0,0 +1,127 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Inference-only OpenAI Privacy Filter model.
4
+
5
+gpt-oss reused as a bidirectional encoder for token classification: every
6
+layer runs non-causal attention with a banded ±sliding_window mask, and
7
+the LM head is replaced with a 33-class BIOES score head.
8
+"""
9
+
10
+from collections.abc import Iterable
11
+
12
+import torch
13
+from torch import nn
14
+
15
+from vllm.config import CacheConfig, VllmConfig
16
+from vllm.model_executor.layers.attention.encoder_only_attention import (
17
+ EncoderOnlyAttention,
18
+)
19
+from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
20
+from vllm.model_executor.layers.quantization import QuantizationConfig
21
+from vllm.sequence import IntermediateTensors
22
+
23
+from .gpt_oss import GptOssForCausalLM, GptOssModel, OAIAttention, TransformerBlock
24
+from .interfaces_base import attn_type, default_pooling_type
25
+from .utils import AutoWeightsLoader, maybe_prefix
26
+
27
+
28
+class OpenAIPrivacyFilterAttention(OAIAttention):
29
+ # Privacy-filter uses GPT-J style RoPE (interleaved pairs), not NeoX.
30
+ rope_is_neox_style = False
31
+
32
+ def _build_attention(
33
+ self,
34
+ config,
35
+ cache_config: CacheConfig | None,
36
+ quant_config: QuantizationConfig | None,
37
+ prefix: str,
38
+ ) -> EncoderOnlyAttention:
39
+ # HF stores sliding_window+1 so each token attends to ±W neighbors;
40
+ # the encoder-only path applies this as a symmetric (W-1, W-1) mask.
41
+ return EncoderOnlyAttention(
42
+ num_heads=self.num_local_attention_heads,
43
+ head_size=self.head_dim,
44
+ scale=self.scaling,
45
+ num_kv_heads=self.num_local_key_value_heads,
46
+ cache_config=cache_config,
47
+ quant_config=quant_config,
48
+ per_layer_sliding_window=config.sliding_window + 1,
49
+ prefix=f"{prefix}.attn",
50
+ sinks=self.sinks,
51
+ )
52
+
53
+
54
+class OpenAIPrivacyFilterDecoderLayer(TransformerBlock):
55
+ attention_cls = OpenAIPrivacyFilterAttention
56
+
57
+
58
+class OpenAIPrivacyFilterModel(GptOssModel):
59
+ block_cls = OpenAIPrivacyFilterDecoderLayer
60
+
61
+
62
+def _interleave_gate_up_concat_to_pairs(
63
+ weights: Iterable[tuple[str, torch.Tensor]],
64
+) -> Iterable[tuple[str, torch.Tensor]]:
65
+ # HF gate_up_proj is concat [gate | up]; swigluoai_and_mul wants
66
+ # gate/up interleaved. MXFP4/quark suffixes are already interleaved.
67
+ for name, weight in weights:
68
+ if name.endswith(".gate_up_proj") or name.endswith(".gate_up_proj_bias"):
69
+ *lead, two_i = weight.shape
70
+ i = two_i // 2
71
+ weight = (
72
+ torch.stack([weight[..., :i], weight[..., i:]], dim=-1)
73
+ .reshape(*lead, two_i)
74
+ .contiguous()
75
+ )
76
+ yield name, weight
77
+
78
+
79
+@attn_type("encoder_only")
80
+@default_pooling_type(tok_pooling_type="ALL")
81
+class OpenAIPrivacyFilterForTokenClassification(nn.Module):
82
+ is_pooling_model = True
83
+ hf_to_vllm_mapper = GptOssForCausalLM.hf_to_vllm_mapper
84
+
85
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
86
+ super().__init__()
87
+ config = vllm_config.model_config.hf_config
88
+ self.head_dtype = vllm_config.model_config.head_dtype
89
+ self.num_labels = config.num_labels
90
+
91
+ self.model = OpenAIPrivacyFilterModel(
92
+ vllm_config=vllm_config,
93
+ prefix=maybe_prefix(prefix, "model"),
94
+ )
95
+ self.score = nn.Linear(
96
+ config.hidden_size, config.num_labels, dtype=self.head_dtype
97
+ )
98
+
99
+ pooler_config = vllm_config.model_config.pooler_config
100
+ assert pooler_config is not None
101
+ self.pooler = pooler_for_token_classify(pooler_config)
102
+
103
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
104
+ return self.model.embed_input_ids(input_ids)
105
+
106
+ def forward(
107
+ self,
108
+ input_ids: torch.Tensor | None,
109
+ positions: torch.Tensor,
110
+ intermediate_tensors: IntermediateTensors | None = None,
111
+ inputs_embeds: torch.Tensor | None = None,
112
+ ) -> torch.Tensor:
113
+ hidden_states = self.model(
114
+ input_ids=input_ids,
115
+ positions=positions,
116
+ intermediate_tensors=intermediate_tensors,
117
+ inputs_embeds=inputs_embeds,
118
+ )
119
+ hidden_states = hidden_states.to(self.head_dtype)
120
+ return self.score(hidden_states)
121
+
122
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
123
+ loader = AutoWeightsLoader(self)
124
+ return loader.load_weights(
125
+ _interleave_gate_up_concat_to_pairs(weights),
126
+ mapper=self.hf_to_vllm_mapper,
127
+ )
vllm/model_executor/models/registry.py
+4
@@ -291,6 +291,10 @@ _TOKEN_CLASSIFICATION_MODELS = {
291
"modernbert",
292
"ModernBertForTokenClassification",
293
),
294
+ "OpenAIPrivacyFilterForTokenClassification": (
295
+ "openai_privacy_filter",
296
+ "OpenAIPrivacyFilterForTokenClassification",
297
+ ),
298
"Qwen3ASRForcedAlignerForTokenClassification": (
299
"qwen3_asr_forced_aligner",
300
"Qwen3ASRForcedAlignerForTokenClassification",
vllm/v1/attention/backends/flash_attn.py
+1
@@ -1164,6 +1164,7 @@ class FlashAttentionImpl(AttentionImpl):
1164
k_descale=layer._k_scale.expand(descale_shape), # type: ignore[operator]
1165
v_descale=layer._v_scale.expand(descale_shape), # type: ignore[operator]
1166
num_splits=1 if self.batch_invariant_enabled else 0,
1167
+ s_aux=self.sinks,
1168
)
1169
1170
return output