fix: keep the refusal-direction projection dynamo-safe

With VLLM_ABLITERATION set, boot crashed: dynamo traces
abliteration.register from Qwen3NextGatedDeltaNet.__init__, and
buf.copy_(r) fails on the meta device during the compile-time model
construction pass. Disable tracing for register and guard the copy so
an unmaterializable buffer turns projection off with a warning.

Log activation at registration instead of inside project_: project_
runs in the compiled region, and a logger call there aborts dynamo.
The width guard uses the same escape hatch.

Also project the embedding that Qwen3NextMultiTokenPredictor builds
for itself, which the first pass missed, and add tests for the
2560-value Flash-Next direction against a 5120 model and for a
half-width activation.

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

tests/models/test_abliteration.py

@@ -16,8 +16,10 @@ HIDDEN = 8
 @pytest.fixture(autouse=True)
 def clear_direction_cache():
     abliteration._read.cache_clear()
+    abliteration.clear_logs()
     yield
     abliteration._read.cache_clear()
+    abliteration.clear_logs()
 
 
 def write_direction(tmp_path, name="r.json", size=HIDDEN):
@@ -97,6 +99,39 @@ def test_hidden_size_mismatch_raises(tmp_path, monkeypatch):
         registered(torch.nn.Module(), size=HIDDEN + 1)
 
 
+def test_first_projection_logs_that_it_is_active(tmp_path, monkeypatch, caplog):
+    """The boot log proves the switch is live, not just loaded."""
+    import logging
+
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path))
+    module = registered(torch.nn.Linear(HIDDEN, HIDDEN))
+    with caplog.at_level(
+        logging.INFO, logger="vllm.model_executor.models.abliteration"
+    ):
+        abliteration.project_(module, torch.randn(2, HIDDEN))
+    assert "active, projected a 8-wide tensor" in caplog.text
+
+
+def test_flash_next_direction_is_rejected_by_the_27b_model(tmp_path, monkeypatch):
+    """Flash-Next stores 2560 values; the dense 27B stream is 5120 wide."""
+    path = write_direction(tmp_path, name="flash_next.json", size=2560)
+    monkeypatch.setenv("VLLM_ABLITERATION", path)
+    with pytest.raises(
+        ValueError, match="holds 2560 values; the model's hidden size is 5120"
+    ):
+        registered(torch.nn.Module(), size=5120)
+
+
+def test_half_width_activation_is_skipped(tmp_path, monkeypatch):
+    """A 2560-wide tensor cannot be projected by the 5120 direction."""
+    monkeypatch.setenv("VLLM_ABLITERATION", write_direction(tmp_path, size=5120))
+    module = registered(torch.nn.Module(), size=5120)
+    x = torch.randn(4, 2560)
+    before = x.clone()
+    abliteration.project_(module, x)
+    assert torch.equal(x, before)
+
+
 def test_zero_direction_raises(tmp_path, monkeypatch):
     path = tmp_path / "zero.json"
     path.write_text(json.dumps({"refusal_direction": [0.0] * HIDDEN}))

vllm/model_executor/models/abliteration.py

@@ -39,6 +39,14 @@ DIRECTION = "_abliteration_r"
 ENABLED = "_abliteration_on"
 
 _shape_warned = False
+_register_logged = False
+
+
+def clear_logs() -> None:
+    """Forget the once-only log flags, for tests."""
+    global _register_logged, _shape_warned
+    _register_logged = False
+    _shape_warned = False
 
 
 def direction_path() -> str | None:
@@ -92,29 +100,44 @@ def load_direction(hidden_size: int) -> torch.Tensor | None:
     return r
 
 
+@torch._dynamo.disable
 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
+    It holds zeros when the switch is off. A construction pass that cannot hold
-    device cannot hold a direction, which turns the projection off for that
+    data, the meta device or a fake-tensor compile pass, turns the projection off
-    module.
+    for that module and says so.
     """
     r = load_direction(hidden_size)
     buf = torch.zeros(hidden_size, dtype=dtype)
     on = r is not None
-    if on and buf.device.type == "meta":
+    if on:
-        logger.warning(
+        try:
-            "%s: %s is built on the meta device, which cannot hold a direction; "
+            buf.copy_(r)
-            "projection is off for this module",
+        except (RuntimeError, NotImplementedError):
-            ENV,
+            # Meta or fake tensor construction: no room for a direction here.
-            type(module).__name__,
+            logger.warning(
-        )
+                "%s: %s is constructed on the %s device, which cannot hold a "
-        on = False
+                "direction; projection is off for this module",
-    elif on:
+                ENV,
-        buf.copy_(r)
+                type(module).__name__,
+                buf.device.type,
+            )
+            on = False
     module.register_buffer(DIRECTION, buf, persistent=False)
     setattr(module, ENABLED, on)
+    # Logging from project_ would sit inside a torch.compile region, which cannot
+    # trace a logger call, so the proof that the switch is live goes here.
+    global _register_logged
+    if on and not _register_logged:
+        _register_logged = True
+        logger.info(
+            "%s: projection is active, first registered on %s (hidden size %d)",
+            ENV,
+            type(module).__name__,
+            hidden_size,
+        )
 
 
 def project_(module: torch.nn.Module, x: torch.Tensor) -> torch.Tensor:
@@ -138,6 +161,7 @@ def project_(module: torch.nn.Module, x: torch.Tensor) -> torch.Tensor:
     return x
 
 
+@torch._dynamo.disable
 def _warn_shape(name: str, width: int, size: int) -> None:
     global _shape_warned
     if not _shape_warned:

vllm/model_executor/models/qwen3_next_mtp.py

@@ -21,6 +21,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
 )
+from vllm.model_executor.models import abliteration
 from vllm.model_executor.models.qwen3_next import (
     Qwen3NextDecoderLayer,
     Qwen3NextModel,
@@ -66,6 +67,7 @@ class Qwen3NextMultiTokenPredictor(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 the checkpoint quant exclude list (its `ignore` glob
@@ -113,7 +115,7 @@ class Qwen3NextMultiTokenPredictor(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,