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

[JadenFiotto-Kaufman]

Rework of #662 against current 0.8.

Reading this diff: the branch is built on current 0.8, which this branch's base is 167 commits behind, so GitHub shows ~190 files. The actual change is 5 files — everything else is 0.8 churn that arrives because of the base. Read the two commits, or diff against 0.8:

git diff 0.8...pr662-rework

What happened to the original eight

Four of #662's items already landed on 0.8 and are dropped:

#	Item	Status on 0.8
3	Surface deferred intervention errors	Landed — vllm.py:965, async_backend.py:153, serve/backend.py:68
6	Restore nnsight-serve packaging	Landed — pyproject.toml serve extra + [project.scripts]
7	Sync the vLLM README	Superseded — the README became docs/developing/vllm-integration.md
8	Fold intervention-gaps/ into the docs	Landed — the directory is gone

Two more are dropped deliberately:

• Dynamic project version #5, enable_prefix_caching=False by default — obsolete. 0.8 solves this per request: _attach_mediators sets param.skip_reading_prefix_cache = True on every traced request (vllm.py:1042). That is strictly better — the cache stays on for untraced traffic while traced requests recompute. The one case it cannot cover, an edit() registration riding a request nnsight did not create, already warns at registration time (registration.py:77). Flipping the engine-wide default now would give back throughput for nothing.
• Proxy torch function proxy argument not necessarily the first argument should loop through arguments to find the proxy #4, multi-invoke async — deferred to its own PR. The gap is real (async_backend.py:81 still refuses more than one prompt) and the design in vllm: follow-ups from the v0.7.0 merges #662 is sound: one engine request per invoke, fan the streams into one queue, merge invoke-shared saves onto the last finished output. But 0.8 moved collection out of the backend into engines/engine.py (acollect/attach), and vllm: follow-ups from the v0.7.0 merges #662's implementation reverts to raw collective_rpc("collect_nnsight", ...). That drops two things: the RequestOutput that acollect passes down so a block parked on tracer.result can be served, and attach's merge of registered values and per-sequence (n > 1) saves. It needs re-plumbing rather than redesigning, and that is better reviewed on its own.

What is here

fix(vllm): narrow and widen on the token axis, not always dim 0

A model vLLM has no definition of its own is served through its Transformers backend, which runs the wrapped HuggingFace module with a leading singleton batch dim. Its decoder layers emit [1, total_tokens, hidden], so tokens sit on dim 1.

VLLMBatcher inherited the base's dim-0-only row math, which gates on shape[0] == total. On a 3-D activation that reads 1 != total, 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 is not batched — so an intervention quietly does nothing as soon as a second request shares the step.

Unlike #662, this does not touch src/nnsight/intervention/batching.py. The base already documents _narrow_tensor/_widen_tensor as the extension point for non-dim-0 layouts, and DiffusionBatcher already uses it that way, so the new _batch_dim hook in core was not needed. VLLMBatcher overrides both and locates the axis with _token_dim.

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 it had no reader.

Known gap, documented not guessed at: Interleaver.replay trims a padded tap tensor with t[:total], still dim-0-only, so taps over a Transformers-backend model can be served the step's padding rows. Deciding the token axis of a padded tensor needs a rule _token_dim cannot supply — its whole test is shape[dim] == total, which a padded tensor fails by construction. Noted in the dev doc and in Limitations.

feat(vllm): NNSIGHT_VLLM_CLONE_READS serves copies instead of views

Replaces #662's clone-inside-tracing.save() approach to #661.

What a block reads 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. VLLMBatcher.narrow now clones what it serves when NNSIGHT_VLLM_CLONE_READS is set.

Two reasons for moving it out of save():

• Coverage. narrow is the single chokepoint every read passes through — .save(), tracer.cache() (which narrows at interleaver.py:839), appends under tracer.iter, and taps, whose replay uses the same handoff. A clone at the save point only ever covered the first: a saved container is marked once while its elements keep aliasing, which is the failure mode docs/models/vllm.md actually warns about.
• Blast radius. It keeps the change inside the vLLM module instead of changing what save() returns for every backend.

