diff --git a/python/sglang/srt/layers/communicator_nsa_cp.py b/python/sglang/srt/layers/communicator_nsa_cp.py index 296d14568..243b18e00 100644 --- a/python/sglang/srt/layers/communicator_nsa_cp.py +++ b/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( - hidden_states_input_mode, output_mode - ) and context.is_same_group_size(residual_input_mode, output_mode): - return NSACPCommunicateSummableTensorPairFn._trivial - + # Check exact enum match first: even if group sizes happen to be equal + # (e.g. tp_size == attn_cp_size makes FULL and SCATTERED both size 1), + # FULL and SCATTERED have different data layouts under CP and require + # 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=}" )