undercurrent 0.1.0__tar.gz → 0.2.0__tar.gz

{undercurrent-0.1.0 → undercurrent-0.2.0}/CHANGELOG.md

@@ -18,6 +18,74 @@ how a release is cut.
 
 ## [Unreleased]
 
+## [0.2.0] - 2026-11-01
+
+Performance and correctness work on the vLLM adapter, and experimental
+graph-mode `residual_stream` capture that keeps CUDA graphs and
+`torch.compile`. There are no breaking changes: code written for 0.1.0 runs
+unchanged.
+
+### Added
+
+- vLLM adapter (experimental): graph-mode `residual_stream` capture.
+  `load_model(graph_capture_layers=[...])` (also accepted by
+  `ProbedModel(..., backend="vllm")`) captures the residual stream at the given
+  decoder layers without `enforce_eager=True`, so the engine keeps CUDA graphs
+  and `torch.compile`. It reuses vLLM's EAGLE-3 auxiliary hidden-state outputs,
+  enabled by the `vllm.general_plugins` entry point before compilation. Other
+  tensor types, other layers, models without that interface, and speculative
+  decoding that uses aux outputs are rejected with an error. The benchmark
+  harness gains `--graph-capture`.
+
+### Changed
+
+- vLLM adapter: activation capture on CUDA no longer blocks the host once per
+  captured row. Each forward hook gathers its matched rows on the GPU and
+  starts one non-blocking copy into pinned host memory; records are built after
+  one wait per step (at the end of the step in-process, or when the driver
+  polls on the cross-process path). Probes see the same records, in the same
+  order, with inline aborts taking effect at the same step as before. See
+  "Copying captures to the host" in the vLLM adapter internals docs.
+- vLLM adapter: on the cross-process polling path (vLLM 0.28's default), the
+  worker sends each poll's activations as one packed batch (float32 `bytes`
+  plus per-field lists) instead of one dict per record with the tensor as a
+  nested Python list, and the driver routes each drain in one executor call
+  instead of one per record. In a CPU microbenchmark at a 3,584-wide hidden
+  size, transport fell from about 350 µs to under 15 µs per record. Probes see
+  the same records in the same order; aborts are sent after the drain is
+  routed.
+- vLLM adapter: forward hooks skip (layer, tensor type) pairs no request
+  captures with one set lookup, and position matching runs once per step
+  instead of once per hooked layer; the captured tensor is only extracted when
+  something matched. Hook-path CPU time per step fell by about a third in a
+  microbenchmark (16 requests x 8 probes, 28 layers).
+
+### Fixed
+
+- vLLM adapter: `ActivationRecord.timestamp` is now the capture time, as
+  documented, on every path. On the cross-process polling path the record was
+  rebuilt in the driver with a fresh timestamp, so a probe (or a latency
+  measurement) saw the time of routing, not of capture, and couldn't tell how
+  stale an activation was. The worker now stamps each capture when its hook
+  runs and the packed batch carries it to the driver.
+- vLLM adapter: with `OverflowPolicy.BLOCK` on the cross-process polling path,
+  async probes ran one at a time instead of in parallel. Each drained backlog
+  was routed one request at a time, so `route()` waited on one request's full
+  queue while every other request's worker sat idle. Drained records are now
+  interleaved round-robin across requests, keeping each request's forward-pass
+  order. On an A100 (Qwen2.5-7B, 16 requests, 76 ms probe) the old order made
+  `block` about 13x slower than the probe's own pace. The docs now also say
+  that on this path `block` delays results rather than slowing the engine.
+- vLLM adapter: activations could be silently dropped under concurrency on the
+  polling path (engine core in a subprocess, the vLLM 0.28 default). One
+  request's drain of the shared worker buffer could still be routing another
+  request's records when that request received its final output and ended, so
+  those records were discarded as late. A request that finishes in the step that
+  captured its only extraction point (e.g. `prompt[-1]` with one new token) lost
+  its capture entirely; longer generations could lose their last activations.
+  Drains are now serialised, so every record captured up to a request's final
+  step is routed before the request ends.
+
 ## [0.1.0] - 2026-10-02
 
 The first public release of **Undercurrent**, Wrynx's activation-probing
@@ -181,5 +249,6 @@ for details and workarounds.
   `state_dict`s, and refuses to load on torch older than 2.6 (the first
   release with a fix for CVE-2025-32434).
 
-[Unreleased]: https://github.com/wrynx/undercurrent/compare/v0.1.0...HEAD
+[Unreleased]: https://github.com/wrynx/undercurrent/compare/v0.2.0...HEAD
+[0.2.0]: https://github.com/wrynx/undercurrent/compare/v0.1.0...v0.2.0
 [0.1.0]: https://github.com/wrynx/undercurrent/releases/tag/v0.1.0