An 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 must be set before VLLM(...). Falsy spellings match the ones NNSIGHT_DISABLE_CPP_BACKTRACE already accepts.

Off by default, and here is the reason to review it. 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 — as silent as the bug it fixes. .output/.input register no transform write-back, so they rely purely on aliasing; an eproperty that does register one is unaffected, and replacement (module.output = out) routes through widen and lands either way. Both halves are pinned by tests and both docs give the form that works under either setting.

The cheap alternative that would keep in-place edits — copying the clone back into the view after the read — needs a hook in core interleaver.py and collides with the eproperty transform slot, so it is not attempted here.

Testing

23 new CPU-only tests in tests/test_batching.py, no engine required. 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.

Full CPU suite: 1054 passed, 16 failed — all 16 are tests/test_tensor_parallel_rules.py and reproduce unchanged on 0.8, so they pre-date this branch.

tests/vllm/ was not run. The box available had a GPU but only vLLM 0.15.1, and the integration targets 0.27.x. The end-to-end vLLM behaviour of both commits is unverified and wants a run before merge.

🤖 Generated with Claude Code

[JadenFiotto-Kaufman]

vLLM suite run — hakone, vLLM 0.27.1, 4x A100-80GB

Closes the "tests/vllm/ was not run" gap in the description. Both commits and the 0.8 merge-base were run through the identical suite, env and GPUs, back to back.

	pr662-rework (28eb29e3)	0.8 baseline (3bed88cc)
passed	216	216
failed	1	1
skipped	17	17
errors	8	8
wall	8:51	8:56

A set-diff of every FAILED and ERROR at setup of line between the two runs is empty — same test names, same counts. The branch is neutral on the vLLM suite.

Both residuals are pre-existing on 0.8:

• 8 errors — tests/vllm/test_deepseek.py. DeepSeek-V2 is not in that box's HF cache and the run used HF_HUB_OFFLINE=1, so every fixture in the file errors at setup (LocalEntryNotFoundError). Environmental.
• 1 failure — test_tensor_parallel.py::TestAdHocCall::test_row_parallel_call_takes_and_returns_the_whole[2]. _min_row_cosine returns 0.126 against assert > 0.99. Reproduces unchanged on the baseline, so it is a real TP=2 ad-hoc-call bug that predates this branch and wants its own issue.

I checked the two mechanisms by which these commits could have touched that TP test before running the baseline: the dropped _nnsight_batch marker has exactly one reader (intervention/backward.py:79, unreachable on the vLLM path), and for a 2-D [total_tokens, hidden] activation the new _narrow_tensor is arithmetically identical to the base's. The baseline confirms it.

CPU suite is unchanged from the description: 1054 passed, and the 16 tests/test_tensor_parallel_rules.py failures reproduce on 0.8 too.

Unrelated finding worth its own issue

The first attempt reported 130 failed, all with RuntimeError: AttributeError: 'Tensor' object has no attribute 'save'. Cause: src/nnsight/_c/*.so is gitignored and built in place, so a fresh clone or git worktree has no compiled extension. mount(save, "save") in src/nnsight/__init__.py:147 is wrapped in a bare except Exception: pass, so the mount fails silently, .save() never exists, and every test that saves anything fails with an error naming neither the extension nor the mount. Copying a matching py_mount.cpython-312-*.so in took it from 130 failures to 1. A one-line warning in that except would have saved an 8-minute diagnostic run.

_{Env: /disk/u/localjadenfk/v27 (vLLM 0.27.1, torch 2.13.0+cu130), CUDA 13 compat libs, VLLM_USE_FLASHINFER_SAMPLER=0, VLLM_ALLOW_INSECURE_SERIALIZATION=1, CUDA_VISIBLE_DEVICES=0,2,6,7.}

🤖 Generated with Claude Code