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>
leavelet
bed4601a20