From 3d312643b9f1035916c6a715b7f0a25b6e1fc484 Mon Sep 17 00:00:00 2001 From: yudian0504 <138860534+yudian0504@users.noreply.github.com> Date: Mon, 23 Mar 2026 15:12:58 +0800 Subject: [PATCH] [BUGFIX] Fix CP residual size mismatch crash when tp_size == attn_cp_size (#21170) --- python/sglang/srt/layers/communicator_nsa_cp.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) 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=}" )