[Qwen3.5] Project the refusal direction out of the residual stream at runtime

orcarouter/Qwen3.8-27B-Uncensored removes refusal by editing weights. The
Swift release we serve is AutoRound W8A16, so re-editing weights means
re-quantizing. Project one unit direction r out of every residual-stream
writer output instead: x <- x - r (r . x), which equals the weight edit
W' = W - r(r^T W) on every writer feeding that stream (Arditi et al. 2024).

Sites: attention output for both the full-attention and GDN branches, the
MLP output, and token embeddings, in Qwen3NextDecoderLayer.forward and
Qwen3NextModel.embed_input_ids, which Qwen3_5DecoderLayer and Qwen3_5Model
inherit but whose base constructors they skip, so both register there too.
The MTP drafter reuses Qwen3_5DecoderLayer and calls its own
embed_input_ids, so the draft path is covered at the same width, and
_maybe_share_embeddings at PP=1 hands the draft the target embedding
module. Sequence parallelism shards dim 0 only, so the last dim is
hidden_size at every site.

VLLM_ABLITERATION names a JSON file with key refusal_direction, a
safetensors tensor, or a bare name resolved under
/opt/qwen38/abliteration then /data/abliteration. Empty or 0 is off.
register() always creates the buffer _abliteration_r, zeros when off, so
one compiled graph serves ON and OFF and toggling the switch
(declared in envs.py, hence in compile_factors) invalidates the compile
cache. project_() is an in-place rank-1 update with no full-size
temporary, and skips tensors whose width is not hidden_size.

Co-authored-by: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

