diff --git a/docs/docs/developer_guide/benchmark_and_profiling.mdx b/docs/docs/developer_guide/benchmark_and_profiling.mdx index a11147873..bd4508b6b 100644 --- a/docs/docs/developer_guide/benchmark_and_profiling.mdx +++ b/docs/docs/developer_guide/benchmark_and_profiling.mdx @@ -296,6 +296,37 @@ The profile merger generates: - Individual rank trace files: `{profile_id}-TP-{tp}-DP-{dp}-PP-{pp}-EP-{ep}.trace.json.gz` - Merged trace file: `merged-{profile_id}.trace.json.gz` +### Profile the CUDA graph capture phase + +The tools above profile the steady-state runtime (prefill / decode). To instead profile the **CUDA graph capture phase** that runs once at server startup, launch the server with `--enable-profile-cuda-graph`. This runs a PyTorch Profiler pass over the decode CUDA-graph capture, which is useful for diagnosing slow or memory-heavy graph capture. + +`--enable-profile-cuda-graph` (server arg) builds the capture profiler and always emits the per-kernel CPU/CUDA time summary tables and a CUDA memory snapshot. Persisting Chrome traces to disk is opt-in via one of two env vars (both no-ops unless `--enable-profile-cuda-graph` is also set): + +- `SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE=1` — writes **one combined trace per tensor-parallel rank** for the whole capture pass, named `cuda_graph_capture-{runner}-TP-{tp_rank}.json.gz`. +- `SGLANG_GRAPH_BATCH_CAPTURE=1` — writes **one trace per captured batch size per rank**, named `{runner}_bs_{bs}_rank{tp_rank}.json.gz`. The profiler runs on a `wait=2, warmup=0, active=1` schedule (the two dummy runs before each capture are skipped) with `record_shapes`, `with_stack`, `with_flops`, and `profile_memory` enabled, giving per-shape kernel identities, input shapes, FLOPs, and memory. + +If both env vars are set, `SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE` (the single combined trace) takes precedence. + +```bash Command +# set trace path +export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log + +# opt in to per-batch-size capture traces (or set +# SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE=1 for a single combined trace per rank) +export SGLANG_GRAPH_BATCH_CAPTURE=1 + +# launch the server with CUDA graph capture profiling enabled +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-profile-cuda-graph +``` + +Behavior and output: + +- All traces are written to `${SGLANG_TORCH_PROFILER_DIR}/graph_capture_profile/` (defaults to `/tmp/graph_capture_profile/` if the variable is unset). Files are namespaced by runner class and TP rank so concurrent capture passes (e.g. EAGLE target/draft/draft-extend) and ranks don't collide. +- A CUDA memory snapshot (`cuda_graph_runner_memory_usage.pickle`) and per-kernel CPU/CUDA time summary tables are always emitted for the capture phase (independent of the env vars above). +- Only the decode CUDA-graph runner is profiled. + +The capture traces are viewed the same way as other PyTorch Profiler traces (see [View traces](#view-traces)). + ### Possible PyTorch bugs If in any cases you encounter the following error (for example, using qwen 2.5 VL): ```bash Command diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index cb68793d1..33cd636ba 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -328,6 +328,10 @@ class Envs: ) SGLANG_RECORD_STEP_TIME = EnvBool(False) SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE = EnvBool(False) + # Opt-in: emit one CUDA-graph capture trace per captured batch size (per-bs). + # SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE (single combined trace) takes + # precedence when both are set. + SGLANG_GRAPH_BATCH_CAPTURE = EnvBool(False) SGLANG_FORCE_SHUTDOWN = EnvBool(False) SGLANG_DEBUG_MEMORY_POOL = EnvBool(False) SGLANG_DSPARK_DEBUG_CONFIDENCE_PREFIX_SCHEDULER = EnvBool(False) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 66bcd6973..338e26b2b 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -28,6 +28,7 @@ from __future__ import annotations import contextlib import inspect import logging +import os from types import SimpleNamespace from typing import TYPE_CHECKING, Callable, Optional, Union @@ -102,7 +103,10 @@ from sglang.srt.utils import ( require_mlp_tp_gather, ) from sglang.srt.utils.device_timer import device_timer_ctx -from sglang.srt.utils.profile_utils import export_cuda_graph_capture_trace +from sglang.srt.utils.profile_utils import ( + export_cuda_graph_capture_trace, + graph_capture_profile_dir, +) try: from kt_kernel import KTMoEWrapper @@ -682,11 +686,63 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): and capture_hidden_mode_matches ) - def _init_profile_context_and_memory_record(self): - profile_context = profile( - activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], - record_shapes=True, + def _graph_batch_capture_active(self) -> bool: + """Whether the per-batch-size capture-trace feature is active. + + Gated by SGLANG_GRAPH_BATCH_CAPTURE. The original single-trace export + (SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE) takes precedence: when both are + set we fall back to the original behavior. + """ + return ( + envs.SGLANG_GRAPH_BATCH_CAPTURE.get() + and not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get() ) + + def _init_profile_context_and_memory_record(self): + if self._graph_batch_capture_active(): + # Per-batch-size capture traces (SGLANG_GRAPH_BATCH_CAPTURE): a + # scheduled profiler is stepped once per batch size (see + # FullCudaGraphBackend.capture_one) and on_trace_ready writes one + # chrome trace per bs. + rank = get_parallel().tp_rank + runner_name = type(self).__name__ + trace_dir = graph_capture_profile_dir() + os.makedirs(trace_dir, exist_ok=True) + + # Track which BS is currently being captured for trace file naming + self._profile_bs_list = list(reversed(self.capture_bs)) + self._profile_bs_idx = 0 + + def on_trace_ready(prof): + bs = self._profile_bs_list[self._profile_bs_idx] + trace_file = os.path.join( + trace_dir, f"{runner_name}_bs_{bs}_rank{rank}.json.gz" + ) + prof.export_chrome_trace(trace_file) + logger.info(f"Saved trace for bs={bs} to {trace_file}") + self._profile_bs_idx += 1 + + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + # Schedule: wait=2 (skip 2 dummy runs), warmup=0, active=1 + # (capture run); repeat=0 repeats the cycle so each batch size + # gets its own trace. + schedule=torch.profiler.schedule(wait=2, warmup=0, active=1, repeat=0), + record_shapes=True, + with_stack=True, + with_flops=True, + profile_memory=True, + on_trace_ready=on_trace_ready, + ) + else: + # a single unscheduled pass over the whole + # capture. The combined trace (if any) is exported in + # _post_process_after_profile via export_cuda_graph_capture_trace, + # gated by SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE. + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + ) torch.cuda.memory._record_memory_history() return profile_context @@ -706,10 +762,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): ) logger.info(log_message) - # Optionally persist the shaped capture trace (record_shapes=True) for - # offline per-kernel analysis -- opt-in via - # SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE; the in-log tables above are - # unchanged. + # single-trace export for the whole capture pass; no-op unless + # SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE is set. In per-bs mode + # (SGLANG_GRAPH_BATCH_CAPTURE) that env is unset, so this stays a no-op + # and the per-bs on_trace_ready handles export instead. export_cuda_graph_capture_trace( prof_context, runner_name=type(self).__name__, @@ -886,8 +942,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): self.model_runner, self.captured_req_width ) profile_context = empty_context() + # Holds the active torch profiler during capture so the backend can + # advance its schedule (profiler.step()) per batch size. Only the + # scheduled per-bs profiler (SGLANG_GRAPH_BATCH_CAPTURE) needs stepping; + # the original unscheduled pass leaves this None. + self._profiler = None if self.enable_profile_cuda_graph: profile_context = self._init_profile_context_and_memory_record() + if self._graph_batch_capture_active(): + self._profiler = profile_context # share_buffers() coalesces seq_lens / seq_lens_cpu through the process- # wide pool, so they may alias a buffer seeded by an earlier runner (the @@ -924,6 +987,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): if self.enable_profile_cuda_graph: self._post_process_after_profile(prof) + self._profiler = None # No pool-side pin to clear: the captured full-physical write loc rides the # backend's `ForwardMetadata.out_cache_loc_full_physical` (-> KVWriteLoc.full_loc). diff --git a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py index e4f13381e..ff2b64910 100644 --- a/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py +++ b/python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py @@ -58,6 +58,7 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): self._graphs: Dict[Any, torch.cuda.CUDAGraph] = {} self._outputs: Dict[Any, Any] = {} self._pool = None + self._cuda_graph_runner = cuda_graph_runner self._device_module = cuda_graph_runner.device_module self._tp_group = cuda_graph_runner.model_runner.tp_group self._capture_stream: Optional[torch.cuda.Stream] = None @@ -84,12 +85,29 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): capture_inputs: Optional[Any] = None, post_warmup_hook: Optional[Callable[[], None]] = None, ) -> None: + # When per-bs capture traces are enabled (--enable-profile-cuda-graph + + # SGLANG_GRAPH_BATCH_CAPTURE), the runner created a scheduled + # torch profiler (wait=2, active=1) and exposed it as _profiler. We step() + # past the two warmup runs so only the capture run is recorded, and each + # batch size produces its own trace via the profiler's on_trace_ready. + # With --enable-profile-cuda-graph alone the runner leaves _profiler None + # (its unscheduled profiler records the whole capture in one pass), so no + # stepping happens here. + runner = self._cuda_graph_runner + profiler = ( + getattr(runner, "_profiler", None) + if getattr(runner, "enable_profile_cuda_graph", False) + else None + ) + # Two warmups so kernels are loaded and one-time setup is paid before capture. # post_warmup_hook lets the attention backend reset state that warmup mutated. for _ in range(2): self._device_module.synchronize() self._tp_group.barrier() forward_fn() + if profiler is not None: + profiler.step() if post_warmup_hook is not None: post_warmup_hook() @@ -110,6 +128,9 @@ class FullCudaGraphBackend(BaseCudaGraphBackend): with graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream): out = forward_fn() + if profiler is not None: + profiler.step() + self._graphs[shape_key] = graph self._outputs[shape_key] = out diff --git a/python/sglang/srt/utils/profile_utils.py b/python/sglang/srt/utils/profile_utils.py index 562ffd28d..705f1010b 100644 --- a/python/sglang/srt/utils/profile_utils.py +++ b/python/sglang/srt/utils/profile_utils.py @@ -30,6 +30,19 @@ if _is_npu: logger = logging.getLogger(__name__) +# Single source of truth for the CUDA-graph capture trace output directory, +# shared by both the original single-trace export and the per-batch-size +# (SGLANG_GRAPH_BATCH_CAPTURE) traces so they land in the same place. +GRAPH_CAPTURE_PROFILE_DIRNAME = "graph_capture_profile" + + +def graph_capture_profile_dir() -> str: + """``/graph_capture_profile`` — the one directory + both capture-trace modes write to. Change the location here only.""" + return os.path.join( + envs.SGLANG_TORCH_PROFILER_DIR.get(), GRAPH_CAPTURE_PROFILE_DIRNAME + ) + def export_cuda_graph_capture_trace(prof_context, *, runner_name: str, tp_rank: int): """Persist a CUDA-graph capture profiler trace (chrome trace) to disk. @@ -37,15 +50,13 @@ def export_cuda_graph_capture_trace(prof_context, *, runner_name: str, tp_rank: Opt-in via ``SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE`` (no-op otherwise). The capture profiler must have run with ``record_shapes=True`` so the trace can be inspected offline as a per-kernel shape/identity record. The file lands in - ``/graph_capture_profile/`` and is namespaced by - runner class and TP rank so concurrent capture passes (e.g. EAGLE3 - target/draft/draft-extend) and ranks don't overwrite each other. + ``graph_capture_profile_dir()`` and is namespaced by runner class and TP rank + so concurrent capture passes (e.g. EAGLE3 target/draft/draft-extend) and + ranks don't overwrite each other. """ if not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get(): return - output_dir = os.path.join( - envs.SGLANG_TORCH_PROFILER_DIR.get(), "graph_capture_profile" - ) + output_dir = graph_capture_profile_dir() os.makedirs(output_dir, exist_ok=True) path = os.path.join( output_dir, f"cuda_graph_capture-{runner_name}-TP-{tp_rank}.json.gz" diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py new file mode 100644 index 000000000..a1f6bb2cd --- /dev/null +++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_runner.py @@ -0,0 +1,294 @@ +"""Unit tests for ``DecodeCudaGraphRunner`` capture-phase profiling — CPU-only. + +Two capture-trace modes plus their precedence: + + * **Original single-trace** (``SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE``): + ``_init_profile_context_and_memory_record`` builds an *unscheduled* profiler + (``record_shapes`` only, no schedule / no ``on_trace_ready``); the combined + trace is exported in ``_post_process_after_profile`` via + ``export_cuda_graph_capture_trace``. + * **Per-batch-size traces** (``SGLANG_GRAPH_BATCH_CAPTURE``): a *scheduled* + profiler (``wait=2, warmup=0, active=1, repeat=0``) with the trace-export + knobs (record_shapes / with_stack / with_flops / profile_memory) and an + ``on_trace_ready`` hook that writes one trace per batch size to + ``/graph_capture_profile/`` named + ``{runner_name}_bs_{bs}_rank{rank}.json.gz``. + * **Precedence**: when both env vars are set, the original single-trace path + wins (no per-bs schedule / dir / bookkeeping). + +The profiler / CUDA-memory APIs are mocked; the directory + naming + schedule +logic is pure-Python and runs on CPU. The method is invoked unbound against a +lightweight stand-in (with the real precedence helper bound) so no model or +server is constructed. +""" + +import os +import tempfile +import unittest +from types import SimpleNamespace +from unittest import mock + +from sglang.srt.model_executor.runner import decode_cuda_graph_runner as mod +from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( + DecodeCudaGraphRunner, +) +from sglang.srt.utils import profile_utils as putils +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +_CAPTURE_TRACE = "SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE" +_BATCH_CAPTURE = "SGLANG_GRAPH_BATCH_CAPTURE" + + +def _make_fake_self(capture_bs): + """Stand-in ``self`` with the real precedence helper bound so the env-var + gating in ``_init_profile_context_and_memory_record`` applies.""" + fake_self = SimpleNamespace(capture_bs=list(capture_bs)) + fake_self._graph_batch_capture_active = ( + DecodeCudaGraphRunner._graph_batch_capture_active.__get__(fake_self) + ) + return fake_self + + +class TestInitProfileBatchMode(CustomTestCase): + """SGLANG_GRAPH_BATCH_CAPTURE -> scheduled per-bs profiler.""" + + def _invoke(self, *, capture_bs, rank=0, profiler_dir=None): + fake_self = _make_fake_self(capture_bs) + env = {_BATCH_CAPTURE: "1"} + if profiler_dir is not None: + env["SGLANG_TORCH_PROFILER_DIR"] = profiler_dir + with mock.patch.dict(os.environ, env, clear=False), mock.patch.object( + mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank) + ), mock.patch.object(mod, "profile") as mock_profile, mock.patch( + "torch.profiler.schedule" + ) as mock_schedule, mock.patch( + "torch.cuda.memory._record_memory_history" + ) as mock_record_history: + os.environ.pop(_CAPTURE_TRACE, None) # original flag off + if profiler_dir is None: + os.environ.pop("SGLANG_TORCH_PROFILER_DIR", None) + ctx = DecodeCudaGraphRunner._init_profile_context_and_memory_record( + fake_self + ) + self.assertIs(ctx, mock_profile.return_value) + return fake_self, mock_profile, mock_schedule, mock_record_history + + def test_creates_graph_capture_profile_dir(self): + with tempfile.TemporaryDirectory() as tmp: + self._invoke(capture_bs=[1, 2, 4], profiler_dir=tmp) + self.assertTrue(os.path.isdir(os.path.join(tmp, "graph_capture_profile"))) + + def test_primes_reversed_bs_list_and_zero_index(self): + with tempfile.TemporaryDirectory() as tmp: + fake_self, *_ = self._invoke(capture_bs=[1, 2, 4, 8], profiler_dir=tmp) + # Capture iterates large -> small, so the bs list is reversed. + self.assertEqual(fake_self._profile_bs_list, [8, 4, 2, 1]) + self.assertEqual(fake_self._profile_bs_idx, 0) + + def test_profiler_built_with_trace_export_knobs(self): + with tempfile.TemporaryDirectory() as tmp: + _, mock_profile, mock_schedule, mock_record_history = self._invoke( + capture_bs=[1, 2], profiler_dir=tmp + ) + self.assertEqual(mock_profile.call_count, 1) + kwargs = mock_profile.call_args.kwargs + self.assertTrue(kwargs["record_shapes"]) + self.assertTrue(kwargs["with_stack"]) + self.assertTrue(kwargs["with_flops"]) + self.assertTrue(kwargs["profile_memory"]) + self.assertTrue(callable(kwargs["on_trace_ready"])) + # Schedule skips the two dummy/warmup runs and records the capture. + mock_schedule.assert_called_once_with(wait=2, warmup=0, active=1, repeat=0) + self.assertIs(kwargs["schedule"], mock_schedule.return_value) + # Memory history recording is armed alongside the profiler. + mock_record_history.assert_called_once() + + def test_default_dir_used_when_profiler_dir_env_unset(self): + # No SGLANG_TORCH_PROFILER_DIR -> falls back to the envs default base dir. + # Patch makedirs so the test never writes to the cwd. + fake_self = _make_fake_self([1]) + with mock.patch.dict( + os.environ, {_BATCH_CAPTURE: "1"}, clear=False + ), mock.patch.object( + mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0) + ), mock.patch.object( + mod, "profile" + ), mock.patch( + "torch.profiler.schedule" + ), mock.patch( + "torch.cuda.memory._record_memory_history" + ), mock.patch.object( + mod.os, "makedirs" + ) as mock_makedirs: + os.environ.pop("SGLANG_TORCH_PROFILER_DIR", None) + os.environ.pop(_CAPTURE_TRACE, None) + DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self) + + mock_makedirs.assert_called_once() + self.assertEqual( + mock_makedirs.call_args.args[0], + os.path.join("/tmp", "graph_capture_profile"), + ) + + +class TestInitProfileOriginalMode(CustomTestCase): + """No flag, original flag only, or both (precedence) -> unscheduled pass with + no per-bs schedule / directory / bookkeeping.""" + + def _invoke_original(self, *, env): + fake_self = _make_fake_self([1, 2]) + with tempfile.TemporaryDirectory() as tmp: + environ = dict(env) + environ["SGLANG_TORCH_PROFILER_DIR"] = tmp + with mock.patch.dict(os.environ, environ, clear=False), mock.patch.object( + mod, "get_parallel", return_value=SimpleNamespace(tp_rank=0) + ), mock.patch.object(mod, "profile") as mock_profile, mock.patch( + "torch.profiler.schedule" + ) as mock_schedule, mock.patch( + "torch.cuda.memory._record_memory_history" + ): + for k in (_CAPTURE_TRACE, _BATCH_CAPTURE): + if k not in environ: + os.environ.pop(k, None) + DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self) + kwargs = mock_profile.call_args.kwargs + # Unscheduled pass: record_shapes only, no schedule / on_trace_ready. + self.assertTrue(kwargs["record_shapes"]) + self.assertIsNone(kwargs.get("schedule")) + self.assertIsNone(kwargs.get("on_trace_ready")) + mock_schedule.assert_not_called() + self.assertFalse(os.path.isdir(os.path.join(tmp, "graph_capture_profile"))) + self.assertFalse(hasattr(fake_self, "_profile_bs_list")) + + def test_no_flags(self): + self._invoke_original(env={}) + + def test_original_flag_only(self): + self._invoke_original(env={_CAPTURE_TRACE: "1"}) + + def test_both_flags_original_takes_precedence(self): + self._invoke_original(env={_CAPTURE_TRACE: "1", _BATCH_CAPTURE: "1"}) + + +class TestOnTraceReadyNaming(CustomTestCase): + def _build_on_trace_ready(self, *, capture_bs, rank, tmp): + fake_self = _make_fake_self(capture_bs) + with mock.patch.dict( + os.environ, + {"SGLANG_TORCH_PROFILER_DIR": tmp, _BATCH_CAPTURE: "1"}, + clear=False, + ), mock.patch.object( + mod, "get_parallel", return_value=SimpleNamespace(tp_rank=rank) + ), mock.patch.object( + mod, "profile" + ) as mock_profile, mock.patch( + "torch.profiler.schedule" + ), mock.patch( + "torch.cuda.memory._record_memory_history" + ): + os.environ.pop(_CAPTURE_TRACE, None) + DecodeCudaGraphRunner._init_profile_context_and_memory_record(fake_self) + on_trace_ready = mock_profile.call_args.kwargs["on_trace_ready"] + return fake_self, on_trace_ready + + def test_exports_one_named_trace_per_bs_and_advances_index(self): + with tempfile.TemporaryDirectory() as tmp: + capture_bs = [1, 2, 4] # reversed -> [4, 2, 1] + fake_self, on_trace_ready = self._build_on_trace_ready( + capture_bs=capture_bs, rank=0, tmp=tmp + ) + trace_dir = os.path.join(tmp, "graph_capture_profile") + runner = type(fake_self).__name__ + + exported = [] + for expected_bs in [4, 2, 1]: + prof = mock.Mock() + prof.export_chrome_trace.side_effect = lambda p: exported.append(p) + on_trace_ready(prof) + prof.export_chrome_trace.assert_called_once_with( + os.path.join(trace_dir, f"{runner}_bs_{expected_bs}_rank0.json.gz") + ) + + self.assertEqual( + exported, + [ + os.path.join(trace_dir, f"{runner}_bs_4_rank0.json.gz"), + os.path.join(trace_dir, f"{runner}_bs_2_rank0.json.gz"), + os.path.join(trace_dir, f"{runner}_bs_1_rank0.json.gz"), + ], + ) + # Index advanced once per flush. + self.assertEqual(fake_self._profile_bs_idx, 3) + + def test_rank_in_trace_filename(self): + with tempfile.TemporaryDirectory() as tmp: + fake_self, on_trace_ready = self._build_on_trace_ready( + capture_bs=[8], rank=3, tmp=tmp + ) + runner = type(fake_self).__name__ + prof = mock.Mock() + on_trace_ready(prof) + prof.export_chrome_trace.assert_called_once_with( + os.path.join( + tmp, "graph_capture_profile", f"{runner}_bs_8_rank3.json.gz" + ) + ) + + +class TestOriginalTraceExport(CustomTestCase): + """export_cuda_graph_capture_trace (original single combined trace per rank), + gated by SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE, and the shared dir helper. + Both trace modes land under graph_capture_profile/.""" + + def test_writes_named_trace_when_flag_set(self): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.dict( + os.environ, + {"SGLANG_TORCH_PROFILER_DIR": tmp, _CAPTURE_TRACE: "1"}, + clear=False, + ): + prof = mock.Mock() + putils.export_cuda_graph_capture_trace( + prof, runner_name="DecodeCudaGraphRunner", tp_rank=2 + ) + expected = os.path.join( + tmp, + "graph_capture_profile", + "cuda_graph_capture-DecodeCudaGraphRunner-TP-2.json.gz", + ) + prof.export_chrome_trace.assert_called_once_with(expected) + self.assertTrue( + os.path.isdir(os.path.join(tmp, "graph_capture_profile")) + ) + + def test_noop_when_flag_unset(self): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.dict( + os.environ, {"SGLANG_TORCH_PROFILER_DIR": tmp}, clear=False + ): + os.environ.pop(_CAPTURE_TRACE, None) + prof = mock.Mock() + putils.export_cuda_graph_capture_trace( + prof, runner_name="DecodeCudaGraphRunner", tp_rank=0 + ) + prof.export_chrome_trace.assert_not_called() + self.assertFalse( + os.path.isdir(os.path.join(tmp, "graph_capture_profile")) + ) + + def test_dir_helper_uses_profiler_dir(self): + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.dict( + os.environ, {"SGLANG_TORCH_PROFILER_DIR": tmp}, clear=False + ): + self.assertEqual( + putils.graph_capture_profile_dir(), + os.path.join(tmp, "graph_capture_profile"), + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py b/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py new file mode 100644 index 000000000..4987bce1f --- /dev/null +++ b/test/registered/unit/model_executor/runner_backend/test_full_cuda_graph_backend.py @@ -0,0 +1,162 @@ +"""Unit tests for ``FullCudaGraphBackend.capture_one`` profiling hooks — CPU-only. + +These cover the changes from the "cuda graph profile traces" PR that wire the +runner's torch profiler into the capture loop: + + * When profiling is disabled, ``capture_one`` runs exactly two warmups + one + capture and never touches a profiler (behavior-identical to before the PR). + * When the runner exposes an active ``_profiler`` (per-bs capture profiling, + ``--enable-profile-cuda-graph`` + ``SGLANG_GRAPH_BATCH_CAPTURE``), + ``capture_one`` calls ``profiler.step()`` past the two warmups and once after + the capture (schedule ``wait=2, warmup=0, active=1``). The captured forward is + NOT wrapped in a ``record_function``; per-bs trace naming is handled by the + profiler's ``on_trace_ready`` callback instead. + * The ``getattr`` guards mean a runner that sets the flag but has no + ``_profiler`` attribute degrades gracefully (no stepping, no crash). + +The real capture path needs CUDA (``torch.cuda.CUDAGraph`` + device graph +context), so those are mocked; the logic under test (call counts, ordering, +profiler stepping) is pure-Python and runs on CPU. +""" + +import contextlib +import unittest +from types import SimpleNamespace +from unittest import mock + +from sglang.srt.model_executor.runner.shape_key import ShapeKey +from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import ( + FullCudaGraphBackend, +) +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=10, suite="base-a-test-cpu") + +# Sentinel: distinguishes "runner has no _profiler attribute" from +# "_profiler is None" in the test fixtures. +_UNSET = object() + + +class _FakeGraphCtx: + """Stand-in for ``device_module.graph(...)`` — a no-op context manager.""" + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +def _make_backend(runner): + """Build a ``FullCudaGraphBackend`` without running ``__init__`` (which would + touch CUDA), wiring just the attributes ``capture_one`` reads.""" + backend = FullCudaGraphBackend.__new__(FullCudaGraphBackend) + backend._graphs = {} + backend._outputs = {} + backend._pool = None + backend._capture_stream = None + backend._memory_saver_adapter = None + backend._cuda_graph_runner = runner + backend._device_module = runner.device_module + backend._tp_group = runner.model_runner.tp_group + return backend + + +def _make_runner(*, enable_profile, profiler, num_tokens_per_bs=1, mode_name="DECODE"): + device_module = SimpleNamespace( + synchronize=mock.Mock(name="synchronize"), + graph=mock.Mock(name="graph", side_effect=lambda **kw: _FakeGraphCtx()), + ) + tp_group = SimpleNamespace(barrier=mock.Mock(name="barrier")) + runner = SimpleNamespace( + device_module=device_module, + model_runner=SimpleNamespace(tp_group=tp_group), + num_tokens_per_bs=num_tokens_per_bs, + capture_forward_mode=SimpleNamespace(name=mode_name), + enable_profile_cuda_graph=enable_profile, + ) + if profiler is not _UNSET: + runner._profiler = profiler + return runner + + +class TestCaptureOneNoProfiling(CustomTestCase): + def test_runs_two_warmups_and_capture_without_stepping(self): + runner = _make_runner(enable_profile=False, profiler=None) + backend = _make_backend(runner) + + sentinel_out = object() + forward_fn = mock.Mock(return_value=sentinel_out) + post_warmup_hook = mock.Mock() + shape_key = ShapeKey(size=4) + + with mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"): + backend.capture_one( + shape_key, forward_fn, post_warmup_hook=post_warmup_hook + ) + + # 2 warmups + 1 capture. + self.assertEqual(forward_fn.call_count, 3) + # post_warmup_hook only runs in the two warmup iterations. + self.assertEqual(post_warmup_hook.call_count, 2) + # Graph + output are recorded against the shape key. + self.assertEqual(backend._graphs[shape_key], "GRAPH") + self.assertIs(backend._outputs[shape_key], sentinel_out) + + def test_enable_flag_set_but_no_profiler_attr_does_not_step(self): + # The runner advertises the flag but never created a profiler; the + # getattr guard must keep capture_one on the non-profiling path. + runner = _make_runner(enable_profile=True, profiler=_UNSET) + backend = _make_backend(runner) + self.assertFalse(hasattr(runner, "_profiler")) + + forward_fn = mock.Mock(return_value=object()) + with mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"): + backend.capture_one(ShapeKey(size=2), forward_fn) + + self.assertEqual(forward_fn.call_count, 3) + + +class TestCaptureOneWithProfiling(CustomTestCase): + def _run(self, *, size, num_tokens_per_bs, mode_name): + profiler = SimpleNamespace(step=mock.Mock(name="step")) + runner = _make_runner( + enable_profile=True, + profiler=profiler, + num_tokens_per_bs=num_tokens_per_bs, + mode_name=mode_name, + ) + backend = _make_backend(runner) + + forward_fn = mock.Mock(return_value=object()) + rf_names = [] + + def _fake_record_function(name): + rf_names.append(name) + return contextlib.nullcontext() + + with mock.patch("torch.cuda.CUDAGraph", return_value="GRAPH"), mock.patch( + "torch.profiler.record_function", side_effect=_fake_record_function + ): + backend.capture_one(ShapeKey(size=size), forward_fn) + + return profiler, forward_fn, rf_names + + def test_steps_twice_in_warmup_and_once_after_capture(self): + profiler, forward_fn, _ = self._run( + size=4, num_tokens_per_bs=1, mode_name="DECODE" + ) + # Schedule wait=2 + active=1 => one step per warmup (x2) + one post-capture. + self.assertEqual(profiler.step.call_count, 3) + self.assertEqual(forward_fn.call_count, 3) + + def test_capture_not_wrapped_in_record_function(self): + # The capture forward is no longer wrapped in a record_function; per-bs + # trace naming is handled by the profiler's on_trace_ready callback. + _, _, rf_names = self._run(size=4, num_tokens_per_bs=1, mode_name="DECODE") + self.assertEqual(rf_names, []) + + +if __name__ == "__main__": + unittest.main()