[diffusion] Fix symbolic replicated-mode counting under torch.compile (#34824)

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -130,6 +130,22 @@ def _kv_gather_unsupported_reason(
     return None
 
 
+def _count_active_replicated_modes(
+    num_replicated_prefix: int,
+    num_replicated_suffix: int,
+    num_replicated_kv_prefix: int,
+) -> int:
+    """Count active replicated-token modes without adding symbolic booleans."""
+    return sum(
+        int(value > 0)
+        for value in (
+            num_replicated_prefix,
+            num_replicated_suffix,
+            num_replicated_kv_prefix,
+        )
+    )
+
+
 def build_varlen_mask_meta(
     key_mask: torch.Tensor,
 ) -> dict:
@@ -822,13 +838,10 @@ class USPAttention(nn.Module):
             and not effective_skip_sp
             and get_sequence_parallel_world_size() > 1
         )
-        replicated_mode_count = sum(
-            value > 0
-            for value in (
-                num_replicated_prefix,
-                num_replicated_suffix,
-                num_replicated_kv_prefix,
-            )
+        replicated_mode_count = _count_active_replicated_modes(
+            num_replicated_prefix,
+            num_replicated_suffix,
+            num_replicated_kv_prefix,
         )
         if (
             self.sp_attention_mode == "kv_gather"

python/sglang/multimodal_gen/test/unit/test_usp_attention_kv_gather.py

@@ -9,10 +9,12 @@ from sglang.multimodal_gen.runtime.layers.attention.layer import (
     UlyssesAttention,
     UlyssesAttention_VSA,
     USPAttention,
+    _count_active_replicated_modes,
     _kv_gather_unsupported_reason,
     _resolve_sp_attention_mode,
 )
 from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.test.test_utils import CustomTestCase
 
 _LAYER = "sglang.multimodal_gen.runtime.layers.attention.layer"
 
@@ -344,5 +346,21 @@ class TestKVGatherCallSupport(unittest.TestCase):
                 attn.forward(q, q, q, qkv_pre_all_to_all=True)
 
 
+class TestReplicatedModeCountCompile(CustomTestCase):
+    def test_symbolic_shape_compiles(self):
+        def add_mode_count(x):
+            count = _count_active_replicated_modes(x.shape[0], 0, 0)
+            return x + count
+
+        compiled = torch.compile(
+            add_mode_count,
+            backend="eager",
+            fullgraph=True,
+            dynamic=True,
+        )
+        actual = compiled(torch.ones(3))
+        torch.testing.assert_close(actual, torch.full((3,), 2.0))
+
+
 if __name__ == "__main__":
     unittest.main()