feat(vllm): instrument the V2 model runner for graph-capture traces

[tkarttunen]

Summary

The graph-capture profiler is wired into capture_model() in the legacy GPUModelRunner only. Since vLLM 0.22, VllmConfig.use_v2_model_runner selects the rewritten runner under vllm/v1/worker/gpu/ by default for many configurations, and that runner contains no reference to capture_torch_profiler at all.

The failure is silent. The patch applies, the server starts and serves, CUDA graphs are captured normally, capture_traces/ stays empty, and neither of the patch's own log lines is emitted. Nothing reports an error, so the only way to notice is to go looking for traces that were never written. The workaround has been to pin VLLM_USE_V2_MODEL_RUNNER=0, which forces users onto a runner vLLM no longer selects by default.

This instruments CudaGraphManager.capture() in vllm/v1/worker/gpu/cudagraph_utils.py, which owns the per-shape capture loop for both PIECEWISE and FULL modes, for 0.22 through 0.25.

What changed

• CudaGraphManager already holds vllm_config, so the profiler is built where it is used. Nothing is threaded through model_runner.py the way the legacy patch threads a profiler= kwarg down two call levels.
• Behaviour matches the legacy path exactly: same rank-0 guard, same trace directory (torch_profiler_dir/capture_traces), same worker_name, and the same capture_{num_tokens}_{mode} annotation that CaptureAnnotation's capture_(\d+)_(.*) pattern parses.
• Warmup passes stay outside the profiled region, matching the legacy runner.
• Additions only: 514 lines across the four patch files, no deletions.

Evidence

Verified on MI350X for 0.22 and 0.25, the two ends of the affected range, with no VLLM_USE_V2_MODEL_RUNNER set anywhere:

INFO [gpu_worker.py] Using V2 Model Runner
INFO [cudagraph_utils.py] Rank 0: Torch profiler enabled for CUDA graph capture,
     traces will be saved to: /tmp/traces/capture_traces

Version	Runner selected	Capture shards	Annotation sample
0.22	V2	102	capture_512_PIECEWISE
0.25	V2	102	capture_512_PIECEWISE

102 shards is the same count the legacy runner produces for the same model, so the V2 path now yields equivalent output rather than merely non-empty output.

Test plan

• All four patches apply to their base images (vllm/vllm-openai-rocm:v0.2{2,3,4,5}.0) with strict git apply, as a single invocation
• Every touched file byte-compiles after patching
• Hunk headers verified: declared line counts and offsets match content
• No lines exceed the 88-column limit
• Capture traces produced on GPU with the V2 runner active and no environment pin
• Annotations parse under CaptureAnnotation

Notes for review

• Indenting the capture body under the profiler pushed two pre-existing comments past 88 columns, so they are rewrapped. No logic in them changed.
• No file overlap with Fix magpie profiler flags #755, which covers 0.14 through 0.21.
• 0.23 and 0.24 were validated statically but not on GPU; their capture() bodies are byte-identical to 0.25's, so the same patch text applies unchanged.

Made with Cursor

[devalshahamd]

Thanks @tkarttunen for the PR!

I am going through the patches, I think we need some additions to cover speculator and encoder models, and also annotations for that. I will try to build on top of this and merge both PRs after testing.