tests/models/test_abliteration.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Tests for the runtime refusal-direction projection."""
+
+import json
+import math
+
+import pytest
+import torch
+
+from vllm import envs
+from vllm.model_executor.models import abliteration
+
+HIDDEN = 8
+
+
+@pytest.fixture(autouse=True)
+def clear_direction_cache():
+    abliteration._read.cache_clear()
+    yield
+    abliteration._read.cache_clear()
+
+
+def write_direction(tmp_path, name="r.json", size=HIDDEN):
+    direction = [math.sin(i) for i in range(size)]
+    path = tmp_path / name
+    path.write_text(json.dumps({"hidden_size": size, "refusal_direction": direction}))
+    return str(path)
+
+
+def registered(module, size=HIDDEN, dtype=torch.float32):
+    abliteration.register(module, size, dtype)
+    return module
+
+
+def test_off_is_a_noop_and_still_registers(tmp_path, monkeypatch):
+    """With the switch off the buffer exists as zeros and x is unchanged."""
+    monkeypatch.delenv("VLLM_ABLITERATION", raising=False)
+    monkeypatch.delenv("QWEN38_ABLITERATION", raising=False)
+    module = registered(torch.nn.Module())
+    assert module._abliteration_on is False
+    assert torch.count_nonzero(module._abliteration_r) == 0
+
+    x = torch.randn(3, HIDDEN)
+    expected = x.clone()
+    assert abliteration.project_(module, x) is x
+    abliteration.project_(module, x)
+    assert torch.equal(x, expected)
+
+
+def test_projection_drops_the_component(tmp_path, monkeypatch):
+    """After projection the tensor keeps no component along the direction."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Module())
+    assert module._abliteration_on is True
+
+    r = module._abliteration_r
+    x = torch.randn(4, HIDDEN)
+    abliteration.project_(module, x)
+    assert torch.allclose(x @ r, torch.zeros(4), atol=1e-5)
+    assert not torch.allclose(x, torch.zeros_like(x))
+
+    # Projecting an already projected tensor changes nothing.
+    projected = x.clone()
+    abliteration.project_(module, x)
+    assert torch.allclose(x, projected, atol=1e-6)
+
+
+def test_registered_direction_is_unit(tmp_path, monkeypatch):
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Module())
+    assert torch.allclose(module._abliteration_r.norm(), torch.tensor(1.0), atol=1e-6)
+
+
+def test_non_contiguous_input_is_projected(tmp_path, monkeypatch):
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Module())
+    x = torch.randn(4, HIDDEN * 2)[:, ::2]
+    assert not x.is_contiguous()
+    abliteration.project_(module, x)
+    assert torch.allclose(x @ module._abliteration_r, torch.zeros(4), atol=1e-5)
+
+
+def test_width_mismatch_is_skipped(tmp_path, monkeypatch):
+    """A tensor without the full hidden dim must not be projected."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Module())
+    x = torch.randn(2, HIDDEN // 2)
+    before = x.clone()
+    abliteration.project_(module, x)
+    assert torch.equal(x, before)
+
+
+def test_hidden_size_mismatch_raises(tmp_path, monkeypatch):
+    """A direction for another model fails at boot, not silently."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path, size=HIDDEN))
+    with pytest.raises(ValueError, match="holds 8 values"):
+        registered(torch.nn.Module(), size=HIDDEN + 1)
+
+
+def test_zero_direction_raises(tmp_path, monkeypatch):
+    path = tmp_path / "zero.json"
+    path.write_text(json.dumps({"refusal_direction": [0.0] * HIDDEN}))
+    monkeypatch.setenv("VLLM_ABLITERATION", str(path))
+    with pytest.raises(ValueError, match="zero or non-finite"):
+        abliteration.load_direction(HIDDEN)
+
+
+def test_bare_name_resolves_in_builtin_dir(tmp_path, monkeypatch):
+    path = write_direction(tmp_path, name="Uncensored.json")
+    monkeypatch.setattr(abliteration, "BUILTIN_DIRS", (str(tmp_path),))
+    monkeypatch.setenv("VLLM_ABLITERATION", "Uncensored")
+    assert abliteration.direction_path() == path
+
+
+def test_missing_direction_file_raises(tmp_path, monkeypatch):
+    monkeypatch.setattr(abliteration, "BUILTIN_DIRS", (str(tmp_path / "nope"),))
+    monkeypatch.setenv("VLLM_ABLITERATION", "Uncensored")
+    with pytest.raises(FileNotFoundError):
+        abliteration.load_direction(HIDDEN)
+
+
+def test_meta_device_turns_projection_off(tmp_path, monkeypatch):
+    """A meta-initialized module cannot keep the direction, so it stays off."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = torch.nn.Module()
+    with torch.device("meta"):
+        abliteration.register(module, HIDDEN, torch.float32)
+    assert module._abliteration_r.device.type == "meta"
+    assert module._abliteration_on is False
+
+
+def test_legacy_env_name_is_accepted(tmp_path, monkeypatch):
+    """QWEN38_ABLITERATION still works, as existing deployments set it."""
+    path = write_direction(tmp_path)
+    monkeypatch.delenv("VLLM_ABLITERATION", raising=False)
+    monkeypatch.setenv("QWEN38_ABLITERATION", path)
+    assert path == envs.VLLM_ABLITERATION
+    assert abliteration.direction_path() == path
+
+
+def test_compile_projects_without_a_graph_break(tmp_path, monkeypatch):
+    """The projection stays traceable, so it does not split the compiled model."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+
+    class Project(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            abliteration.register(self, HIDDEN, torch.float32)
+
+        def forward(self, x):
+            return abliteration.project_(self, x)
+
+    module = Project().eval()
+    x = torch.randn(16, HIDDEN)
+    expected = abliteration.project_(module, x.clone())
+    got = torch.compile(module, fullgraph=True)(x.clone())
+    assert torch.allclose(got, expected, atol=1e-6)
+
+
+def qwen3_5_config(dense: bool):
+    """A two-layer Qwen3.5 config, one full-attention and one linear layer."""
+    if dense:
+        from transformers import Qwen3_5TextConfig as config_cls
+
+        extra = {"intermediate_size": 256, "hidden_act": "silu"}
+    else:
+        from transformers import Qwen3_5MoeTextConfig as config_cls
+
+        extra = {
+            "moe_intermediate_size": 32,
+            "shared_expert_intermediate_size": 32,
+            "num_experts": 2,
+            "num_experts_per_tok": 1,
+            "n_group": 1,
+            "topk_group": 1,
+            "norm_topk_prob": False,
+            "decoder_sparse_step": 1,
+            "mlp_only_layers": [],
+        }
+    return config_cls(
+        vocab_size=64,
+        hidden_size=HIDDEN,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=32,
+        rms_norm_eps=1e-6,
+        layer_types=["full_attention", "linear_attention"],
+        linear_key_head_dim=32,
+        linear_value_head_dim=32,
+        linear_num_key_heads=2,
+        linear_num_value_heads=2,
+        num_nextn_predict_layers=1,
+        **extra,
+    )
+
+
+@pytest.mark.parametrize("dense", [True, False])
+def test_qwen3_5_model_projects_embeddings(tmp_path, monkeypatch, dense, dist_init):
+    """A real Qwen3.5 model registers the direction and drops it from embeddings."""
+    from types import SimpleNamespace
+
+    from vllm.config import VllmConfig, set_current_vllm_config
+    from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
+
+    config = qwen3_5_config(dense)
+    path = write_direction(tmp_path, size=HIDDEN)
+    monkeypatch.setenv("VLLM_ABLITERATION", path)
+
+    vllm_config = VllmConfig()
+    vllm_config.model_config = SimpleNamespace(
+        hf_config=config,
+        hf_text_config=config,
+        dtype=torch.bfloat16,
+        max_model_len=128,
+        is_diffusion=False,
+        head_dtype=None,
+        is_mm_prefix_lm=False,
+        multimodal_config=None,
+        quantization_config=None,
+        runner_type="generate",
+        is_moe=not dense,
+        logits_processors=None,
+        rswa_window=None,
+    )
+    vllm_config.quant_config = None
+
+    with set_current_vllm_config(vllm_config):
+        model = Qwen3_5Model(vllm_config=vllm_config)
+    with torch.no_grad():
+        for param in model.parameters():
+            param.data.normal_(0, 0.02)
+
+    assert model._abliteration_on is True
+    for layer in model.layers:
+        assert layer._abliteration_on is True
+
+    r = model._abliteration_r.float()
+    # The buffer carries the model dtype, so bf16 rounding keeps it near unit norm.
+    assert torch.allclose(r.norm(), torch.tensor(1.0), atol=2e-2)
+
+    embeddings = model.embed_input_ids(torch.arange(6)).float()
+    residual = float((embeddings @ r).abs().max())
+    scale = float(embeddings.abs().max())
+    assert residual < 1e-2 * scale

vllm/envs.py

@@ -227,6 +227,7 @@ if TYPE_CHECKING:
     VLLM_REGEX_COMPILATION_TIMEOUT_S: int = 5
     VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256
     VLLM_ALLOW_INSECURE_SERIALIZATION: bool = False
+    VLLM_ABLITERATION: str = ""
     VLLM_DISABLE_REQUEST_ID_RANDOMIZATION: bool = False
     VLLM_NIXL_SIDE_CHANNEL_HOST: str = "localhost"
     VLLM_NIXL_SIDE_CHANNEL_PORT: int = 5600
@@ -1189,6 +1190,15 @@ environment_variables: dict[str, Callable[[], Any]] = {
         os.environ.get("VLLM_ALLOW_RUNTIME_LORA_UPDATING", "0").strip().lower()
         in ("1", "true")
     ),
+    # Path to a refusal direction (a JSON file with a "refusal_direction" list,
+    # or a safetensors file with that tensor) projected out of every
+    # residual-stream write of the served model, or a bare name resolved under
+    # /opt/qwen38/abliteration and /data/abliteration. Empty or "0" is off. See
+    # vllm/model_executor/models/abliteration.py. QWEN38_ABLITERATION is the
+    # older name some deployments still set.
+    "VLLM_ABLITERATION": lambda: (
+        os.getenv("VLLM_ABLITERATION") or os.getenv("QWEN38_ABLITERATION", "")
+    ),
     # We assume drivers can report p2p status correctly.
     # If the program hangs when using custom allreduce,
     # potantially caused by a bug in the driver (535 series),

vllm/model_executor/models/abliteration.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Optional runtime refusal-direction removal ("abliteration").
+
+``VLLM_ABLITERATION`` (or the legacy ``QWEN38_ABLITERATION``) names a JSON file
+with a ``refusal_direction`` list, or a safetensors file with that tensor, of
+length ``hidden_size``. A bare name resolves to ``<name>.json`` in
+``/opt/qwen38/abliteration``, then in ``/data/abliteration``. When set, every
+write into the residual stream loses its component along the unit direction r:
+
+    x <- x - r (r . x)
+
+This runs on the attention block outputs (full and linear attention), the MLP
+output and the text token embeddings. For one direction it equals the weight
+edit ``W' = W - r (r^T W)`` on ``self_attn.o_proj``,
+``linear_attn.out_proj``, ``mlp.down_proj`` and the rows of ``embed_tokens``
+(Arditi et al., 2024), so the served weights, quantized or not, stay untouched.
+
+Unset, no projection kernel runs. The buffer always exists, holding zeros while
+the switch is off, so the attribute a compiled or captured graph refers to is
+always present, and the switch is part of the torch.compile cache key, so
+toggling it never reuses a stale graph.
+"""
+
+import json
+import math
+import os
+from functools import lru_cache
+
+import torch
+
+from vllm import envs
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+ENV = "VLLM_ABLITERATION"
+BUILTIN_DIRS = ("/opt/qwen38/abliteration", "/data/abliteration")
+DIRECTION = "_abliteration_r"
+ENABLED = "_abliteration_on"
+
+_shape_warned = False
+
+
+def direction_path() -> str | None:
+    """The configured direction file, or None when the switch is off."""
+    value = (envs.VLLM_ABLITERATION or "").strip()
+    if value in ("", "0"):
+        return None
+    if os.sep in value or value.endswith((".json", ".safetensors")):
+        return value
+    for directory in BUILTIN_DIRS:
+        path = os.path.join(directory, f"{value}.json")
+        if os.path.exists(path):
+            return path
+    return os.path.join(BUILTIN_DIRS[0], f"{value}.json")
+
+
+@lru_cache(maxsize=4)
+def _read(path: str) -> torch.Tensor:
+    """The direction as a unit float32 CPU tensor, read once per path."""
+    if path.endswith(".json"):
+        with open(path) as f:
+            values = json.load(f)["refusal_direction"]
+        # A model built on the meta device would otherwise make this CPU tensor
+        # a meta one, and a meta tensor cannot be inspected or copied.
+        r = torch.tensor(values, dtype=torch.float32, device="cpu")
+    else:
+        from safetensors import safe_open
+
+        with safe_open(path, framework="pt", device="cpu") as f:
+            r = f.get_tensor("refusal_direction")
+        r = r.to(dtype=torch.float32, device="cpu")
+    r = r.flatten()
+    norm = float(r.norm())
+    if not math.isfinite(norm) or norm == 0.0:
+        raise ValueError(f"{ENV}: {path} holds a zero or non-finite direction")
+    logger.info("%s: refusal direction (%d values) from %s", ENV, r.numel(), path)
+    return r / norm
+
+
+def load_direction(hidden_size: int) -> torch.Tensor | None:
+    """The unit refusal direction on CPU, or None when the switch is off."""
+    path = direction_path()
+    if path is None:
+        return None
+    r = _read(path)
+    if r.numel() != hidden_size:
+        raise ValueError(
+            f"{ENV}: {path} holds {r.numel()} values; the model's hidden size "
+            f"is {hidden_size}"
+        )
+    return r
+
+
+def register(module: torch.nn.Module, hidden_size: int, dtype: torch.dtype) -> None:
+    """Attach the direction as a non-persistent buffer, from a model constructor.
+
+    Call it before CUDA-graph capture so the buffer lands on the model's device.
+    It holds zeros when the switch is off. A buffer that lands on the meta
+    device cannot hold a direction, which turns the projection off for that
+    module.
+    """
+    r = load_direction(hidden_size)
+    buf = torch.zeros(hidden_size, dtype=dtype)
+    on = r is not None
+    if on and buf.device.type == "meta":
+        logger.warning(
+            "%s: %s is built on the meta device, which cannot hold a direction; "
+            "projection is off for this module",
+            ENV,
+            type(module).__name__,
+        )
+        on = False
+    elif on:
+        buf.copy_(r)
+    module.register_buffer(DIRECTION, buf, persistent=False)
+    setattr(module, ENABLED, on)
+
+
+def project_(module: torch.nn.Module, x: torch.Tensor) -> torch.Tensor:
+    """Remove the component of x along the direction, in place. No-op when off."""
+    if not getattr(module, ENABLED, False):
+        return x
+    r = getattr(module, DIRECTION)
+    width = x.shape[-1]
+    if width != r.numel():
+        _warn_shape(type(module).__name__, width, r.numel())
+        return x
+    if r.dtype != x.dtype:
+        # Activations do not always carry the model dtype, and the rank-1 update
+        # needs one dtype throughout.
+        r = r.to(x.dtype)
+    if x.is_contiguous():
+        flat = x.view(-1, width)
+        flat.addr_(flat @ r, r, alpha=-1.0)
+    else:
+        x.sub_(torch.outer(x.reshape(-1, width) @ r, r).view(x.shape))
+    return x
+
+
+def _warn_shape(name: str, width: int, size: int) -> None:
+    global _shape_warned
+    if not _shape_warned:
+        _shape_warned = True
+        logger.warning(
+            "%s: skipped a %d-wide tensor from %s; the direction has %d values",
+            ENV,
+            width,
+            name,
+            size,
+        )

vllm/model_executor/models/qwen3_5.py

@@ -66,6 +66,7 @@ from vllm.multimodal.inputs import MultiModalFeatureSpec
 from vllm.sequence import IntermediateTensors
 from vllm.tokenizers.registry import cached_tokenizer_from_config
 
+from . import abliteration
 from .interfaces import (
     HasInnerState,
     IsHybrid,
@@ -191,6 +192,7 @@ class Qwen3_5DecoderLayer(Qwen3NextDecoderLayer):
         self.post_attention_layernorm = Qwen3_5RMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
+        abliteration.register(self, config.hidden_size, model_config.dtype)
 
         self.layer_scale = getattr(config, "layer_scale", False)
         if self.layer_scale:
@@ -254,6 +256,7 @@ class Qwen3_5Model(Qwen3NextModel):
             quant_config=self.quant_config,
             prefix=f"{prefix}.embed_tokens",
         )
+        abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype)
 
         def get_layer(prefix: str):
             return Qwen3_5DecoderLayer(

vllm/model_executor/models/qwen3_5_mtp.py

@@ -37,6 +37,7 @@ from vllm.model_executor.models.utils import sequence_parallel_chunk
 from vllm.sequence import IntermediateTensors
 from vllm.transformers_utils.repo_utils import try_get_local_file
 
+from . import abliteration
 from .interfaces import (
     MultiModalEmbeddings,
     SupportsMultiModal,
@@ -146,6 +147,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
             self.vocab_size,
             config.hidden_size,
         )
+        abliteration.register(self, config.hidden_size, model_config.dtype)
 
         # Workaround: mtp.fc is stored as BF16 in NVFP4 checkpoints but is
         # missing from hf_quant_config.json exclude_modules. Force unquantized.
@@ -217,7 +219,7 @@ class Qwen3_5MultiTokenPredictor(nn.Module):
         )
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
+        return abliteration.project_(self, self.embed_tokens(input_ids))
 
     def forward(
         self,

vllm/model_executor/models/qwen3_next.py

@@ -61,6 +61,7 @@ from vllm.platforms import current_platform
 from vllm.sequence import IntermediateTensors
 from vllm.v1.attention.backend import AttentionType
 
+from . import abliteration
 from .interfaces import (
     EagleModelMixin,
     HasInnerState,
@@ -537,6 +538,7 @@ class Qwen3NextDecoderLayer(nn.Module):
         self.post_attention_layernorm = Qwen3NextRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps
         )
+        abliteration.register(self, config.hidden_size, model_config.dtype)
 
         self.layer_scale = getattr(config, "layer_scale", False)
         if self.layer_scale:
@@ -584,6 +586,8 @@ class Qwen3NextDecoderLayer(nn.Module):
         else:
             raise ValueError("Invalid layer_type")
 
+        abliteration.project_(self, hidden_states)
+
         if self.layer_scale:
             if len(hidden_states.shape) == 2:
                 hidden_states = hidden_states * (
@@ -612,6 +616,8 @@ class Qwen3NextDecoderLayer(nn.Module):
         else:
             hidden_states = self.mlp(hidden_states)
 
+        abliteration.project_(self, hidden_states)
+
         if self.layer_scale:
             if len(hidden_states.shape) == 2:
                 hidden_states = hidden_states * (
@@ -661,6 +667,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
             self.vocab_size,
             config.hidden_size,
         )
+        abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype)
 
         def get_layer(prefix: str):
             return Qwen3NextDecoderLayer(
@@ -689,7 +696,7 @@ class Qwen3NextModel(nn.Module, EagleModelMixin):
         self.aux_hidden_state_layers: tuple[int, ...] = ()
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
+        return abliteration.project_(self, self.embed_tokens(input_ids))
 
     @property
     def use_sequence_parallel(self) -> bool: