sgl-project--sglang/python/sglang/srt/layers/communicator.py

# Copyright 2023-2024 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import logging
from contextlib import contextmanager
from dataclasses import dataclass
from enum import Enum, auto
from functools import partial
from typing import Callable, Dict, List, Optional, Tuple, Union

import torch

from sglang.srt.distributed import (
    attention_tensor_model_parallel_all_reduce,
    attention_tensor_model_parallel_quant_all_reduce,
    get_tp_group,
    moe_tensor_model_parallel_all_reduce,
    tensor_model_parallel_all_reduce,
)
from sglang.srt.distributed.device_communicators.pynccl_allocator import (
    use_symmetric_memory,
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.utils import (
    dsa_use_prefill_cp,
    is_dsa_enable_prefill_cp,
)
from sglang.srt.layers.dp_attention import (
    attn_tp_all_gather_into_tensor,
    attn_tp_reduce_scatter_tensor,
    dp_gather_partial,
    dp_gather_replicate,
    dp_reduce_scatter_tensor,
    dp_scatter,
    get_dp_global_num_tokens,
    get_global_dp_buffer,
    get_local_dp_buffer,
    get_moe_cp_rank,
    get_moe_cp_size,
    is_allocation_symmetric,
    is_dp_attention_enabled,
    is_enable_moe_cp_allgather,
    moe_cp_all_gather_into_tensor,
)
from sglang.srt.layers.flashinfer_comm_fusion import is_flashinfer_allreduce_unavailable
from sglang.srt.layers.moe import (
    get_moe_a2a_backend,
    should_use_dp_reduce_scatterv,
    should_use_flashinfer_cutlass_moe_fp4_allgather,
)
from sglang.srt.layers.quantization.fp8_utils import (
    _use_aiter_bpreshuffle_gfx95,
    materialize_bpreshuffle_fp8_scale_tuple,
)
from sglang.srt.layers.utils.cp_utils import (
    is_mla_prefill_cp_enabled,
    mla_use_prefill_cp,
)
from sglang.srt.model_executor.cuda_graph_config import (
    Backend,
    Phase,
    check_cuda_graph_backend,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.runtime_context import get_forward, get_parallel, get_server_args
from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
from sglang.srt.utils import (
    get_bool_env_var,
    is_cuda,
    is_flashinfer_available,
    is_gfx95_supported,
    is_hip,
    is_npu,
    is_sm90_supported,
    is_sm100_supported,
)

_is_cuda = is_cuda()
_is_flashinfer_available = is_flashinfer_available()
_is_sm90_supported = _is_cuda and is_sm90_supported()
_is_sm100_supported = _is_cuda and is_sm100_supported()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
_is_gfx95_supported = is_gfx95_supported()
_is_npu = is_npu()
_use_ag_after_qlora = envs.SGLANG_USE_AG_AFTER_QLORA.get()

if _use_aiter:
    from aiter.ops.rmsnorm import add_rmsnorm_quant as _aiter_add_rmsnorm_quant
    from aiter.ops.rmsnorm import rmsnorm_quant as _aiter_rmsnorm_quant

    from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype as _aiter_fp8_dtype

    if _is_gfx95_supported:
        from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant

        from sglang.srt.layers.quantization.rocm_mxfp4_utils import (
            fused_rms_mxfp4_quant,
        )
elif _is_npu:
    from sglang.srt.hardware_backend.npu.cmo import prepare_weight_cache


def _fused_rmsnorm_fp8_per_token_quant(
    hidden_states: torch.Tensor,
    weight: torch.Tensor,
    epsilon: float,
    residual: Optional[torch.Tensor] = None,
):
    """Fused (optional residual-add +) RMSNorm + FP8 per-token quantization.

    Only used with the aiter (ROCm) backend.

    Args:
        residual: if provided, computes hidden_states + residual before RMSNorm
                  and returns updated residual_out as second element.

    Returns:
        If residual is None:  (out_fp8, scale)
        If residual provided: ((out_fp8, scale), residual_out)
    """
    M, N = hidden_states.shape
    out_fp8 = torch.empty((M, N), dtype=_aiter_fp8_dtype, device=hidden_states.device)
    scale = torch.empty(M, dtype=torch.float32, device=hidden_states.device)
    if residual is not None:
        residual_out = torch.empty_like(hidden_states)
        _aiter_add_rmsnorm_quant(
            out_fp8,
            hidden_states,
            residual,
            residual_out,
            scale,
            weight,
            epsilon,
            0,  # group_size=0 → per-token
        )
        return (out_fp8, scale.unsqueeze(1)), residual_out
    else:
        _aiter_rmsnorm_quant(
            out_fp8,
            hidden_states,
            scale,
            weight,
            epsilon,
            0,  # group_size=0 → per-token
        )
        return (out_fp8, scale.unsqueeze(1))


# TODO: According to the discussion in https://github.com/flashinfer-ai/flashinfer/issues/1223#issuecomment-3047256465
# We set the max token num to 128 for allreduce fusion with min-latency case(use_oneshot=True).
FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048


def apply_flashinfer_allreduce_fusion(batch_size: int):
    return (
        # NOTE: flashinfer 0.6.1 caused performance regression on sm100 for allreduce fusion
        # Ref: https://github.com/sgl-project/sglang/issues/17237
        (_is_sm90_supported or _is_sm100_supported)
        and _is_flashinfer_available
        and batch_size > 0
        and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE
        and not is_dp_attention_enabled()
        and get_server_args().flashinfer_allreduce_fusion_backend is not None
        and not is_flashinfer_allreduce_unavailable()
    )


def apply_aiter_all_reduce_fusion(input_tensor: torch.Tensor):
    n = input_tensor.shape[-1]
    total_bytes = input_tensor.numel() * input_tensor.element_size()
    # Aiter's should_custom_ar uses <= max_size/2 (64 MB); match that boundary.
    return (
        _use_aiter
        and total_bytes > 0
        and n <= 16384
        and total_bytes <= 8 * 1024 * 8192
        and get_parallel().tp_size != 6
        and not is_dp_attention_enabled()
        and get_server_args().enable_aiter_allreduce_fusion
    )


class ScatterMode(Enum):
    """
    Suppose we have TP=4, DP=2, enable-dp-attention, and the system handles seq a,b,c,d
    Model input/output: [ab, ab, cd, cd] for four ranks respectively
    SCATTERED: [a, b, c, d]
    TP_ATTN_FULL: [ab, ab, cd, cd], i.e. all ranks inside a TP attn group have full data of the group
    FULL: [abcd, abcd, abcd, abcd]
    MOE_FULL: full within the MoE group (cp_per_moe CP chunks), used when moe_dp_size < attn_cp_size
    """

    SCATTERED = auto()
    TP_ATTN_FULL = auto()
    FULL = auto()
    MOE_FULL = auto()

    @staticmethod
    def model_input_output():
        """The scatter mode for model forward pass input and output data"""
        if is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled():
            return ScatterMode.SCATTERED

        return ScatterMode.TP_ATTN_FULL


class AttentionInputs:
    def __init__(
        self,
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        qkv_latent_func: Callable,
    ):
        self.hidden_states_local = hidden_states
        self.forward_batch = forward_batch
        self.qkv_latent_func = qkv_latent_func
        self.hidden_states_ = None
        self.qkv_latent_ = None

    def tp_all_gather_hidden_states(self, hidden_states, forward_batch):
        total_tokens = forward_batch.input_ids.shape[0]
        output = hidden_states.new_empty((total_tokens, hidden_states.shape[-1]))
        get_tp_group().all_gather_into_tensor(output, hidden_states)
        return output

    def fetch_qkv_latent(self):
        if self.qkv_latent_ is not None:
            return self.qkv_latent_
        assert self.qkv_latent_func is not None
        self.qkv_latent_ = self.qkv_latent_func(
            self.hidden_states_local, self.forward_batch
        )
        if get_attn_tp_context().input_scattered:
            self.qkv_latent_ = self.tp_all_gather_hidden_states(
                self.qkv_latent_, self.forward_batch
            )
        return self.qkv_latent_

    def fetch_hidden_states(self):
        if self.hidden_states_ is not None:
            return self.hidden_states_
        self.hidden_states_ = self.hidden_states_local
        if get_attn_tp_context().input_scattered:
            self.hidden_states_ = self.tp_all_gather_hidden_states(
                self.hidden_states_, self.forward_batch
            )
        return self.hidden_states_


class AttnTpContext:
    def __init__(self):
        self.allow_input_scattered = False
        self.is_dsa = False

    def init_context(self, q_lora_rank, is_dsa):
        self.is_dsa = is_dsa
        self.allow_input_scattered = (
            get_server_args().enable_attn_tp_input_scattered
            and (_is_cuda or _is_npu)
            and q_lora_rank is not None
            and not is_dsa
            and get_parallel().tp_size > 1
            and not is_dp_attention_enabled()
            and get_moe_a2a_backend().is_none()
            and not enable_moe_dense_fully_dp()
            and not check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE)
            and get_server_args().speculative_algorithm != "EAGLE3"
        )
        if get_server_args().enable_attn_tp_input_scattered:
            if not self.allow_input_scattered:
                logging.info(
                    "attn_tp_input_scattered is not enabled while other conditions are not met"
                )
            else:
                logging.info("attn_tp_input_scattered is enabled")

    def use_input_scattered(self, forward_batch: ForwardBatch):
        return (
            self.allow_input_scattered
            and forward_batch.forward_mode.is_extend()
            and not forward_batch.forward_mode.is_target_verify()
            and forward_batch.input_ids is not None
            and not forward_batch.can_run_tbo
        )

    @property
    def input_scattered(self):
        return get_forward().attn_input_scattered

    def set_attn_inputs(self, attn_inputs: AttentionInputs):
        get_forward().set("attn_inputs", attn_inputs)

    def fetch_qkv_latent(self):
        attn_inputs = get_forward().attn_inputs
        assert attn_inputs is not None
        return attn_inputs.fetch_qkv_latent()

    def fetch_hidden_states(self):
        attn_inputs = get_forward().attn_inputs
        assert attn_inputs is not None
        return attn_inputs.fetch_hidden_states()

    def clear_attn_inputs(self) -> None:
        get_forward().set("attn_inputs", None)

    @contextmanager
    def maybe_input_scattered(self, forward_batch: ForwardBatch):
        flag = self.use_input_scattered(forward_batch)
        forward = get_forward()
        # scoped() also restores when the forward raises — the old in-place
        # swap leaked the flag on exceptions.
        with forward.scoped(attn_input_scattered=flag):
            try:
                yield
            finally:
                forward.set("attn_inputs", None)


ATTN_TP_CONTEXT = AttnTpContext()


def get_attn_tp_context():
    return ATTN_TP_CONTEXT


@dataclass
class _LayerModeComputationContext:
    num_layers: int
    layer_id: int
    is_layer_sparse: bool
    is_previous_layer_sparse: Optional[bool]
    is_next_layer_sparse: Optional[bool]

    def previous_layer(self):
        assert self.is_previous_layer_sparse is not None
        return _LayerModeComputationContext(
            num_layers=self.num_layers,
            layer_id=self.layer_id - 1,
            is_layer_sparse=self.is_previous_layer_sparse,
            is_previous_layer_sparse=None,
            is_next_layer_sparse=self.is_layer_sparse,
        )


@dataclass
class LayerScatterModes:
    layer_input_mode: ScatterMode
    attn_mode: ScatterMode
    # Can be further split into e.g. mlp_input_mode and mlp_output_mode if needed
    mlp_mode: ScatterMode
    middle_residual_mode: ScatterMode
    layer_output_mode: ScatterMode

    @classmethod
    def init_new(cls, **kwargs):
        context = _LayerModeComputationContext(**kwargs)
        return cls(
            layer_input_mode=cls._compute_layer_input_mode(context),
            attn_mode=ScatterMode.TP_ATTN_FULL,
            mlp_mode=cls._compute_mlp_mode(context),
            middle_residual_mode=cls._compute_middle_residual_mode(context),
            layer_output_mode=cls._compute_layer_output_mode(context),
        )

    @classmethod
    def _compute_layer_input_mode(cls, context: _LayerModeComputationContext):
        if context.layer_id == 0:
            return ScatterMode.model_input_output()
        return cls._compute_layer_output_mode(context.previous_layer())

    @classmethod
    def _compute_mlp_mode(cls, context: _LayerModeComputationContext):
        if context.is_layer_sparse:
            if (
                # Token dispatch/combine will be handled outside of LayerCommunicator for these modes.
                not get_moe_a2a_backend().is_none()
                or should_use_flashinfer_cutlass_moe_fp4_allgather()
            ):
                return ScatterMode.SCATTERED
            # DSA CP and MLA CP both don't support MOE_FULL yet; fall back to FULL.
            if is_enable_moe_cp_allgather() and not (
                is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled()
            ):
                return ScatterMode.MOE_FULL
            return ScatterMode.FULL
        else:
            return (
                ScatterMode.SCATTERED
                if enable_moe_dense_fully_dp()
                else ScatterMode.FULL
            )

    @classmethod
    def _should_gather_for_tbo(cls, context: _LayerModeComputationContext):
        return (
            not context.is_layer_sparse
            and context.is_next_layer_sparse
            and enable_moe_dense_fully_dp()
            and get_server_args().enable_two_batch_overlap
        )

    @classmethod
    def _compute_middle_residual_mode(cls, context: _LayerModeComputationContext):
        mlp_mode = cls._compute_mlp_mode(context)
        if mlp_mode == ScatterMode.SCATTERED:
            return ScatterMode.SCATTERED
        if mlp_mode in (ScatterMode.FULL, ScatterMode.MOE_FULL):
            return ScatterMode.TP_ATTN_FULL
        raise NotImplementedError

    @classmethod
    def _compute_layer_output_mode(cls, context: _LayerModeComputationContext):
        mlp_mode = cls._compute_mlp_mode(context)
        if context.layer_id == context.num_layers - 1:
            return ScatterMode.model_input_output()
        if mlp_mode == ScatterMode.SCATTERED:
            if cls._should_gather_for_tbo(context):
                return ScatterMode.TP_ATTN_FULL
            return ScatterMode.SCATTERED
        if mlp_mode in (ScatterMode.FULL, ScatterMode.MOE_FULL):
            return ScatterMode.TP_ATTN_FULL
        raise NotImplementedError


def enable_moe_dense_fully_dp():
    return get_server_args().moe_dense_tp_size == 1


class LayerCommunicator:
    def __init__(
        self,
        layer_scatter_modes: LayerScatterModes,
        input_layernorm: torch.nn.Module,
        post_attention_layernorm: torch.nn.Module,
        # Reduce scatter requires skipping all-reduce in model code after MoE/MLP, so only enable for models which have that implemented. Remove flag once done for all models that use LayerCommunicator.
        allow_reduce_scatter: bool = False,
        is_last_layer: bool = False,
        qkv_latent_func: Optional[Callable] = None,
        force_layernorm_before_dp_gather: bool = False,
    ):
        self.layer_scatter_modes = layer_scatter_modes
        self.input_layernorm = input_layernorm
        self.post_attention_layernorm = post_attention_layernorm
        self.allow_reduce_scatter = allow_reduce_scatter
        self.is_last_layer = is_last_layer
        self.qkv_latent_func = qkv_latent_func
        self.force_layernorm_before_dp_gather = force_layernorm_before_dp_gather

        self._context = CommunicateContext.init_new()
        self._context.force_layernorm_before_dp_gather = (
            force_layernorm_before_dp_gather
        )
        self._post_init_communicate()
        self._speculative_algo = SpeculativeAlgorithm.from_string(
            get_server_args().speculative_algorithm
        )

    def _post_init_communicate(self):
        self._communicate_simple_fn = CommunicateSimpleFn.get_fn(
            input_mode=self.layer_scatter_modes.layer_input_mode,
            output_mode=self.layer_scatter_modes.attn_mode,
            context=self._context,
        )
        self._communicate_with_all_reduce_and_layer_norm_fn = (
            CommunicateWithAllReduceAndLayerNormFn.get_fn(
                hidden_states_input_mode=self.layer_scatter_modes.attn_mode,
                residual_input_mode=self.layer_scatter_modes.layer_input_mode,
                hidden_states_output_mode=self.layer_scatter_modes.mlp_mode,
                residual_output_mode=self.layer_scatter_modes.middle_residual_mode,
                context=self._context,
            )
        )
        self._communicate_summable_tensor_pair_fn = (
            CommunicateSummableTensorPairFn.get_fn(
                hidden_states_input_mode=self.layer_scatter_modes.mlp_mode,
                residual_input_mode=self.layer_scatter_modes.middle_residual_mode,
                output_mode=self.layer_scatter_modes.layer_output_mode,
                context=self._context,
            )
        )

    def prepare_attn_and_capture_last_layer_outputs(
        self,
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        captured_last_layer_outputs: Optional[List[torch.Tensor]] = None,
        post_residual_addition: Optional[torch.Tensor] = None,
        quant_format: str = "",
    ):
        hidden_states, residual = self.prepare_attn(
            hidden_states,
            residual,
            forward_batch,
            quant_format=quant_format,
            post_residual_addition=post_residual_addition,
        )
        if captured_last_layer_outputs is not None:
            gathered_last_layer_output = self._communicate_simple_fn(
                hidden_states=residual,
                forward_batch=forward_batch,
                context=self._context,
            )
            if (
                gathered_last_layer_output is residual
                and not self._post_attn_residual_is_read_only(residual)
            ):
                gathered_last_layer_output = residual.clone()
            captured_last_layer_outputs.append(gathered_last_layer_output)
        return hidden_states, residual

    def _post_attn_residual_is_read_only(self, residual: torch.Tensor) -> bool:
        """True if ``prepare_mlp``'s post-attention RMSNorm leaves ``residual``
        untouched, so Eagle3 aux capture can keep its reference and skip the clone.

        Only the flashinfer all-reduce-fusion path writes a fresh ``residual_out``
        (see ``flashinfer_allreduce_residual_rmsnorm``); the aiter fused kernel and
        every plain norm fold into ``residual`` in place. That path is reachable
        only from the ``_gather_*`` communicate-fns, and only when they fall past
        their input-scattered branch.
        """
        norm_fn = getattr(
            self._communicate_with_all_reduce_and_layer_norm_fn,
            "func",
            self._communicate_with_all_reduce_and_layer_norm_fn,
        )
        uses_gather_norm = norm_fn in (
            CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual,
            CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual_moe,
        )
        return (
            uses_gather_norm
            and not get_attn_tp_context().input_scattered
            and apply_flashinfer_allreduce_fusion(residual.shape[0])
        )

    def prepare_attn(
        self,
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        quant_format: str = "",
        post_residual_addition: Optional[torch.Tensor] = None,
    ):
        if get_attn_tp_context().input_scattered:
            hidden_states, residual = self._tp_reduce_scatter(
                hidden_states,
                residual,
            )
        if hidden_states.shape[0] == 0:
            residual = hidden_states
        else:
            if (
                residual is not None
                and hasattr(hidden_states, "_sglang_needs_allreduce_fusion")
                and hidden_states._sglang_needs_allreduce_fusion
            ):
                if (
                    apply_aiter_all_reduce_fusion(hidden_states)
                    or apply_flashinfer_allreduce_fusion(hidden_states.shape[0])
                ) and hasattr(self.input_layernorm, "forward_with_allreduce_fusion"):
                    hidden_states, residual = (
                        self.input_layernorm.forward_with_allreduce_fusion(
                            hidden_states, residual, use_attn_tp_group=False
                        )
                    )
                else:
                    hidden_states = moe_tensor_model_parallel_all_reduce(hidden_states)
                    hidden_states, residual = self.input_layernorm(
                        hidden_states, residual
                    )
            else:
                if residual is None:
                    residual = hidden_states

                    if _use_aiter and _is_gfx95_supported and ("mxfp4" in quant_format):
                        hidden_states, *_, _ = fused_rms_mxfp4_quant(
                            hidden_states,
                            self.input_layernorm.weight,
                            self.input_layernorm.variance_epsilon,
                            None,
                            None,
                            None,
                            None,
                        )
                    elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"):
                        # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant.
                        # When DSA is active, also preserve the unquantized bf16
                        # output as a 3-tuple (fp8, scale, bf16) so the DSA
                        # indexer can skip redundant FP8 dequantization.
                        _dsa_needs_bf16 = get_attn_tp_context().is_dsa
                        hidden_states, _unq_bf16, _, _res = fused_rms_fp8_group_quant(
                            hidden_states,
                            self.input_layernorm.weight,
                            self.input_layernorm.variance_epsilon,
                            inp2=None,
                            inp2_weight=None,
                            inp2_epsilon=None,
                            group_size=128,
                            dtype_quant=torch.float8_e4m3fn,
                            res1=None,
                            output_unquantized_inp1=_dsa_needs_bf16,
                            transpose_scale=False,
                        )
                        if _use_aiter_bpreshuffle_gfx95:
                            hidden_states = materialize_bpreshuffle_fp8_scale_tuple(
                                hidden_states
                            )
                        if _dsa_needs_bf16:
                            hidden_states = (
                                hidden_states[0],
                                hidden_states[1],
                                _unq_bf16,
                            )

                    elif _use_aiter and (quant_format == "fp8_per_token"):
                        hidden_states = _fused_rmsnorm_fp8_per_token_quant(
                            hidden_states,
                            self.input_layernorm.weight.data,
                            self.input_layernorm.variance_epsilon,
                        )

                    else:
                        hidden_states = self.input_layernorm(hidden_states)
                else:
                    if _use_aiter and _is_gfx95_supported and ("mxfp4" in quant_format):
                        hidden_states, *_, residual = fused_rms_mxfp4_quant(
                            hidden_states,
                            self.input_layernorm.weight,
                            self.input_layernorm.variance_epsilon,
                            None,
                            None,
                            None,
                            residual,
                        )
                    elif _use_aiter and _is_gfx95_supported and (quant_format == "fp8"):
                        # aiter (ROCm gfx95) fused RMSNorm + FP8 group quant
                        # with residual addition. When DSA is active, pack
                        # the unquantized bf16 as a 3-tuple (fp8, scale, bf16).
                        _dsa_needs_bf16 = get_attn_tp_context().is_dsa
                        hidden_states, _unq_bf16, _, residual = (
                            fused_rms_fp8_group_quant(
                                hidden_states,
                                self.input_layernorm.weight,
                                self.input_layernorm.variance_epsilon,
                                inp2=None,
                                inp2_weight=None,
                                inp2_epsilon=None,
                                group_size=128,
                                dtype_quant=torch.float8_e4m3fn,
                                res1=residual,
                                output_unquantized_inp1=_dsa_needs_bf16,
                                transpose_scale=False,
                            )
                        )
                        if _use_aiter_bpreshuffle_gfx95:
                            hidden_states = materialize_bpreshuffle_fp8_scale_tuple(
                                hidden_states
                            )
                        if _dsa_needs_bf16:
                            hidden_states = (
                                hidden_states[0],
                                hidden_states[1],
                                _unq_bf16,
                            )
                    elif _use_aiter and (quant_format == "fp8_per_token"):
                        if post_residual_addition is not None:
                            residual = residual + post_residual_addition
                        hidden_states, residual = _fused_rmsnorm_fp8_per_token_quant(
                            hidden_states,
                            self.input_layernorm.weight.data,
                            self.input_layernorm.variance_epsilon,
                            residual=residual,
                        )
                    else:
                        hidden_states, residual = self.input_layernorm(
                            hidden_states,
                            residual,
                            post_residual_addition,
                        )

        hidden_states = self._communicate_simple_fn(
            hidden_states=hidden_states,
            forward_batch=forward_batch,
            context=self._context,
        )
        if self.qkv_latent_func is not None:
            attn_inputs = AttentionInputs(
                hidden_states, forward_batch, self.qkv_latent_func
            )
            get_attn_tp_context().set_attn_inputs(attn_inputs)
        return hidden_states, residual

    def _tp_reduce_scatter(
        self,
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        if hidden_states.shape[0] == 0:
            return hidden_states, hidden_states
        assert (
            hidden_states.shape[0] % self._context.tp_size == 0
        ), f"Expected total tokens {hidden_states.shape[0]} % tp_size {self._context.tp_size} to be 0"
        local_tokens = hidden_states.shape[0] // self._context.tp_size
        output = hidden_states.new_empty(local_tokens, *hidden_states.shape[1:])
        get_tp_group().reduce_scatter_tensor(output, hidden_states)
        if residual is not None:
            residual = residual.tensor_split(self._context.tp_size)[
                self._context.tp_rank
            ]
        return output, residual

    def prepare_mlp(
        self,
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        cache=None,
    ):
        if cache is not None:
            self._context.cache = cache

        return self._communicate_with_all_reduce_and_layer_norm_fn(
            hidden_states=hidden_states,
            residual=residual,
            forward_batch=forward_batch,
            layernorm=self.post_attention_layernorm,
            context=self._context,
        )

    def postprocess_layer(
        self,
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
    ):
        return self._communicate_summable_tensor_pair_fn(
            hidden_states=hidden_states,
            residual=residual,
            forward_batch=forward_batch,
            context=self._context,
            allow_reduce_scatter=self.allow_reduce_scatter,
        )

    def should_use_reduce_scatter(self, forward_batch: ForwardBatch):
        if not self.allow_reduce_scatter:
            return False
        if (
            self._communicate_summable_tensor_pair_fn
            is CommunicateSummableTensorPairFn._scatter_hidden_states
        ):
            if should_use_dp_reduce_scatterv():
                return True
            if forward_batch.dp_padding_mode.is_max_len():
                return True
        if dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch):
            return True
        if get_attn_tp_context().input_scattered and not self.is_last_layer:
            return True
        return False

    # NOTE: This function will cause torch recompilation
    def should_fuse_mlp_allreduce_with_next_layer(
        self, forward_batch: ForwardBatch
    ) -> bool:
        # When MOE_FULL is active (moe_cp allgather), fusion must be disabled because
        # the fusion path skips postprocess_layer which contains the moe_cp scatter.
        # Without scatter, hidden_states remain at MOE_FULL size while residual is at
        # TP_ATTN_FULL size, causing a shape mismatch.
        if is_enable_moe_cp_allgather():
            return False

        if (
            is_dp_attention_enabled()
            and self._speculative_algo is not None
            and self._speculative_algo.is_eagle()
        ):
            return False

        if get_attn_tp_context().input_scattered:
            return False

        batch_size = (
            forward_batch.input_ids.shape[0]
            if hasattr(forward_batch, "input_ids")
            else 0
        )

        # When mlp_mode is SCATTERED, the MLP runs on scattered data with no TP
        # all-reduce, so there is nothing to fuse with the next layer.
        if self.layer_scatter_modes.mlp_mode == ScatterMode.SCATTERED:
            return False

        return (
            (
                apply_flashinfer_allreduce_fusion(batch_size)
                or (
                    _use_aiter
                    and batch_size > 0
                    and get_parallel().tp_size != 6
                    and not is_dp_attention_enabled()
                    and get_moe_a2a_backend().is_none()
                    and get_server_args().enable_aiter_allreduce_fusion
                )
            )
            and (not self.is_last_layer)
            and (self._context.tp_size > 1)
        )


@dataclass
class CommunicateContext:
    process_group_sizes: Dict[ScatterMode, int]
    attn_tp_rank: int
    attn_tp_size: int
    attn_dp_size: int
    attn_cp_rank: int
    attn_cp_size: int
    tp_size: int
    cache = None
    tp_rank: int
    force_layernorm_before_dp_gather: bool = False

    def is_same_group_size(self, a: ScatterMode, b: ScatterMode):
        return self.process_group_sizes[a] == self.process_group_sizes[b]

    @classmethod
    def init_new(cls):
        attn_tp_rank = get_parallel().attn_tp_rank
        attn_tp_size = get_parallel().attn_tp_size
        attn_dp_size = get_parallel().attn_dp_size
        attn_cp_size = get_parallel().attn_cp_size
        attn_cp_rank = get_parallel().attn_cp_rank
        tp_size = get_parallel().tp_size
        tp_rank = get_parallel().tp_rank
        moe_cp_size = get_moe_cp_size()
        process_group_sizes = {
            ScatterMode.SCATTERED: 1,
            ScatterMode.TP_ATTN_FULL: attn_tp_size,
            # TODO: support --moe-dense-tp-size > 1
            # With context parallel enabled, we should exclude
            # the attn_cp_size from the total tp_size
            ScatterMode.FULL: tp_size // attn_cp_size,
            ScatterMode.MOE_FULL: tp_size // (attn_cp_size // moe_cp_size),
        }
        return cls(
            process_group_sizes=process_group_sizes,
            attn_tp_rank=attn_tp_rank,
            attn_tp_size=attn_tp_size,
            attn_dp_size=attn_dp_size,
            attn_cp_rank=attn_cp_rank,
            attn_cp_size=attn_cp_size,
            tp_size=tp_size,
            tp_rank=tp_rank,
        )


class CommunicateSimpleFn:
    @staticmethod
    def get_fn(
        input_mode: ScatterMode,
        output_mode: ScatterMode,
        context: CommunicateContext,
    ):
        if context.is_same_group_size(input_mode, output_mode):
            return CommunicateSimpleFn._trivial

        if (input_mode == ScatterMode.SCATTERED) and (
            output_mode == ScatterMode.TP_ATTN_FULL
        ):
            if _use_ag_after_qlora:
                return CommunicateSimpleFn._trivial
            return CommunicateSimpleFn._scattered_to_tp_attn_full

        raise NotImplementedError(f"{input_mode=} {output_mode=}")

    @staticmethod
    def _trivial(
        hidden_states: torch.Tensor,
        forward_batch: ForwardBatch,
        context: CommunicateContext,
    ) -> torch.Tensor:
        return hidden_states

    @staticmethod
    def _scattered_to_tp_attn_full(
        hidden_states: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
        forward_batch: ForwardBatch,
        context: CommunicateContext,
    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        if isinstance(hidden_states, tuple):
            gathered_hidden_states = []
            for local_hidden_states in hidden_states:
                with use_symmetric_memory(
                    get_tp_group(),
                    disabled=not is_allocation_symmetric(),
                ):
                    output = torch.empty(
                        (
                            local_hidden_states.shape[0] * context.attn_tp_size,
                            *local_hidden_states.shape[1:],
                        ),
                        dtype=local_hidden_states.dtype,
                        device=local_hidden_states.device,
                    )
                attn_tp_all_gather_into_tensor(
                    output,
                    local_hidden_states,
                )
                gathered_hidden_states.append(output)
            return tuple(gathered_hidden_states)

        hidden_states, local_hidden_states = (
            get_local_dp_buffer(get_parallel().attn_tp_group),
            hidden_states,
        )
        attn_tp_all_gather_into_tensor(
            hidden_states,
            local_hidden_states,
        )
        return hidden_states


class CommunicateWithAllReduceAndLayerNormFn:
    """Besides communication, needs to
    1. All reduce in tp_attn_group on hidden_states
    2. Apply layer norm
    """

    @staticmethod
    def get_fn(
        hidden_states_input_mode: ScatterMode,
        residual_input_mode: ScatterMode,
        hidden_states_output_mode: ScatterMode,
        residual_output_mode: ScatterMode,
        context: CommunicateContext,
    ):

        if (
            context.is_same_group_size(
                hidden_states_input_mode, hidden_states_output_mode
            )
            and context.is_same_group_size(residual_input_mode, residual_output_mode)
            and context.attn_tp_size == 1
        ):
            return CommunicateWithAllReduceAndLayerNormFn._simple

        if (
            (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
            and (
                residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL]
            )
            and (hidden_states_output_mode == ScatterMode.FULL)
            and (residual_output_mode == ScatterMode.TP_ATTN_FULL)
        ):
            return partial(
                CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual,
                residual_input_mode=residual_input_mode,
            )

        if (
            (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
            and (
                residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL]
            )
            and (hidden_states_output_mode == ScatterMode.MOE_FULL)
            and (residual_output_mode == ScatterMode.TP_ATTN_FULL)
        ):
            return partial(
                CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual_moe,
                residual_input_mode=residual_input_mode,
            )

        if (
            (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
            and (
                residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL]
            )
            and (hidden_states_output_mode == ScatterMode.SCATTERED)
            and (residual_output_mode == ScatterMode.SCATTERED)
        ):
            return partial(
                CommunicateWithAllReduceAndLayerNormFn._scatter_hidden_states_and_residual,
                residual_input_mode=residual_input_mode,
            )

        if (
            (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
            and (
                residual_input_mode in [ScatterMode.SCATTERED, ScatterMode.TP_ATTN_FULL]
            )
            and (hidden_states_output_mode == ScatterMode.TP_ATTN_FULL)
            and (residual_output_mode == ScatterMode.TP_ATTN_FULL)
            and context.attn_tp_size > 1
        ):
            # Used when the dense MLP is tensor-parallelized along the
            # attention TP group (``moe_dense_tp_size > 1``): hidden states
            # need an all-reduce inside the attention TP group before the
            # next layernorm, while staying in TP_ATTN_FULL on both sides.
            return (
                CommunicateWithAllReduceAndLayerNormFn._tp_attn_all_reduce_and_layernorm
            )

        raise NotImplementedError(
            f"{hidden_states_input_mode=} {residual_input_mode=} {hidden_states_output_mode=} {residual_output_mode=}"
        )

    @staticmethod
    def _simple(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        layernorm: torch.nn.Module,
        context: CommunicateContext,
    ):
        # TODO move these `if shape != 0` into LayerNorm itself
        if hidden_states.shape[0] != 0:
            hidden_states, residual = layernorm(hidden_states, residual)
        return hidden_states, residual

    @staticmethod
    def _tp_attn_all_reduce_and_layernorm(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        layernorm: torch.nn.Module,
        context: CommunicateContext,
    ):
        """All-reduce hidden states inside the attention TP group, then layernorm.

        Used when the dense MLP shares the attention TP group
        (``moe_dense_tp_size > 1``): both hidden states and residual stay in
        ``TP_ATTN_FULL`` across the boundary.
        """
        hidden_states = get_parallel().attn_tp_group.all_reduce(hidden_states)
        if hidden_states.shape[0] != 0:
            hidden_states, residual = layernorm(hidden_states, residual)
        return hidden_states, residual

    @staticmethod
    def _gather_hidden_states_and_residual(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        layernorm: torch.nn.Module,
        context: CommunicateContext,
        *,
        residual_input_mode,
    ):
        if get_attn_tp_context().input_scattered:
            return CommunicateWithAllReduceAndLayerNormFn._tp_all_reduce_with_scattered_residual(
                hidden_states,
                residual,
                layernorm,
                context,
            )

        if residual_input_mode == ScatterMode.SCATTERED and context.attn_tp_size > 1:
            residual, local_residual = (
                get_local_dp_buffer(get_parallel().attn_tp_group),
                residual,
            )
            attn_tp_all_gather_into_tensor(residual, local_residual)
        if context.attn_dp_size != 1:
            use_layer_norm_before_gather = (
                context.force_layernorm_before_dp_gather or context.attn_tp_size == 1
            )
            if use_layer_norm_before_gather and hidden_states.shape[0] != 0:
                if context.attn_tp_size > 1:
                    hidden_states = attention_tensor_model_parallel_all_reduce(
                        hidden_states
                    )
                with use_symmetric_memory(
                    get_tp_group(),
                    disabled=not is_allocation_symmetric(),
                ):
                    hidden_states, residual = layernorm(hidden_states, residual)
            elif context.attn_tp_rank == 0:
                hidden_states += residual

            hidden_states, local_hidden_states = (
                get_global_dp_buffer(get_tp_group()),
                hidden_states,
            )
            if use_layer_norm_before_gather:
                dp_gather_replicate(hidden_states, local_hidden_states, forward_batch)
            else:
                dp_gather_partial(hidden_states, local_hidden_states, forward_batch)

            if not use_layer_norm_before_gather:
                dp_scatter(residual, hidden_states, forward_batch)
                if hidden_states.shape[0] != 0:
                    hidden_states = layernorm(hidden_states)
        else:
            handled = False
            if (
                apply_aiter_all_reduce_fusion(hidden_states)
                or apply_flashinfer_allreduce_fusion(hidden_states.shape[0])
            ) and hasattr(layernorm, "forward_with_allreduce_fusion"):
                hidden_states, residual = layernorm.forward_with_allreduce_fusion(
                    hidden_states, residual, use_attn_tp_group=True
                )
                handled = True

            if not handled:
                quantize_communications = (
                    not forward_batch.forward_mode.is_decode_or_idle()
                    and get_server_args().enable_quant_communications
                )
                if quantize_communications:
                    hidden_states = attention_tensor_model_parallel_quant_all_reduce(
                        hidden_states
                    )
                else:
                    hidden_states = attention_tensor_model_parallel_all_reduce(
                        hidden_states
                    )
                if _is_npu and context.cache is not None:
                    _ = prepare_weight_cache(hidden_states, context.cache)
                hidden_states, residual = layernorm(hidden_states, residual)
        return hidden_states, residual

    @staticmethod
    def _scatter_hidden_states_and_residual(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        layernorm: torch.nn.Module,
        context: CommunicateContext,
        *,
        residual_input_mode,
    ):
        input_hidden_states = hidden_states
        hidden_states = hidden_states.tensor_split(context.attn_tp_size)[
            context.attn_tp_rank
        ]
        attn_tp_reduce_scatter_tensor(hidden_states, input_hidden_states)
        if residual_input_mode == ScatterMode.TP_ATTN_FULL:
            residual = residual.tensor_split(context.attn_tp_size)[context.attn_tp_rank]
        if hidden_states.shape[0] != 0:
            hidden_states, residual = layernorm(hidden_states, residual)
        return hidden_states, residual

    @staticmethod
    def _tp_all_reduce_with_scattered_residual(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        layernorm: torch.nn.Module,
        context: CommunicateContext,
    ):
        if hidden_states.shape[0] == 0:
            return hidden_states, hidden_states

        scattered_states = hidden_states.tensor_split(context.tp_size)[context.tp_rank]
        scattered_states += residual
        residual = tensor_model_parallel_all_reduce(hidden_states)
        hidden_states = layernorm(residual)
        return hidden_states, residual

    @staticmethod
    def _gather_hidden_states_and_residual_moe(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch,
        layernorm: torch.nn.Module,
        context: CommunicateContext,
        *,
        residual_input_mode,
    ):
        """Allgather tokens for MoE when moe_dp_size < attn_cp_size.

        Steps:
          1. Standard attn-TP all-reduce + optional DP allgather + layernorm (same as
             _gather_hidden_states_and_residual for the dp>1 case, or simple all-reduce
             + layernorm for dp==1).
          2. moe_cp allgather: gather tokens from cp_per_moe CP ranks so each rank holds
             all tokens for its MoE group.

        Residual is left at TP_ATTN_FULL throughout.
        """
        # Early return on empty tensor is safe for MOE_CP because:
        # - During CP extend: zigzag split guarantees all CP ranks have non-zero tokens,
        #   so no rank hits this path while others proceed to the allgather.
        # - During decode: moe_cp allgather is skipped (guarded by is_context_parallel_extend).
        # - CUDA graph warmup: not applicable when --disable-piecewise-cuda-graph is used.
        if hidden_states.shape[0] == 0:
            return hidden_states, residual

        # Step 1: Standard all-reduce/DP-allgather + layernorm (reuse existing logic).
        hidden_states, residual = (
            CommunicateWithAllReduceAndLayerNormFn._gather_hidden_states_and_residual(
                hidden_states=hidden_states,
                residual=residual,
                forward_batch=forward_batch,
                layernorm=layernorm,
                context=context,
                residual_input_mode=residual_input_mode,
            )
        )

        # Step 2: moe_cp allgather — gather across cp_per_moe CP ranks.
        # Only active during prefill (context-parallel extend); decode keeps existing path.
        moe_cp_size = get_moe_cp_size()
        if (
            moe_cp_size > 1
            and hidden_states.shape[0] > 0
            and forward_batch.forward_mode.is_context_parallel_extend()
            and forward_batch.attn_cp_metadata is not None
        ):
            # Zigzag split can produce unequal token counts across CP ranks
            # (when seq_len % (cp_size * 2) != 0). NCCL allgather requires
            # equal input sizes, so pad to the max per-rank token count.
            per_rank_tokens = forward_batch.attn_cp_metadata.per_rank_actual_token
            max_tokens = max(per_rank_tokens)
            pad_size = max_tokens - hidden_states.shape[0]
            if pad_size > 0:
                hidden_states = torch.nn.functional.pad(
                    hidden_states, [0, 0, 0, pad_size]
                )

            output = torch.empty(
                (max_tokens * moe_cp_size, hidden_states.shape[1]),
                dtype=hidden_states.dtype,
                device=hidden_states.device,
            )
            moe_cp_all_gather_into_tensor(output, hidden_states)
            hidden_states = output

        return hidden_states, residual


class CommunicateSummableTensorPairFn:
    """It is allowed to make (hidden_states, residual) := (hidden_states + residual, None) if needed."""

    @classmethod
    def execute(
        cls,
        hidden_states_input_mode,
        residual_input_mode,
        output_mode,
        context,
        **kwargs,
    ):
        return cls.get_fn(
            hidden_states_input_mode=hidden_states_input_mode,
            residual_input_mode=residual_input_mode,
            output_mode=output_mode,
            context=context,
        )(context=context, **kwargs)

    @staticmethod
    def get_fn(
        hidden_states_input_mode: ScatterMode,
        residual_input_mode: ScatterMode,
        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 CommunicateSummableTensorPairFn._trivial

        if (
            (hidden_states_input_mode == ScatterMode.FULL)
            and (residual_input_mode == ScatterMode.TP_ATTN_FULL)
            and (output_mode == ScatterMode.TP_ATTN_FULL)
        ):
            return CommunicateSummableTensorPairFn._scatter_hidden_states

        if (
            (hidden_states_input_mode == ScatterMode.SCATTERED)
            and (residual_input_mode == ScatterMode.SCATTERED)
            and (output_mode == ScatterMode.TP_ATTN_FULL)
        ):
            return CommunicateSummableTensorPairFn._gather

        if (
            (hidden_states_input_mode == ScatterMode.TP_ATTN_FULL)
            and (residual_input_mode == ScatterMode.TP_ATTN_FULL)
            and (output_mode == ScatterMode.SCATTERED)
        ):
            return CommunicateSummableTensorPairFn._scatter

        if (
            (hidden_states_input_mode == ScatterMode.MOE_FULL)
            and (residual_input_mode == ScatterMode.TP_ATTN_FULL)
            and (output_mode == ScatterMode.TP_ATTN_FULL)
        ):
            return CommunicateSummableTensorPairFn._scatter_hidden_states_moe

        raise NotImplementedError(
            f"{hidden_states_input_mode=} {residual_input_mode=} {output_mode=}"
        )

    @staticmethod
    def _trivial(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        context: CommunicateContext,
        **kwargs,
    ):
        return hidden_states, residual

    @staticmethod
    def _scatter_hidden_states(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        context: CommunicateContext,
        allow_reduce_scatter: bool = False,
    ):
        if get_parallel().tp_size == get_parallel().attn_dp_size:
            group = get_tp_group()
        else:
            group = get_parallel().attn_tp_group
        hidden_states, global_hidden_states = (
            get_local_dp_buffer(group),
            hidden_states,
        )
        if should_use_dp_reduce_scatterv():
            get_tp_group().reduce_scatterv(
                global_hidden_states,
                output=hidden_states,
                sizes=get_dp_global_num_tokens(),
            )
        elif allow_reduce_scatter and forward_batch.dp_padding_mode.is_max_len():
            dp_reduce_scatter_tensor(hidden_states, global_hidden_states)
        else:
            dp_scatter(hidden_states, global_hidden_states, forward_batch)
        return hidden_states, residual

    @staticmethod
    def _gather(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        context: CommunicateContext,
        **kwargs,
    ):
        hidden_states += residual
        residual = None
        hidden_states, local_hidden_states = (
            get_local_dp_buffer(get_parallel().attn_tp_group),
            hidden_states,
        )
        attn_tp_all_gather_into_tensor(
            hidden_states,
            local_hidden_states,
        )
        return hidden_states, residual

    @staticmethod
    def _scatter(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        context: CommunicateContext,
    ):
        assert residual is None, "not yet handled residual!=None"
        tensor_list = list(hidden_states.tensor_split(context.attn_tp_size))
        hidden_states = tensor_list[context.attn_tp_rank]
        return hidden_states, residual

    @staticmethod
    def _scatter_hidden_states_moe(
        hidden_states: torch.Tensor,
        residual: torch.Tensor,
        forward_batch: ForwardBatch,
        context: CommunicateContext,
        **kwargs,
    ):
        """Scatter MoE output back to TP_ATTN_FULL after MOE_FULL computation.

        After moe_tensor_model_parallel_all_reduce (which runs unconditionally since
        mlp_reduce_scatter=False for this path), all ranks in the moe_cp group hold the
        full MoE result for all cp_per_moe token chunks. We simply slice out this rank's
        CP-local portion.

        If DP>1, further scatter back to the local DP slice.
        """
        # Only scatter back during prefill; decode was never allgathered so no-op.
        # Safe w.r.t. empty tensors: same reasoning as _gather_hidden_states_and_residual_moe
        # — CP extend always has non-zero tokens per rank, and decode skips this path.
        moe_cp_size = get_moe_cp_size()
        if (
            moe_cp_size > 1
            and forward_batch.forward_mode.is_context_parallel_extend()
            and forward_batch.attn_cp_metadata is not None
        ):
            moe_cp_rank = get_moe_cp_rank()
            # The allgather was padded to max_tokens_per_rank (equal chunks).
            # Extract this rank's actual (non-padded) tokens from its chunk.
            per_rank_tokens = forward_batch.attn_cp_metadata.per_rank_actual_token
            max_tokens_per_rank = max(per_rank_tokens)
            actual_local_tokens = per_rank_tokens[moe_cp_rank]
            hidden_states = hidden_states.narrow(
                0, moe_cp_rank * max_tokens_per_rank, actual_local_tokens
            ).contiguous()

        # DP scatter (if DP attention is enabled)
        if context.attn_dp_size > 1:
            if get_parallel().tp_size == get_parallel().attn_dp_size:
                group = get_tp_group()
            else:
                group = get_parallel().attn_tp_group
            hidden_states_output, global_hidden_states = (
                get_local_dp_buffer(group),
                hidden_states,
            )
            dp_scatter(hidden_states_output, global_hidden_states, forward_batch)
            hidden_states = hidden_states_output

        return hidden_states, residual