[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>

docs/models/pooling_models/token_classify.md

@@ -45,6 +45,7 @@ The BAAI/bge-m3 model leverages token classification for sparse retrieval. For m
 | ------------ | ------ | ----------------- | --------------------------- | --------------------------------------- |
 | `BertForTokenClassification` | bert-based | `boltuix/NeuroBERT-NER` (see note), etc. | | |
 | `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | |
+| `OpenAIPrivacyFilterForTokenClassification` | gpt-oss-based encoder | `openai/privacy-filter` | | |
 | `Qwen3ForTokenClassification`<sup>C</sup> | Qwen3-based | `bd2lcco/Qwen3-0.6B-finetuned` | | |
 | `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* |
 

tests/models/language/pooling/test_token_classification.py

@@ -116,6 +116,53 @@ def test_modernbert_models(
         torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)
 
 
+PRIVACY_FILTER_PROMPTS = [
+    "My name is Harry Potter.",
+    "Email me at harry.potter@hogwarts.edu.",
+    "Call me on +44 20 7946 0958 tomorrow.",
+    "My account number is 12345678 and the API key is sk-live-abc123def456.",
+    "I live at 4 Privet Drive, Little Whinging.",
+    "Visit https://example.com/profile/harry for more info.",
+    "We met on 12 January 2024.",
+]
+
+
+@pytest.mark.parametrize("model", ["openai/privacy-filter"])
+@pytest.mark.parametrize("dtype", ["bfloat16"])
+@torch.inference_mode
+def test_openai_privacy_filter(
+    hf_runner,
+    vllm_runner,
+    model: str,
+    dtype: str,
+) -> None:
+    with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
+        vllm_outputs = vllm_model.token_classify(PRIVACY_FILTER_PROMPTS)
+
+    hf_model_kwargs = {}
+    if current_platform.is_rocm():
+        hf_model_kwargs["attn_implementation"] = "eager"
+
+    with hf_runner(
+        model,
+        dtype=dtype,
+        auto_cls=AutoModelForTokenClassification,
+        model_kwargs=hf_model_kwargs,
+    ) as hf_model:
+        tokenizer = hf_model.tokenizer
+        hf_outputs = []
+        for prompt in PRIVACY_FILTER_PROMPTS:
+            inputs = tokenizer([prompt], return_tensors="pt")
+            inputs = hf_model.wrap_device(inputs)
+            output = hf_model.model(**inputs)
+            hf_outputs.append(softmax(output.logits[0]))
+
+    for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
+        hf_output = hf_output.detach().clone().cpu().float()
+        vllm_output = vllm_output.detach().clone().cpu().float()
+        torch.testing.assert_close(hf_output, vllm_output, atol=0.1, rtol=1e-2)
+
+
 @pytest.mark.parametrize("model", ["bd2lcco/Qwen3-0.6B-finetuned"])
 @pytest.mark.parametrize("dtype", ["float"])
 @torch.inference_mode

tests/models/registry.py

@@ -707,6 +707,10 @@ _TOKEN_CLASSIFICATION_EXAMPLE_MODELS = {
     "ModernBertForTokenClassification": _HfExamplesInfo(
         "disham993/electrical-ner-ModernBERT-base"
     ),
+    "OpenAIPrivacyFilterForTokenClassification": _HfExamplesInfo(
+        "openai/privacy-filter",
+        min_transformers_version="5.6.0.dev0",
+    ),
 }
 
 _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

@@ -11,6 +11,7 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm._aiter_ops import rocm_aiter_ops
 from vllm.config.kernel import MoEBackend
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.all2all_utils import (
     maybe_make_prepare_finalize,
 )
@@ -83,6 +84,14 @@ def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBac
         if moe_config.moe_parallel_config.dp_size > 1:
             _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)
 
+        # HACK: unquantized FlashInfer aliases SWIGLUOAI to plain Swiglu
+        # (swiglu_alpha/limit only set on the MXFP4 branch). Route to
+        # Triton's swigluoai_and_mul until that's plumbed through. Same
+        # demotion pattern as the Qwen3.5/dp_size hack above.
+        if moe_config.activation == MoEActivation.SWIGLUOAI:
+            _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_TRTLLM)
+            _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)
+
     elif current_platform.is_xpu():
         _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.XPU]
     elif current_platform.is_cpu():

