@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 0a5747d410

[Profiler] Add Proton CUDA graph attribution for MRV2 (#51084)

Signed-off-by: Luossu <zhluosuu@outlook.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

鐘天楽 committed Sep 14, 2026 at 15:02 UTC 0a5747d410da9b08145dba2b6f3e6614bab35753

8 files changed +657 -45

docs/contributing/profiling.md

+32 -14

@@ -90,18 +90,18 @@ $ curl -X POST http://localhost:8000/stop_profile

90 is Triton's GPU profiler. It can collect a low-overhead aggregate tree or a

91 Chrome trace and works through the same vLLM profiling controls as the PyTorch

92 and CUDA profilers. Proton currently supports NVIDIA GPUs through CUPTI and

93 -requires eager execution.

93 +supports CUDA graph attribution.

95 -Start a server with a local output directory:

95 +Start a server with a local output directory and graph attribution:

97 ```bash

98 vllm serve meta-llama/Llama-3.1-8B-Instruct \

99 - --enforce-eager \

99 --profiler-config '{

100 "profiler": "proton",

101 "proton_profiler_dir": "./proton_profile",

102 "proton_output_format": "hatchet",

104 - "proton_hook": "triton"

103 + "proton_hook": "triton",

104 + "proton_graph_attribution": true

105 }'

106 ```

107

@@ -111,7 +111,7 @@ rank-qualified output name, such as

111 `proton_dp0_pp0_tp0_dcp0_ep0_rank0_pid1234_0123456789abcdef0123456789abcdef_run0.hatchet`,

112 so distributed workers, restarted servers, and repeated profiling runs do not

113 overwrite one another. A `profile_prefix` is included when supplied. Each

114 -profile is finalized by `/stop_profile` and is ready to inspect immediately.

114 +profile is written by `/stop_profile` and is ready to inspect immediately.

115

116 The Proton-specific options are:

117

@@ -121,21 +121,39 @@ The Proton-specific options are:

121 - `proton_mode`: an optional backend mode string

122 - `proton_hook`: `triton` to record Triton launch metadata, or unset

123 - `proton_output_format`: `hatchet`, `hatchet_msgpack`, `chrome_trace`, or unset

124 +- `proton_graph_attribution`: observe CUDA graph capture for replay attribution;

125 + disabled by default and requires `proton_data: "tree"`

126

127 `hatchet` and `hatchet_msgpack` require `proton_data: "tree"`, while

126 -`chrome_trace` requires `proton_data: "trace"`. CUDA graph profiling is not yet

127 -supported, so `--enforce-eager` is required for every Proton mode.

128 +`chrome_trace` requires `proton_data: "trace"`.

129

130 Automatic backend selection is recommended. vLLM currently supports Proton's

131 `cupti` backend on NVIDIA GPUs. ROCm support is not yet available. vLLM does not

132 expose Proton's experimental instrumentation backend because current upstream

132 -Triton builds can produce profiles without timing metrics. Backend-specific

133 -modes, including `pcsampling`, can be selected with `proton_mode`.

134 -

135 -Triton 3.6 supports explicit `hatchet` and `chrome_trace` output. The

136 -`hatchet_msgpack` format and `periodic_flushing` mode require Triton 3.7 or

137 -newer; vLLM rejects these options with the detected version before starting a

138 -session.

133 +can produce profiles without timing metrics. When `proton_graph_attribution` is

134 +enabled, Proton observes vLLM's CUDA graph capture with the configured profiling

135 +session active, then deactivates that same session until profiling starts. This

136 +lets later profiles attribute replayed kernels without retaining model-startup

137 +activity. Backend-specific modes can be selected with `proton_mode`;

138 +`pcsampling` synchronizes the CUDA context and therefore requires

139 +`--enforce-eager`. When CUDA graphs are enabled (including encoder graphs),

140 +Proton requires `proton_graph_attribution: true` to collect replayed kernels.

141 +For Chrome traces, disable CUDA graphs with `--enforce-eager`.

142 +

143 +CUDA graph-attributed profiles support repeated `start_profile`/`stop_profile`

144 +runs. Each stop flushes and writes one tree-data phase while preserving the

145 +graph-aware session. Without `proton_graph_attribution`, each `stop_profile`

146 +instead finalizes and writes an independent Proton session.

147 +

148 +CUDA graph attribution requires Triton 3.7 or newer. It uses the phase data API

149 +to discard graph-capture activity and separate profiling runs. The `hatchet_msgpack`

150 +output format and `periodic_flushing` mode also require Triton 3.7 or newer.

151 +`periodic_flushing` cannot be combined with graph attribution because both

152 +manage the session's data phases. Ordinary Proton profiling remains available

153 +with Triton 3.6.

154 +

155 +Graph attribution retains capture metadata for the worker lifetime. With eager

156 +execution or no graphs to capture, profiling uses ordinary independent sessions.

157

158 Inspect tree profiles with:

159

tests/v1/worker/test_gpu_profiler.py

+393 -5

@@ -350,17 +350,24 @@ def test_profiler_entered_during_capture():

350

351

352 def make_proton(session_id: int | None = 7):

353 + data = SimpleNamespace(

354 + advance_phase=Mock(side_effect=range(1, 100)),

355 + clear=Mock(),

356 + get=Mock(return_value={"traceEvents": []}),

357 + get_msgpack=Mock(return_value=b"profile"),

358 + )

359 return SimpleNamespace(

360 start=Mock(return_value=session_id),

361 activate=Mock(),

362 deactivate=Mock(),

363 finalize=Mock(),

364 scope=Mock(return_value=nullcontext()),

365 + data=data,

366 )

367

368

369 def make_proton_wrapper(

363 - tmp_path, proton=None, triton_version="3.6.0", **config_overrides

370 + tmp_path, proton=None, triton_version="3.7.0", **config_overrides

371 ):

372 proton = proton or make_proton()

373 config = ProfilerConfig(

@@ -415,6 +422,13 @@ class TestProtonConfig:

422 },

423 "requires proton_data",

424 ),

425 + (

426 + {

427 + "proton_data": "trace",

428 + "proton_graph_attribution": True,

429 + },

430 + "requires proton_data='tree'",

431 + ),

432 ],

433 )

434 def test_rejects_invalid_option_combinations(self, tmp_path, options, message):

@@ -450,6 +464,7 @@ class TestProtonConfig:

464 ("proton_mode", "pcsampling"),

465 ("proton_hook", "triton"),

466 ("proton_output_format", "chrome_trace"),

467 + ("proton_graph_attribution", True),

468 ],

469 )

470 def test_rejects_proton_options_for_other_profilers(self, field, value):

@@ -478,17 +493,103 @@ class TestProtonConfig:

493 compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.NONE),

494 )

495

481 - def test_rejects_proton_when_cuda_graphs_are_enabled(self, tmp_path):

482 - with pytest.raises(ValueError, match="requires CUDA graphs to be disabled"):

496 + @pytest.mark.parametrize("encoder_only", [False, True])

497 + @pytest.mark.parametrize("attribution", [False, True])

498 + def test_cuda_graphs_require_attribution(self, tmp_path, encoder_only, attribution):

499 + # Encoder graphs are independent of the decoder's cudagraph_mode.

500 + expected = (

501 + nullcontext()

502 + if attribution

503 + else pytest.raises(

504 + ValueError, match="requires proton_graph_attribution=True"

505 + )

506 + )

507 + with expected:

508 + VllmConfig(

509 + profiler_config=ProfilerConfig(

510 + profiler="proton",

511 + proton_profiler_dir=str(tmp_path),

512 + proton_graph_attribution=attribution,

513 + ),

514 + compilation_config=CompilationConfig(

515 + cudagraph_mode=(

516 + CUDAGraphMode.NONE if encoder_only else CUDAGraphMode.FULL

517 + ),

518 + cudagraph_mm_encoder=encoder_only,

519 + ),

520 + )

521 +

522 + def test_validates_default_cuda_graph_mode_after_resolution(self, tmp_path):

523 + with pytest.raises(ValueError, match="requires proton_graph_attribution=True"):

524 VllmConfig(

525 profiler_config=ProfilerConfig(

526 profiler="proton", proton_profiler_dir=str(tmp_path)

527 ),

528 + )

529 +

530 + @pytest.mark.parametrize(

531 + "mode",

532 + [

533 + "periodic_flushing",

534 + "periodic_flushing:format=hatchet",

535 + "PERIODIC_FLUSHING:format=hatchet",

536 + ],

537 + )

538 + def test_graph_attribution_rejects_periodic_flushing(self, tmp_path, mode):

539 + # Reject before Proton's native phase manager can abort the worker.

540 + with pytest.raises(ValueError, match="incompatible with periodic_flushing"):

541 + ProfilerConfig(

542 + profiler="proton",

543 + proton_profiler_dir=str(tmp_path),

544 + proton_graph_attribution=True,

545 + proton_mode=mode,

546 + )

547 +

548 + @pytest.mark.parametrize("attribution", [False, True])

549 + @pytest.mark.parametrize(

550 + "mode", ["pcsampling", "pcsampling:interval=100", "PcSampling:interval=100"]

551 + )

552 + @pytest.mark.parametrize("encoder_only", [False, True])

553 + def test_pcsampling_requires_graphs_disabled(

554 + self, tmp_path, attribution, mode, encoder_only

555 + ):

556 + with pytest.raises(ValueError, match="PC sampling requires CUDA graphs"):

557 + VllmConfig(

558 + profiler_config=ProfilerConfig(

559 + profiler="proton",

560 + proton_profiler_dir=str(tmp_path),

561 + proton_graph_attribution=attribution,

562 + proton_mode=mode,

563 + ),

564 compilation_config=CompilationConfig(

488 - cudagraph_mode=CUDAGraphMode.PIECEWISE

565 + cudagraph_mode=CUDAGraphMode.NONE

566 + if encoder_only

567 + else CUDAGraphMode.FULL,

568 + cudagraph_mm_encoder=encoder_only,

569 ),

570 )

571

572 + def test_ordinary_proton_keeps_mrv1_graph_restriction(self, tmp_path, monkeypatch):

573 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")

574 + with pytest.raises(ValueError, match="requires proton_graph_attribution=True"):

575 + VllmConfig(

576 + profiler_config=ProfilerConfig(

577 + profiler="proton", proton_profiler_dir=str(tmp_path)

578 + ),

579 + compilation_config=CompilationConfig(cudagraph_mode=CUDAGraphMode.FULL),

580 + )

581 +

582 + def test_graph_attribution_rejects_mrv1(self, tmp_path, monkeypatch):

583 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")

584 + with pytest.raises(ValueError, match="requires the V2 model runner"):

585 + VllmConfig(

586 + profiler_config=ProfilerConfig(

587 + profiler="proton",

588 + proton_profiler_dir=str(tmp_path),

589 + proton_graph_attribution=True,

590 + )

591 + )

592 +

593

594 @_requires_cuda_for_proton

595 class TestProtonProfilerWrapper:

@@ -564,7 +665,20 @@ class TestProtonProfilerWrapper:

665 )

666 def test_newer_features_reject_triton_3_6(self, tmp_path, option, value, feature):

667 with pytest.raises(RuntimeError, match=feature):

567 - make_proton_wrapper(tmp_path, **{option: value})

668 + make_proton_wrapper(tmp_path, triton_version="3.6.0", **{option: value})

669 +

670 + @pytest.mark.parametrize("version", ["3.6.0", "unknown"])

671 + def test_graph_attribution_requires_phase_api(self, tmp_path, version):

672 + with pytest.raises(RuntimeError, match="requires Triton >= 3.7"):

673 + make_proton_wrapper(

674 + tmp_path, triton_version=version, proton_graph_attribution=True

675 + )

676 +

677 + def test_ordinary_profiling_still_supports_triton_3_6(self, tmp_path):

678 + wrapper, proton = make_proton_wrapper(tmp_path, triton_version="3.6.0")

679 + wrapper.start()

680 + wrapper.stop()

681 + proton.finalize.assert_called_once_with(session=7)

682

683 @pytest.mark.parametrize(

684 ("option", "value"),

@@ -583,6 +697,125 @@ class TestProtonProfilerWrapper:

697 with pytest.raises(RuntimeError, match="does not support selecting"):

698 make_proton_wrapper(tmp_path, proton, proton_output_format="hatchet")

699

700 + def test_cuda_graph_tree_phase_is_written_at_stop(self, tmp_path):

701 + wrapper, proton = make_proton_wrapper(

702 + tmp_path,

703 + proton_context="python",

704 + proton_graph_attribution=True,

705 + )

706 + with wrapper.capture_cuda_graphs():

707 + proton.start.assert_called_once()

708 + proton.deactivate.assert_not_called()

709 +

710 + capture_args = proton.start.call_args.kwargs

711 + assert capture_args["context"] == "python"

712 + assert capture_args["data"] == "tree"

713 + proton.data.advance_phase.assert_called_once_with(7)

714 + proton.deactivate.assert_called_once_with(session=7, flushing=True)

715 + proton.data.clear.assert_called_once_with(7, 0)

716 +

717 + wrapper.start()

718 + wrapper.stop()

719 + wrapper.start()

720 + wrapper.stop()

721 +

722 + assert proton.data.get.call_args_list == [call(7, 1), call(7, 2)]

723 + proton.data.clear.assert_has_calls([call(7, 0), call(7, 1), call(7, 2)])

724 + output_names = sorted(tmp_path.glob("proton_rank_3_*.hatchet"))

725 + assert len(output_names) == 2

726 + assert output_names[0].name.endswith("_run0.hatchet")

727 + assert output_names[1].name.endswith("_run1.hatchet")

728 +

729 + def test_capture_is_noop_without_opt_in(self, tmp_path):

730 + wrapper, proton = make_proton_wrapper(tmp_path)

731 + with wrapper.capture_cuda_graphs():

732 + pass

733 + proton.start.assert_not_called()

734 +

735 + @pytest.mark.parametrize("delay", [0, 2])

736 + def test_duplicate_start_preserves_output_prefix(self, tmp_path, delay):

737 + wrapper, proton = make_proton_wrapper(

738 + tmp_path, proton_graph_attribution=True, delay_iterations=delay

739 + )

740 + with wrapper.capture_cuda_graphs():

741 + pass

742 + wrapper.set_output_name("first")

743 + wrapper.start()

744 + wrapper.set_output_name("duplicate")

745 + wrapper.start()

746 + for _ in range(delay):

747 + wrapper.step()

748 + wrapper.stop()

749 + assert len(list(tmp_path.glob("proton_first_*.hatchet"))) == 1

750 + assert not list(tmp_path.glob("proton_duplicate_*"))

751 + wrapper.set_output_name("second")

752 + wrapper.start()

753 + for _ in range(delay):

754 + wrapper.step()

755 + wrapper.stop()

756 + assert len(list(tmp_path.glob("proton_second_*.hatchet"))) == 1

757 +

758 + def test_failed_export_clears_activity_and_allows_next_interval(self, tmp_path):

759 + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)

760 + with wrapper.capture_cuda_graphs():

761 + pass

762 + proton.data.get.side_effect = [OSError("disk full"), []]

763 + wrapper.start()

764 + wrapper.stop()

765 + proton.data.clear.assert_called_with(7, 1)

766 + wrapper.start()

767 + wrapper.stop()

768 + proton.data.clear.assert_called_with(7, 2)

769 + assert len(list(tmp_path.glob("*_run1.hatchet"))) == 1

770 +

771 + def test_cuda_graph_context_deactivates_after_capture_error(self, tmp_path):

772 + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)

773 +

774 + with (

775 + pytest.raises(RuntimeError, match="capture failed"),

776 + wrapper.capture_cuda_graphs(),

777 + ):

778 + raise RuntimeError("capture failed")

779 +

780 + proton.deactivate.assert_called_once_with(session=7, flushing=True)

781 +

782 + def test_cuda_graph_capture_deactivates_when_phase_advance_fails(self, tmp_path):

783 + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)

784 + proton.data.advance_phase.side_effect = RuntimeError("advance failed")

785 +

786 + with (

787 + pytest.raises(RuntimeError, match="advance failed"),

788 + wrapper.capture_cuda_graphs(),

789 + ):

790 + pass

791 +

792 + proton.deactivate.assert_called_once_with(session=7, flushing=True)

793 + proton.data.clear.assert_not_called()

794 +

795 + def test_cuda_graph_stop_deactivates_when_phase_advance_fails(self, tmp_path):

796 + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)

797 + with wrapper.capture_cuda_graphs():

798 + pass

799 + proton.data.clear.reset_mock()

800 + proton.data.advance_phase.side_effect = RuntimeError("advance failed")

801 +

802 + wrapper._start()

803 + with pytest.raises(RuntimeError, match="advance failed"):

804 + wrapper._stop()

805 +

806 + proton.deactivate.assert_called_with(session=7, flushing=True)

807 + proton.data.clear.assert_not_called()

808 +

809 + def test_shutdown_finalizes_cuda_graph_capture_session(self, tmp_path):

810 + wrapper, proton = make_proton_wrapper(tmp_path, proton_graph_attribution=True)

811 + with wrapper.capture_cuda_graphs():

812 + pass

813 +

814 + wrapper.shutdown()

815 + wrapper.shutdown()

816 +

817 + proton.finalize.assert_called_once_with(session=7)

818 +

819 def test_missing_proton_has_actionable_error(self, tmp_path):

820 config = ProfilerConfig(profiler="proton", proton_profiler_dir=str(tmp_path))

821 with (

@@ -638,6 +871,7 @@ def test_gpu_worker_recreates_proton_profiler_for_each_run():

871 ),

872 patch("vllm.v1.worker.gpu_worker.ProtonProfilerWrapper") as wrapper,

873 ):

874 + wrapper.return_value.has_cuda_graph_session = False

875 Worker.profile(worker, profile_prefix="first")

876 Worker.profile(worker, is_start=False)

877 Worker.profile(worker, profile_prefix="second")

@@ -647,3 +881,157 @@ def test_gpu_worker_recreates_proton_profiler_for_each_run():

881 call(worker.profiler_config, worker_name="second_rank1"),

882 ]

883 assert wrapper.return_value.start.call_count == 2

884 +

885 +

886 +@_requires_cuda_for_proton

887 +def test_gpu_worker_reuses_cuda_graph_proton_session():

888 + worker = MagicMock()

889 + worker.rank = 1

890 + worker.profiler = MagicMock(spec=ProtonProfilerWrapper)

891 + worker.profiler.has_cuda_graph_session = True

892 + worker.profiler_config.profiler = "proton"

893 +

894 + with patch(

895 + "vllm.distributed.utils.get_worker_rank_suffix",

896 + return_value="rank1",

897 + ):

898 + Worker.profile(worker, profile_prefix="first")

899 + Worker.profile(worker, is_start=False)

900 +

901 + worker.profiler.set_output_name.assert_called_once_with("first_rank1")

902 + worker.profiler.start.assert_called_once_with()

903 + worker.profiler.stop.assert_called_once_with()

904 +

905 +

906 +@_requires_cuda_for_proton

907 +@pytest.mark.parametrize("runner", ["attribution_off", "v1", "no_capture"])

908 +def test_proton_not_initialized_without_capture(runner):

909 + worker = MagicMock()

910 + worker.profiler = None

911 + worker.profiler_config.profiler = "proton"

912 + worker.profiler_config.proton_graph_attribution = runner != "attribution_off"

913 + worker.use_v2_model_runner = runner != "v1"

914 + worker.model_runner.needs_cudagraph_capture.return_value = runner != "no_capture"

915 +

916 + context = Worker._get_cudagraph_capture_context(worker)

917 +

918 + assert worker.profiler is None

919 + with context:

920 + pass

921 +

922 +

923 +@_requires_cuda_for_proton

924 +def test_proton_initializes_before_cuda_graph_capture():

925 + class FakeProtonProfiler:

926 + def __init__(self, config, worker_name):

927 + self.config = config

928 + self.worker_name = worker_name

929 + self.capture_context = nullcontext()

930 +

931 + def capture_cuda_graphs(self):

932 + return self.capture_context

933 +

934 + worker = MagicMock()

935 + worker.rank = 2

936 + worker.profiler = None

937 + worker.profiler_config.profiler = "proton"

938 + worker.profiler_config.proton_graph_attribution = True

939 + worker.use_v2_model_runner = True

940 + worker.model_runner.needs_cudagraph_capture.return_value = True

941 +

942 + with (

943 + patch(

944 + "vllm.distributed.utils.get_worker_rank_suffix",

945 + return_value="rank2",

946 + ),

947 + patch(

948 + "vllm.profiler.wrapper.ProtonProfilerWrapper",

949 + FakeProtonProfiler,

950 + ),

951 + ):

