# 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