From 169783d42f3145cf00e6daa6692580da3685c37c Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 10 Aug 2026 10:22:16 +0800 Subject: [PATCH] [diffusion] chore: make torch.compile opt-in for speed mode (#34173) --- docs/docs/sglang-diffusion/api/cli.mdx | 2 +- .../sglang-diffusion/deployment_cookbook.mdx | 4 +- .../model_deployment_config.py | 6 +-- .../runtime/server_args/auto_tune.py | 4 +- .../test/unit/test_server_args.py | 49 ++++++++++++++----- 5 files changed, 43 insertions(+), 22 deletions(-) diff --git a/docs/docs/sglang-diffusion/api/cli.mdx b/docs/docs/sglang-diffusion/api/cli.mdx index 0e66cd459..214c2cf42 100644 --- a/docs/docs/sglang-diffusion/api/cli.mdx +++ b/docs/docs/sglang-diffusion/api/cli.mdx @@ -80,7 +80,7 @@ Use `sglang generate --help` and `sglang serve --help` for the full argument lis - `--lora-path {PATH}` and `--lora-nickname {NAME}`: load a LoRA adapter - `--lora-merge-mode {auto|merge|dynamic}`: choose how LoRA is applied. `auto` statically merges regular weights and uses dynamic LoRA for FSDP-sharded weights to avoid full-gather peaks. - `--num-gpus {N}`: number of GPUs to use -- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default and keeps safe offload defaults, using FSDP only for validated DiT-offload replacement paths; `speed` also enables `--enable-torch-compile` unless the model-specific deployment config opts out or you explicitly disable it. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes. +- `--performance-mode {manual|auto|speed|memory}` / `--mode`: preset for latency/throughput and memory defaults. `auto` is the default and keeps safe offload defaults, using FSDP only for validated DiT-offload replacement paths. `speed` keeps `torch.compile` disabled unless a model-specific deployment config opts in after validation; pass `--enable-torch-compile true` to enable it explicitly. Use `manual` to keep performance-related server args under explicit user control. Explicit offload, FSDP, and parallelism flags take precedence in all modes. - `--tp-size {N}`: tensor parallelism size. Depending on the pipeline, it can shard the DiT, one or more encoders, or both. - `--sp-degree {N}`: sequence parallelism size - `--dp-size {N}` (alias `--data-parallel-size`): number of data-parallel replicas. Each replica is a full copy of the engine on `num_gpus / N` GPUs with its own ingress; generation requests round-robin across replicas, realtime sessions stick to the replica holding their state, and control operations (weights, LoRA, memory occupation, shutdown) apply to every replica. Combines with the other parallelism axes (`num_gpus = dp × cfg × tp × sp`); monolithic serving only. diff --git a/docs/docs/sglang-diffusion/deployment_cookbook.mdx b/docs/docs/sglang-diffusion/deployment_cookbook.mdx index 7582853c6..bc0d333a2 100644 --- a/docs/docs/sglang-diffusion/deployment_cookbook.mdx +++ b/docs/docs/sglang-diffusion/deployment_cookbook.mdx @@ -103,7 +103,7 @@ status-code contract and warmup-mode behavior. `speed` - Favors GPU-resident execution for lower latency and higher throughput. Disables CPU offload when unset and enables `torch.compile` by default unless explicitly disabled; may OOM. + Favors GPU-resident execution for lower latency and higher throughput. Disables CPU offload when unset. `torch.compile` stays off unless the model has a validated default or it is enabled explicitly; may OOM. `memory` @@ -112,7 +112,7 @@ status-code contract and warmup-mode behavior. -`auto` checks selected GPU memory before applying FSDP. In multi-GPU runs it uses the least available memory across selected GPUs, and only turns on FSDP automatically when doing so can replace DiT offload. Text encoder, image encoder, and other component residency still follow the offload policy unless the model marks a high-memory resident path as safe. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `speed` intentionally does not check memory; it is the mode for users who prefer latency/throughput and accept OOM risk, and it will turn on `torch.compile` automatically unless `--enable-torch-compile false` is set explicitly. +`auto` checks selected GPU memory before applying FSDP. In multi-GPU runs it uses the least available memory across selected GPUs, and only turns on FSDP automatically when doing so can replace DiT offload. Text encoder, image encoder, and other component residency still follow the offload policy unless the model marks a high-memory resident path as safe. When the model default uses CFG and the user did not set a parallelism policy, `auto` may also enable CFG parallelism. `speed` intentionally does not check memory; it is the mode for users who prefer latency/throughput and accept OOM risk. It keeps `torch.compile` disabled by default because its effect varies by model and workload. A model-specific deployment config may enable a validated compile path, and `--enable-torch-compile true` always opts in explicitly. The modes tune residency for native pipeline components declared to the component residency manager. Today this covers the major DiT, text/image encoder, VAE, vocoder, and upsampler components; DiT can use layerwise offload when supported, while text encoders use either resident execution or component CPU offload. Do not assume text-encoder layerwise offload unless a model implements and validates it. diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py index 189f79dc9..c73d9243b 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/model_deployment_config.py @@ -25,9 +25,9 @@ class ModelDeploymentConfig: auto_enable_cfg_parallel: bool = True # degree 1 keeps CFG parallel disabled and leaves GPUs available for SP auto_cfg_parallel_degree_by_num_gpus: tuple[tuple[int, int], ...] = () - # Let performance_mode=speed opt into torch.compile unless the model has - # established that the compiled path changes its numerical contract. - speed_mode_enable_torch_compile_by_default: bool = True + # torch.compile is model opt-in because it can be slower than eager for + # diffusion workloads dominated by already-optimized kernels + speed_mode_enable_torch_compile_by_default: bool = False supports_cfg_parallel: bool = True def get_auto_cfg_parallel_degree(self, num_gpus: int) -> int: diff --git a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py index 2458fa7dc..d2228463d 100644 --- a/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py +++ b/python/sglang/multimodal_gen/runtime/server_args/auto_tune.py @@ -78,9 +78,7 @@ class ServerArgsAutoTuner: and not args.enable_torch_compile and not args.is_arg_explicitly_set("enable_torch_compile") ): - # speed means fastest: compile by default. An explicit - # --enable-torch-compile false still wins (e.g. models where - # compile is slower or changes the numerical contract). + # only models with a validated compile win opt in by default args.enable_torch_compile = True logger.info( "performance_mode=speed enables torch.compile " diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 6ce48bd6f..8a20f399a 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -24,6 +24,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( from sglang.multimodal_gen.configs.pipeline_configs.minimax_h3 import ( MiniMaxH3PipelineConfig, ) +from sglang.multimodal_gen.configs.pipeline_configs.model_deployment_config import ( + ModelDeploymentConfig, +) from sglang.multimodal_gen.configs.pipeline_configs.mova import MOVAPipelineConfig from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import ( QwenImagePipelineConfig, @@ -1590,7 +1593,7 @@ class TestOffloadDefaults(unittest.TestCase): self.assertFalse(args.text_encoder_cpu_offload) self.assertFalse(args.image_encoder_cpu_offload) - def test_speed_mode_enables_torch_compile_by_default(self): + def test_speed_mode_keeps_torch_compile_off_by_default(self): args = self._from_dict_with_pipeline_config( QwenImagePipelineConfig(), kwargs={ @@ -1599,20 +1602,40 @@ class TestOffloadDefaults(unittest.TestCase): }, ) - self.assertTrue(args.enable_torch_compile) - - def test_speed_mode_preserves_explicit_torch_compile_off(self): - args = self._from_dict_with_pipeline_config( - QwenImagePipelineConfig(), - kwargs={ - "model_path": "Qwen/Qwen-Image", - "performance_mode": "speed", - "enable_torch_compile": False, - }, - ) - self.assertFalse(args.enable_torch_compile) + def test_speed_mode_preserves_explicit_torch_compile_setting(self): + for enabled in (False, True): + with self.subTest(enabled=enabled): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "model_path": "Qwen/Qwen-Image", + "performance_mode": "speed", + "enable_torch_compile": enabled, + }, + ) + + self.assertEqual(args.enable_torch_compile, enabled) + + def test_speed_mode_honors_model_torch_compile_opt_in(self): + with patch.object( + QwenImagePipelineConfig, + "get_model_deployment_config", + return_value=ModelDeploymentConfig( + speed_mode_enable_torch_compile_by_default=True + ), + ): + args = self._from_dict_with_pipeline_config( + QwenImagePipelineConfig(), + kwargs={ + "model_path": "Qwen/Qwen-Image", + "performance_mode": "speed", + }, + ) + + self.assertTrue(args.enable_torch_compile) + def test_speed_mode_uses_minimax_h3_compile_policy(self): for explicit, expected in ((None, False), (True, True)): kwargs = {