952 + context = Worker._get_cudagraph_capture_context(worker)

953 +

954 + assert worker.profiler.config is worker.profiler_config

955 + assert worker.profiler.worker_name == "rank2"

956 + assert context is worker.profiler.capture_context

957 +

958 +

959 +@_requires_cuda_for_proton

960 +@pytest.mark.parametrize("context", ["shadow", "python"])

961 +@pytest.mark.parametrize("output_format", ["hatchet", "hatchet_msgpack"])

962 +def test_proton_cuda_graph_replay_attribution_on_gpu(tmp_path, context, output_format):

963 + """Both intervals contain replay kernels, without capture-only activity."""

964 + import json

965 +

966 + import torch

967 + import triton

968 + import triton.profiler as proton

969 + from packaging.version import Version

970 +

971 + if Version(triton.__version__) < Version("3.7"):

972 + pytest.skip("Graph attribution requires Triton >= 3.7")

973 +

974 + wrapper = ProtonProfilerWrapper(

975 + ProfilerConfig(

976 + profiler="proton",

977 + proton_profiler_dir=str(tmp_path),

978 + proton_graph_attribution=True,

979 + proton_context=context,

980 + proton_output_format=output_format,

981 + ),

982 + worker_name="gpu",

983 + )

984 + worker = SimpleNamespace(

985 + profiler=wrapper,

986 + use_v2_model_runner=True,

987 + )

988 + x = torch.ones(1024, device="cuda")

989 + graph = torch.cuda.CUDAGraph()

990 +

991 + def capture_only():

992 + x.add_(1)

993 +

994 + def captured_add():

995 + x.add_(1)

996 +

997 + def kernel_metrics(value):

998 + if isinstance(value, list):

999 + return [metric for child in value for metric in kernel_metrics(child)]

1000 + if isinstance(value, dict):

1001 + metrics = value.get("metrics", {})

1002 + return ([metrics] if metrics.get("count", 0) else []) + [

1003 + metric for child in value.values() for metric in kernel_metrics(child)

1004 + ]

1005 + return []

1006 +

1007 + try:

1008 + with Worker._get_cudagraph_capture_context(worker):

1009 + with proton.scope("capture_only"):

1010 + capture_only()

1011 + torch.accelerator.synchronize()

1012 + with torch.cuda.graph(graph), proton.scope("captured_add"):

1013 + captured_add()

1014 + for run in range(2):

1015 + wrapper.start()

1016 + with wrapper.annotate_context_manager(f"replay_{run}"):

1017 + graph.replay()

1018 + wrapper.stop()

1019 + (path,) = tmp_path.glob(f"*_run{run}.{output_format}")

1020 + if output_format == "hatchet_msgpack":

1021 + import msgpack

1022 +

1023 + data = msgpack.unpackb(path.read_bytes())

1024 + else:

1025 + data = json.loads(path.read_text())

1026 + serialized = json.dumps(data)

1027 + assert "<captured_at>" in serialized

1028 + assert "captured_add" in serialized

1029 + assert "capture_only" not in serialized

1030 + assert f"replay_{1 - run}" not in serialized

1031 + metrics = kernel_metrics(data)

1032 + assert sum(metric["count"] for metric in metrics) == 1

1033 + assert sum(metric["time (ns)"] for metric in metrics) > 0

1034 + torch.testing.assert_close(x, torch.full_like(x, 4))

1035 + finally:

1036 + wrapper.shutdown()

1037 + assert not list(tmp_path.glob(".proton_cuda_graph_session*"))

vllm/config/profiler.py

+16

@@ -76,6 +76,10 @@ class ProfilerConfig:

76 """Optional format passed to Proton when finalizing a profile. ``None``

77 uses the default format for ``proton_data``."""

79 + proton_graph_attribution: bool = False

80 + """Observe CUDA graph capture so replayed kernels can be attributed.

81 + Requires Triton >= 3.7 and ``proton_data='tree'``."""

82 +

83 torch_profiler_with_stack: bool = True

