[BUGFIX] Fix CP residual size mismatch crash when tp_size == attn_cp_size (#21170)

python/sglang/srt/layers/communicator_nsa_cp.py

@@ -174,11 +174,10 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
         output_mode: ScatterMode,
         context: CommunicateContext,
     ):
-        if context.is_same_group_size(
+        # Check exact enum match first: even if group sizes happen to be equal
-            hidden_states_input_mode, output_mode
+        # (e.g. tp_size == attn_cp_size makes FULL and SCATTERED both size 1),
-        ) and context.is_same_group_size(residual_input_mode, output_mode):
+        # FULL and SCATTERED have different data layouts under CP and require
-            return NSACPCommunicateSummableTensorPairFn._trivial
+        # an explicit scatter operation.
-
         if (
             (hidden_states_input_mode == ScatterMode.FULL)
             and (residual_input_mode == ScatterMode.SCATTERED)
@@ -186,6 +185,11 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn):
         ):
             return NSACPCommunicateSummableTensorPairFn._scatter_hidden_states
 
+        if context.is_same_group_size(
+            hidden_states_input_mode, output_mode
+        ) and context.is_same_group_size(residual_input_mode, output_mode):
+            return NSACPCommunicateSummableTensorPairFn._trivial
+
         raise NotImplementedError(
             f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}"
         )