{undercurrent-0.1.0 → undercurrent-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: undercurrent
-Version: 0.1.0
+Version: 0.2.0
 Summary: Undercurrent: production-grade LLM activation probing and intervention. Works with Hugging Face Transformers and vLLM.
 Project-URL: Homepage, https://github.com/wrynx/undercurrent
 Project-URL: Source, https://github.com/wrynx/undercurrent
@@ -43,7 +43,7 @@ Requires-Dist: pre-commit>=4; extra == 'dev'
 Requires-Dist: pytest-cov>=5; extra == 'dev'
 Requires-Dist: pytest-timeout>=2.3; extra == 'dev'
 Requires-Dist: pytest>=7; extra == 'dev'
-Requires-Dist: ruff<0.17,>=0.16; extra == 'dev'
+Requires-Dist: ruff<0.18,>=0.16; extra == 'dev'
 Requires-Dist: types-pyyaml; extra == 'dev'
 Provides-Extra: docs
 Requires-Dist: mkdocs-material; extra == 'docs'
@@ -94,13 +94,13 @@ without adding latency to the response.
   [Deploy with vLLM](https://wrynx.github.io/undercurrent/guides/vllm-deployment/)),
   and with Hugging Face Transformers on a laptop CPU.
 
-> **Known issues (v0.1):** the vLLM adapter is experimental: its GPU tests
+> **Known issues (v0.2):** the vLLM adapter is experimental: its GPU tests
 > pass on an NVIDIA L4 with vLLM 0.28.0, run by hand (there's no GPU CI yet),
 > and only vLLM 0.28 is supported. Fixed in 0.1.0: on vLLM, `residual_stream`
 > on Llama-family (fused-residual) models is captured as
 > `hidden_states + residual`, the same residual stream as the HF backend
 > (GPU-verified). See
-> [Known issues](https://wrynx.github.io/undercurrent/compatibility/#known-issues-v01).
+> [Known issues](https://wrynx.github.io/undercurrent/compatibility/#known-issues-v02).
 
 ## How it works
 

{undercurrent-0.1.0 → undercurrent-0.2.0}/README.md

@@ -37,13 +37,13 @@ without adding latency to the response.
   [Deploy with vLLM](https://wrynx.github.io/undercurrent/guides/vllm-deployment/)),
   and with Hugging Face Transformers on a laptop CPU.
 
-> **Known issues (v0.1):** the vLLM adapter is experimental: its GPU tests
+> **Known issues (v0.2):** the vLLM adapter is experimental: its GPU tests
 > pass on an NVIDIA L4 with vLLM 0.28.0, run by hand (there's no GPU CI yet),
 > and only vLLM 0.28 is supported. Fixed in 0.1.0: on vLLM, `residual_stream`
 > on Llama-family (fused-residual) models is captured as
 > `hidden_states + residual`, the same residual stream as the HF backend
 > (GPU-verified). See
-> [Known issues](https://wrynx.github.io/undercurrent/compatibility/#known-issues-v01).
+> [Known issues](https://wrynx.github.io/undercurrent/compatibility/#known-issues-v02).
 
 ## How it works
 

{undercurrent-0.1.0 → undercurrent-0.2.0}/benchmarks/README.md

@@ -134,6 +134,7 @@ or surprising runs too: the page must show what was measured. Don't commit
 | `--dtype {auto,float32,float16,bfloat16}` | `auto` | Model dtype (hf: `from_pretrained`; vllm: engine `dtype`) |
 | `--vllm-arg KEY=VALUE` | none | vllm only, repeatable: engine arg passed to both the baseline `vllm.LLM` and the probed engine. VALUE is parsed as JSON when it can be |
 | `--isolate {auto,process,none}` | `auto` | Run each scenario in its own subprocess (auto: yes for vllm, no for hf) |
+| `--graph-capture` | off | vllm only: run probed scenarios in graph mode (CUDA graphs, no `enforce_eager`), capturing `residual_stream` at the scenario's layers via `load_model(graph_capture_layers=...)`. Passing `--vllm-arg enforce_eager=true` still applies to `baseline`, so run the graph baseline separately |
 
 ## Scenarios
 
@@ -268,7 +269,7 @@ One JSON object.
 | `run_id` | string | The run directory's name |
 | `created_utc` | string | ISO 8601 UTC start time |
 | `cli_argv` | list of strings | The exact command-line arguments (after `run.py`) |
-| `args` | object | Every resolved option (after `--dry-run` and backend defaults), keyed by argparse destination: `backend`, `model`, `num_prompts`, `batch_size`, `max_new_tokens`, `input_len`, `scenarios`, `warmup`, `repeats`, `out`, `dry_run`, `seed`, `probe_delay_ms`, `slow_queue_depth`, `flat_tolerance_pct`, `worker_pool_size`, `device`, `dtype`, `vllm_engine_kwargs`, `isolate` |
+| `args` | object | Every resolved option (after `--dry-run` and backend defaults), keyed by argparse destination: `backend`, `model`, `num_prompts`, `batch_size`, `max_new_tokens`, `input_len`, `scenarios`, `warmup`, `repeats`, `out`, `dry_run`, `seed`, `probe_delay_ms`, `slow_queue_depth`, `flat_tolerance_pct`, `worker_pool_size`, `device`, `dtype`, `vllm_engine_kwargs`, `isolate`, `graph_capture` |
 | `scenarios` | list of strings | Expanded scenario names, in run order |
 | `dry_run` | bool | Whether `--dry-run` was given |
 | `device` | string or null | hf: the torch device; vllm: null |

{undercurrent-0.1.0 → undercurrent-0.2.0}/benchmarks/run.py

@@ -670,15 +670,22 @@ class VLLMRunner:
 
         from undercurrent.model import ProbedModel
 
+        spec = build_spec(scenario, self.num_layers)
+        engine_kwargs = dict(self.engine_kwargs)
+        if self.args.graph_capture:
+            # Graph mode: the probed engine runs with CUDA graphs and captures residual_stream
+            # through the model's aux outputs at exactly the layers this scenario probes.
+            engine_kwargs.pop("enforce_eager", None)
+            engine_kwargs["graph_capture_layers"] = sorted({p["layers"] for p in spec["extraction_points"]})
         pm = ProbedModel(
             self.args.model,
             backend="vllm",
-            spec=build_spec(scenario, self.num_layers),
+            spec=spec,
             probes=probes,
             router_kwargs=router_kwargs,
             metrics_sink=sink,
             max_concurrency=self.args.batch_size,
-            **self.engine_kwargs,
+            **engine_kwargs,
         )
 
         class Probed(Session):
@@ -1195,6 +1202,12 @@ def build_parser() -> argparse.ArgumentParser:
         default="auto",
         help="run each scenario in its own subprocess (auto: yes for vllm, no for hf)",
     )
+    p.add_argument(
+        "--graph-capture",
+        action="store_true",
+        help="vllm only: run probed scenarios in graph mode (CUDA graphs + torch.compile, no enforce_eager), "
+        "capturing residual_stream via load_model(graph_capture_layers=...); the baseline is unaffected",
+    )
     p.add_argument("--worker-output", help=argparse.SUPPRESS)
     return p
 
@@ -1226,6 +1239,8 @@ def resolve_args(args: argparse.Namespace) -> argparse.Namespace:
             )
         if args.vllm_engine_kwargs:
             raise HarnessError("--vllm-arg only applies to --backend vllm")
+        if args.graph_capture:
+            raise HarnessError("--graph-capture only applies to --backend vllm")
         if args.device is None:
             import torch
 
@@ -1259,6 +1274,11 @@ def check_vllm_environment() -> None:
             "--backend vllm needs vLLM, which undercurrent never installs by default. Run the harness in your "
             'vLLM environment or image, or `pip install "undercurrent[vllm]"`.'
         )
+    # torch.cuda.is_available() below poisons fork() for CUDA in this process,
+    # but doesn't mark CUDA initialised, so vLLM would still fork its engine
+    # core and that child fails with "Cannot re-initialize CUDA in forked
+    # subprocess". Make vLLM spawn it instead (unless the caller chose).
+    os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
     import torch
 
     if not torch.cuda.is_available():
@@ -1303,6 +1323,8 @@ def _child_argv(args: argparse.Namespace, scenario: Scenario, output: str) -> li
         argv += ["--device", args.device]
     if args.worker_pool_size is not None:
         argv += ["--worker-pool-size", str(args.worker_pool_size)]
+    if args.graph_capture:
+        argv.append("--graph-capture")
     for key, value in args.vllm_engine_kwargs.items():
         argv += ["--vllm-arg", f"{key}={json.dumps(value)}"]
     return argv

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/about/intended-use.md

@@ -98,7 +98,7 @@ What the project does:
 What the project does not do:
 
 - It doesn't modify activations or weights. There is no steering, ablation or
-  "refusal removal" feature, and none is planned for v0.1.
+  "refusal removal" feature, and none is planned for v0.2.
 - It doesn't ship trained probes, jailbreaks or recipes for disabling a
   model's safeguards.
 - It doesn't give access to a model you couldn't already run.
@@ -133,4 +133,4 @@ would document a model (a "probe card"). At a minimum, include:
   Undercurrent and engine versions you tested with.
 
 Publishing pretrained probes on the Hugging Face Hub is on the roadmap. It is
-not part of v0.1.
+not part of v0.2.

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/compatibility.md

@@ -38,10 +38,11 @@ Only combinations that have actually been run are listed here.
 | Python | torch | transformers | vllm | CUDA | Status |
 |--------|-------|--------------|------|------|--------|
 | 3.14.4 | 2.13.0 (`+cu130` build) | 5.16.1 | 0.28.0 | As provided by the vLLM wheel; no GPU present | **Tested locally**: full unit test suite, CPU only. GPU and real vLLM generation tests were skipped. |
+| 3.13.16 | 2.13.0 (`+cu130` build) | 5.21.0 | 0.28.0 | 13.0, NVIDIA L4 (driver 580.82.07) | **GPU-verified, 2026-11-01** (0.2.0, commit `5ded89e`): `scripts/gpu_check.sh` passed, covering every GPU test (vLLM integration, `ProbedModel` with the vLLM backend, and the `residual_stream` HF-vs-vLLM equivalence check) and the vLLM adapter suite against the real vLLM. Run on Google Colab. |
 | 3.13.15 | 2.13.0 (`+cu130` build) | 5.17.0 | 0.28.0 | 13.0, NVIDIA L4 (driver 580.82.07) | **GPU-verified, 2026-10-02**: `scripts/gpu_check.sh` passed, covering every GPU test (vLLM integration, `ProbedModel` with the vLLM backend, and the `residual_stream` HF-vs-vLLM equivalence check) and the vLLM adapter suite against the real vLLM. Run on Google Colab. |
 | 3.10.21 | 2.1.2 (`+cpu` build) | 4.40.2 | not installed | n/a | **The declared lower bounds** (with pydantic 2.0.3, PyYAML 6.0.3, numpy 1.23.5, packaging 22.0): CPU unit test suite passes locally. The CI `floors` job installs the same versions and runs the suite. Tests tied to newer releases skip themselves (the byte-for-byte JSON Schema comparisons need pydantic 2.11+; the content-safety demo needs torch 2.6+ to load checkpoints). |
 
-There is no GPU CI for v0.1. Instead, before every release a maintainer runs
+There is no GPU CI for v0.2. Instead, before every release a maintainer runs
 `scripts/gpu_check.sh` on a GPU machine with vLLM installed. It runs the GPU
 and real-vLLM tests and prints the combination it ran with, which is then
 added to the table above (see [RELEASING.md](https://github.com/wrynx/undercurrent/blob/main/RELEASING.md)).
@@ -53,7 +54,7 @@ pip install -e ".[dev]"     # from a checkout, into your vLLM environment
 scripts/gpu_check.sh
 ```
 
-## Known issues (v0.1)
+## Known issues (v0.2)
 
 These affect the vLLM adapter only. The Hugging Face backend is not affected.
 
@@ -104,15 +105,25 @@ torch 2.13.0+cu130) with `scripts/gpu_check.sh`.**
 
 ### No GPU CI for the vLLM adapter
 
-There is no GPU CI for v0.1. The GPU and real-vLLM tests run by hand: a
+There is no GPU CI for v0.2. The GPU and real-vLLM tests run by hand: a
 maintainer runs
 [`scripts/gpu_check.sh`](https://github.com/wrynx/undercurrent/blob/main/scripts/gpu_check.sh)
 on a GPU machine before each release, and the combination it ran with goes
-into [Tested combinations](#tested-combinations). For 0.1.0 that's an NVIDIA
+into [Tested combinations](#tested-combinations). For 0.2.0 that's an NVIDIA
 L4 with vLLM 0.28.0. Other GPUs, larger models and multi-GPU topologies
 haven't been run, and the vLLM adapter is still experimental. You can run the
 check on your own GPU machine too (see above).
 
+### Graph-mode `residual_stream` capture is experimental
+
+New in 0.2.0, `load_model(graph_capture_layers=[...])` captures
+`residual_stream` without `enforce_eager=True` by reusing vLLM's EAGLE-3
+auxiliary hidden-state outputs. It's limited to `residual_stream` at the layers
+given at load time, on models that implement that interface, and it patches
+more vLLM internals than hook mode does. Every other configuration still needs
+`enforce_eager=True`. See
+[Keep CUDA graphs](guides/vllm-deployment.md#keep-cuda-graphs-graph-mode-residual-capture-experimental).
+
 ### vLLM 0.29 and 0.30 aren't supported yet
 
 The supported range is `vllm>=0.28,<0.29`. vLLM 0.30 is already out, so a

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/concepts/execution-modes.md

@@ -138,7 +138,7 @@ Each async extraction point of each request gets its own **binding**:
 | --- | --- | --- |
 | `drop_oldest` (default) | Evict the oldest waiting activation, accept the new one | Generation never waits; the probe loses older activations |
 | `drop_newest` | Discard the incoming activation | Generation never waits; the probe loses the latest activations |
-| `block` | `route()` waits until there is room | Nothing is dropped; generation slows down to the probe's pace |
+| `block` | `route()` waits until there is room | Nothing is dropped; generation slows down to the probe's pace (on vLLM's default cross-process path, results arrive late instead; see [Async execution](../production/async-execution.md)) |
 
 Every drop is counted in the router's [metrics](../production/metrics.md), so
 you can see when a probe falls behind.

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/concepts/interventions.md

@@ -3,7 +3,7 @@
 <!-- owner: p3-concepts -->
 
 An **intervention** is an inline probe changing the course of generation. In
-v0.1 there is one intervention: **abort**. When an inline probe returns
+v0.2 there is one intervention: **abort**. When an inline probe returns
 `ProbeSignal(action=ProbeAction.ABORT)`, the engine adapter stops generating
 for that request. The router still finalizes every probe, so you get each
 probe's `ProbeResult` for the truncated generation.

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/guides/custom-probe.md

@@ -623,7 +623,7 @@ and probes; see [Deploy with vLLM](vllm-deployment.md).
 ### Probes with trained weights
 
 A probe that runs a trained classifier needs its weights. There is no probe
-save/load API in v0.1. Load the weights yourself, with `safetensors`, which
+save/load API in v0.2. Load the weights yourself, with `safetensors`, which
 `transformers` already depends on. Loading a file on every `spawn` would add
 disk I/O to every request, so load once and share the read-only result, for
 example:

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/guides/vllm-deployment.md

@@ -15,13 +15,14 @@ server and want to drive the engine adapter and `Router` yourself, read
     GPU, use the default Hugging Face backend from the
     [Quickstart](../getting-started/quickstart.md).
 
-!!! warning "Known issues (v0.1)"
-    The vLLM adapter hasn't been verified on a GPU yet. Fixed in 0.1.0: on
-    Llama-family models (and other fused-residual architectures),
-    `residual_stream` is captured as `hidden_states + residual`, the same
-    residual stream as the HF backend; earlier builds captured the layer's
-    MLP output. GPU verification of the fix is pending. See
-    [Known issues](../compatibility.md#known-issues-v01) for details.
+!!! warning "Known issues (v0.2)"
+    The vLLM adapter is experimental. Its GPU tests are run by hand before
+    each release (there's no GPU CI yet), and only vLLM 0.28 is supported.
+    Fixed in 0.1.0 and GPU-verified: on Llama-family models (and other
+    fused-residual architectures), `residual_stream` is captured as
+    `hidden_states + residual`, the same residual stream as the HF backend;
+    earlier builds captured the layer's MLP output. See
+    [Known issues](../compatibility.md#known-issues-v02) for details.
 
 ## Install into your vLLM environment
 
@@ -260,7 +261,7 @@ Everything below `ProbedModel` stays reachable:
 
 ## Serving over HTTP
 
-Undercurrent v0.1 doesn't include an HTTP server. A supported, OpenAI-compatible
+Undercurrent v0.2 doesn't include an HTTP server. A supported, OpenAI-compatible
 `serve` command is [on the roadmap](https://github.com/wrynx/undercurrent/blob/main/ROADMAP.md).
 Until then, put `VLLMEngineAdapter` (or `ProbedModel`) behind the web server
 you already run.
@@ -378,6 +379,9 @@ flagged stream (`choices[0].finish_reason == "content_flagged"`, empty
   per-topology support status.
 - **Use `enforce_eager=True`.** With CUDA graphs, replayed steps don't run
   Python, so the forward hooks can't fire. The adapter doesn't set it for you.
+  Eager mode is expensive (on an A100, Qwen2.5-7B at batch 16 decodes at
+  about half its CUDA-graph throughput). If you only need `residual_stream`,
+  see [Keep CUDA graphs: graph-mode residual capture](#keep-cuda-graphs-graph-mode-residual-capture-experimental).
 - **`block_until_signal` stalls the whole batch.** Under continuous batching,
   a blocking inline probe holds the scheduler step that every co-batched
   request shares, for up to `timeout_ms`. The adapter logs a one-time warning
@@ -395,6 +399,43 @@ flagged stream (`choices[0].finish_reason == "content_flagged"`, empty
   point, but a bare `vllm serve` process doesn't drive probing. See
   [Running under `vllm serve`](../internals/vllm-adapter.md#running-under-vllm-serve).
 
+### Keep CUDA graphs: graph-mode residual capture (experimental)
+
+*New in 0.2.0.* If every extraction point is a `residual_stream`, you can
+leave `enforce_eager` off and keep vLLM's CUDA graphs and `torch.compile`.
+Pass the decoder layers to capture as `graph_capture_layers`:
+
+```py
+with ProbedModel.from_pretrained(
+    "Qwen/Qwen2.5-0.5B-Instruct",
+    backend="vllm",
+    spec=SPEC,  # residual_stream at layer 6 only
+    probes={"norm_gate": norm_gate},
+    graph_capture_layers=[6],  # instead of enforce_eager=True
+) as model:
+    outputs = model.generate(prompts, max_new_tokens=32)
+```
+
+`VLLMEngineAdapter.load_model(..., graph_capture_layers=[...])` works the same
+way. Instead of forward hooks, the adapter reads the auxiliary hidden-state
+outputs vLLM already computes inside the graph for EAGLE-3 speculative
+decoding. Probes see the same `residual_stream` tensor as in hook mode.
+
+Limits, each enforced with an error rather than a silent miss:
+
+- `residual_stream` only, and only at the layers given at load time. An
+  extraction point with another tensor type or layer is rejected when it's
+  registered.
+- The model must implement vLLM's EAGLE-3 aux-output interface (Llama and
+  Qwen2, among others; GPT-2 doesn't).
+- It can't be combined with speculative decoding that uses aux outputs
+  itself.
+- vLLM's compile cache is disabled for the engine, so start-up recompiles.
+
+Benchmark it on your own model with `benchmarks/run.py --graph-capture`. How
+it works is covered in
+[Graph-mode residual capture](../internals/vllm-adapter.md#graph-mode-residual-capture-experimental).
+
 [The vLLM adapter](../internals/vllm-adapter.md) explains how capture,
 position mapping and aborts work.
 
@@ -414,7 +455,7 @@ change in any minor release, so the range stops at the next one.
   best-effort for one more Undercurrent release, then it's dropped.
 
 !!! warning "GPU CI status"
-    There is no GPU CI for v0.1: the GPU tests run by hand before each release
+    There is no GPU CI for v0.2: the GPU tests run by hand before each release
     (`scripts/gpu_check.sh`). The vLLM rows on the
     [Compatibility](../compatibility.md#tested-combinations) page say which
     combinations have had a GPU run.
@@ -489,7 +530,8 @@ doesn't look (it checks `model.model.layers`, `model.layers` and
 
 **No probe results, or `aborted` is `False` although a probe asked to abort**
 
-Check that you passed `enforce_eager=True`, that the extraction point's
+Check that you passed `enforce_eager=True` (or, in graph mode, that the
+extraction point is a `residual_stream` at one of the `graph_capture_layers`), that the extraction point's
 `tensor_type` and `position` are captured on your model (see
 [Known limitations](../internals/vllm-adapter.md#known-limitations)), and
 that a class probe puts its `ABORT` signals in `signal_history` (see

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/internals/vllm-adapter.md

@@ -228,6 +228,64 @@ wants that (layer, tensor type), so unused hooks cost little.
 Which tensor each hook sees under tensor and pipeline parallelism is covered in
 [Tensor & pipeline parallelism](vllm-parallelism.md).
 
+### Graph-mode residual capture (experimental)
+
+Forward hooks are Python, and neither CUDA-graph replay nor vLLM's
+`torch.compile`d forward runs them, so hook-based capture needs
+`enforce_eager=True`. That is expensive: on an A100, eager decoding of
+Qwen2.5-7B at batch 16 runs at about half the CUDA-graph throughput, and
+Qwen2.5-0.5B at about an eighth, before any probe runs.
+
+For `residual_stream`, `load_model(graph_capture_layers=[...])` avoids eager
+mode. It reuses the auxiliary hidden-state outputs vLLM already has for EAGLE-3
+speculative decoding: models implementing that interface (Llama and Qwen2 among
+them) output `hidden_states + residual` after the chosen decoder layers, inside
+the compiled and captured graph. Aux index `k + 1` is decoder layer `k`'s
+output, the same tensor hook mode captures at layer `k`.
+
+- The adapter sets `UNDERCURRENT_VLLM_GRAPH_CAPTURE_LAYERS` for the engine
+  process and disables vLLM's compile cache (its key doesn't cover the extra
+  outputs).
+- In the engine process, the `vllm.general_plugins` entry point patches
+  `GPUModelRunner.load_model`: after loading and before compilation, it calls
+  `set_aux_hidden_state_layers` and sets `use_aux_hidden_state_outputs`, as
+  vLLM does for EAGLE-3, and wraps `_model_forward` to hand the aux outputs to
+  the worker extension after each forward pass.
+- The worker extension installs no forward hooks. It stages each probed
+  layer's matched rows from the aux outputs with the same batched,
+  non-blocking copy as hook mode (see below), before the next replay overwrites
+  them.
+
+Limits, each enforced with an error rather than silently capturing nothing:
+`residual_stream` only, at the layers given at load time; the model must
+implement vLLM's EAGLE-3 aux-output interface; speculative decoding that uses
+aux outputs itself can't be combined with it. `graph_capture_layers=[]` runs a
+graph-mode engine with nothing captured (the benchmark's `probed-0`).
+
+### Copying captures to the host
+
+Probes run on the host, so every captured row is copied from the GPU. A forward
+hook runs while the GPU is still working through earlier kernels, and a
+blocking copy inside the hook makes the host wait for the GPU. While the host
+waits, it can't queue the next layers' kernels, so the GPU then idles while the
+host catches up. Doing that once per matched row (as before) cost one stall per
+(row, layer) per step.
+
+For CUDA tensors, each hook instead gathers all of its matched rows on the GPU
+and starts one non-blocking copy into pinned host memory, recording a CUDA
+event after it. The records are built after one wait on the latest event, once
+per step:
+
+- in-process (`bind_router()` succeeded): at the end of the `execute_model`
+  call, before the driver polls for aborts, so inline aborts take effect at the
+  same step as before;
+- cross-process polling: when the driver calls `pop_pending_activations()`, or
+  at the start of the next step, whichever comes first. The copies have
+  finished by then, so the wait costs nothing.
+
+Records still arrive in capture order (hook order, then row order). On CPU each
+row is copied synchronously, since there's no device to overlap with.
+
 ### Engine settings the adapter forces
 
 `load_model()` sets two environment variables before it builds the engine. Each
@@ -262,14 +320,21 @@ the `bind_router` call fails. The adapter logs this and falls back to polling:
 
 - `_emit()` buffers each `ActivationRecord` in the worker.
 - After every `RequestOutput`, the adapter calls `pop_pending_activations` over
-  `collective_rpc`. Records travel as plain dicts with the tensor as a nested
-  list, because vLLM's RPC encoder doesn't handle arbitrary dataclasses or
-  tensors nested in ad-hoc structures.
-- `_drain_pending_activations()` rebuilds the `ActivationRecord`s in the driver
-  and routes them **in capture order**, one at a time, through
-  `run_in_executor`. Order matters for trajectory probes, and running routing
-  off the event loop stops a slow probe from stalling every other in-flight
-  request's traffic with the engine.
+  `collective_rpc`. Each worker returns one packed batch
+  (`activation_batch.py`): the tensors stacked into a float32 `bytes` buffer
+  per tensor width, and the other fields as parallel lists. vLLM's RPC encoder
+  carries `bytes` and lists natively, but not arbitrary dataclasses or tensors
+  nested in ad-hoc structures. (Earlier versions sent each tensor as a nested
+  Python list, which cost a few hundred microseconds per record at a
+  3,584-wide hidden size; a batch costs a few microseconds per record.)
+- `_drain_pending_activations()` unpacks the batches, rebuilds the
+  `ActivationRecord`s in the driver and routes them **in capture order**, one
+  at a time, in a single `run_in_executor` call per drain (one thread hand-off
+  per record cost about as much again as the old encoding). Order matters for
+  trajectory probes, and running routing off the event loop stops a slow probe
+  from stalling every other in-flight request's traffic with the engine.
+  Aborts requested by those records are sent to the engine after the batch is
+  routed.
 
 Capture, position matching and row mapping are the same on both paths; only the
 place where `route()` runs differs. On the polling path, activations (and the
@@ -367,7 +432,7 @@ engine core and each worker). Here's what the current code does on that route:
   requests registered through its `register_extraction` RPC, and only hands
   records over when someone calls `pop_pending_activations`.
   `VLLMEngineAdapter.generate()` makes those calls; a bare `vllm serve` process
-  has nothing that does. In Undercurrent v0.1 this route gives you a
+  has nothing that does. In Undercurrent v0.2 this route gives you a
   correctly configured engine with the extension loaded, but not end-to-end
   probing.
 
@@ -382,9 +447,9 @@ These all fail loudly unless the table says otherwise.
 | Limitation | Impact | Workaround | Tracked in |
 | --- | --- | --- | --- |
 | More than one worker (tensor or pipeline parallelism, Ray) isn't validated. | `load_model()` raises `VLLMAdapterLimitationError`. | `allow_unsupported_executor=True`. Results from all workers are merged and deduplicated, but this is unvalidated. | [Tensor & pipeline parallelism](vllm-parallelism.md) |
-| `residual_stream` needs to recognise the decoder layer's output: a single tensor, `(hidden_states, None)`, or a fused-residual `(hidden_states, residual)` pair of matching shape and dtype (captured as their sum). | Any other output (for example HunYuan's 3-tuple) raises `VLLMAdapterLimitationError` when a `residual_stream` extraction point fires on it. GPU verification of the fused-residual sum against the HF backend is pending `scripts/gpu_check.sh`. | Use `attn_out`, `mlp_out` or `final_norm`, or the HF backend, on such models. | [Compatibility](../compatibility.md#fixed-in-010-residual_stream-on-fused-residual-vllm-layers) |
+| `residual_stream` needs to recognise the decoder layer's output: a single tensor, `(hidden_states, None)`, or a fused-residual `(hidden_states, residual)` pair of matching shape and dtype (captured as their sum). | Any other output (for example HunYuan's 3-tuple) raises `VLLMAdapterLimitationError` when a `residual_stream` extraction point fires on it. The fused-residual sum is GPU-verified against the HF backend (`scripts/gpu_check.sh`). | Use `attn_out`, `mlp_out` or `final_norm`, or the HF backend, on such models. | [Compatibility](../compatibility.md#fixed-in-010-residual_stream-on-fused-residual-vllm-layers) |
 | `attn_out` is only hooked when the layer has a `self_attn` submodule. vLLM's GPT-2 implementation names it `attn`. | **Silent.** `attn_out` extraction points never fire on GPT-2-style models. | Use `residual_stream` or `mlp_out` on those models. | Future work |
-| Only tested with `enforce_eager=True`. The adapter doesn't set it. | With CUDA graphs, replayed steps run recorded GPU kernels without executing Python, so forward hooks can't fire during them. | Pass `enforce_eager=True` to `load_model()`, as the tests and examples do. | Future work |
+| Hook-based capture needs `enforce_eager=True`. The adapter doesn't set it. | With CUDA graphs, replayed steps run recorded GPU kernels without executing Python, so forward hooks can't fire during them. | Pass `enforce_eager=True` to `load_model()`, as the tests and examples do; or, for `residual_stream` only, use `graph_capture_layers=[...]` (experimental, see [Graph-mode residual capture](#graph-mode-residual-capture-experimental)). | Other tensor types: future work |
 | `tensor: kv` is rejected. The KV cache lives in paged blocks addressed through a slot mapping, not in a per-step tensor with one row per token. | `register_extraction` raises `WorkerExtractionError`. | Use `residual_stream`, `attn_out`, `mlp_out` or `final_norm`. | Future work |
 | Negative-indexed single-point selectors (`generated[-1]`) never match. Resolving them needs the request's final length, which vLLM only knows once the request finishes, and by then the step's activations are gone. | **Silent.** No records for those extraction points. | `generated[*]` and slices (`generated[5:]`, `generated[2:8]`) work. Keep the last record from a trajectory probe yourself. | Future work |
 | `block_until_signal` stalls every request in the batch, not only the triggering one. | Higher latency for co-batched requests, up to `timeout_ms` per blocked step. | Use `mode: reject` or `execution_mode: async` under concurrency. | By design (see [above](#block_until_signal-and-continuous-batching)) |

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/internals/vllm-parallelism.md

@@ -10,7 +10,7 @@ The vLLM source details below were checked against the vLLM version listed in
 
 ## Summary
 
-Multi-worker topologies are **not validated** in Undercurrent v0.1:
+Multi-worker topologies are **not validated** in Undercurrent v0.2:
 `VLLMEngineAdapter.load_model()` raises `VLLMAdapterLimitationError` when vLLM
 starts more than one worker, unless you pass `allow_unsupported_executor=True`.
 The capture side is sound for stock Llama-style models. Under TP, every rank's

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/production/async-execution.md

@@ -67,7 +67,7 @@ When `route()` finds a binding's queue full, the overflow policy decides:
 | --- | --- | --- | --- |
 | `drop_oldest` (default) | Evict the oldest queued activation, accept the new one | Older activations; the probe sees the most recent window | Nothing on the generation path |
 | `drop_newest` | Reject the incoming activation; the queue is unchanged | Newer activations; the probe sees the earliest ones | Nothing on the generation path |
-| `block` | `route()` waits until the worker frees a slot | Nothing is dropped | Generation slows to the probe's speed |
+| `block` | `route()` waits until the worker frees a slot | Nothing is dropped | Generation slows to the probe's speed (HF, and vLLM in-process); on vLLM's cross-process path, time to results instead |
 
 How to choose:
 
@@ -79,10 +79,24 @@ How to choose:
 - **`block`** suits offline evaluation and data collection, where every
   activation must be processed and throughput matters less than completeness.
   Avoid it in latency-sensitive serving: a slow probe becomes slow generation.
-  On vLLM, the blocked `route()` call runs on the shared engine thread, so it
-  stalls every request in the batch, not just this one. And if the binding is
-  waiting for a pool thread (see above), `block` can stall generation until
-  another request ends.
+  On vLLM, what waits depends on where `route()` runs:
+  - **In-process** (the router reached the worker): the blocked `route()` call
+    runs on the shared engine thread, so it stalls every request in the batch,
+    not just this one.
+  - **Cross-process polling** (the default with vLLM 0.28, whose engine core
+    runs in its own process): `route()` runs in the driver, after the engine
+    has already produced the tokens. The engine is **not** slowed; it keeps
+    generating while the driver falls behind, captured activations accumulate
+    in the engine process until the driver catches up, and each `generate()`
+    call returns only once all of its activations have been processed. So
+    `block` there means "lossless, results delayed", not backpressure. The
+    adapter routes each drained backlog round-robin across requests, so all
+    bindings work in parallel and the delay is about the probe's own pace;
+    measured on an A100 with Qwen2.5-7B, 16 requests and a 76 ms probe, routing
+    one request's backlog at a time had made it about 13x slower than that.
+
+  And if the binding is waiting for a pool thread (see above), `block` can
+  stall until another request ends.
 
 The policy is set router-wide:
 

{undercurrent-0.1.0 → undercurrent-0.2.0}/docs/production/embedding.md

@@ -341,7 +341,7 @@ don't set one. When a queue is full, the router-wide
 | --- | --- | --- |
 | `OverflowPolicy.DROP_OLDEST` (default) | Unaffected | The oldest queued activations; the probe sees the most recent ones |
 | `OverflowPolicy.DROP_NEWEST` | Unaffected | The incoming activations; the probe sees the earliest ones |
-| `OverflowPolicy.BLOCK` | `route()` waits for space, so generation slows to the probe's pace (on vLLM, the whole batch does) | Nothing |
+| `OverflowPolicy.BLOCK` | `route()` waits for space, so generation slows to the probe's pace (on vLLM in-process, the whole batch does; on vLLM's cross-process path the engine isn't slowed and results arrive late instead) | Nothing |
 
 Use a drop policy for live serving and watch the drop count; use `BLOCK` for
 offline evaluation where every activation must be processed. A bigger queue

{undercurrent-0.1.0 → undercurrent-0.2.0}/pyproject.toml

@@ -98,7 +98,7 @@ dev = [
     # Hard per-test time limit for `gpu` tests (tests/conftest.py): a stuck vLLM
     # engine dumps every thread's stack and fails instead of hanging the run.
     "pytest-timeout>=2.3",
-    "ruff>=0.16,<0.17",
+    "ruff>=0.16,<0.18",
     "pre-commit>=4",
     # Validates example specs against schema/probe-spec.schema.json in tests.
     "jsonschema>=4.18,<5",

{undercurrent-0.1.0 → undercurrent-0.2.0}/src/undercurrent/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.1.0'
-__version_tuple__ = version_tuple = (0, 1, 0)
+__version__ = version = '0.2.0'
+__version_tuple__ = version_tuple = (0, 2, 0)
 
 __commit_id__ = commit_id = None