84 """If `True`, enables stack tracing in the torch profiler. Enabled by default

85 as it is useful for debugging. Can be disabled via

@@ -199,6 +203,7 @@ class ProfilerConfig:

203 ("proton_mode", self.proton_mode, None),

204 ("proton_hook", self.proton_hook, None),

205 ("proton_output_format", self.proton_output_format, None),

206 + ("proton_graph_attribution", self.proton_graph_attribution, False),

207 )

208 if value != default

209 ]

@@ -220,6 +225,17 @@ class ProfilerConfig:

225

226 if self.profiler == "proton":

227 output_format = self.proton_output_format

228 + if self.proton_graph_attribution and self.proton_data != "tree":

229 + raise ValueError("proton_graph_attribution requires proton_data='tree'")

230 + if (

231 + self.proton_graph_attribution

232 + and self.proton_mode

233 + and self.proton_mode.split(":", 1)[0].lower() == "periodic_flushing"

234 + ):

235 + raise ValueError(

236 + "proton_graph_attribution is incompatible with periodic_flushing: "

237 + "both manage Proton data phases."

238 + )

239 if output_format == "chrome_trace" and self.proton_data != "trace":

240 raise ValueError("chrome_trace output requires proton_data='trace'")

241 if (

vllm/config/vllm.py

+39 -12

@@ -1551,18 +1551,6 @@ class VllmConfig:

1551 self.compilation_config.mode = CompilationMode.NONE

1552 self.compilation_config.cudagraph_mode = CUDAGraphMode.NONE

1553

1554 - if self.profiler_config.profiler == "proton":

1555 - if not current_platform.is_cuda():

1556 - raise ValueError(

1557 - "The Proton profiler currently supports NVIDIA CUDA only"

1558 - )

1559 - if self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:

1560 - raise ValueError(

1561 - "The Proton profiler requires CUDA graphs to be disabled. "

1562 - "Use --enforce-eager or set "

1563 - "--compilation-config.cudagraph_mode=none."

1564 - )

1565 -

1554 if os.environ.get("TORCH_COMPILE_DISABLE") == "1":

1555 logger.warning_once(

1556 "TORCH_COMPILE_DISABLE is set, disabling torch.compile. "

@@ -1860,6 +1848,7 @@ class VllmConfig:

1848 else:

1849 self._validate_v1_model_runner()

1850

1851 + self._validate_profiler_config()

1852 self._validate_batch_sharded_sampling()

1853 self._validate_adaptive_verification()

1854

@@ -3042,6 +3031,44 @@ class VllmConfig:

3031

3032 return unsupported

3033

3034 + def _validate_profiler_config(self) -> None:

3035 + if self.profiler_config.profiler != "proton":

3036 + return

3037 +

3038 + from vllm.platforms import current_platform

3039 +

3040 + if not current_platform.is_cuda():

3041 + raise ValueError("The Proton profiler currently supports NVIDIA CUDA only")

3042 + if (

3043 + self.profiler_config.proton_graph_attribution

3044 + and not self.use_v2_model_runner

3045 + ):

3046 + raise ValueError(

3047 + "Proton CUDA graph attribution requires the V2 model runner."

3048 + )

3049 +

3050 + has_cuda_graphs = (

3051 + self.compilation_config.cudagraph_mode != CUDAGraphMode.NONE

3052 + or (

3053 + self.compilation_config.cudagraph_mm_encoder

3054 + and not (self.model_config and self.model_config.enforce_eager)

3055 + )

3056 + )

3057 + if not has_cuda_graphs:

3058 + return

3059 + mode = self.profiler_config.proton_mode

3060 + if mode and mode.split(":", 1)[0].lower() == "pcsampling":

3061 + raise ValueError(

3062 + "Proton PC sampling requires CUDA graphs to be disabled. "

3063 + "Use --enforce-eager."

3064 + )

3065 + if not self.profiler_config.proton_graph_attribution:

3066 + raise ValueError(

3067 + "Proton profiling with CUDA graphs requires "

3068 + "proton_graph_attribution=True to capture replayed kernels. "

3069 + "Enable attribution or use --enforce-eager."

3070 + )

3071 +

3072 def _validate_v2_model_runner(self) -> None:

3073 """Check for features not yet supported by the V2 model runner."""

3074 if not HAS_TRITON:

vllm/profiler/wrapper.py

+138 -11

@@ -6,8 +6,9 @@ import inspect

6 import json

7 import os

8 from abc import ABC, abstractmethod

9 -from collections.abc import Callable

10 -from contextlib import nullcontext

9 +from collections.abc import Callable, Iterator

10 +from contextlib import AbstractContextManager, contextmanager, nullcontext, suppress

11 +from glob import glob

12 from typing import Literal

13 from uuid import uuid4

@@ -160,6 +161,15 @@ class WorkerProfiler(ABC):

161 if self._running:

162 self.stop()

163

164 + @property

165 + def has_cuda_graph_session(self) -> bool:

166 + """Whether a capture-time session is retained to attribute replays."""

167 + return False

168 +

169 + def capture_cuda_graphs(self) -> AbstractContextManager[None]:

170 + """Observe graph creation for backends that attribute replay activity."""

171 + return nullcontext()

172 +

173 def annotate_context_manager(self, name: str):

174 """Return a context manager to annotate profiler traces."""

175 return nullcontext()

@@ -387,6 +397,7 @@ class ProtonProfilerWrapper(WorkerProfiler):

397 self._mode = profiler_config.proton_mode

398 self._hook = profiler_config.proton_hook

399 self._output_format = profiler_config.proton_output_format

400 + self._graph_attribution = profiler_config.proton_graph_attribution

401 self._triton_version_string = getattr(triton, "__version__", "unknown")

402 try:

403 self._triton_version = Version(self._triton_version_string)

@@ -398,6 +409,13 @@ class ProtonProfilerWrapper(WorkerProfiler):

409 # worker cannot overwrite profiles left by an earlier server process.

410 self._instance_id = f"pid{os.getpid()}_{uuid4().hex}"

411 self._run_id = 0

412 + self._graph_session = False

413 + self._phase = 0

414 + self._active_output_path: str | None = None

415 + self._session_storage_path = os.path.join(

416 + self._output_dir,

417 + f".proton_cuda_graph_session_{worker_name}_{self._instance_id}",

418 + )

419

420 logger.info_once(

421 "Proton profiling enabled. Output will be saved under: %s",

@@ -412,6 +430,10 @@ class ProtonProfilerWrapper(WorkerProfiler):

430 )

431

432 def _validate_capabilities(self) -> None:

433 + if self._graph_attribution:

434 + self._require_triton_version(

435 + "CUDA graph attribution", _TRITON_PROTON_3_7_VERSION

436 + )

437 if self._output_format is not None:

438 parameters = inspect.signature(self._proton.finalize).parameters

439 supports_output_format = "output_format" in parameters or any(

@@ -428,7 +450,7 @@ class ProtonProfilerWrapper(WorkerProfiler):

450 self._require_triton_version(

451 "hatchet_msgpack output", _TRITON_PROTON_3_7_VERSION

452 )

431 - if self._mode and self._mode.split(":", 1)[0] == "periodic_flushing":

453 + if self._mode and self._mode.split(":", 1)[0].lower() == "periodic_flushing":

454 self._require_triton_version(

455 "periodic flushing", _TRITON_PROTON_3_7_VERSION

456 )

@@ -447,28 +469,112 @@ class ProtonProfilerWrapper(WorkerProfiler):

469 raise RuntimeError("Proton did not create a profiling session")

470 return session_id

471

472 + @property

473 + def has_cuda_graph_session(self) -> bool:

474 + return self._graph_session

475 +

476 + def set_output_name(self, worker_name: str) -> None:

477 + """Set the next run's output name after startup graph capture."""

