vLLM-4x-RTX-3090/patches/0102-qwen4_exp-add-QWEN38_ABLITERATION-runtime-refusal-di.patch

From c0d73eb28886bf2119856a9e66637a8eb532906e Mon Sep 17 00:00:00 2001
From: Hinata <hinata@idihore.id>
Date: Sun, 11 Oct 2026 03:47:14 +0000
Subject: [PATCH 102/103] qwen4_exp: add QWEN38_ABLITERATION runtime
 refusal-direction removal

---
 vllm/models/qwen4_exp/nvidia/abliteration.py | 88 ++++++++++++++++++++
 vllm/models/qwen4_exp/nvidia/model.py        |  7 +-
 vllm/models/qwen4_exp/nvidia/mtp.py          |  4 +-
 vllm/models/qwen4_exp/nvidia/ple_layer.py    |  3 +
 4 files changed, 100 insertions(+), 2 deletions(-)
 create mode 100644 vllm/models/qwen4_exp/nvidia/abliteration.py

diff --git a/vllm/models/qwen4_exp/nvidia/abliteration.py b/vllm/models/qwen4_exp/nvidia/abliteration.py
new file mode 100644
index 0000000..840168b
--- /dev/null
+++ b/vllm/models/qwen4_exp/nvidia/abliteration.py
@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Optional refusal-direction removal ("abliteration") applied at runtime.
+
+``QWEN38_ABLITERATION`` names a JSON file with a ``refusal_direction`` list, or a safetensors file with a
+``refusal_direction`` tensor, of length ``hidden_size``. A bare name such as ``Qwen3.8-Flash-Next-Uncensored`` resolves to
+``/opt/qwen38/abliteration/<name>.json``.
+With it set, every residual-stream write loses its component along the unit direction r:
+
+    x <- x - r (r . x)
+
+applied to the attention and linear-attention block outputs, the MoE block output (routed plus shared experts),
+the PLE value projection and the text token embeddings. For one direction this equals the weight edit
+``W' = W - r (r^T W)`` on ``self_attn.o_proj``, ``linear_attn.out_proj``, every expert's and the shared expert's
+``down_proj``, ``ple.value_proj`` and the rows of ``embed_tokens`` (Arditi et al., 2024): the routed and shared
+expert outputs are summed with scalar weights before the residual write, so projecting their sum is the same as
+editing each down projection. The INT4 experts and every other weight stay untouched. Unset, nothing is loaded and
+no work is added.
+"""
+
+import json
+import math
+import os
+
+import torch
+
+ENV = "QWEN38_ABLITERATION"
+BUILTIN_DIR = "/opt/qwen38/abliteration"
+
+
+def direction_path() -> str | None:
+    value = os.environ.get(ENV, "").strip()
+    if value in ("", "0"):
+        return None
+    if os.sep not in value and not value.endswith((".json", ".safetensors")):
+        value = os.path.join(BUILTIN_DIR, f"{value}.json")
+    return value
+
+
+def load_direction(hidden_size: int) -> torch.Tensor | None:
+    """The unit refusal direction as a float32 CPU tensor, or None when the switch is off."""
+    path = direction_path()
+    if path is None:
+        return None
+    if path.endswith(".json"):
+        with open(path) as f:
+            r = torch.tensor(json.load(f)["refusal_direction"], dtype=torch.float32)
+    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(torch.float32).flatten()
+    if r.numel() != hidden_size:
+        raise ValueError(f"{ENV}: {path} holds {r.numel()} values; the model's hidden size is {hidden_size}")
+    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")
+    return r / norm
+
+
+def register(module: torch.nn.Module, hidden_size: int, dtype: torch.dtype) -> None:
+    """Attach the direction as a non-persistent buffer ``_abliteration_r`` (None when the switch is off).
+
+    Call inside the model constructor so the buffer lands on the model's device before CUDA-graph capture. The
+    CPU-only PLE worker process never runs these forwards and gets no buffer.
+    """
+    from vllm.model_executor.layers.ple_offload_layer import is_offload_process
+
+    r = None if is_offload_process() else load_direction(hidden_size)
+    if r is None:
+        module._abliteration_r = None
+        return
+    buf = torch.empty(hidden_size, dtype=dtype)
+    buf.copy_(r)
+    module.register_buffer("_abliteration_r", buf, persistent=False)
+
+
+def project_(x: torch.Tensor, r: torch.Tensor | None) -> torch.Tensor:
+    """In place: remove the component along r from x's last dimension. No-op when r is None."""
+    if r is None:
+        return x
+    width = x.shape[-1]
+    if x.is_contiguous():
+        flat = x.view(-1, width)
+        flat.addr_(flat @ r, r, alpha=-1.0)  # rank-1 update, no full-size temporary
+    else:
+        x.sub_(torch.outer(x.reshape(-1, width) @ r, r).view(x.shape))
+    return x
diff --git a/vllm/models/qwen4_exp/nvidia/model.py b/vllm/models/qwen4_exp/nvidia/model.py
index fe416d4..b4fe373 100644
--- a/vllm/models/qwen4_exp/nvidia/model.py
+++ b/vllm/models/qwen4_exp/nvidia/model.py
@@ -89,6 +89,7 @@ from vllm.v1.e2_guards import (
 from vllm.v1.kv_cache_interface import MambaSpec
 
 from ..config import Qwen4ExpConfig
+from . import abliteration
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
 from .ple_layer import Qwen4ExpPLELayer
@@ -384,6 +385,7 @@ class Qwen4ExpDecoderLayer(nn.Module):
             hc_config,
             prefix=maybe_prefix(prefix, "mlp_hyper_connection"),
         )
