{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cell-0",
   "metadata": {},
   "source": [
    "# Extract Residual Stream Activations\n",
    "\n",
    "This notebook demonstrates how to capture residual-stream activations during\n",
    "vLLM inference using the async engine and vllm-lens."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-1",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "cell-2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from vllm import AsyncEngineArgs, AsyncLLMEngine, SamplingParams"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-3",
   "metadata": {},
   "source": [
    "## Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "cell-4",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL = \"facebook/opt-125m\"\n",
    "PROMPT = \"The future of AI is\"\n",
    "LAYERS = [0, 2, 5]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-5",
   "metadata": {},
   "source": [
    "## Initialize the Async Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cell-6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b0314366850e4b35b06579ca676ff1ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/651 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 03-14 09:58:52 [model.py:531] Resolved architecture: OPTForCausalLM\n",
      "INFO 03-14 09:58:52 [model.py:1554] Using max model len 2048\n",
      "INFO 03-14 09:58:52 [scheduler.py:231] Chunked prefill is enabled with max_num_batched_tokens=2048.\n",
      "INFO 03-14 09:58:52 [vllm.py:747] Asynchronous scheduling is enabled.\n",
      "WARNING 03-14 09:58:52 [vllm.py:781] Enforce eager set, disabling torch.compile and CUDAGraphs. This is equivalent to setting -cc.mode=none -cc.cudagraph_mode=none\n",
      "WARNING 03-14 09:58:52 [vllm.py:792] Inductor compilation was disabled by user settings, optimizations settings that are only active during inductor compilation will be ignored.\n",
      "INFO 03-14 09:58:52 [vllm.py:957] Cudagraph is disabled under eager mode\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "538c064752844ee1a9b81684bfb50490",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/685 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c57384c7470d47e3a5f1a59780b71f97",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.json: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc14443ba2ae4d6ba04db25906cd7b5b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "merges.txt: 0.00B [00:00, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d9bd8e05ea26490ea8b46861eb1740c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "special_tokens_map.json:   0%|          | 0.00/441 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "029fa4e43c064139b4986790a2faf8ee",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "generation_config.json:   0%|          | 0.00/137 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:58:55 [core.py:101] Initializing a V1 LLM engine (v0.17.0) with config: model='facebook/opt-125m', speculative_config=None, tokenizer='facebook/opt-125m', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=2048, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, data_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=True, enable_return_routed_experts=False, kv_cache_dtype=auto, device_config=cuda, structured_outputs_config=StructuredOutputsConfig(backend='auto', disable_fallback=False, disable_any_whitespace=False, disable_additional_properties=False, reasoning_parser='', reasoning_parser_plugin='', enable_in_reasoning=False), observability_config=ObservabilityConfig(show_hidden_metrics_for_version=None, otlp_traces_endpoint=None, collect_detailed_traces=None, kv_cache_metrics=False, kv_cache_metrics_sample=0.01, cudagraph_metrics=False, enable_layerwise_nvtx_tracing=False, enable_mfu_metrics=False, enable_mm_processor_stats=False, enable_logging_iteration_details=False), seed=0, served_model_name=facebook/opt-125m, enable_prefix_caching=True, enable_chunked_prefill=True, pooler_config=None, compilation_config={'level': None, 'mode': <CompilationMode.NONE: 0>, 'debug_dump_path': None, 'cache_dir': '', 'compile_cache_save_format': 'binary', 'backend': 'inductor', 'custom_ops': ['all'], 'splitting_ops': [], 'compile_mm_encoder': False, 'compile_sizes': [], 'compile_ranges_split_points': [2048], 'inductor_compile_config': {'enable_auto_functionalized_v2': False, 'combo_kernels': True, 'benchmark_combo_kernel': True}, 'inductor_passes': {}, 'cudagraph_mode': <CUDAGraphMode.NONE: 0>, 'cudagraph_num_of_warmups': 0, 'cudagraph_capture_sizes': [], 'cudagraph_copy_inputs': False, 'cudagraph_specialize_lora': True, 'use_inductor_graph_partition': False, 'pass_config': {'fuse_norm_quant': True, 'fuse_act_quant': True, 'fuse_attn_quant': False, 'enable_sp': False, 'fuse_gemm_comms': False, 'fuse_allreduce_rms': False}, 'max_cudagraph_capture_size': 0, 'dynamic_shapes_config': {'type': <DynamicShapesType.BACKED: 'backed'>, 'evaluate_guards': False, 'assume_32_bit_indexing': False}, 'local_cache_dir': None, 'fast_moe_cold_start': True, 'static_all_moe_layers': []}\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:58:55 [worker_base.py:283] Injected <class 'vllm_lens._worker_ext.HiddenStatesExtension'> into <class 'vllm.v1.worker.gpu_worker.Worker'> for extended collective_rpc calls ['clear_captured_states', 'clear_steering_data', 'get_captured_states', 'install_hooks', 'set_steering_data']\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:58:58 [parallel_state.py:1393] world_size=1 rank=0 local_rank=0 distributed_init_method=tcp://10.242.3.185:60453 backend=nccl\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:58:58 [parallel_state.py:1715] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, PCP rank 0, TP rank 0, EP rank N/A, EPLB rank N/A\n",
      "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n",
      "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n",
      "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n",
      "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n",
      "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n",
      "[Gloo] Rank 0 is connected to 0 peer ranks. Expected number of connected peer ranks is : 0\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:58:58 [base.py:106] Offloader set to NoopOffloader\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:58:58 [gpu_model_runner.py:4255] Starting to load model facebook/opt-125m...\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:58:59 [cuda.py:405] Using FLASH_ATTN attention backend out of potential backends: ['FLASH_ATTN', 'FLASHINFER', 'TRITON_ATTN', 'FLEX_ATTENTION'].\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:58:59 [flash_attn.py:587] Using FlashAttention version 3\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:59:03 [weight_utils.py:561] Time spent downloading weights for facebook/opt-125m: 3.814851 seconds\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7342843d498648ba873e20a6e6f4326d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading pt checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:59:04 [default_loader.py:293] Loading weights took 0.14 seconds\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:59:04 [gpu_model_runner.py:4338] Model loading took 0.24 GiB memory and 5.151460 seconds\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:59:09 [gpu_worker.py:424] Available KV cache memory: 80.23 GiB\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:59:09 [kv_cache_utils.py:1314] GPU KV cache size: 2,336,944 tokens\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:59:09 [kv_cache_utils.py:1319] Maximum concurrency for 2,048 tokens per request: 1141.09x\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:59:10 [core.py:282] init engine (profile, create kv cache, warmup model) took 5.54 seconds\n",
      "(EngineCore_DP0 pid=278813) INFO 03-14 09:59:10 [vllm.py:957] Cudagraph is disabled under eager mode\n",
      "WARNING 03-14 09:59:10 [vllm.py:781] Enforce eager set, disabling torch.compile and CUDAGraphs. This is equivalent to setting -cc.mode=none -cc.cudagraph_mode=none\n",
      "(EngineCore_DP0 pid=278813) WARNING 03-14 09:59:10 [vllm.py:792] Inductor compilation was disabled by user settings, optimizations settings that are only active during inductor compilation will be ignored.\n",
      "(EngineCore_DP0 pid=278813) "
     ]
    }
   ],
   "source": [
    "engine_args = AsyncEngineArgs(\n",
    "    model=MODEL,\n",
    "    gpu_memory_utilization=0.85,\n",
    ")\n",
    "engine = AsyncLLMEngine.from_engine_args(engine_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-7",
   "metadata": {},
   "source": [
    "## Generate with Activation Capture\n",
    "\n",
    "Pass `output_residual_stream` via `extra_args` to request residual-stream\n",
    "activations at the specified layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cell-8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING 03-14 09:59:10 [input_processor.py:254] Passing raw prompts to InputProcessor is deprecated and will be removed in v0.18. You should instead pass the outputs of Renderer.render_cmpl() or Renderer.render_chat().\n"
     ]
    }
   ],
   "source": [
    "sampling_params = SamplingParams(\n",
    "    temperature=0.7,\n",
    "    max_tokens=64,\n",
    "    extra_args={\"output_residual_stream\": LAYERS},\n",
    ")\n",
    "\n",
    "final = None\n",
    "async for output in engine.generate(PROMPT, sampling_params, \"req-0\"):\n",
    "    final = output"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cell-9",
   "metadata": {},
   "source": [
    "## Inspect Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cell-10",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Prompt: The future of AI is\n",
      "Generated:  cloud computing\n",
      "\n",
      "Jose Antonia\n",
      "\n",
      "AI is destined to become a future of computing. The future of computing will be cloud computing, where it is no longer an urban fantasy.\n",
      "\n",
      "The future of computing will be cloud computing, where it is no longer an urban fantasy. The future of computing will be cloud\n",
      "\n",
      "Residual stream shape: torch.Size([3, 69, 768])\n",
      "Sample activations (layer 0, first token): tensor([-0.1010, -1.0039,  0.1498,  0.1503,  0.0785], dtype=torch.float16)\n"
     ]
    }
   ],
   "source": [
    "print(f\"Prompt: {PROMPT}\")\n",
    "print(f\"Generated: {final.outputs[0].text}\")\n",
    "\n",
    "activation_data = getattr(final, \"activations\", None)\n",
    "if activation_data is not None:\n",
    "    residual_stream = activation_data[\"residual_stream\"]\n",
    "    print(f\"\\nResidual stream shape: {residual_stream.shape}\")\n",
    "    print(f\"Sample activations (layer 0, first token): {residual_stream[0, 0, :5]}\")\n",
    "else:\n",
    "    print(\"Activations: NOT FOUND (ERROR)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45324e92",
   "metadata": {},
   "source": "## Persistent Hooks (Garçon-style)\n\nInstead of requesting activations per-call, you can register a persistent hook\nthat captures data across many requests and collect results in bulk afterwards."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "943a1a1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from vllm_lens import Hook\n",
    "\n",
    "\n",
    "# Define a hook that accumulates activations across forward passes\n",
    "def capture_residual(ctx, h):\n",
    "    if \"parts\" not in ctx.saved:\n",
    "        ctx.saved[\"parts\"] = []\n",
    "    ctx.saved[\"parts\"].append(h.cpu())\n",
    "    return None  # don't modify hidden states\n",
    "\n",
    "\n",
    "hook = Hook(fn=capture_residual, layer_indices=LAYERS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30e05d4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "import cloudpickle\n",
    "\n",
    "# Register the hook — it will fire on every subsequent request\n",
    "await engine.collective_rpc(\"install_hooks\")\n",
    "await engine.collective_rpc(\n",
    "    \"set_persistent_hooks\",\n",
    "    args=(cloudpickle.dumps([hook]),),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10e87b52",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run several prompts — hooks fire automatically, results stay server-side\n",
    "prompts = [\n",
    "    \"The future of AI is\",\n",
    "    \"In the beginning, there was\",\n",
    "    \"The meaning of life is\",\n",
    "]\n",
    "sp = SamplingParams(temperature=0.7, max_tokens=16)\n",
    "\n",
    "for i, prompt in enumerate(prompts):\n",
    "    async for output in engine.generate(prompt, sp, f\"persistent-{i}\"):\n",
    "        if output.finished:\n",
    "            print(f\"[{i}] {prompt} → {output.outputs[0].text!r}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc1c7f86",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# Collect all results at once — one bulk transfer\n",
    "raw_results = await engine.collective_rpc(\"get_all_hook_results\")\n",
    "for raw in raw_results or ():\n",
    "    if raw is not None:\n",
    "        all_results = pickle.loads(raw)\n",
    "        break\n",
    "\n",
    "print(f\"Collected results for {len(all_results)} requests\")\n",
    "for req_id, hook_data in all_results.items():\n",
    "    parts = hook_data[\"0\"][\"parts\"]\n",
    "    combined = torch.cat(parts, dim=0)\n",
    "    print(f\"  {req_id}: {len(parts)} forward passes, {combined.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d3c05cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up\n",
    "await engine.collective_rpc(\"clear_persistent_hooks\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vllm-lens (3.12.13)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}