478 + if self._active:

479 + return

480 + self._output_path = os.path.join(self._output_dir, f"proton_{worker_name}")

481 +

482 + @override

483 + @contextmanager

484 + def capture_cuda_graphs(self) -> Iterator[None]:

485 + """Keep a Proton session active while vLLM captures CUDA graphs."""

486 + if not self._graph_attribution:

487 + yield

488 + return

489 + if self._session_id is None:

490 + self._session_id = self._create_session(self._session_storage_path)

491 + else:

492 + self._proton.activate(session=self._session_id)

493 + self._graph_session = True

494 +

495 + try:

496 + yield

497 + finally:

498 + captured_phase = self._phase

499 + try:

500 + self._phase = self._proton.data.advance_phase(self._session_id)

501 + finally:

502 + self._proton.deactivate(session=self._session_id, flushing=True)

503 + self._proton.data.clear(self._session_id, captured_phase)

504 +

505 @override

506 def _start(self) -> None:

452 - output_path = f"{self._output_path}_{self._instance_id}_run{self._run_id}"

453 - self._session_id = self._create_session(output_path)

507 + self._active_output_path = (

508 + f"{self._output_path}_{self._instance_id}_run{self._run_id}"

509 + )

510 self._run_id += 1

511 + if self._graph_session:

512 + assert self._session_id is not None

513 + self._proton.activate(session=self._session_id)

514 + else:

515 + self._session_id = self._create_session(self._active_output_path)

516 +

517 + def _write_graph_phase(self, phase: int) -> None:

518 + assert self._active_output_path is not None

519 + output_format = self._output_format or "hatchet"

520 + output_path = f"{self._active_output_path}.{output_format}"

521 + if output_format == "hatchet_msgpack":

522 + with open(output_path, "wb") as output_file:

523 + output_file.write(

524 + self._proton.data.get_msgpack(self._session_id, phase)

525 + )

526 + else:

527 + with open(output_path, "w", encoding="utf-8") as output_file:

528 + json.dump(self._proton.data.get(self._session_id, phase), output_file)

529 +

530 + def _finalize_session(self, session_id: int) -> None:

531 + if self._output_format is None:

532 + self._proton.finalize(session=session_id)

533 + else:

534 + self._proton.finalize(session=session_id, output_format=self._output_format)

535

536 @override

537 def _stop(self) -> None:

538 assert self._session_id is not None

539 session_id = self._session_id

540 + if self._graph_session:

541 + completed_phase = self._phase

542 + try:

543 + self._phase = self._proton.data.advance_phase(session_id)

544 + finally:

545 + self._proton.deactivate(session=session_id, flushing=True)

546 + try:

547 + self._write_graph_phase(completed_phase)

548 + finally:

549 + self._proton.data.clear(session_id, completed_phase)

550 + self._active_output_path = None

551 + return

552 +

553 try:

554 self._proton.deactivate(session=session_id)

555 finally:

556 try:

464 - if self._output_format is None:

465 - self._proton.finalize(session=session_id)

466 - else:

467 - self._proton.finalize(

468 - session=session_id, output_format=self._output_format

469 - )

557 + self._finalize_session(session_id)

558 finally:

559 self._session_id = None

560 + self._active_output_path = None

561 +

562 + @override

563 + def shutdown(self) -> None:

564 + super().shutdown()

565 + if self._graph_session and self._session_id is not None:

566 + session_id = self._session_id

567 + self._session_id = None

568 + try:

569 + self._finalize_session(session_id)

570 + except Exception:

571 + logger.exception(

572 + "Failed to finalize Proton CUDA graph session during shutdown."

573 + )

574 + finally:

575 + for output_path in glob(f"{self._session_storage_path}.*"):

576 + with suppress(FileNotFoundError):

577 + os.remove(output_path)

578

579 @override

580 def annotate_context_manager(self, name: str):

@@ -496,3 +602,24 @@ class CudaProfilerWrapper(WorkerProfiler):

602 @override

603 def annotate_context_manager(self, name: str):

604 return torch.cuda.nvtx.range(name)

605 +

606 +

