From 7c3b5a6732fbacc4d63beb0185141e7a2ec62496 Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:55:34 -0700 Subject: [PATCH] config: every handler declares its cuda-graph decisions (#36725) Co-authored-by: Claude Opus 5 --- .../sglang/srt/hardware_backend/npu/utils.py | 46 ++- .../srt/model_executor/cuda_graph_config.py | 21 +- python/sglang/srt/server_args.py | 292 +++++++++++++++--- .../test_resolution_declarations.py | 98 ++++++ .../unit/server_args/test_server_args.py | 5 +- 5 files changed, 403 insertions(+), 59 deletions(-) diff --git a/python/sglang/srt/hardware_backend/npu/utils.py b/python/sglang/srt/hardware_backend/npu/utils.py index 857ae8634..5bcc8906f 100644 --- a/python/sglang/srt/hardware_backend/npu/utils.py +++ b/python/sglang/srt/hardware_backend/npu/utils.py @@ -8,6 +8,7 @@ import torch from sglang.srt.arg_groups.overrides import declare_resolution from sglang.srt.environ import envs +from sglang.srt.model_executor.cuda_graph_config import Phase, with_phase from sglang.srt.utils import get_npu_memory_capacity, is_npu if TYPE_CHECKING: @@ -71,7 +72,6 @@ def set_default_server_args(args: "ServerArgs"): ) # NPU memory settings - decode = cfg.cuda_graph_config.decode npu_mem = get_npu_memory_capacity() if npu_mem <= 32 * 1024: # Ascend 910B4,910B4_1 @@ -82,11 +82,27 @@ def set_default_server_args(args: "ServerArgs"): "set_default_server_args", chunked_prefill_size=4 * 1024, ) - if decode.max_bs is None: + if cfg.cuda_graph_config.decode.max_bs is None: if cfg.tp_size < 4: - decode.max_bs = 16 + declare_resolution( + args, + "set_default_server_args", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, + Phase.DECODE, + max_bs=16, + ), + ) else: - decode.max_bs = 64 + declare_resolution( + args, + "set_default_server_args", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, + Phase.DECODE, + max_bs=64, + ), + ) elif npu_mem <= 64 * 1024: # Ascend 910B1,910B2,910B2C,910B3,910_9391,910_9392,910_9381,910_9382,910_9372,910_9362 # (chunked_prefill_size 8k, max_bs 64 if tp < 4 else 256) @@ -96,11 +112,27 @@ def set_default_server_args(args: "ServerArgs"): "set_default_server_args", chunked_prefill_size=8 * 1024, ) - if decode.max_bs is None: + if cfg.cuda_graph_config.decode.max_bs is None: if cfg.tp_size < 4: - decode.max_bs = 64 + declare_resolution( + args, + "set_default_server_args", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, + Phase.DECODE, + max_bs=64, + ), + ) else: - decode.max_bs = 256 + declare_resolution( + args, + "set_default_server_args", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, + Phase.DECODE, + max_bs=256, + ), + ) # NPU does not support CustomAllReduce declare_resolution( diff --git a/python/sglang/srt/model_executor/cuda_graph_config.py b/python/sglang/srt/model_executor/cuda_graph_config.py index 04db53ce7..8148ef14b 100644 --- a/python/sglang/srt/model_executor/cuda_graph_config.py +++ b/python/sglang/srt/model_executor/cuda_graph_config.py @@ -23,7 +23,7 @@ inside the function body to preserve that invariant. import argparse import dataclasses import json -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any, Dict, List, Optional @@ -119,6 +119,25 @@ def default_prefill_backend() -> str: return Backend.BREAKABLE if is_cuda() else Backend.TC_PIECEWISE +def with_phase(config: "CudaGraphConfig", phase: str, **changes) -> "CudaGraphConfig": + """A copy of ``config`` with ``changes`` applied to one phase. + + Resolution declares values, so a handler that decides a graph setting hands + the stash a new config instead of editing the one an earlier handler + declared. + """ + if phase not in Phase.ALL: + raise KeyError(phase) + # Not a deep copy: `dataclasses.replace` copies field references, so a + # list-valued `bs` is shared. Rebind `bs`, never mutate it in place. + return CudaGraphConfig( + **{ + name: replace(getattr(config, name), **(changes if name == phase else {})) + for name in Phase.ALL + } + ) + + @dataclass class CudaGraphConfig: """Top-level CUDA graph config: one PhaseConfig per phase.""" diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 66909b9e4..dae349f7b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -85,6 +85,7 @@ from sglang.srt.model_executor.cuda_graph_config import ( Phase, default_cuda_graph_config, parse_cuda_graph_config_arg, + with_phase, ) from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.platforms import current_platform @@ -4055,8 +4056,18 @@ class ServerArgs: ) # cuda_graph_config was already parsed from the legacy boolean, so # flipping the boolean alone would not stop graph capture. - cfg.cuda_graph_config.decode.backend = Backend.DISABLED - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_model_capability_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) + self._declare( + "_handle_model_capability_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) logger.warning( "HRM-Text (prefix_lm) detected: forcing --attention-backend " "triton, --chunked-prefill-size -1, --disable-radix-cache, and " @@ -4139,9 +4150,19 @@ class ServerArgs: prefill_only_disable_kv_cache=True, ) self._validate_prefill_only_disable_kv_cache_args() - cfg.cuda_graph_config.decode.backend = Backend.DISABLED + self._declare( + "_handle_model_capability_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) if is_cuda() and cfg.cuda_graph_config.prefill.backend != Backend.DISABLED: - cfg.cuda_graph_config.prefill.backend = Backend.BREAKABLE + self._declare( + "_handle_model_capability_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.BREAKABLE + ), + ) # CUDA-graph sizing has already run by this point and derives # its generic maximum from the 8K chunked-prefill default. # On the Hopper/Blackwell FA raw-K/V path, raise the unlocked @@ -4157,21 +4178,32 @@ class ServerArgs: self, "_cuda_graph_config_locked", set() ) if (Phase.PREFILL, "max_bs") not in cuda_graph_config_locked: - prefill_config.max_bs = max( - prefill_config.max_bs or 0, - model_config.context_len, - 16384, - ) - if (Phase.PREFILL, "bs") not in cuda_graph_config_locked: - prefill_config.bs = ( - self._generate_prefill_cuda_graph_batch_sizes( - prefill_config.max_bs - ) + sizing = { + "max_bs": max( + prefill_config.max_bs or 0, + model_config.context_len, + 16384, ) + } + if (Phase.PREFILL, "bs") not in cuda_graph_config_locked: + sizing["bs"] = self._generate_prefill_cuda_graph_batch_sizes( + sizing["max_bs"] + ) + self._declare( + "_handle_model_capability_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, **sizing + ), + ) elif not is_cuda(): # BCG is CUDA-only. Other graph backends do not support this # encoder-style prefill, so retain the eager Triton path. - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_model_capability_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) logger.info( "EmbeddingGemma detected: disabling radix cache and chunked " "prefill; using breakable CUDA graph for CUDA prefill." @@ -4716,7 +4748,12 @@ class ServerArgs: "At this moment Ascend platform only support prefill graph compilation with " "cuda_graph_config[prefill].tc_compiler='eager'." ) - cfg.cuda_graph_config.prefill.tc_compiler = "eager" + self._declare( + "_handle_npu_backends", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, tc_compiler="eager" + ), + ) def _handle_mps_backends(self): cfg = resolving_view(self) @@ -4734,7 +4771,12 @@ class ServerArgs: # --cuda-graph-backend-decode (or --cuda-graph-config), keep it # disabled so the default startup doesn't require graph capture. if (Phase.DECODE, "backend") not in self._cuda_graph_config_locked: - cfg.cuda_graph_config.decode.backend = Backend.DISABLED + self._declare( + "_handle_xpu_backends", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) elif cfg.cuda_graph_config.decode.backend not in ( Backend.DISABLED, Backend.FULL, @@ -4744,7 +4786,12 @@ class ServerArgs: "disabling unsupported decode backend '%s'.", cfg.cuda_graph_config.decode.backend, ) - cfg.cuda_graph_config.decode.backend = Backend.DISABLED + self._declare( + "_handle_xpu_backends", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) # ------------------------------------------------------------------ # CUDA graph configuration resolution @@ -4829,8 +4876,15 @@ class ServerArgs: sorted(bs), aligned, ) - cfg.cuda_graph_config.prefill.bs = aligned - cfg.cuda_graph_config.prefill.max_bs = aligned[-1] + self._declare( + "_apply_deepep_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, + Phase.PREFILL, + bs=aligned, + max_bs=aligned[-1], + ), + ) def _parse_cuda_graph_config(self): """Resolve cuda_graph_config from explicit JSON, per-phase @@ -4926,7 +4980,12 @@ class ServerArgs: "Using tc_piecewise CUDA graph for validated multimodal " "decoder prefill." ) - cfg.cuda_graph_config.prefill.backend = Backend.TC_PIECEWISE + self._declare( + "_apply_cuda_graph_compatibility", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.TC_PIECEWISE + ), + ) if cfg.cuda_graph_config.prefill.backend == Backend.TC_PIECEWISE: self._disable_tc_piecewise_cudagraph_if_incompatible() @@ -4939,10 +4998,20 @@ class ServerArgs: cfg = resolving_view(self) if cfg.disaggregation_mode == "prefill": if (Phase.DECODE, "backend") not in self._cuda_graph_config_locked: - cfg.cuda_graph_config.decode.backend = Backend.DISABLED + self._declare( + "_apply_cuda_graph_disaggregation_roles", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) elif cfg.disaggregation_mode == "decode": if (Phase.PREFILL, "backend") not in self._cuda_graph_config_locked: - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_apply_cuda_graph_disaggregation_roles", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) def _disable_tc_piecewise_cudagraph_if_incompatible(self): """TcPiecewise (torch.compile + piecewise) is incompatible with @@ -5018,7 +5087,15 @@ class ServerArgs: ] for _name, predicate in rules: if predicate(): - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_disable_tc_piecewise_cudagraph_if_incompatible", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) + # One decision, one declaration: every rule declares the same + # value, so a later match would only append a duplicate entry. + break def _disable_breakable_cudagraph_if_incompatible(self): """Breakable (segmented capture, no torch.compile). Breakable enforces @@ -5071,7 +5148,12 @@ class ServerArgs: "disabling prefill CUDA graph.", name, ) - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_disable_breakable_cudagraph_if_incompatible", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) return def _disable_full_prefill_cudagraph_if_incompatible(self): @@ -5085,7 +5167,12 @@ class ServerArgs: "disabling prefill CUDA graph.", name, ) - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_disable_full_prefill_cudagraph_if_incompatible", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) return def _disable_prefill_cuda_graph_for_deepseek_trtllm_mla(self): @@ -5115,7 +5202,12 @@ class ServerArgs: "backend explicitly (e.g. --cuda-graph-backend-prefill tc_piecewise) to override.", cfg.cuda_graph_config.prefill.backend, ) - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_disable_prefill_cuda_graph_for_deepseek_trtllm_mla", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) def _validate_cuda_graph_config(self): cfg = resolving_view(self) @@ -5143,8 +5235,18 @@ class ServerArgs: if cfg.cuda_graph_config.decode.backend != Backend.DISABLED: logger.warning("CUDA graph is disabled because --enable-mis is set.") - cfg.cuda_graph_config.decode.backend = Backend.DISABLED - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_multi_item_scoring", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) + self._declare( + "_handle_multi_item_scoring", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) if not cfg.disable_radix_cache: logger.warning("Radix cache is disabled because --enable-mis is set.") @@ -5192,8 +5294,10 @@ class ServerArgs: The coefficient 1.5 is a heuristic value, in the future, we can do better estimation by looking at the model types, hidden sizes or even do a dummy run. """ cfg = resolving_view(self) - decode_cuda_graph_config = cfg.cuda_graph_config.decode - prefill_cuda_graph_config = cfg.cuda_graph_config.prefill + # A copy, so an earlier declaration keeps the value it recorded. + cuda_graph_config = copy.deepcopy(cfg.cuda_graph_config) + decode_cuda_graph_config = cuda_graph_config.decode + prefill_cuda_graph_config = cuda_graph_config.prefill if gpu_mem is not None: if gpu_mem < 20 * 1024: @@ -5340,6 +5444,11 @@ class ServerArgs: ) ) + if cuda_graph_config != cfg.cuda_graph_config: + self._declare( + "_handle_gpu_memory_settings", cuda_graph_config=cuda_graph_config + ) + if cfg.mem_fraction_static is None: if self.post_capture_kv_sizing_planned(): # Post-capture sizing measures free memory after graph capture, so @@ -5779,7 +5888,14 @@ class ServerArgs: # The DSA CP field declarations moved to the override # registry (arg_groups/overrides.py: # _deepseek_family_overrides). - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_model_specific_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, + Phase.PREFILL, + backend=Backend.DISABLED, + ), + ) else: # Pure TP and partial DP Attention mode is active for DSA, logging a warning if cfg.dp_size < cfg.tp_size: @@ -5862,7 +5978,14 @@ class ServerArgs: # the override registry (arg_groups/overrides.py: # _deepseek_family_overrides). if cfg.enable_prefill_cp and self.use_mla_backend(): - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_model_specific_adjustments", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, + Phase.PREFILL, + backend=Backend.DISABLED, + ), + ) # Set moe backend for DeepSeek: the sm100 quant/moe resolution # moved to the resolution pipeline (arg_groups/overrides.py: @@ -6358,15 +6481,35 @@ class ServerArgs: logger.warning( "Cuda graph is disabled because of using torch native attention backend" ) - cfg.cuda_graph_config.decode.backend = Backend.DISABLED - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_attention_backend_compatibility", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) + self._declare( + "_handle_attention_backend_compatibility", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) if attention_backend == "flex_attention": logger.warning( "Cuda graph is disabled because of using torch Flex Attention backend" ) - cfg.cuda_graph_config.decode.backend = Backend.DISABLED - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_attention_backend_compatibility", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) + self._declare( + "_handle_attention_backend_compatibility", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) assert ( cfg.speculative_algorithm is None ), "Speculative decoding is currently not supported with Flex Attention backend" @@ -7212,11 +7355,17 @@ class ServerArgs: and prefill_cfg.max_bs > cfg.chunked_prefill_size and (Phase.PREFILL, "max_bs") not in self._cuda_graph_config_locked ): - prefill_cfg.max_bs = cfg.chunked_prefill_size + clamped = {"max_bs": cfg.chunked_prefill_size} if (Phase.PREFILL, "bs") not in self._cuda_graph_config_locked: - prefill_cfg.bs = self._generate_prefill_cuda_graph_batch_sizes( - prefill_cfg.max_bs + clamped["bs"] = self._generate_prefill_cuda_graph_batch_sizes( + clamped["max_bs"] ) + self._declare( + "_handle_data_parallelism", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, **clamped + ), + ) # Resolve the phase-aware TP LM-head default before validating the # resulting DP/TP LM-head configuration. @@ -7533,8 +7682,18 @@ class ServerArgs: ) if cfg.deepep_mode == "normal": logger.warning("Cuda graph is disabled because deepep_mode=`normal`") - cfg.cuda_graph_config.decode.backend = Backend.DISABLED - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_a2a_moe", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) + self._declare( + "_handle_a2a_moe", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) if a2a_backend == "deepep_v2": self._validate_deepep_v2_model_architecture() @@ -7574,7 +7733,12 @@ class ServerArgs: "--moe-a2a-backend deepep_v2." ) # Prefill reads host counts and is not graph-capturable. - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_a2a_moe", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) logger.warning( f"DeepEP v2 MoE is enabled. The expert parallel size is adjusted to be the same as the tensor parallel size[{cfg.tp_size}]." ) @@ -9299,8 +9463,18 @@ class ServerArgs: logger.warning( "Cuda graph is disabled for diffusion LLM inference on AMD GPUs" ) - cfg.cuda_graph_config.decode.backend = Backend.DISABLED - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_dllm_inference", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) + self._declare( + "_handle_dllm_inference", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) from sglang.srt.arg_groups.overrides import ( _dllm_attention_backend, @@ -9436,8 +9610,18 @@ class ServerArgs: logger.warning( "Cuda graph and server warmup are disabled because of using tensor dump mode" ) - cfg.cuda_graph_config.decode.backend = Backend.DISABLED - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_other_validations", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) + self._declare( + "_handle_other_validations", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) self._declare("_handle_other_validations", skip_server_warmup=True) if cfg.msprobe_dump_config is not None: @@ -9446,8 +9630,18 @@ class ServerArgs: "cuda graph is disabled because msProbe only supports dump in eager mode, " "warmup is disabled(skip_server_warmup=True) because there is no need to dump data for this stage." ) - cfg.cuda_graph_config.decode.backend = Backend.DISABLED - cfg.cuda_graph_config.prefill.backend = Backend.DISABLED + self._declare( + "_handle_other_validations", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.DECODE, backend=Backend.DISABLED + ), + ) + self._declare( + "_handle_other_validations", + cuda_graph_config=with_phase( + cfg.cuda_graph_config, Phase.PREFILL, backend=Backend.DISABLED + ), + ) self._declare("_handle_other_validations", skip_server_warmup=True) # Validate limit_mm_per_prompt modalities diff --git a/test/registered/unit/server_args/test_resolution_declarations.py b/test/registered/unit/server_args/test_resolution_declarations.py index 583f26701..32e3e18ca 100644 --- a/test/registered/unit/server_args/test_resolution_declarations.py +++ b/test/registered/unit/server_args/test_resolution_declarations.py @@ -920,5 +920,103 @@ class TestResolutionDeclarations(CustomTestCase): ) +class TestDeclaredValuesAreNotEditedLater(CustomTestCase): + """A declaration records a value, not a handle on one. + + The stash keeps whatever object the declaring handler passed, so a handler + that declares a mutable and then edits it in place rewrites an entry that + already went into the log. The projection still answers with the end state, + which is why nothing else notices: what is lost is *which* handler decided + what, and `validate_declarations` never sees the later change at all. + """ + + def setUp(self): + super().setUp() + environment = dict(os.environ) + + def restore(): + os.environ.clear() + os.environ.update(environment) + + self.addCleanup(restore) + + def _resolve_recording_each_entry(self, **supplied): + """Resolve, deep-copying every stash entry the moment it is appended.""" + from sglang.srt.arg_groups import overrides + + recorded = [] + + def watch(name): + original = getattr(overrides, name) + + def wrapper(server_args, *args, **kwargs): + result = original(server_args, *args, **kwargs) + stash = getattr(server_args, "_resolved_overrides", None) or [] + while len(recorded) < len(stash): + index = len(recorded) + recorded.append((index, copy.deepcopy(stash[index]))) + return result + + return original, wrapper + + # Every path that appends to the stash. + patched = {} + for name in ( + "declare_resolution", + "declare_late_resolution", + "declare_direct_writes", + "run_post_process_pass", + ): + original, wrapper = watch(name) + patched[name] = original + setattr(overrides, name, wrapper) + try: + path = tempfile.mkdtemp(prefix="declared_values_") + self.addCleanup(shutil.rmtree, path, ignore_errors=True) + with open(os.path.join(path, "config.json"), "w") as handle: + json.dump(_MINI_CONFIG, handle) + server_args = ServerArgs( + model_path=path, device="cuda", random_seed=42, **supplied + ) + server_args.resolve_once() + finally: + for name, original in patched.items(): + setattr(overrides, name, original) + return server_args, recorded + + def test_no_entry_changes_after_it_is_recorded(self): + # One shape per family of handlers that decides a graph setting. + for label, supplied in ( + ("plain", {}), + ("cuda_graph_knobs", {"cuda_graph_max_bs_decode": 16}), + ("chunked_prefill", {"chunked_prefill_size": 1024}), + ("explicit_json", {"cuda_graph_config": {"decode": {"max_bs": 12}}}), + ("disaggregation", {"disaggregation_mode": "prefill"}), + ("deterministic", {"enable_deterministic_inference": True}), + ("speculative", {"speculative_algorithm": "EAGLE"}), + ("dp_attention", {"tp_size": 2, "dp_size": 2, "enable_dp_attention": True}), + ): + with self.subTest(shape=label): + server_args, recorded = self._resolve_recording_each_entry(**supplied) + stash = server_args._resolved_overrides + self.assertGreater( + len(recorded), + 0, + "nothing was recorded, so this case is not watching the " + "declaration paths it thinks it is", + ) + drifted = [ + (index, was, stash[index]) + for index, was in recorded + if stash[index] != was + ] + self.assertEqual( + [], + drifted, + "these entries changed after they were declared, so the log " + f"credits the wrong handler for the end state: {drifted}", + ) + + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/server_args/test_server_args.py b/test/registered/unit/server_args/test_server_args.py index ed9ce7ed3..f43ff6c64 100644 --- a/test/registered/unit/server_args/test_server_args.py +++ b/test/registered/unit/server_args/test_server_args.py @@ -2245,8 +2245,9 @@ class TestDeepEPv2Args(CustomTestCase): for mode in ("direct", "hybrid"): args = self._args(moe_runner_backend="deep_gemm", deepep_v2_mode=mode) args._handle_a2a_moe() - self.assertEqual(args.cuda_graph_config.decode.backend, Backend.FULL) - self.assertEqual(args.cuda_graph_config.prefill.backend, Backend.DISABLED) + declared = resolution_result(args, "cuda_graph_config") + self.assertEqual(declared.decode.backend, Backend.FULL) + self.assertEqual(declared.prefill.backend, Backend.DISABLED) def test_two_batch_overlap_rejected(self): args = self._args(moe_runner_backend="deep_gemm", enable_two_batch_overlap=True)