[diffusion] amd: fix SGLANG_DIFFUSION_ATTENTION_BACKEND env var for diffusion attention backend selection (#16325)

Co-authored-by: root <root@mi300x8-008.atl1.do.cpe.ice.amd.com>

.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 \

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.
+    # previously forced via global_force_attn_backend() or --attention-backend CLI arg.
-    #
-    # THIS SELECTION OVERRIDES THE SGLANG_DIFFUSION_ATTENTION_BACKEND
-    # ENVIRONMENT VARIABLE.
     from sglang.multimodal_gen.runtime.platforms import current_platform
 
     supported_attention_backends = set(supported_attention_backends)

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

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)

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. "
+                    "AITer backend works best with fp16/bf16 inputs but got dtype=%s. "
-                    "Falling back to Torch SDPA backend.",
+                    "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"