diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index f7e7ba813..352dfb853 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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" diff --git a/python/sglang/multimodal_gen/test/unit/test_usp_attention_kv_gather.py b/python/sglang/multimodal_gen/test/unit/test_usp_attention_kv_gather.py index a04bb5b8a..f913310f8 100644 --- a/python/sglang/multimodal_gen/test/unit/test_usp_attention_kv_gather.py +++ b/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()