[diffusion] fix: tiny fix _templated_ring_attention bug (#15053)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Mick <mickjagger19@icloud.com>

python/sglang/multimodal_gen/runtime/layers/usp.py

@@ -229,13 +229,13 @@ def ring_attn(
     use_segment_id = parse(torch.__version__).release >= parse("2.6.0").release
 
     attn_kwargs = dict(
-        mesh=ring_pg,
         op=attn_callable_adapter,
         dropout_p=dropout_p,
         is_causal=is_causal,
         query=query,
         key=key,
         value=value,
+        group=ring_pg,  # https://github.com/pytorch/pytorch/blob/c907c778f42ba2fdaf25b733dd25baf9779c6a12/torch/distributed/tensor/experimental/_context_parallel/_attention.py#L309
     )
 
     if use_segment_id:

python/sglang/multimodal_gen/test/run_suite.py

@@ -114,7 +114,9 @@ def run_pytest(files):
             and "AssertionError" in full_output
         )
 
-        if not is_perf_assertion:
+        is_flaky_ci_assertion = "SafetensorError" in full_output
+
+        if not (is_perf_assertion or is_flaky_ci_assertion):
             return returncode
 
     logger.info(f"Max retry exceeded")

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -304,7 +304,7 @@
                 "LatentPreparationStage": 12.4,
                 "TimestepPreparationStage": 2.71,
                 "DenoisingStage": 54705.41,
-                "DecodingStage": 311.13
+                "DecodingStage": 469.47
             },
             "denoise_step_ms": {
                 "0": 1067.03,
@@ -682,7 +682,7 @@
                 "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 44.28,
                 "LatentPreparationStage": 0.29,
-                "DenoisingStage": 9054.39,
+                "DenoisingStage": 9154.39,
                 "DecodingStage": 5995.09,
                 "per_frame_generation": null
             },
@@ -694,8 +694,8 @@
                 "4": 1489.93,
                 "5": 1487.02
             },
-            "expected_e2e_ms": 15672.15,
-            "expected_avg_denoise_ms": 1508.46,
+            "expected_e2e_ms": 16672.15,
+            "expected_avg_denoise_ms": 1608.46,
             "expected_median_denoise_ms": 1488.48
         },
         "wan2_2_i2v_a14b_2gpu": {

python/sglang/multimodal_gen/test/server/test_server_common.py

@@ -52,9 +52,14 @@ def diffusion_server(case: DiffusionTestCase) -> ServerContext:
     server_args = case.server_args
     sampling_params = case.sampling_params
     extra_args = os.environ.get("SGLANG_TEST_SERVE_ARGS", "")
-    extra_args += (
-        f" --num-gpus {server_args.num_gpus} --ulysses-degree {server_args.num_gpus}"
-    )
+    extra_args += f" --num-gpus {server_args.num_gpus}"
+
+    if server_args.ulysses_degree is not None:
+        extra_args += f" --ulysses-degree {server_args.ulysses_degree}"
+
+    if server_args.ring_degree is not None:
+        extra_args += f" --ring-degree {server_args.ring_degree}"
+
     # LoRA support
     if server_args.lora_path:
         extra_args += f" --lora-path {server_args.lora_path}"

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -119,6 +119,8 @@ class DiffusionServerArgs:
     custom_validator: str | None = None  # optional custom validator name
     # resources
     num_gpus: int = 1
+    ulysses_degree: int | None = None
+    ring_degree: int | None = None
     # LoRA
     lora_path: str | None = None  # LoRA adapter path (HF repo or local path)
 
@@ -485,6 +487,9 @@ TWO_GPU_CASES_B = [
             warmup_text=1,
             warmup_edit=0,
             num_gpus=2,
+            # test ring attn
+            ulysses_degree=1,
+            ring_degree=2,
         ),
         T2I_sampling_params,
     ),