[hmellor]

Overview

The cudagraph "Assigning / modifying buffers" check from #11596 raises on any update in Dynamo's residual bytecode, so it also rejects unrelated dict updates. This PR only flags <module>._buffers.update, which is what Dynamo emits for a buffer assigned during forward.

Claims

• Transformers backend models whose config sets is_causal (e.g. Ettin / ModernBERT decoder) compile with cudagraphs again: Transformers' is_causal handling leaves kwargs.update({'is_causal': ...}) in the residual bytecode.
• Buffer assignment during forward is still rejected.
• Needed by [Transformers backend] Run the output head of causal models #61276, whose Ettin test hits this.

Validation

tests/compile/test_wrapper.py::test_cudagraph_buffer_assignment_check covers both cases on CPU; the kwargs case fails without the fix.

pytest tests/compile/test_wrapper.py
pytest tests/models/transformers/test_backend.py::test_output_head_transform  # with #61276

Tests: test_wrapper.py passes (CPU and B200); with #61276, Ettin test_output_head_transform goes from failing to passing and all 58 test_backend.py tests pass on B200.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (Mandatory for agents, optional for humans).

• AI assistance was used during the creation of this PR.

• Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.

• Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.

• Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.

• Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.

🤖 Generated with Claude Code

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[SIDDARTHAREDDY8]

I read vllm/compilation/wrapper.py and the new test in tests/compile/test_wrapper.py. The narrowing from "update" in new_code.co_names to _updates_module_buffers is exactly right for the reported false positive. I verified the matcher with dis on Python 3.12: x._buffers.update(...) compiles to consecutive LOAD_ATTR _buffers / LOAD_ATTR update, so the pairwise check keys on the attribute names rather than the opcode, and a plain self.model.buffer = ... does not match on its own — detection relies on Dynamo rewriting the assignment into _buffers.update, which the new BufferAssignWrapper test covers, and KwargsUpdateWrapper covers the kwargs.update case that used to be flagged.

Two things worth considering:

1. Scope: bytecode_hook returns early unless old_code is the wrapper's own forward code object, so a buffer assignment inside a helper method called by forward — compiled by Dynamo as a separate code object — would not be inspected. Is that the intended scope, or should the check apply to any code object the hook sees? If it is intended, a short comment at the early return would make that clear.

2. Nit: the error message still says to search for uses of the function update. Now that the check is specific to _buffers.update, naming that in the message would point at the culprit more directly.

Otherwise this looks good — the test pins both sides of the behavior change.

[hmellor]

/ci run

[github-actions]

✅ Triggered Buildkite CI #101480 for commit a3e674527208.

ℹ️ This PR is 589 commits behind upstream main. Merged upstream main into this branch as a3e674527208 before starting CI. Pull the branch before pushing more changes.

[ZJY0516]

LGTM