diff --git a/python/sglang/srt/layers/communicator_nsa_cp.py b/python/sglang/srt/layers/communicator_nsa_cp.py index 243b18e00..5a89066b4 100644 --- a/python/sglang/srt/layers/communicator_nsa_cp.py +++ b/python/sglang/srt/layers/communicator_nsa_cp.py @@ -19,6 +19,7 @@ from typing import Callable, Optional import torch from sglang.srt.layers.attention.nsa.utils import ( + _cp_attn_tp_all_gather_padded_tensor, is_nsa_enable_prefill_cp, nsa_use_prefill_cp, ) @@ -32,9 +33,7 @@ from sglang.srt.layers.communicator import ( ScatterMode, ) from sglang.srt.layers.dp_attention import ( - attn_cp_all_gather_into_tensor, attn_cp_reduce_scatter_tensor, - get_local_dp_buffer, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -149,17 +148,26 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn( ): if hidden_states.shape[0] != 0: hidden_states, residual = layernorm(hidden_states, residual) - # for prefill: attn tp scattered -> full - # for decode: attn tp full -> full + # for prefill: attn cp scattered -> full + # for decode: attn cp full -> full if nsa_use_prefill_cp(forward_batch): assert context.attn_dp_size == 1 - hidden_states, local_hidden_states = ( - get_local_dp_buffer(), + # The per-rank CP shards are uneven under shared-KV (the page-aligned + # in-seq split page-rounds each request's 2*cp_size segments), so a raw + # even all_gather_into_tensor is invalid. Pad this rank's shard to + # max_rank_len and all-gather into a [max_rank_len * cp_size, H] + # rank-major buffer (same idiom as the NSA index/KV gathers). The + # trailing padding rows are dropped on the scatter side. When shards are + # already even (e.g. non-shared-KV) the padding is zero, so this reduces + # to the previous behavior. NOTE: do NOT use get_local_dp_buffer() here; + # its length (global_num_tokens[dp_rank]) can be smaller than + # max_rank_len * cp_size once page-padding inflates the per-rank max. + hidden_states = _cp_attn_tp_all_gather_padded_tensor( hidden_states, - ) - attn_cp_all_gather_into_tensor( - hidden_states, - local_hidden_states, + forward_batch.nsa_cp_metadata.total_seq_lens, + context.attn_cp_size, + forward_batch, + torch.cuda.current_stream(), ) return hidden_states, residual @@ -202,13 +210,25 @@ class NSACPCommunicateSummableTensorPairFn(CommunicateSummableTensorPairFn): context: CommunicateContext, allow_reduce_scatter: bool = False, ): - # for prefill: full -> attn tp scattered - # for decode: full -> attn tp full + # for prefill: full -> attn cp scattered + # for decode: full -> attn cp full if nsa_use_prefill_cp(forward_batch): assert context.attn_dp_size == 1 + # Inverse of the padded rank-major gather. The FULL is + # [max_rank_len * cp_size, H]; tensor_split(cp_size) gives equal + # max_rank_len-row chunks (the canonical in-place reduce-scatter + # layout), and reduce-scatter sum-combines this rank's EP-local MoE + # partials into its chunk. Then drop the trailing padding rows to + # recover the rank's actual tokens, aligning with the SCATTERED + # residual (which was never gathered). input_hidden_states = hidden_states hidden_states = hidden_states.tensor_split(context.attn_cp_size)[ context.attn_cp_rank ] attn_cp_reduce_scatter_tensor(hidden_states, input_hidden_states) + hidden_states = hidden_states[ + : forward_batch.nsa_cp_metadata.per_rank_actual_token[ + context.attn_cp_rank + ] + ] return hidden_states, residual