From 64a31d4b75fac33840156479b1c9e9a62c0e5ae8 Mon Sep 17 00:00:00 2001 From: sunxxuns <126995791+sunxxuns@users.noreply.github.com> Date: Fri, 9 Jan 2026 04:16:58 -0800 Subject: [PATCH] [diffusion] amd: fix SGLANG_DIFFUSION_ATTENTION_BACKEND env var for diffusion attention backend selection (#16325) Co-authored-by: root --- .github/workflows/pr-test-amd.yml | 2 -- .../runtime/layers/attention/selector.py | 5 +---- .../runtime/managers/scheduler.py | 19 ++++++++++++++++++- .../runtime/pipelines_core/lora_pipeline.py | 16 +++++++++++++++- .../multimodal_gen/runtime/platforms/rocm.py | 13 ++----------- 5 files changed, 36 insertions(+), 19 deletions(-) diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml index 187aa4dc5..e1b7bc2e0 100644 --- a/.github/workflows/pr-test-amd.yml +++ b/.github/workflows/pr-test-amd.yml @@ -409,7 +409,6 @@ jobs: -e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \ -e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \ -e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \ - -e SGLANG_DIFFUSION_ATTENTION_BACKEND=AITER \ -e AITER_JIT_DIR=/sgl-data/aiter-kernels \ -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \ -e HF_HUB_ENABLE_HF_TRANSFER=1 \ @@ -529,7 +528,6 @@ jobs: -e SGLANG_DENOISE_STEP_TOLERANCE=0.6 \ -e SGLANG_DENOISE_AGG_TOLERANCE=0.3 \ -e SGLANG_TEST_NUM_INFERENCE_STEPS=5 \ - -e SGLANG_DIFFUSION_ATTENTION_BACKEND=AITER \ -e AITER_JIT_DIR=/sgl-data/aiter-kernels \ -e MIOPEN_USER_DB_PATH=/sgl-data/miopen-cache \ -e HF_HUB_ENABLE_HF_TRANSFER=1 \ diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py index 159088d65..fe526d69c 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/selector.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/selector.py @@ -108,10 +108,7 @@ def _cached_get_attn_backend( supported_attention_backends: tuple[AttentionBackendEnum], ) -> type[AttentionBackend]: # Check whether a particular choice of backend was - # previously forced. - # - # THIS SELECTION OVERRIDES THE SGLANG_DIFFUSION_ATTENTION_BACKEND - # ENVIRONMENT VARIABLE. + # previously forced via global_force_attn_backend() or --attention-backend CLI arg. from sglang.multimodal_gen.runtime.platforms import current_platform supported_attention_backends = set(supported_attention_backends) diff --git a/python/sglang/multimodal_gen/runtime/managers/scheduler.py b/python/sglang/multimodal_gen/runtime/managers/scheduler.py index 620de251f..af579840b 100644 --- a/python/sglang/multimodal_gen/runtime/managers/scheduler.py +++ b/python/sglang/multimodal_gen/runtime/managers/scheduler.py @@ -91,6 +91,10 @@ class Scheduler: self.prepare_server_warmup_reqs() + # Maximum consecutive errors before terminating the event loop + self._max_consecutive_errors = 3 + self._consecutive_error_count = 0 + def _handle_set_lora(self, reqs: List[Any]) -> OutputBatch: # TODO: return set status # TODO: return with SetLoRAResponse or something more appropriate @@ -252,11 +256,24 @@ class Scheduler: new_reqs = self.recv_reqs() new_reqs = self.process_received_reqs_with_req_based_warmup(new_reqs) self.waiting_queue.extend(new_reqs) + # Reset error count on success + self._consecutive_error_count = 0 except Exception as e: + self._consecutive_error_count += 1 logger.error( - f"Error receiving requests in scheduler event loop: {e}", + f"Error receiving requests in scheduler event loop " + f"(attempt {self._consecutive_error_count}/{self._max_consecutive_errors}): {e}", exc_info=True, ) + if self._consecutive_error_count >= self._max_consecutive_errors: + logger.error( + f"Maximum consecutive errors ({self._max_consecutive_errors}) reached. " + "Terminating scheduler event loop." + ) + raise RuntimeError( + f"Scheduler terminated after {self._max_consecutive_errors} " + f"consecutive errors. Last error: {e}" + ) from e continue # 2: execute, make sure a reply is always sent diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py index 3b2fbbdd7..f611cf7f9 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py @@ -305,7 +305,21 @@ class LoRAPipeline(ComposedPipelineBase): Load the LoRA, and setup the lora_adapters for later weight replacement """ assert lora_path is not None - lora_local_path = maybe_download_lora(lora_path) + + # Only rank 0 downloads to avoid race conditions where other ranks + # try to load incomplete downloads + if rank == 0: + lora_local_path = maybe_download_lora(lora_path) + else: + lora_local_path = None + + # Synchronize all ranks after download completes + if dist.is_initialized(): + dist.barrier() + + # Non-rank-0 workers now download (will hit cache since rank 0 completed) + if rank != 0: + lora_local_path = maybe_download_lora(lora_path) raw_state_dict = load_file(lora_local_path) lora_state_dict = normalize_lora_state_dict(raw_state_dict, logger=logger) diff --git a/python/sglang/multimodal_gen/runtime/platforms/rocm.py b/python/sglang/multimodal_gen/runtime/platforms/rocm.py index 7741cb0b9..5a2c5e844 100644 --- a/python/sglang/multimodal_gen/runtime/platforms/rocm.py +++ b/python/sglang/multimodal_gen/runtime/platforms/rocm.py @@ -11,7 +11,6 @@ from typing import Any import torch -import sglang.multimodal_gen.envs as envs from sglang.multimodal_gen.runtime.platforms.interface import ( AttentionBackendEnum, DeviceCapability, @@ -93,11 +92,6 @@ class RocmPlatform(Platform): head_size: int, dtype: torch.dtype, ) -> str: - logger.info( - "Trying SGLANG_DIFFUSION_ATTENTION_BACKEND=%s", - envs.SGLANG_DIFFUSION_ATTENTION_BACKEND, - ) - if selected_backend == AttentionBackendEnum.TORCH_SDPA: logger.info("Using Torch SDPA backend.") return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" @@ -108,13 +102,10 @@ class RocmPlatform(Platform): elif selected_backend == AttentionBackendEnum.AITER: if dtype not in (torch.float16, torch.bfloat16): logger.warning( - "AITer backend only supports fp16/bf16 inputs but got dtype=%s. " - "Falling back to Torch SDPA backend.", + "AITer backend works best with fp16/bf16 inputs but got dtype=%s. " + "Proceeding with AITer anyway.", dtype, ) - # TODO: need to compare triton with sdpa as an alternative backend - return "sglang.multimodal_gen.runtime.layers.attention.backends.sdpa.SDPABackend" - logger.info("Using AITer backend on ROCm.") return "sglang.multimodal_gen.runtime.layers.attention.backends.aiter.AITerBackend"