+        abliteration.register(self, config.hidden_size, model_config.dtype)
 
     def forward(
         self,
@@ -434,12 +436,14 @@ class Qwen4ExpDecoderLayer(nn.Module):
             )
         else:
             raise ValueError("Invalid layer_type")
+        abliteration.project_(attn_out, self._abliteration_r)
 
         mlp_hc = self.mlp_hyper_connection
         hidden_states, block_input, injection = mlp_hc.combine_and_mix(
             hidden_states, attn_out, injection
         )
         mlp_out = self.mlp(block_input)
+        abliteration.project_(mlp_out, self._abliteration_r)
         return hidden_states, mlp_out, injection
 
 
@@ -554,6 +558,7 @@ class Qwen4ExpModel(nn.Module):
             self.embed_tokens = embedding_cls(self.vocab_size, config.hidden_size)
         else:
             self.embed_tokens = PPMissingLayer()
+        abliteration.register(self, config.hidden_size, vllm_config.model_config.dtype)
 
         def get_layer(prefix: str) -> Qwen4ExpDecoderLayer:
             layer_idx = extract_layer_index(prefix)
@@ -615,7 +620,7 @@ class Qwen4ExpModel(nn.Module):
             self._mtp_hidden_buffer = None
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
+        return abliteration.project_(self.embed_tokens(input_ids), self._abliteration_r)
 
     @staticmethod
     def _start_layer_ple_prefetch(
diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py
index 01fb69c..39a0ff4 100644
--- a/vllm/models/qwen4_exp/nvidia/mtp.py
+++ b/vllm/models/qwen4_exp/nvidia/mtp.py
@@ -60,6 +60,7 @@ from vllm.v1.e2_guards import (
     assert_ple_mode,
 )
 
+from . import abliteration
 from .hyperconnection import GatedResidual, HyperConnectionConfig
 from .low_latency_gemm import enable_qwen4_exp_low_latency_gemm
 from .model import (
@@ -220,6 +221,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):
             else VocabParallelEmbedding
         )
         self.embed_tokens = embedding_cls(self.vocab_size, self.hidden_size)
+        abliteration.register(self, self.hidden_size, model_config.dtype)
         draft_vllm_config = _make_draft_vllm_config(
             vllm_config,
             self.mtp_start_layer_idx,
@@ -312,7 +314,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):
             buffer[:num_rows].copy_(selected)
 
     def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
-        return self.embed_tokens(input_ids)
+        return abliteration.project_(self.embed_tokens(input_ids), self._abliteration_r)
 
     def forward(
         self,
diff --git a/vllm/models/qwen4_exp/nvidia/ple_layer.py b/vllm/models/qwen4_exp/nvidia/ple_layer.py
index af8f528..443e9dd 100644
--- a/vllm/models/qwen4_exp/nvidia/ple_layer.py
+++ b/vllm/models/qwen4_exp/nvidia/ple_layer.py
@@ -28,6 +28,7 @@ from vllm.v1.attention.backends.short_conv_attn import (
     PleShortConvAttentionMetadata,
 )
 
+from . import abliteration
 from .ngram_embedding import Qwen4ExpNGramEmbedding
 from .ops.ple import ple_conv, ple_gate
 
@@ -136,6 +137,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
             prefix=f"{prefix}.kv_proj",
             disable_tp=True,
         )
+        abliteration.register(self, self.hidden_size, model_config.dtype)
         norm_args = (
             self.hc_hidden_size,
             config.rms_norm_eps,
@@ -449,6 +451,7 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
         embeddings = self._dequantize_embeddings(embeddings, hidden_states.dtype)
         kv, _ = self.kv_proj(embeddings)
         key, value = kv.split(self.kv_proj.output_sizes, dim=-1)
+        abliteration.project_(value, self._abliteration_r)
         gated_output, conv_input = ple_gate(
             key,
             value,
-- 
2.43.0