B300 CP decouple: pad-aware hidden-states gather/scatter for standard a2a under shared-KV

With DeepEP the sparse-MoE layer keeps tokens SCATTERED (DeepEP dispatch/combine
routes tokens to expert-owning ranks), so the NSA layer communicator never
materializes the FULL hidden states. Decoupling CP onto the standard/allgather
a2a (flashinfer_trtllm, moe_a2a_backend=none) flips mlp_mode to FULL, which makes
NSACPLayerCommunicator gather hidden states across CP ranks before MoE and scatter
back after.

The gather used a raw even attn_cp_all_gather_into_tensor into get_local_dp_buffer(),
and the scatter used tensor_split(cp_size) + reduce_scatter. Both require equal
per-rank shards. Under --enable-nsa-prefill-cp-shared-kv the per-rank shards are
UNEVEN: the page-aligned in-seq split page-rounds each request's 2*cp_size segments,
so per_rank_actual_token differs across ranks. Result: "output tensor size must be
equal to world_size times input tensor size" in the CP all-gather.

Fix: make the round-trip pad-aware, mirroring the NSA index/KV gather idiom that
already handles uneven CP shards via max_rank_len.
- Gather: _cp_attn_tp_all_gather_padded_tensor pads this rank's shard to
  max_rank_len and all-gathers into a [max_rank_len * cp_size, H] rank-major
  buffer (allocates its own buffer; get_local_dp_buffer()'s length can be smaller
  than max_rank_len*cp_size once page-padding inflates the per-rank max).
- Scatter: tensor_split(cp_size) now yields equal max_rank_len chunks (canonical
  in-place reduce-scatter layout); reduce-scatter sum-combines this rank's EP-local
  MoE partials, then slice to per_rank_actual_token to drop the padding rows,
  realigning with the SCATTERED residual.

Strict generalization: when shards are even (non-shared-KV) the padding is zero and
this reduces to the previous behavior. Gated by nsa_use_prefill_cp, so decode/non-CP
and the DeepEP baseline (mlp_mode SCATTERED) are untouched. No global_num_tokens
2*cp_size padding needed: max(per_rank_actual_token)*cp_size == ceil_align(total,
cp_size), which the existing cp_size alignment already provides.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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