vllm/model_executor/models/gpt_oss.py

@@ -70,6 +70,10 @@ from .utils import (
 
 
 class OAIAttention(nn.Module):
+    # Override to switch RoPE convention. gpt-oss uses NeoX (chunk halves);
+    # privacy-filter and similar derivatives use GPT-J (interleaved pairs).
+    rope_is_neox_style: bool = True
+
     def __init__(
         self,
         config: GptOssConfig,
@@ -99,7 +103,7 @@ class OAIAttention(nn.Module):
                 "beta_slow": config.rope_parameters["beta_slow"],
                 "truncate": config.rope_parameters.get("truncate", True),
             },
-            is_neox_style=True,
+            is_neox_style=self.rope_is_neox_style,
         )
 
         tp_size = get_tensor_model_parallel_world_size()
@@ -133,9 +137,25 @@ class OAIAttention(nn.Module):
         self.num_local_attention_heads = config.num_attention_heads // tp_size
         self.num_local_key_value_heads = config.num_key_value_heads // tp_size
 
+        self.attn = self._build_attention(
+            config=config,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=prefix,
+        )
+
+    def _build_attention(
+        self,
+        config: GptOssConfig,
+        cache_config: CacheConfig | None,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> Attention:
+        # Override to swap in an encoder-only attention or alter the
+        # per-layer sliding-window policy.
         # Only apply sliding window to every other layer
         sliding_window = config.sliding_window if self.layer_idx % 2 == 0 else None
-        self.attn = Attention(
+        return Attention(
             self.num_local_attention_heads,
             self.head_dim,
             self.scaling,
@@ -222,6 +242,10 @@ class MLPBlock(torch.nn.Module):
 
 
 class TransformerBlock(torch.nn.Module):
+    # Override to swap attention/MLP without re-implementing the block.
+    attention_cls: type[nn.Module] = OAIAttention
+    mlp_cls: type[nn.Module] = MLPBlock
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -234,13 +258,13 @@ class TransformerBlock(torch.nn.Module):
         cache_config = vllm_config.cache_config
 
         self.layer_idx = extract_layer_index(prefix)
-        self.attn = OAIAttention(
+        self.attn = self.attention_cls(
             config,
             prefix=f"{prefix}.attn",
             quant_config=quant_config,
             cache_config=cache_config,
         )
-        self.mlp = MLPBlock(vllm_config, self.layer_idx, prefix=f"{prefix}.mlp")
+        self.mlp = self.mlp_cls(vllm_config, self.layer_idx, prefix=f"{prefix}.mlp")
         self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
         self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
 
@@ -266,6 +290,9 @@ class TransformerBlock(torch.nn.Module):
 
 @support_torch_compile
 class GptOssModel(nn.Module, EagleModelMixin):
+    # Override to swap in an alternative TransformerBlock subclass.
+    block_cls: type[nn.Module] = TransformerBlock
+
     def __init__(
         self,
         *,
@@ -282,7 +309,7 @@ class GptOssModel(nn.Module, EagleModelMixin):
         )
         self.start_layer, self.end_layer, self.layers = make_layers(
             self.config.num_hidden_layers,
-            lambda prefix: TransformerBlock(
+            lambda prefix: self.block_cls(
                 vllm_config,
                 prefix=prefix,
                 quant_config=self.quant_config,

vllm/model_executor/models/openai_privacy_filter.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Inference-only OpenAI Privacy Filter model.
+
+gpt-oss reused as a bidirectional encoder for token classification: every
+layer runs non-causal attention with a banded ±sliding_window mask, and
+the LM head is replaced with a 33-class BIOES score head.
+"""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from vllm.config import CacheConfig, VllmConfig
+from vllm.model_executor.layers.attention.encoder_only_attention import (
+    EncoderOnlyAttention,
+)
+from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_classify
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.sequence import IntermediateTensors
+
+from .gpt_oss import GptOssForCausalLM, GptOssModel, OAIAttention, TransformerBlock
+from .interfaces_base import attn_type, default_pooling_type
+from .utils import AutoWeightsLoader, maybe_prefix
+
+
+class OpenAIPrivacyFilterAttention(OAIAttention):
+    # Privacy-filter uses GPT-J style RoPE (interleaved pairs), not NeoX.
+    rope_is_neox_style = False
+
+    def _build_attention(
+        self,
+        config,
+        cache_config: CacheConfig | None,
+        quant_config: QuantizationConfig | None,
+        prefix: str,
+    ) -> EncoderOnlyAttention:
+        # HF stores sliding_window+1 so each token attends to ±W neighbors;
+        # the encoder-only path applies this as a symmetric (W-1, W-1) mask.
+        return EncoderOnlyAttention(
+            num_heads=self.num_local_attention_heads,
+            head_size=self.head_dim,
+            scale=self.scaling,
+            num_kv_heads=self.num_local_key_value_heads,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            per_layer_sliding_window=config.sliding_window + 1,
+            prefix=f"{prefix}.attn",
+            sinks=self.sinks,
+        )
+
+
+class OpenAIPrivacyFilterDecoderLayer(TransformerBlock):
+    attention_cls = OpenAIPrivacyFilterAttention
+
+
+class OpenAIPrivacyFilterModel(GptOssModel):
+    block_cls = OpenAIPrivacyFilterDecoderLayer
+
+
+def _interleave_gate_up_concat_to_pairs(
+    weights: Iterable[tuple[str, torch.Tensor]],
+) -> Iterable[tuple[str, torch.Tensor]]:
+    # HF gate_up_proj is concat [gate | up]; swigluoai_and_mul wants
+    # gate/up interleaved. MXFP4/quark suffixes are already interleaved.
+    for name, weight in weights:
+        if name.endswith(".gate_up_proj") or name.endswith(".gate_up_proj_bias"):
+            *lead, two_i = weight.shape
+            i = two_i // 2
+            weight = (
+                torch.stack([weight[..., :i], weight[..., i:]], dim=-1)
+                .reshape(*lead, two_i)
+                .contiguous()
+            )
+        yield name, weight
+
+
+@attn_type("encoder_only")
+@default_pooling_type(tok_pooling_type="ALL")
+class OpenAIPrivacyFilterForTokenClassification(nn.Module):
+    is_pooling_model = True
+    hf_to_vllm_mapper = GptOssForCausalLM.hf_to_vllm_mapper
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        self.head_dtype = vllm_config.model_config.head_dtype
+        self.num_labels = config.num_labels
+
+        self.model = OpenAIPrivacyFilterModel(
+            vllm_config=vllm_config,
+            prefix=maybe_prefix(prefix, "model"),
+        )
+        self.score = nn.Linear(
+            config.hidden_size, config.num_labels, dtype=self.head_dtype
+        )
+
+        pooler_config = vllm_config.model_config.pooler_config
+        assert pooler_config is not None
+        self.pooler = pooler_for_token_classify(pooler_config)
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_input_ids(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        hidden_states = self.model(
+            input_ids=input_ids,
+            positions=positions,
+            intermediate_tensors=intermediate_tensors,
+            inputs_embeds=inputs_embeds,
+        )
+        hidden_states = hidden_states.to(self.head_dtype)
+        return self.score(hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(
+            _interleave_gate_up_concat_to_pairs(weights),
+            mapper=self.hf_to_vllm_mapper,
+        )

vllm/model_executor/models/registry.py

@@ -291,6 +291,10 @@ _TOKEN_CLASSIFICATION_MODELS = {
         "modernbert",
         "ModernBertForTokenClassification",
     ),
+    "OpenAIPrivacyFilterForTokenClassification": (
+        "openai_privacy_filter",
+        "OpenAIPrivacyFilterForTokenClassification",
+    ),
     "Qwen3ASRForcedAlignerForTokenClassification": (
         "qwen3_asr_forced_aligner",
         "Qwen3ASRForcedAlignerForTokenClassification",

vllm/v1/attention/backends/flash_attn.py

@@ -1164,6 +1164,7 @@ class FlashAttentionImpl(AttentionImpl):
             k_descale=layer._k_scale.expand(descale_shape),  # type: ignore[operator]
             v_descale=layer._v_scale.expand(descale_shape),  # type: ignore[operator]
             num_splits=1 if self.batch_invariant_enabled else 0,
+            s_aux=self.sinks,
         )
 
         return output