chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
yudian0504
2026-03-23 15:12:58 +08:00
committed by GitHub
parent 7757a9ddd0
commit 3d312643b9
1 changed files with 9 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/sglang/srt/layers/communicator_nsa_cp.py
View File

@@ -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=}"
)