607 +def create_graph_capture_profiler(

608 + profiler_config: ProfilerConfig, global_rank: int

609 +) -> WorkerProfiler | None:

610 + """Create a profiler to observe CUDA graph capture, if configured.

611 +

612 + Applies only to backends that attribute graph replay activity which

613 + need a session around capture (Proton with ``proton_graph_attribution``).

614 + """

615 + if (

616 + profiler_config.profiler == "proton"

617 + and profiler_config.proton_graph_attribution

618 + ):

619 + from vllm.distributed.utils import get_worker_rank_suffix

620 +

621 + return ProtonProfilerWrapper(

622 + profiler_config,

623 + worker_name=get_worker_rank_suffix(global_rank=global_rank),

624 + )

625 + return None

vllm/v1/worker/gpu/model_runner.py

+10

@@ -987,6 +987,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):

987 """Estimate the GPU memory required to capture CUDA graphs."""

988 return _profile_cudagraph_memory(self)

989

990 + def needs_cudagraph_capture(self) -> bool:

991 + """Whether capture_model() has any CUDA graphs left to capture."""

992 + return (

993 + self.model_state.supports_mm_inputs

994 + and self.model_state.encoder_runner.has_cudagraph()

995 + ) or (

996 + self.cudagraph_manager is not None

997 + and self.cudagraph_manager.needs_capture()

998 + )

999 +

1000 @torch.inference_mode()

1001 def capture_model(self, *, profile_only: bool = False) -> int:

1002 assert self.cudagraph_manager is not None

vllm/v1/worker/gpu_worker.py

+26 -3

@@ -9,7 +9,7 @@ from collections.abc import Callable

9 from contextlib import AbstractContextManager, contextmanager, nullcontext

10 from datetime import timedelta

11 from types import NoneType

12 -from typing import TYPE_CHECKING, Any

12 +from typing import TYPE_CHECKING, Any, cast

14 import numpy as np

15 import regex as re

@@ -64,6 +64,7 @@ from vllm.profiler.wrapper import (

64 CudaProfilerWrapper,

65 ProtonProfilerWrapper,

66 TorchProfilerWrapper,

67 + create_graph_capture_profiler,

68 )

69 from vllm.sequence import IntermediateTensors

70 from vllm.tasks import SupportedTask

@@ -142,6 +143,7 @@ def maybe_rocm_profiling_fallback(profile_result: MemoryProfilingResult) -> int

143 if TYPE_CHECKING:

144 from vllm.device_allocator.sleep_mode_backend import SleepModeBackend

145 from vllm.model_executor.model_loader.tensorizer import TensorizerConfig

146 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2

147 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

148

149

@@ -810,7 +812,8 @@ class Worker(WorkerBase):

812

813 cuda_graph_memory_bytes = 0

814 if not self.model_config.enforce_eager:

813 - cuda_graph_memory_bytes = self.model_runner.capture_model()

815 + with self._get_cudagraph_capture_context():

816 + cuda_graph_memory_bytes = self.model_runner.capture_model()

817

818 # Compare actual vs estimated CUDA graph memory (if we did profiling)

819 if (

@@ -949,6 +952,21 @@ class Worker(WorkerBase):

952 encoder=self.compilation_config.encoder_compilation_time,

953 )

954

955 + def _get_cudagraph_capture_context(self) -> AbstractContextManager[None]:

956 + """Let the configured profiler observe CUDA graph capture."""

957 + if not self.use_v2_model_runner:

958 + return nullcontext()

959 + if self.profiler is None:

960 + model_runner = cast("GPUModelRunnerV2", self.model_runner)

961 + if not model_runner.needs_cudagraph_capture():

962 + return nullcontext()

963 + self.profiler = create_graph_capture_profiler(

964 + self.profiler_config, global_rank=self.rank

965 + )

966 + if self.profiler is None:

967 + return nullcontext()

968 + return self.profiler.capture_cuda_graphs()

969 +

970 def reset_mm_cache(self) -> None:

971 self.model_runner.reset_mm_cache()

972

@@ -1248,6 +1266,9 @@ class Worker(WorkerBase):

1266 else:

1267 trace_name = rank_suffix

1268

1269 + if profiler_type == "proton" and self.profiler is not None:

1270 + self.profiler.set_output_name(trace_name)

1271 +

1272 # Create the profiler wrapper only on the first start call

1273 if self.profiler is None:

1274 if profiler_type == "torch":

@@ -1284,7 +1305,9 @@ class Worker(WorkerBase):

1305 try:

1306 self.profiler.stop()

1307 finally:

1287 - if self.profiler_config.profiler == "proton":

1308 + if self.profiler_config.profiler == "proton" and not (

1309 + self.profiler.has_cuda_graph_session

1310 + ):

1311 # Proton output names are fixed when the wrapper is constructed.

1312 # Recreate it so the next profile_prefix is honored.

1313 self.profiler = None

vllm/v1/worker/mm_encoder_model_runner.py

@@ -70,6 +70,9 @@ class MMEncoderModelRunner(GPUModelRunner):

70 def capture_model(self, *, profile_only: bool = False) -> int:

71 return 0

73 + def needs_cudagraph_capture(self) -> bool:

74 + return False

75 +

76 def _dummy_run(

77 self, *args: Any, **kwargs: Any

78 ) -> tuple[torch.Tensor, torch.Tensor]: