[diffusion] Fuse Cosmos3 Nano T2I attention on Hopper (#36571)

python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py

@@ -66,6 +66,22 @@ def is_cosmos_layer(name: str, _module: object) -> bool:
     return is_module_list_entry_in(name, ("layers", "gen_layers"))
 
 
+def _can_enable_t1_fused_qk_norm_rope(
+    *,
+    is_blackwell: bool,
+    is_hopper: bool,
+    hidden_act: str,
+    tp_size: int,
+    sp_size: int,
+    is_compiled: bool,
+) -> bool:
+    if is_compiled:
+        return False
+    if is_blackwell:
+        return True
+    return is_hopper and hidden_act != "relu2" and tp_size == 1 and sp_size == 1
+
+
 # -----------------------------------------------------------------------------
 # mRoPE position ID computation (Qwen3VL-style)
 # -----------------------------------------------------------------------------
@@ -1602,12 +1618,17 @@ class Cosmos3OmniTransformer(CachableDiT, LayerwiseOffloadableModuleMixin):
 
         self._ensure_cache_dicts()
 
-        # The T=1 fused path is faster on Blackwell, but regresses the Hopper
-        # Cosmos3-Super two-GPU workload. Keep Hopper on the original split path.
-        enable_t1_fused_qk_norm_rope = (
-            T == 1
-            and current_platform.is_blackwell()
-            and not self._gen_layers_torch_compiled
+        # The T=1 fused path is faster on Blackwell. It also benefits the
+        # single-GPU Hopper Nano (SwiGLU) workload, while the Hopper
+        # Cosmos3-Super (dense MLP) multi-GPU workload remains on the split
+        # path because that shape regresses with the fusion.
+        enable_t1_fused_qk_norm_rope = T == 1 and _can_enable_t1_fused_qk_norm_rope(
+            is_blackwell=current_platform.is_blackwell(),
+            is_hopper=current_platform.is_hopper(),
+            hidden_act=self.hidden_act,
+            tp_size=get_tp_world_size(),
+            sp_size=get_sp_world_size(),
+            is_compiled=self._gen_layers_torch_compiled,
         )
 
         # Compute UND K/V cache for this cache_key if not already cached

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

@@ -46,6 +46,7 @@ from sglang.multimodal_gen.runtime.loader.component_loaders.scheduler_loader imp
 from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping
 from sglang.multimodal_gen.runtime.models.dits.cosmos3video import (
     DomainAwareLinear,
+    _can_enable_t1_fused_qk_norm_rope,
     compute_mrope_position_ids_action,
     compute_mrope_position_ids_sound,
     compute_mrope_position_ids_vision,
@@ -91,6 +92,55 @@ def _cosmos3_server_args(config=None, batching_max_size=1):
     )
 
 
+class TestCosmos3T1FusedQKNormRoPE(unittest.TestCase):
+    def _can_enable(
+        self,
+        *,
+        is_blackwell=False,
+        is_hopper=False,
+        hidden_act="silu",
+        tp_size=1,
+        sp_size=1,
+        is_compiled=False,
+    ):
+        return _can_enable_t1_fused_qk_norm_rope(
+            is_blackwell=is_blackwell,
+            is_hopper=is_hopper,
+            hidden_act=hidden_act,
+            tp_size=tp_size,
+            sp_size=sp_size,
+            is_compiled=is_compiled,
+        )
+
+    def test_blackwell_remains_enabled(self):
+        self.assertTrue(
+            self._can_enable(
+                is_blackwell=True,
+                hidden_act="relu2",
+                tp_size=2,
+                sp_size=2,
+            )
+        )
+
+    def test_hopper_swiglu_single_gpu_enabled(self):
+        self.assertTrue(self._can_enable(is_hopper=True))
+
+    def test_hopper_dense_mlp_disabled(self):
+        self.assertFalse(self._can_enable(is_hopper=True, hidden_act="relu2"))
+
+    def test_hopper_tensor_parallel_disabled(self):
+        self.assertFalse(self._can_enable(is_hopper=True, tp_size=2))
+
+    def test_hopper_sequence_parallel_disabled(self):
+        self.assertFalse(self._can_enable(is_hopper=True, sp_size=2))
+
+    def test_compiled_path_disabled(self):
+        self.assertFalse(self._can_enable(is_hopper=True, is_compiled=True))
+
+    def test_other_architecture_disabled(self):
+        self.assertFalse(self._can_enable())
+
+
 class TestCosmos3ParamNamesMapping(unittest.TestCase):
     """Verify diffusers → sglang weight key translations."""