vllm: follow-ups from the v0.7.0 merges by khaiwang · Pull Request #662 · ndif-team/nnsight

khaiwang · 2026-05-07T03:11:33Z

Follow-ups from the vLLM-integration merges that landed into dev for v0.7.0. This PR merges a mix of correctness fixes, one default change, packaging, and docs. Each is independent but small.

#	Change	Kind	Commits
1	Clone inference-mode tensors on `.save()` (fixes #661)	correctness	`a40a5f2`, `fe7c072`, `cc56513`
2	Narrow/swap on the token axis for Transformers-backend models	correctness	`0bc57e3`
3	Surface deferred intervention errors on local sync/async paths	correctness	`5ef4466`
4	Submit every invoke in async traces, not just the first	correctness	`18f0b41`
5	Disable vLLM prefix caching by default	default	`ee2022d`
6	Restore `nnsight-serve` install machinery	packaging	`bf2949d`
7	Sync vLLM integration README to the v0.7.0 async path	docs	`c8bd79e`
8	Fold `intervention-gaps/` into `docs/models/vllm.md`; move probes to `tests/`	docs	`fc6c10b`

1. Clone inference-mode tensors on `.save()` (fixes #661)

vLLM runs forward passes inside torch.inference_mode(), and several of its kernels (e.g. fused_add_rms_norm) mutate buffers in place. Without a clone, references returned by Envoy.output / Envoy.inputs alias those buffers, so values surviving past the trace reflect post-mutation state, not what the user asked to save.

intervention/tracing/globals.save() now clones when the saved object is an inference-mode tensor. The clone allocates a fresh, non-inference tensor so downstream fused ops mutate the original buffer rather than the user's saved reference. No-op for normal tensors — HF / vanilla PyTorch paths are unaffected.

Object.save() returns save(self) instead of self, so the cloned tensor (not the original) is what the trace's local-frame filter retains via Globals.saves.

Verification (SmolLM2-135M on vLLM 0.19.1, one A100): max reference-vs-clone diff across saved residual / attention / MLP tensors dropped from 4064.23 → 0.00.

Regression coverage in TestSaveCloning pins the invariants:

Inference-mode tensor save returns a clone (mutation of source doesn't corrupt save).
Normal tensor save returns the original (no-overhead contract for HF path).
End-to-end: module.output.save() inside torch.inference_mode() returns a non-inference tensor.

cc56513 makes the end-to-end inference-mode test deterministic: the old assert x.std() > 50 depended on an unseeded 8-element randn sample and flaked ~2.4% of runs; it's replaced with an exact, RNG-free check (the saved clone holds the pre-mutation values; the in-place-mutated source is exactly saved * 1000).

Commits: a40a5f2, fe7c072, cc56513.

2. Narrow/swap on the token axis for Transformers-backend models

Models without a native vLLM definition (e.g. SmolLM3) run through vLLM's Transformers backend, which wraps the HuggingFace model and adds a leading singleton batch dim (inputs_embeds[None, ...]). Their decoder-layer activations are 3D [1, total_tokens, hidden], so the batched token axis is dim 1, not dim 0.

Batcher._narrow / _swap hard-assumed the token axis was dim 0 and gated on shape[0] == total_batch_size. For the 3D case that gate is 1 == total_tokens (False), so reads returned the full batch (all prompts) and writes were silently discarded — every intervention became a no-op once batching was active (2+ prompts). Native vLLM models (2D [total_tokens, hidden]) were unaffected, which is why Qwen3 worked but SmolLM3 did not.

The base Batcher now narrows/swaps along an axis reported by a new _batch_dim() hook (still dim 0 by default; preserves the existing in-place, concat-for-view/grad-leaf, and passthrough paths). VLLMBatcher overrides _batch_dim to recognize the Transformers-backend [1, total_tokens, hidden] shape and select dim 1. CPU-only regression tests (TestVLLMBatcherAxis): the 3D swap/narrow tests fail on the unpatched batcher and pass after; the 2D native test pins no regression.

Commit: 0bc57e3.

3. Surface deferred intervention errors on local sync/async paths

The vLLM interleaver always runs in defer mode (GPUModelRunner.load_model sets defer_exceptions=True) so a single bad intervention can't crash the engine that's serving other requests. The serve path re-raised captured errors via surface_server_errors, but the local VLLM.trace() sync and async paths only collected output.saves and never read the __nnsight_exceptions__ envelope — so an intervention that errored on the worker (e.g. an in-place write on an inference-mode tensor) failed silently: no exception, and every .save() after the failing line was dropped.

VLLM.__call__ (sync) and AsyncVLLMBackend.__aiter__ (async) now read the envelope and re-raise via surface_server_errors, mirroring the serve path. The error surfaces at the trace boundary while the engine stays alive for the next trace (verified: a clean trace and a clone-based intervention both work in the same process after the surfaced error). Envelopes are merged per request so a multi-invoke trace doesn't clobber one request's error with another's. Adds test_inplace_inference_write_surfaces_error covering the sync path.

Commit: 5ef4466.

4. Submit every invoke in async traces, not just the first

AsyncVLLMBackend.__call__ serialized every invoke's mediator but then submitted only prompts[0] / params[0] under a single request_id, so a multi-invoke async trace ran only the first prompt: invokes past the first never reached the engine, their per-invoke saves came back empty, and trace-shared saves were never collected (the worker's received_count never reached expected_count). Every prior async test used a single prompt, so this code path had no coverage.

The backend now submits one request per invoke (mirroring the fan-out in serve/server.py) and merges the per-request generators into a single stream via an asyncio queue, collecting saves per finished request. The single-invoke streaming path is behavior-preserving, and the deferred-error surfacing from #3 is reused via a shared _attach_saves helper.

Adds test_async_multi_invoke_runs_all_invokes — the async counterpart of test_shared_list_across_invokes: a two-invoke async trace must produce two finished requests and collect the trace-shared list from both invokes.

Verification (gpt2, mode="async", vLLM 0.19.1, one A100): the new test fails on the old code (Expected 2 finished requests, got 1) and passes after; the full TestAsyncEngine suite is 6/6 green (the 5 pre-existing single-invoke async tests unaffected).

Commit: 18f0b41.

5. Disable vLLM prefix caching by default

vLLM's prefix caching reuses KV values from previously-seen sequences. When the next request shares a prefix, those tokens skip the forward pass — hooks don't fire and interventions on those tokens are silently skipped, with no error.

VLLM(...) now defaults to enable_prefix_caching=False so interventions consistently see every token. Users who explicitly opt in (e.g. for throughput on workloads that don't need to hook prefill tokens) can still pass enable_prefix_caching=True. Matches what docs/models/vllm.md documents as the integration's default.

Commit: ee2022d.

6. Restore `nnsight-serve` install machinery

PR #656 merged the nnsight-serve sources (cli.py, server.py, LocalServeBackend, ServeInterleavingTracer, …) onto dev, but the pyproject.toml changes were dropped during conflict resolution. Result: pip install "nnsight[serve]" returns "no matching distribution" and the nnsight-serve CLI shim isn't on PATH for a fresh install.

Restored:

serve optional-dependency that pulls vllm + FastAPI + uvicorn.
[project.scripts] entry registering nnsight-serve → nnsight.modeling.vllm.serve.cli:main.
all extended to include serve.

After this, the documented pip install "nnsight[serve]" / nnsight-serve … workflow works without the python -m nnsight.modeling.vllm.serve.cli workaround.

Commit: bf2949d.

7. Sync `src/nnsight/modeling/vllm/README.md` to the v0.7.0 async path

Two earlier refactors landed without README updates; the drift has been live through v0.7.0:

d124cc5 (2026-03-12, "refactor vLLM input processing") eliminated AsyncInterleavingTracer entirely. AsyncVLLMBackend now calls _setup_interleaver() directly; the async path uses the default RemoteInterleavingTracer.
bb61efa (2026-03-28, "refactor async backend") collapsed the dual-call __call__(tracer) / __call__() pattern into a single required __call__(self, tracer) that submits to AsyncLLM.generate() immediately on trace exit. _stream() was removed; iteration moved to __aiter__/__await__. tracer.backend() (with parens) now raises TypeError — the correct iteration form is async for output in tracer.backend (no parens).

This commit syncs the README accordingly:

Drop async_tracer.py file listing and all AsyncInterleavingTracer references.
Rewrite the AsyncVLLMBackend description (file responsibilities + Key Classes) to enumerate the current __call__/__aiter__/__await__ surface.
Redraw the async architecture diagram: VLLM.trace() injects only the backend → default tracer applies → __call__(tracer) submits and parks the generator → __aiter__ streams.
Replace 5 _stream() mentions with __aiter__().
Replace 4 tracer.backend() parens-form mentions with tracer.backend.
Replace the "Why AsyncInterleavingTracer Bypasses RemoteableMixin" section with "How VLLM.trace() Routes the Async Path", showing current setdefault-based routing (RemoteableMixin.trace() never hard-coded tracer_cls — the previous prose was also wrong on that point).
Usage example: drop parens; wrap in async def main() under an if __name__ == "__main__": guard (AsyncLLM uses multiprocessing spawn); note that output.saves is only set on output.finished.

Verified the corrected usage example runs end-to-end on this branch (gpt2, mode="async"): 8 RequestOutputs streamed, finished=True on the last, output.saves == {'logits': Tensor[1, 50257]}.

Note: the README's AsyncVLLMBackend surface section is accurate as of bb61efa; the multi-invoke fan-out added in #4 (18f0b41) does not change the public __call__/__aiter__/__await__ shape it describes.

Commit: c8bd79e.

8. Fold `intervention-gaps/` into `docs/models/vllm.md`; move probes to `tests/`

Migrate the durable content of intervention-gaps/{REPORT,VLLM_GUIDE}.md into the maintained user doc and delete the two stale docs (vLLM 0.15.1-era; their in-place-write recipes no longer work).

docs/models/vllm.md:

New "Intervention recipes" section: clone-and-replace writes, ablation, steering, logit lens, two-trace activation patching, tracer.cache().
"What each module returns" table: dual-stream (hidden, residual) output, int64 position-id .input, fused-RMSNorm/RowParallel tuples, merged qkv_proj/gate_up_proj, flat [total_tokens, hidden] layout.
New gotchas: in-place writes raise (replace instead), clone-on-save, enable_prefix_caching=False default, deferred errors keep the engine alive, no attention weights, vLLM ≠ transformers numerics.
Drop stale claims: tracer.cache() is supported; version 0.15.1 → 0.19.1.

tests/vllm_intervention_gaps/:

git mv run_all.py + test_*.py here (executable vLLM-vs-HF diagnostic suite) and add a README.

Recipes verified on vLLM 0.19.1 (Qwen2.5-0.5B): in-place writes raise, replacement works, logit-lens matmul (norm(hs) @ lm_head.weight.T) bitwise-matches model.logits at the last layer, and TP≥2 sub-module access works (the old "crashes at tp≥2" claim was stale).

Commit: fc6c10b.

Models without a native vLLM definition (e.g. SmolLM3) run through vLLM's Transformers backend, which wraps the HuggingFace model and adds a leading singleton batch dim (inputs_embeds[None, ...]). Their decoder-layer activations are 3D [1, total_tokens, hidden], so the batched token axis is dim 1, not dim 0. Batcher._narrow/_swap hard-assumed the token axis was dim 0 and gated on shape[0] == total_batch_size. For the 3D case that gate is 1 == total_tokens (False), so reads returned the full batch (all prompts) and writes were silently discarded -- every intervention became a no-op once needs_batching (2+ prompts) was active. Native vLLM models (2D [total_tokens, hidden]) were unaffected, which is why Qwen3 worked but SmolLM3 did not. Generalize the base Batcher to narrow/swap along an axis reported by a new _batch_dim() hook (still dim 0 by default; preserves the existing in-place, concat-for-view/grad-leaf, and passthrough paths). VLLMBatcher overrides _batch_dim to recognize the Transformers-backend's [1, total_tokens, hidden] shape and select dim 1. Adds CPU-only regression tests (TestVLLMBatcherAxis): the 3D swap/narrow tests fail on the unpatched batcher and pass after; the 2D native test pins no regression. Related to #661/#662. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

Butanium · 2026-07-20T21:19:01Z

@khaiwang please ping me when this is fixed so that i can go back to nnterp new version

A model vLLM has no definition of its own is served through vLLM's Transformers backend, which runs the wrapped HuggingFace module with a leading singleton batch dim (`inputs_embeds[None, ...]`). Its decoder layers emit `[1, total_tokens, hidden]`, so tokens sit on dim 1, not dim 0. `VLLMBatcher` inherited the base `Batcher`'s dim-0-only row math, which gates on `shape[0] == total`. On a 3-D activation that reads `1 != total`, so the tensor is called unbatched and passes through whole: every read hands the block every in-flight request's tokens, and every write is discarded. Nothing surfaces — a passthrough is indistinguishable from a tensor that legitimately isn't batched — so an intervention just quietly does nothing as soon as a second request shares the step. Native models (2-D slabs) were never affected, which is why this went unnoticed. `VLLMBatcher` now overrides `_narrow_tensor`/`_widen_tensor`, the extension point the base documents for non-dim-0 layouts (as `DiffusionBatcher` already does), and locates the token axis with a new `_token_dim`. Nothing in `intervention/batching.py` changes: the layout is a vLLM concern, and core batching should not grow a hook for it. The override drops the base's `_nnsight_batch` view marker, which exists for backward to redirect a batch slice through its storage-owning base. Backward is not supported on the vLLM path, so the marker had no reader. Not covered: `Interleaver.replay` trims a padded tap tensor with `t[:total]`, which is still dim-0-only, so taps over a Transformers-backend model can be served padding. Deciding the token axis of a *padded* tensor needs a rule this one can't supply, and taps on such a model are untested besides; noted in the dev doc rather than guessed at here. Tests are CPU-only and need no engine. The three Transformers-backend cases fail on the old batcher (the in-place one writes all 8 tokens instead of its 5); the three native cases pass either way and pin no regression.

Replaces the clone-in-`tracing.save()` approach for #661. What a block reads on vLLM is a view into engine memory, and vLLM's fused kernels (`fused_add_rms_norm`, MLA's in-place rotation of the `q_proj` output) overwrite those buffers a few ops later, so a value kept past its read point comes back holding a later layer's data with nothing to indicate it. The documented fix is `.clone()` at each site, which is silent when forgotten. `VLLMBatcher.narrow` now clones what it serves when the env var is set. Doing it here rather than in `save()` covers every way a value is kept — `.save()`, `tracer.cache()` (which narrows through the same call), appends under `tracer.iter`, and taps, whose replay goes through the same handoff — where a clone at the save point only ever covered the first, since a saved *container* is marked once while its elements keep aliasing. It also keeps the change inside the vLLM module instead of changing what `save()` returns for every backend. Env var rather than a `CONFIG.APP` field because the batcher that narrows is built in the engine's worker process (`GPUModelRunner.load_model`); a field set client-side would never reach it. Read once per batcher, so it has to be set before `VLLM(...)`. Falsy spellings match the ones `NNSIGHT_DISABLE_CPP_BACKTRACE` already accepts. Off by default, because it costs in-place edits: with nothing aliasing engine memory, `layers[10].output[0][:] += v` writes to the copy and the model never sees it. That is as silent as the problem it fixes, so it is pinned by a test, called out in both docs, and paired with the form that works either way (edit the copy, then assign it back — a swap routes through `widen`). An `eproperty` with a registered `transform` write-back is unaffected; that path already splices through `widen`. The copy is of the request's own span, not the whole slab, but it is still one allocation per module read per step.

`TestAdHocCall` saved `module.output[0]` and compared an ad-hoc call's result against it. A served value is the engine's live buffer, and for `mlp.down_proj` that buffer *is* the decoder layer's returned hidden_states: the next layer opens with `input_layernorm(hidden_states, residual)` -> fused_add_rms_norm, which rewrites it in place. So by the time the trace exits, the reference holds post-mutation state while the ad-hoc result — computed by a forward that allocated its own output — holds the real value, and `_min_row_cosine` comes back 0.126 against an assert of > 0.99. The test was blaming the ad-hoc call for the corruption of its own reference. `test_row_parallel_call_takes_and_returns_the_whole[2]` has been red on 0.8 for this reason, and it reads as a `Fragments.split` /`whole` bug, which is what its docstring says it guards — nobody would look at aliasing. This is the anti-pattern docs/models/vllm.md:109 names outright ("a served value is the model's live buffer, and the next layer's fused add+norm rewrites it in place"), so the fix is the documented idiom: clone before saving. The column-parallel case gets the same clone. It passes un-cloned today because act_fn allocates rather than writing back over gate_up's output, but that is luck, not a property the test should rest on. Found while verifying this branch on hardware: the failure disappears under NNSIGHT_VLLM_CLONE_READS=1 with no test change (1 failed/1 passed -> 2 passed), which is what pointed at the buffer rather than the gather.

khaiwang requested a review from JadenFiotto-Kaufman May 7, 2026 18:36

khaiwang changed the title ~~fix(vllm): clone inference-mode tensors on .save() to prevent in-plac…~~ vllm: follow-ups from the v0.7.0 merges May 21, 2026

khaiwang force-pushed the zikai/vllm-clone-on-save branch from 18f0b41 to 701e3d6 Compare August 18, 2026 02:38

JadenFiotto-Kaufman changed the base branch from dev to 0.8 August 19, 2026 17:24

JadenFiotto-Kaufman added 2 commits September 9, 2026 14:23

JadenFiotto-Kaufman mentioned this pull request Sep 9, 2026

vllm: rework #662 against 0.8 — token axis + NNSIGHT_VLLM_CLONE_READS #727

Merged

JadenFiotto-Kaufman force-pushed the zikai/vllm-clone-on-save branch from 701e3d6 to f1de399 Compare September 9, 2026 20:28

JadenFiotto-Kaufman merged commit cb5a1ac into 0.8 Sep 9, 2026
2 checks passed

JadenFiotto-Kaufman deleted the zikai/vllm-clone-on-save branch September 10, 2026 01:56

Navigation Menu

vllm: follow-ups from the v0.7.0 merges - #662

vllm: follow-ups from the v0.7.0 merges#662
JadenFiotto-Kaufman merged 3 commits into
0.8from
zikai/vllm-clone-on-save

khaiwang commented May 7, 2026 •

edited

Loading

Uh oh!

Butanium commented Jul 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

khaiwang commented May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

1. Clone inference-mode tensors on .save() (fixes #661)

2. Narrow/swap on the token axis for Transformers-backend models

3. Surface deferred intervention errors on local sync/async paths

4. Submit every invoke in async traces, not just the first

5. Disable vLLM prefix caching by default

6. Restore nnsight-serve install machinery

7. Sync src/nnsight/modeling/vllm/README.md to the v0.7.0 async path

8. Fold intervention-gaps/ into docs/models/vllm.md; move probes to tests/

Uh oh!

Butanium commented Jul 20, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

khaiwang commented May 7, 2026 •

edited

Loading

1. Clone inference-mode tensors on `.save()` (fixes #661)

6. Restore `nnsight-serve` install machinery

7. Sync `src/nnsight/modeling/vllm/README.md` to the v0.7.0 async path

8. Fold `intervention-gaps/` into `docs/models/vllm.md`; move probes to `tests/`