From 4f574fcb5e0d0c753a54dba99cc3bd39c51582df Mon Sep 17 00:00:00 2001
From: "cyko@ibm.com;6J3007897;Irene Ko" <cyko1@ccc-login5.pok.ibm.com>
Date: Tue, 7 Apr 2026 12:23:26 -0400
Subject: [PATCH 1/5] Expand the support for locating hidden state id; device
 types

Signed-off-by: cyko@ibm.com;6J3007897;Irene Ko <cyko1@ccc-login5.pok.ibm.com>
---
 .../state_control/act_add/control.py          |  3 +-
 .../algorithms/state_control/caa/control.py   |  4 +-
 .../state_control/common/hook_utils.py        | 44 +++++++++++++++++--
 .../algorithms/state_control/iti/control.py   |  4 +-
 4 files changed, 46 insertions(+), 9 deletions(-)

diff --git a/aisteer360/algorithms/state_control/act_add/control.py b/aisteer360/algorithms/state_control/act_add/control.py
index 25f054bc..4de2de79 100644
--- a/aisteer360/algorithms/state_control/act_add/control.py
+++ b/aisteer360/algorithms/state_control/act_add/control.py
@@ -10,6 +10,7 @@
 from aisteer360.algorithms.state_control.common.estimators import SinglePairEstimator
 from aisteer360.algorithms.state_control.common.gates import AlwaysOpenGate
 from aisteer360.algorithms.state_control.common.hook_utils import (
+    get_model_dtype,
     get_model_layer_list,
     extract_hidden_states,
     replace_hidden_states,
@@ -78,7 +79,7 @@ def steer(
             )
 
         device = next(model.parameters()).device
-        sv = sv.to(device, dtype=model.dtype)
+        sv = sv.to(device, dtype=get_model_dtype(model))
 
         # resolve layer_id via selector
         if self.layer_id is not None:
diff --git a/aisteer360/algorithms/state_control/caa/control.py b/aisteer360/algorithms/state_control/caa/control.py
index ea6e75f7..5830b732 100644
--- a/aisteer360/algorithms/state_control/caa/control.py
+++ b/aisteer360/algorithms/state_control/caa/control.py
@@ -7,7 +7,7 @@
 
 from aisteer360.algorithms.state_control.base import StateControl
 from aisteer360.algorithms.state_control.common.gates import AlwaysOpenGate
-from aisteer360.algorithms.state_control.common.hook_utils import get_model_layer_list
+from aisteer360.algorithms.state_control.common.hook_utils import get_model_dtype, get_model_layer_list
 from aisteer360.algorithms.state_control.common.selectors import FixedLayerSelector, FractionalDepthSelector
 from aisteer360.algorithms.state_control.common.token_scope import compute_prompt_lens, make_token_mask
 from aisteer360.algorithms.state_control.common.transforms import AdditiveTransform, NormPreservingTransform
@@ -88,7 +88,7 @@ def steer(
             sv = estimator.fit(model, tokenizer, data=self.data, spec=self.train_spec)
 
         # move to device
-        sv = sv.to(device, dtype=model.dtype)
+        sv = sv.to(device, dtype=get_model_dtype(model))
 
         # optionally normalize the vector
         if self.normalize_vector:
diff --git a/aisteer360/algorithms/state_control/common/hook_utils.py b/aisteer360/algorithms/state_control/common/hook_utils.py
index 6f9be4d6..9ef89192 100644
--- a/aisteer360/algorithms/state_control/common/hook_utils.py
+++ b/aisteer360/algorithms/state_control/common/hook_utils.py
@@ -1,5 +1,6 @@
 """Utilities for hook registration and model inspection."""
 import torch
+import torch.nn as nn
 from transformers import PreTrainedModel
 
 
@@ -33,11 +34,38 @@ def get_model_layer_list(model: PreTrainedModel) -> tuple[list, list[str]]:
     return modules, names
 
 
+def get_model_dtype(model: nn.Module) -> torch.dtype:
+    """Return the dtype of a model's parameters.
+
+    Works with both HuggingFace ``PreTrainedModel`` (which exposes
+    ``model.dtype``) and vLLM ``nn.Module`` by falling back to the 
+    first parameter's dtype.
+    """
+    if hasattr(model, "dtype"):
+        return model.dtype
+    return next(model.parameters()).dtype
+
+
+def _hidden_states_index(input_args: tuple) -> int | None:
+    """Return the index of the hidden_states tensor in positional args.
+
+    HF layers: ``forward(hidden_states, ...)`` → index 0.
+    vLLM layers: ``forward(positions, hidden_states, residual)`` → index 1.
+
+    Uses ``ndim >= 2`` to skip 1-D tensors like positions.
+    Returns ``None`` if no suitable tensor is found.
+    """
+    for i, arg in enumerate(input_args):
+        if isinstance(arg, torch.Tensor) and arg.ndim >= 2:
+            return i
+    return None
+
+
 def extract_hidden_states(input_args: tuple, input_kwargs: dict) -> torch.Tensor | None:
     """Extract hidden_states tensor from a pre-hook's arguments.
 
-    HuggingFace transformer layers receive hidden_states either as the
-    first positional argument or as a keyword argument.
+    Works with both HuggingFace layers (hidden_states as first arg)
+    and vLLM layers (hidden_states as second arg after positions).
 
     Args:
         input_args: Positional args from the pre-hook.
@@ -47,7 +75,9 @@ def extract_hidden_states(input_args: tuple, input_kwargs: dict) -> torch.Tensor
         The hidden_states tensor, or None if not found.
     """
     if input_args:
-        return input_args[0]
+        idx = _hidden_states_index(input_args)
+        if idx is not None:
+            return input_args[idx]
     return input_kwargs.get("hidden_states")
 
 
@@ -58,6 +88,8 @@ def replace_hidden_states(
 ) -> tuple[tuple, dict]:
     """Return modified (input_args, input_kwargs) with hidden_states replaced.
 
+    Works with both HuggingFace and vLLM layer argument patterns.
+
     Args:
         input_args: Original positional args.
         input_kwargs: Original keyword args.
@@ -67,7 +99,11 @@ def replace_hidden_states(
         Tuple of (new_input_args, new_input_kwargs).
     """
     if input_args:
-        return (new_hidden, *input_args[1:]), input_kwargs
+        idx = _hidden_states_index(input_args)
+        if idx is not None:
+            args_list = list(input_args)
+            args_list[idx] = new_hidden
+            return tuple(args_list), input_kwargs
     input_kwargs = dict(input_kwargs)
     input_kwargs["hidden_states"] = new_hidden
     return input_args, input_kwargs
diff --git a/aisteer360/algorithms/state_control/iti/control.py b/aisteer360/algorithms/state_control/iti/control.py
index f53cd323..a4e3a7e2 100644
--- a/aisteer360/algorithms/state_control/iti/control.py
+++ b/aisteer360/algorithms/state_control/iti/control.py
@@ -8,7 +8,7 @@
 
 from aisteer360.algorithms.state_control.base import StateControl
 from aisteer360.algorithms.state_control.common.gates import AlwaysOpenGate
-from aisteer360.algorithms.state_control.common.hook_utils import get_model_layer_list
+from aisteer360.algorithms.state_control.common.hook_utils import get_model_dtype, get_model_layer_list
 from aisteer360.algorithms.state_control.common.selectors import TopKHeadSelector
 from aisteer360.algorithms.state_control.common.steering_vector import SteeringVector
 from aisteer360.algorithms.state_control.common.token_scope import compute_prompt_lens, make_token_mask
@@ -102,7 +102,7 @@ def steer(
             sv = estimator.fit(model, tokenizer, data=self.data, spec=self.train_spec)
 
         # move to device
-        sv = sv.to(device, dtype=model.dtype)
+        sv = sv.to(device, dtype=get_model_dtype(model))
         self._steering_vector = sv
 
         # resolve head selection

From f3c0afee7309f8a367116fcefa839bc83f3763e5 Mon Sep 17 00:00:00 2001
From: "cyko@ibm.com;6J3007897;Irene Ko" <cyko1@ccc-login5.pok.ibm.com>
Date: Tue, 7 Apr 2026 13:19:15 -0400
Subject: [PATCH 2/5] added gating and interfacing with vllm; sampling param
 interpreter for vllm

Signed-off-by: cyko@ibm.com;6J3007897;Irene Ko <cyko1@ccc-login5.pok.ibm.com>
---
 .../algorithms/core/steering_pipeline.py      | 150 +++++++++++++++++-
 1 file changed, 146 insertions(+), 4 deletions(-)

diff --git a/aisteer360/algorithms/core/steering_pipeline.py b/aisteer360/algorithms/core/steering_pipeline.py
index 07b72970..bd0b7b79 100644
--- a/aisteer360/algorithms/core/steering_pipeline.py
+++ b/aisteer360/algorithms/core/steering_pipeline.py
@@ -67,6 +67,8 @@ class SteeringPipeline:
     device: torch.device | str | None = None
     hf_model_kwargs: dict = field(default_factory=dict)
     lazy_init: bool = False
+    backend: str = "hf"
+    vllm_kwargs: dict = field(default_factory=dict)
 
     # lazy‑filled fields
     model: PreTrainedModel | None = field(init=False, default=None)
@@ -78,6 +80,7 @@ class SteeringPipeline:
     output_control: OutputControl = field(init=False)
 
     _is_steered: bool = field(default=False, init=False, repr=False)
+    _vllm_engine: object = field(default=None, init=False, repr=False)
 
     def __post_init__(self) -> None:
 
@@ -88,8 +91,16 @@ def __post_init__(self) -> None:
         self.state_control = controls_merged["state_control"]
         self.output_control = controls_merged["output_control"]
 
+        if self.backend == "vllm":
+            # vLLM backend: skip model loading, only load tokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.tokenizer_name_or_path or self.model_name_or_path,
+                trust_remote_code=True,
+            )
+            self.tokenizer = ensure_pad_token(self.tokenizer)
+
         # load HF artifacts
-        if not self.lazy_init:
+        elif not self.lazy_init:
             if self.model_name_or_path is None:
                 raise ValueError("`model_name_or_path` must be provided when lazy_init=False")
 
@@ -164,6 +175,10 @@ def steer(self, **steer_kwargs) -> None:
         if self._is_steered:
             return
 
+        if self.backend == "vllm":
+            self._steer_vllm(**steer_kwargs)
+            return
+
         # steer each control (bottom-up order: structural -> input -> state -> output)
         for control in (self.structural_control, self.input_control, self.state_control, self.output_control):
             steer_fn = getattr(control, "steer", None)
@@ -198,12 +213,47 @@ def steer(self, **steer_kwargs) -> None:
         # return steered pipeline
         self._is_steered = True
 
+    def _steer_vllm(self, **steer_kwargs) -> None:
+        """vLLM backend: steer on a temp HF model, then boot vLLM engine."""
+        import gc
+
+        from aisteer360.adapters.vllm_hook.backend import boot_vllm_engine
+
+        # Load a temporary HF model to run steer(), 
+        # same loop as the regular steer() for preparing the control
+        temp_model = AutoModelForCausalLM.from_pretrained(
+            self.model_name_or_path,
+            **self.hf_model_kwargs,
+        )
+
+        for control in (self.structural_control, self.input_control, self.state_control, self.output_control):
+            steer_fn = getattr(control, "steer", None)
+            if callable(steer_fn):
+                maybe_new = steer_fn(temp_model, tokenizer=self.tokenizer, **steer_kwargs)
+                if isinstance(maybe_new, nn.Module):
+                    temp_model = maybe_new
+
+        # Unload temp model to free up the space before booting vLLM; 
+        # This does not work with structural control
+        del temp_model
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+        # Boot vLLM engine with the state control
+        self._vllm_engine = boot_vllm_engine(
+            model_name_or_path=str(self.model_name_or_path),
+            state_control=self.state_control,
+            vllm_kwargs=self.vllm_kwargs,
+        )
+        self._is_steered = True
+
     def _prepare_inputs(
             self,
             input_ids: list[int] | torch.LongTensor,
             attention_mask: torch.Tensor | None,
             runtime_kwargs: dict | None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor] | list[str]:
         """Apply input control and normalize input tensors.
 
         Transforms the prompt via the input control's adapter and ensures both input_ids and attention_mask are
@@ -215,10 +265,10 @@ def _prepare_inputs(
             runtime_kwargs: Per-call parameters for input control
 
         Returns:
-            tuple[torch.Tensor, torch.Tensor]: (steered_input_ids, attention_mask), both as 2D tensors on model device
+            HF backend: tuple[torch.Tensor, torch.Tensor] — (steered_input_ids, attention_mask), both as 2D tensors on model device
+            vLLM backend: list[str] — decoded prompt strings
         """
         runtime_kwargs = runtime_kwargs or {}
-        device = self.model.device
 
         # apply input control adapter
         adapter = self.input_control.get_prompt_adapter(runtime_kwargs)
@@ -229,6 +279,12 @@ def _prepare_inputs(
             steered_input_ids = torch.tensor(steered_input_ids, dtype=torch.long)
         if steered_input_ids.ndim == 1:
             steered_input_ids = steered_input_ids.unsqueeze(0)
+
+        # decode to strings and return 
+        if self.backend == "vllm":
+            return self.tokenizer.batch_decode(steered_input_ids, skip_special_tokens=False)
+
+        device = self.model.device
         steered_input_ids = steered_input_ids.to(device)
 
         # normalize attention_mask
@@ -302,6 +358,9 @@ def generate(
         if not self._is_steered:
             raise RuntimeError("Must call `.steer()` before `.generate()`.")
 
+        if self.backend == "vllm":
+            return self._generate_vllm(input_ids, attention_mask, runtime_kwargs, **gen_kwargs)
+
         runtime_kwargs = runtime_kwargs or {}
         return_full_sequence = bool(gen_kwargs.pop("return_full_sequence", False))
 
@@ -342,11 +401,61 @@ def generate_text(self, *args, **kwargs) -> str | list[str]:
         Returns:
             Decoded text string (single prompt) or list of strings (batch)
         """
+        if self.backend == "vllm":
+            return self._generate_text_vllm(*args, **kwargs)
         ids = self.generate(*args, **kwargs)
         if ids.ndim == 1:
             return self.tokenizer.decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
         return self.tokenizer.batch_decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
 
+    
+    def _generate_vllm(
+            self,
+            input_ids: list[int] | torch.LongTensor,
+            attention_mask: torch.Tensor | None = None,
+            runtime_kwargs: dict | None = None,
+            **gen_kwargs,
+    ) -> torch.Tensor:
+        """Generate via vLLM backend.
+
+        Returns:
+            Generated token IDs
+        """
+        prompts = self._prepare_inputs(input_ids, None, runtime_kwargs)
+
+        sampling_params = _gen_kwargs_to_sampling_params(gen_kwargs)
+        use_hook = gen_kwargs.pop("use_hook", True)
+        outputs = self._vllm_engine.generate(prompts, sampling_params, use_hook=use_hook)
+
+        # convert RequestOutput → padded token ID tensor
+        all_ids = [torch.tensor(out.outputs[0].token_ids, dtype=torch.long) for out in outputs]
+        max_len = max(t.size(0) for t in all_ids)
+        pad_id = self.tokenizer.pad_token_id or 0
+        padded = torch.full((len(all_ids), max_len), pad_id, dtype=torch.long)
+        for i, t in enumerate(all_ids):
+            padded[i, : t.size(0)] = t
+        return padded
+
+    def _generate_text_vllm(
+            self,
+            input_ids: list[int] | torch.LongTensor,
+            attention_mask: torch.Tensor | None = None,
+            runtime_kwargs: dict | None = None,
+            **gen_kwargs,
+    ) -> str | list[str]:
+
+        if not self._is_steered:
+            raise RuntimeError("Must call `.steer()` before `.generate_text()`.")
+
+        prompts = self._prepare_inputs(input_ids, None, runtime_kwargs)
+
+        sampling_params = _gen_kwargs_to_sampling_params(gen_kwargs)
+        use_hook = gen_kwargs.pop("use_hook", True)
+        outputs = self._vllm_engine.generate(prompts, sampling_params, use_hook=use_hook)
+
+        texts = [out.outputs[0].text for out in outputs]
+        return texts[0] if len(texts) == 1 else texts
+
     def compute_logprobs(
             self,
             input_ids: list[int] | torch.LongTensor,
@@ -381,6 +490,11 @@ def compute_logprobs(
         """
         if not self._is_steered:
             raise RuntimeError("Must call `.steer()` before `.compute_logprobs()`.")
+        if self.backend == "vllm":
+            raise NotImplementedError(
+                "compute_logprobs() is not supported with backend='vllm'. "
+                "Use backend='hf' for logprob computation."
+            )
         if ref_output_ids is None:
             raise ValueError("`ref_output_ids` is required for `compute_logprobs()`.")
 
@@ -537,3 +651,31 @@ def compute_logprobs(
             all_logprobs.append(token_logprobs)
 
         return torch.cat(all_logprobs, dim=0)
+
+
+# ---------------------------------------------------------------------------
+# Module-level helpers
+# ---------------------------------------------------------------------------
+
+def _gen_kwargs_to_sampling_params(gen_kwargs: dict):
+    """Map HuggingFace-style generation kwargs to vLLM SamplingParams."""
+    from vllm import SamplingParams
+
+    mapping = {
+        "max_new_tokens": "max_tokens",
+        "max_length": "max_tokens",
+        "temperature": "temperature",
+        "top_p": "top_p",
+        "top_k": "top_k",
+        "repetition_penalty": "repetition_penalty",
+    }
+    sp_kwargs: dict = {}
+    for hf_key, vllm_key in mapping.items():
+        if hf_key in gen_kwargs:
+            sp_kwargs[vllm_key] = gen_kwargs.pop(hf_key)
+
+    # sensible defaults
+    sp_kwargs.setdefault("max_tokens", 256)
+    sp_kwargs.setdefault("temperature", 0.7)
+
+    return SamplingParams(**sp_kwargs)

From 234b741181dbe9cf1e0f72a89599054315d55b1a Mon Sep 17 00:00:00 2001
From: "cyko@ibm.com;6J3007897;Irene Ko" <cyko1@ccc-login5.pok.ibm.com>
Date: Tue, 7 Apr 2026 14:12:57 -0400
Subject: [PATCH 3/5] debug the path for vllm engine boot

Signed-off-by: cyko@ibm.com;6J3007897;Irene Ko <cyko1@ccc-login5.pok.ibm.com>
---
 aisteer360/algorithms/core/steering_pipeline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/aisteer360/algorithms/core/steering_pipeline.py b/aisteer360/algorithms/core/steering_pipeline.py
index bd0b7b79..7479d712 100644
--- a/aisteer360/algorithms/core/steering_pipeline.py
+++ b/aisteer360/algorithms/core/steering_pipeline.py
@@ -217,7 +217,7 @@ def _steer_vllm(self, **steer_kwargs) -> None:
         """vLLM backend: steer on a temp HF model, then boot vLLM engine."""
         import gc
 
-        from aisteer360.adapters.vllm_hook.backend import boot_vllm_engine
+        from aisteer360.adapter_vllm_hook.backend import boot_vllm_engine
 
         # Load a temporary HF model to run steer(), 
         # same loop as the regular steer() for preparing the control

From ea5ac6a2e35a71d56ed49fcbd4b0db8af0752d44 Mon Sep 17 00:00:00 2001
From: "cyko@ibm.com;6J3007897;Irene Ko" <cyko1@ccc-login5.pok.ibm.com>
Date: Tue, 7 Apr 2026 15:25:24 -0400
Subject: [PATCH 4/5] core adapter for vllm hook and dummy test

Signed-off-by: cyko@ibm.com;6J3007897;Irene Ko <cyko1@ccc-login5.pok.ibm.com>
---
 aisteer360/adapter_vllm_hook/__init__.py |   0
 aisteer360/adapter_vllm_hook/backend.py  |  77 ++++++++++
 aisteer360/adapter_vllm_hook/recipe.py   | 123 +++++++++++++++
 aisteer360/adapter_vllm_hook/worker.py   | 187 +++++++++++++++++++++++
 test_vllm_state.py                       |  71 +++++++++
 5 files changed, 458 insertions(+)
 create mode 100644 aisteer360/adapter_vllm_hook/__init__.py
 create mode 100644 aisteer360/adapter_vllm_hook/backend.py
 create mode 100644 aisteer360/adapter_vllm_hook/recipe.py
 create mode 100644 aisteer360/adapter_vllm_hook/worker.py
 create mode 100644 test_vllm_state.py

diff --git a/aisteer360/adapter_vllm_hook/__init__.py b/aisteer360/adapter_vllm_hook/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/aisteer360/adapter_vllm_hook/backend.py b/aisteer360/adapter_vllm_hook/backend.py
new file mode 100644
index 00000000..4d3f628a
--- /dev/null
+++ b/aisteer360/adapter_vllm_hook/backend.py
@@ -0,0 +1,77 @@
+"""Boot a vLLM engine with AISteer360 steering applied via a custom worker."""
+from __future__ import annotations
+
+import json
+import logging
+import multiprocessing as mp
+import os
+import tempfile
+
+logger = logging.getLogger(__name__)
+
+
+def boot_vllm_engine(
+    model_name_or_path: str,
+    state_control,
+    vllm_kwargs: dict | None = None,
+) -> object:
+    """Create a HookLLM engine with the AISteer360 generic worker.
+
+    Serializes the post-steer ``state_control`` into a recipe, 
+    registers the :class:`AISteer360Worker` in vLLM-Hook's plugin registry, 
+    and boots a ``HookLLM`` instance.
+
+    Args:
+        model_name_or_path: HuggingFace model identifier or local path.
+        state_control: A post-steer StateControl instance.
+        vllm_kwargs: Extra kwargs forwarded to ``HookLLM``.
+
+    Returns:
+        A ``HookLLM`` instance ready for generation.
+    """
+    from .recipe import serialize_state_control
+
+    vllm_kwargs = dict(vllm_kwargs or {})
+    work_dir = tempfile.mkdtemp(prefix="aisteer360_vllm_")
+
+    # serialize state control
+    recipe = serialize_state_control(state_control, work_dir)
+    recipe["tokenizer_name_or_path"] = str(model_name_or_path)
+
+    config_path = os.path.join(work_dir, "aisteer360_recipe.json")
+    with open(config_path, "w") as f:
+        json.dump(recipe, f, indent=2)
+
+    os.environ["VLLM_AISTEER360_CONFIG"] = config_path
+
+    # vLLM env 
+    mp.set_start_method("spawn", force=True)
+    os.environ.setdefault("VLLM_USE_V1", "1")
+    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
+
+    # register worker in vLLM-Hook's plugin registry 
+    import vllm.plugins
+    from vllm_hook_plugins import PluginRegistry
+    vllm.plugins.load_general_plugins()
+
+    from .worker import AISteer360Worker
+    PluginRegistry.register_worker("aisteer360", AISteer360Worker)
+
+    # boot HookLLM 
+    from vllm_hook_plugins import HookLLM
+
+    engine_kwargs = {
+        "gpu_memory_utilization": 0.9,
+        "enforce_eager": True,
+        "enable_prefix_caching": True,
+    }
+    engine_kwargs.update(vllm_kwargs)
+
+    logger.info("Booting vLLM engine with AISteer360Worker for %s", model_name_or_path)
+    return HookLLM(
+        model=str(model_name_or_path),
+        worker_name="aisteer360",
+        config_file=config_path,
+        enable_hook=True,
+        **engine_kwargs,
+    )
diff --git a/aisteer360/adapter_vllm_hook/recipe.py b/aisteer360/adapter_vllm_hook/recipe.py
new file mode 100644
index 00000000..d9628bfc
--- /dev/null
+++ b/aisteer360/adapter_vllm_hook/recipe.py
@@ -0,0 +1,123 @@
+"""Serialize and reconstruct StateControl objects across process boundaries.
+
+When using the vLLM backend, the StateControl must cross from the main process into vLLM-Hook's worker subprocess.  
+We serialize a lightweight recipe (the class path, constructor arguments, and pre-computed steering vectors)
+then reconstruct and re-steer inside the worker where the real model lives.
+"""
+from __future__ import annotations
+
+import importlib
+import json
+import logging
+import os
+from dataclasses import fields as dc_fields
+
+from aisteer360.algorithms.state_control.common.steering_vector import SteeringVector
+
+logger = logging.getLogger(__name__)
+
+def serialize_state_control(state_control, work_dir: str) -> dict:
+    """Serialize a post-steer StateControl into a JSON-safe recipe.
+
+    Args:
+        state_control: A steered StateControl instance.
+        work_dir: Directory for temporary ``.svec`` files.
+
+    Returns:
+        Recipe dict suitable for ``json.dump()``.
+    """
+    recipe: dict = {
+        "class": f"{type(state_control).__module__}.{type(state_control).__qualname__}",
+        "constructor_args": {},
+        "steering_vectors": {},
+        "post_steer_state": {},
+    }
+
+    # serialize constructor args from the Args dataclass
+    if hasattr(state_control, "args") and state_control.args is not None:
+        for f in dc_fields(state_control.args):
+            value = getattr(state_control.args, f.name)
+            if isinstance(value, SteeringVector):
+                svec_path = os.path.join(work_dir, f"arg_{f.name}.svec")
+                value.save(svec_path)
+                recipe["steering_vectors"][f.name] = svec_path
+            elif _is_json_safe(value):
+                recipe["constructor_args"][f.name] = value
+            else:
+                recipe["constructor_args"][f.name] = None
+
+    # save post-steer computed steering vector
+    sv = getattr(state_control, "_steering_vector", None)
+    if isinstance(sv, SteeringVector):
+        svec_path = os.path.join(work_dir, "post_steer_vector.svec")
+        sv.save(svec_path)
+        recipe["post_steer_state"]["steering_vector_path"] = svec_path
+
+    # save resolved layer info
+    for attr in ("_layer_id", "_layer_names"):
+        val = getattr(state_control, attr, None)
+        if val is not None:
+            recipe["post_steer_state"][attr] = val
+
+    return recipe
+
+
+def reconstruct_state_control(recipe: dict):
+    """Reconstruct a StateControl from a recipe dict.
+
+    The returned control has its steering vector pre-loaded so that
+    ``steer()`` in the worker only needs to resolve layer names and
+    build the transform (no prompt-pair extraction needed).
+
+    Args:
+        recipe: Recipe dict produced by :func:`serialize_state_control`.
+
+    Returns:
+        A StateControl instance ready for ``steer(model)``.
+    """
+    # import the class
+    class_path = recipe["class"]
+    module_path, class_name = class_path.rsplit(".", 1)
+    module = importlib.import_module(module_path)
+    cls = getattr(module, class_name)
+
+    # rebuild constructor kwargs 
+    kwargs = dict(recipe.get("constructor_args", {}))
+
+    # Load steering vectors saved from constructor args
+    for field_name, svec_path in recipe.get("steering_vectors", {}).items():
+        kwargs[field_name] = SteeringVector.load(svec_path)
+
+    # If prompts were used originally but we have a post-steer vector,
+    # swap to vector mode (skip extraction in the worker)
+    post_steer = recipe.get("post_steer_state", {})
+    sv_path = post_steer.get("steering_vector_path")
+    if sv_path and kwargs.get("steering_vector") is None:
+        kwargs["steering_vector"] = SteeringVector.load(sv_path)
+        # Remove prompt args to satisfy the "exactly one source" validation
+        kwargs.pop("positive_prompt", None)
+        kwargs.pop("negative_prompt", None)
+
+    # Filter out None values for optional args that weren't serializable
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+
+    control = cls(**kwargs)
+
+    # Restore post-steer state so the worker's steer() is lightweight
+    if "_layer_id" in post_steer:
+        control._layer_id = post_steer["_layer_id"]
+    if "_layer_names" in post_steer:
+        control._layer_names = post_steer["_layer_names"]
+
+    return control
+
+
+def _is_json_safe(value) -> bool:
+    """Check if a value can be JSON-serialized."""
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return True
+    if isinstance(value, (list, tuple)):
+        return all(_is_json_safe(v) for v in value)
+    if isinstance(value, dict):
+        return all(isinstance(k, str) and _is_json_safe(v) for k, v in value.items())
+    return False
diff --git a/aisteer360/adapter_vllm_hook/worker.py b/aisteer360/adapter_vllm_hook/worker.py
new file mode 100644
index 00000000..ea9b2cc8
--- /dev/null
+++ b/aisteer360/adapter_vllm_hook/worker.py
@@ -0,0 +1,187 @@
+"""AISteer360 generic worker for vLLM-Hook.
+
+Reconstructs any AISteer360 :class:`StateControl` inside the vLLM worker
+subprocess and registers hooks on the PyTorch model. 
+"""
+from __future__ import annotations
+
+import json
+import logging
+import os
+
+import torch
+from vllm.v1.worker.gpu_worker import Worker as V1Worker
+
+logger = logging.getLogger(__name__)
+
+
+class AISteer360Worker(V1Worker):
+    """vLLM-Hook worker that applies any AISteer360 StateControl."""
+
+    def load_model(self, *args, **kwargs):
+        """Load the model, then reconstruct and apply the StateControl."""
+        result = super().load_model(*args, **kwargs)
+
+        config_path = os.environ.get("VLLM_AISTEER360_CONFIG")
+        if not config_path:
+            logger.warning("VLLM_AISTEER360_CONFIG not set; no steering applied.")
+            self._state_control = None
+            return result
+
+        try:
+            self._install_aisteer360_hooks(config_path)
+            logger.info("AISteer360 hooks installed successfully")
+        except Exception:
+            logger.exception("AISteer360 hook installation failed")
+            self._state_control = None
+
+        return result
+
+    def _install_aisteer360_hooks(self, config_path: str) -> None:
+        """Read recipe, reconstruct control, steer, and register hooks."""
+        from aisteer360.adapter_vllm_hook.recipe import reconstruct_state_control
+
+        with open(config_path) as f:
+            recipe = json.load(f)
+
+        # Reconstruct the StateControl (with pre-computed vectors)
+        state_control = reconstruct_state_control(recipe)
+
+        # Access vLLM's internal PyTorch model
+        model = getattr(self.model_runner, "model", None)
+        if model is None:
+            raise RuntimeError("Could not access model_runner.model")
+
+        # Load tokenizer if needed for steer()
+        tokenizer = None
+        tokenizer_path = recipe.get("tokenizer_name_or_path")
+        if tokenizer_path:
+            from transformers import AutoTokenizer
+            tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
+
+        # Run steer() on the real model — reuses AISteer360's full logic
+        state_control.steer(model, tokenizer=tokenizer)
+
+        # Get hooks from the state control
+        dummy_ids = torch.zeros(1, 1, dtype=torch.long, device="cuda")
+        hooks = state_control.get_hooks(dummy_ids, runtime_kwargs={})
+
+        # Wrap each hook function with the flag file gate so that
+        # HookLLM's use_hook=False (which removes the flag file) disables steering.
+        hook_flag = os.environ.get("VLLM_HOOK_FLAG")
+        for phase in ("pre", "forward", "backward"):
+            for spec in hooks.get(phase, []):
+                spec["hook_func"] = _gated_hook(spec["hook_func"], hook_flag, phase)
+
+        state_control.set_hooks(hooks)
+        state_control._model_ref = model
+        state_control.register_hooks(model)
+
+        # Store reference
+        self._state_control = state_control
+        self.hook_flag = hook_flag
+
+        logger.info(
+            "Registered %d hooks from %s",
+            len(state_control.registered),
+            type(state_control).__name__,
+        )
+
+    def execute_model(self, *args, **kwargs):
+        return super().execute_model(*args, **kwargs)
+
+
+def _is_gated_out(flag_path: str | None) -> bool:
+    """Return True if steering is disabled by a missing flag file."""
+    return bool(flag_path) and not os.path.exists(flag_path)
+
+
+def _reshape_hidden_args(args, kwargs, *, squeeze: bool):
+    """Reshape hidden_states between 2-D ``[N, H]`` and 3-D ``[B, T, H]``.
+
+    When *squeeze* is False, unsqueezes 2-D → 3-D.
+    When *squeeze* is True, squeezes 3-D (batch=1) → 2-D.
+    """
+    from aisteer360.algorithms.state_control.common.hook_utils import (
+        _hidden_states_index,
+        extract_hidden_states,
+    )
+    hidden = extract_hidden_states(args, kwargs)
+    if hidden is None:
+        return args, kwargs
+
+    if squeeze:
+        needs_reshape = hidden.ndim == 3 and hidden.size(0) == 1
+        op = torch.Tensor.squeeze
+    else:
+        needs_reshape = hidden.ndim == 2
+        op = torch.Tensor.unsqueeze
+
+    if not needs_reshape:
+        return args, kwargs
+
+    new_hidden = op(hidden, 0)
+    if args:
+        idx = _hidden_states_index(args)
+        if idx is not None:
+            args = (*args[:idx], new_hidden, *args[idx + 1:])
+    elif "hidden_states" in kwargs:
+        kwargs = {**kwargs, "hidden_states": new_hidden}
+    return args, kwargs
+
+
+def _reshape_output(output, *, squeeze: bool):
+    """Reshape tensors in forward-hook output between 2-D and 3-D."""
+    if output is None:
+        return None
+
+    if squeeze:
+        check = lambda t: t.ndim == 3 and t.size(0) == 1
+        op = lambda t: t.squeeze(0)
+    else:
+        check = lambda t: t.ndim == 2
+        op = lambda t: t.unsqueeze(0)
+
+    if isinstance(output, torch.Tensor):
+        return op(output) if check(output) else output
+    if isinstance(output, tuple):
+        return tuple(op(t) if isinstance(t, torch.Tensor) and check(t) else t for t in output)
+    return output
+
+
+def _gated_hook(original_hook, flag_path: str | None, phase: str):
+    """Wrap a hook with flag-file gating and vLLM tensor shape normalization.
+
+    AISteer360 hooks expect 3-D hidden states ``[B, T, H]`` but vLLM
+    passes 2-D ``[N, H]``.  This wrapper unsqueezes inputs before the
+    original hook and squeezes outputs back, keeping the original hooks
+    untouched.
+    """
+    if flag_path is None and phase not in ("pre", "forward"):
+        return original_hook
+
+    if phase == "pre":
+        def wrapper(module, args, kwargs):
+            if _is_gated_out(flag_path):
+                return None
+            args, kwargs = _reshape_hidden_args(args, kwargs, squeeze=False)
+            result = original_hook(module, args, kwargs)
+            if result is None:
+                return None
+            return _reshape_hidden_args(*result, squeeze=True)
+        return wrapper
+
+    if phase == "forward":
+        def wrapper(module, args, kwargs, output):
+            if _is_gated_out(flag_path):
+                return None
+            output = _reshape_output(output, squeeze=False)
+            result = original_hook(module, args, kwargs, output)
+            return _reshape_output(result, squeeze=True)
+        return wrapper
+
+    def wrapper(*args, **kwargs):
+        if _is_gated_out(flag_path):
+            return None
+        return original_hook(*args, **kwargs)
+    return wrapper
diff --git a/test_vllm_state.py b/test_vllm_state.py
new file mode 100644
index 00000000..84aa86b2
--- /dev/null
+++ b/test_vllm_state.py
@@ -0,0 +1,71 @@
+import multiprocessing as mp
+
+MODEL = "/dccstor/pyrite/irene/models--meta-llama--Llama-3.1-8B-Instruct/snapshots/0e9e39f249a16976918f6564b8830bc894c89659"
+
+
+def main():
+    from aisteer360.algorithms.core.steering_pipeline import SteeringPipeline
+    from aisteer360.algorithms.state_control.act_add import ActAdd
+    from aisteer360.algorithms.state_control.caa import CAA
+    from aisteer360.algorithms.state_control.iti import ITI
+
+    # add backend="vllm"
+    pipe = SteeringPipeline(
+        model_name_or_path=MODEL,
+        controls=[ActAdd(positive_prompt="Love", negative_prompt="Hate", multiplier=5.0)],
+        backend="vllm",
+        vllm_kwargs={"gpu_memory_utilization": 0.9},
+    )
+    # pipe = SteeringPipeline(
+    #     model_name_or_path=MODEL,
+    #     controls=[CAA(                                                                                         
+    #       data={  
+    #           "positives": ["Love", "Kindness", "Joy", "Compassion"],                                         
+    #           "negatives": ["Hate", "Cruelty", "Sadness", "Indifference"],                                    
+    #       },                                                                                                 
+    #       multiplier=5.0,                                                                                    
+    #       token_scope="all",  # important for vLLM — "after_prompt" won't work                               
+    #   )],
+    #   backend="vllm",
+    #     vllm_kwargs={"gpu_memory_utilization": 0.9},
+    # )
+    # pipe = SteeringPipeline(
+    #     model_name_or_path=MODEL,
+    #     controls=[ITI(                                                                                             
+    #         data={                                                                                                 
+    #             "positives": [                                                                                     
+    #                 "The sky is blue.",
+    #                 "Water boils at 100 degrees Celsius.",                                                         
+    #                 "The Earth orbits the Sun.",                                                                   
+    #                 "Humans need oxygen to breathe.",                                                              
+    #             ],                                                                                                 
+    #             "negatives": [                                                                                     
+    #                 "The sky is green.",
+    #                 "Water boils at 50 degrees Celsius.",                                                          
+    #                 "The Sun orbits the Earth.",                                                                   
+    #             ],
+    #         },                                                                                                     
+    #         alpha=15.0, 
+    #         num_heads=48,                                                                                          
+    #         token_scope="all",  # important for vLLM
+    #     )],
+    #     backend="vllm",
+    #     vllm_kwargs={"gpu_memory_utilization": 0.9},
+    # )
+    pipe.steer()
+
+    # Compare steered vs baseline
+    ids = pipe.tokenizer(["Tell me a story."], return_tensors="pt")["input_ids"]
+    steered = pipe.generate_text(ids, use_hook=True)
+    pipe._vllm_engine.llm_engine.reset_prefix_cache()
+    baseline = pipe.generate_text(ids, use_hook=False)
+
+    print("\n=== Steered ===")
+    print(steered)
+    print("\n=== Baseline ===")
+    print(baseline)
+
+
+if __name__ == "__main__":
+    mp.set_start_method("spawn", force=True)
+    main()

From 8cf598b57a49ca80937aae5a9093a3ceaef300a2 Mon Sep 17 00:00:00 2001
From: "cyko@ibm.com;6J3007897;Irene Ko" <cyko1@ccc-login5.pok.ibm.com>
Date: Tue, 7 Apr 2026 15:53:53 -0400
Subject: [PATCH 5/5] Added dependency; remove dummy absolute path

Signed-off-by: cyko@ibm.com;6J3007897;Irene Ko <cyko1@ccc-login5.pok.ibm.com>
---
 pyproject.toml     | 5 +++++
 test_vllm_state.py | 2 +-
 2 files changed, 6 insertions(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index de7c57d1..5b933bd5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -71,6 +71,11 @@ docs = [
     "mkdocs-bibtex",
 ]
 
+vllm = [
+    "vllm>=0.8.0",
+    "vllm-hook-plugins",
+]
+
 dev = [
     "pytest>=8.3.2,<9.0.0",
     "pre-commit>=4.3.0"
diff --git a/test_vllm_state.py b/test_vllm_state.py
index 84aa86b2..1a1df3f7 100644
--- a/test_vllm_state.py
+++ b/test_vllm_state.py
@@ -1,6 +1,6 @@
 import multiprocessing as mp
 
-MODEL = "/dccstor/pyrite/irene/models--meta-llama--Llama-3.1-8B-Instruct/snapshots/0e9e39f249a16976918f6564b8830bc894c89659"
+MODEL = ""
 
 
 def main():