diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index fa3fccf1c..e08d90480 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -1013,10 +1013,14 @@ class GroupCoordinator: self, input_: Union[torch.Tensor, List[torch.Tensor]], sizes: Optional[List[int]] = None, + output: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, List[torch.Tensor]]: """ Supports varying sizes per rank and input tensor list. `sizes`: a list of len(world_size) with the number of items per rank to gather. + `output`: optional pre-allocated destination buffer (single-tensor input only). + When given, NCCL writes the gathered result directly into it, avoiding an + extra output allocation + caller-side copy. """ world_size = self.world_size pynccl_comm = self.pynccl_comm @@ -1027,7 +1031,9 @@ class GroupCoordinator: ), "pynccl is required for all_gatherv" def _all_gather_allocate_output( - input_: torch.Tensor, sizes: Optional[List[int]] = None + input_: torch.Tensor, + sizes: Optional[List[int]] = None, + output: Optional[torch.Tensor] = None, ): input_size = input_.size() if sizes is not None: @@ -1039,6 +1045,12 @@ class GroupCoordinator: sizes = None else: output_size = (input_size[0] * world_size,) + input_size[1:] + if output is not None: + assert tuple(output.shape) == tuple(output_size), ( + f"all_gatherv output buffer shape {tuple(output.shape)} " + f"!= expected {tuple(output_size)}" + ) + return output, sizes # Allocate output tensor. with self.use_symmetric_memory(self, disabled=sizes is not None): output_tensor = torch.empty( @@ -1046,13 +1058,18 @@ class GroupCoordinator: ) return output_tensor, sizes - if isinstance(input_, torch.Tensor): + single_input = isinstance(input_, torch.Tensor) + if single_input: input_ = [input_] + elif output is not None: + raise ValueError("all_gatherv `output` requires a single-tensor input") output_list = [] size_list = [] for inp in input_: - output_tensor, s = _all_gather_allocate_output(inp, sizes=sizes) + output_tensor, s = _all_gather_allocate_output( + inp, sizes=sizes, output=output + ) output_list.append(output_tensor) size_list.append(s) diff --git a/python/sglang/srt/layers/deepseek_v4_rope.py b/python/sglang/srt/layers/deepseek_v4_rope.py index 6e1c755f7..c3c15ab8f 100644 --- a/python/sglang/srt/layers/deepseek_v4_rope.py +++ b/python/sglang/srt/layers/deepseek_v4_rope.py @@ -161,7 +161,7 @@ def apply_rotary_emb_triton_kernel_batched( BLOCK_M: tl.constexpr, BLOCK_P: tl.constexpr, ): - # Batched variant: BLOCK_M tokens per program (mirrors ATOM's inverse_rope_gptj + # Batched variant: BLOCK_M tokens per program # which batches 32 tokens/program) to cut the per-token launch granularity of # the original (one program per token). pid_m = tl.program_id(0) @@ -210,12 +210,12 @@ def apply_rotary_emb_triton_kernel_batched( @triton.jit -def apply_rotary_emb_contig_kernel( +def apply_rotary_emb_flat_kernel( x_ptr, fr_ptr, pos_ptr, - rope_dim, - n_tokens, + n_rows, + n_heads, sx_tok, sx_head, sx_d, @@ -223,53 +223,54 @@ def apply_rotary_emb_contig_kernel( sfr_d, USE_POS: tl.constexpr, IS_INVERSE: tl.constexpr, - BLOCK_M: tl.constexpr, RD: tl.constexpr, RDH: tl.constexpr, + BLOCK_ROWS: tl.constexpr, ): - # CONTIGUOUS-load GPT-J rope (mirrors ATOM's inverse_rope_gptj): load the rope - # slice as a contiguous [BLOCK_M, RD] tile (coalesced, vs the strided 2i/2i+1 - # interleaved loads), and do the pair rotation via reshape+flip. RD tokens of - # one head per program, BLOCK_M tokens batched. - pid_m = tl.program_id(0) - pid_h = tl.program_id(1) - tok = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - tok_mask = tok < n_tokens + # FLAT-row GPT-J rope: iterate over (token, head) pairs flattened as + # row = token * n_heads + head, BLOCK_ROWS *consecutive* rows per program. + # Consecutive rows are sx_head apart in memory (vs sx_tok == n_heads*sx_head + # for the per-head contig kernel), so the read/write is far less scattered -> + # ~2x higher achieved HBM bandwidth (cold) on the 128-head attention output + # (production rope ~168us -> ~59us). + pid = tl.program_id(0) + row = pid * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) + rmask = row < n_rows + tok = row // n_heads + head = row % n_heads d = tl.arange(0, RD) - dmask = d < rope_dim - m = tok_mask[:, None] & dmask[None, :] - - xo = tok[:, None] * sx_tok + pid_h * sx_head + d[None, :] * sx_d - x = tl.load(x_ptr + xo, mask=m, other=0.0).to(tl.float32) + base = tok[:, None] * sx_tok + head[:, None] * sx_head + xo = base + d[None, :] * sx_d + x = tl.load(x_ptr + xo, mask=rmask[:, None], other=0.0).to(tl.float32) if USE_POS: - pos = tl.load(pos_ptr + tok, mask=tok_mask, other=0) + pos = tl.load(pos_ptr + tok, mask=rmask, other=0) else: pos = tok - # element d uses cos/sin of pair (d//2): freqs_real interleaved [cos0,sin0,...] cos_idx = (d // 2) * 2 cos = tl.load( - fr_ptr + pos[:, None] * sfr_pos + cos_idx[None, :] * sfr_d, mask=m, other=0.0 + fr_ptr + pos[:, None] * sfr_pos + cos_idx[None, :] * sfr_d, + mask=rmask[:, None], + other=0.0, ) sin = tl.load( fr_ptr + pos[:, None] * sfr_pos + (cos_idx[None, :] + 1) * sfr_d, - mask=m, + mask=rmask[:, None], other=0.0, ) x_sin = x * sin even = (d % 2 == 0)[None, :] - # inverse: negate evens; forward: negate odds (then flip pairs) if IS_INVERSE: x_neg = tl.where(even, -x_sin, x_sin) else: x_neg = tl.where(even, x_sin, -x_sin) - x_neg = tl.reshape(x_neg, (BLOCK_M, RDH, 2)) + x_neg = tl.reshape(x_neg, (BLOCK_ROWS, RDH, 2)) x_neg = tl.flip(x_neg, 2) - x_rot = tl.reshape(x_neg, (BLOCK_M, RD)) + x_rot = tl.reshape(x_neg, (BLOCK_ROWS, RD)) out = x * cos + x_rot - tl.store(x_ptr + xo, out.to(x_ptr.dtype.element_ty), mask=m) + tl.store(x_ptr + xo, out.to(x_ptr.dtype.element_ty), mask=rmask[:, None]) # Use the batched / contiguous-load rope kernels (faster, coalesced) instead of the @@ -303,16 +304,25 @@ def apply_rotary_emb_triton( else: assert freqs_real.shape[0] == batch_size BLOCK_M = 32 - # 3D (attention-output / q-k rope): contiguous-load kernel (ATOM-style). + # 3D (attention-output / q-k rope): contiguous-load kernel. if is_3d: RD = max(triton.next_power_of_2(rope_dim), 2) - grid = (triton.cdiv(batch_size, BLOCK_M), n_heads) - apply_rotary_emb_contig_kernel[grid]( + # FLAT-row kernel: process (token, head) pairs flattened as + # row = token*n_heads + head, BLOCK_ROWS consecutive rows per program. + # The per-head contig kernel reads BLOCK_M tokens strided by + # n_heads*head_dim (very scattered on the 128-head attention output) and + # only reaches ~2.2 TB/s cold; the flat kernel's rows are head_dim apart + # -> ~4.5 TB/s cold (~2x). Microbench (MI300, 8192x128x64, + # cold): BLOCK_ROWS=16 + num_warps=1. Numerically bit-exact vs contig. + FLAT_BLOCK_ROWS = 16 + n_rows = batch_size * n_heads + grid = (triton.cdiv(n_rows, FLAT_BLOCK_ROWS),) + apply_rotary_emb_flat_kernel[grid]( x, freqs_real, positions, - rope_dim, - batch_size, + n_rows, + n_heads, x.stride(0), x.stride(1), x.stride(2), @@ -320,9 +330,10 @@ def apply_rotary_emb_triton( freqs_real.stride(1), USE_POS=(positions is not None), IS_INVERSE=inverse, - BLOCK_M=BLOCK_M, RD=RD, RDH=RD // 2, + BLOCK_ROWS=FLAT_BLOCK_ROWS, + num_warps=1, ) return x BLOCK_P = max(triton.next_power_of_2(rope_dim // 2), 1) diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index 2776faf2f..fc0d2d8b4 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -580,13 +580,11 @@ def _dp_gather_via_all_gatherv( else: local_real = local_tokens.new_zeros((local_rows, *local_tokens.shape[1:])) local_real[: local_tokens.shape[0]].copy_(local_tokens) - gathered = get_tp_group().all_gatherv(local_real, sizes=sizes) - if isinstance(gathered, list): - # all_gatherv may return a list of per-rank tensors; concatenate them - # along the token dim (taking [0] would drop all but rank 0's tokens). - gathered = torch.cat(gathered, dim=0) - # gathered rows == sum(sizes); must equal the buffer length. - global_tokens[: gathered.shape[0]].copy_(gathered) + # sum(sizes) == global_tokens.shape[0] is guaranteed by the caller (else it + # falls back to all_reduce). Pass global_tokens as the NCCL output buffer so + # the gather writes directly into it -- avoids the previous extra full-buffer + # torch.cat + copy_ (two ~sum(sizes)*hidden DtoD copies, ~700us/layer at c512). + get_tp_group().all_gatherv(local_real, sizes=sizes, output=global_tokens) def _dp_gather( diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 17c8ca2d8..69814b2bf 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -830,6 +830,7 @@ class DeepseekV2MoE(nn.Module): gemm_output_zero_allocator: BumpAllocator = None, input_ids: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None, + skip_shared_experts: bool = False, ) -> torch.Tensor: from sglang.srt.layers.moe.mega_moe import forward_mega_moe, should_use_mega_moe @@ -870,6 +871,7 @@ class DeepseekV2MoE(nn.Module): gemm_output_zero_allocator, input_ids, input_ids_global=input_ids_global, + skip_shared_experts=skip_shared_experts, ) else: return self.forward_deepep( @@ -969,6 +971,7 @@ class DeepseekV2MoE(nn.Module): gemm_output_zero_allocator: BumpAllocator = None, input_ids: Optional[torch.Tensor] = None, input_ids_global: Optional[torch.Tensor] = None, + skip_shared_experts: bool = False, ) -> torch.Tensor: if hasattr(self, "shared_experts") and use_intel_amx_backend( self.shared_experts.gate_up_proj @@ -981,8 +984,16 @@ class DeepseekV2MoE(nn.Module): else None ) defer_shared = not self.experts.moe_runner_config.inplace + # PoC (SGLANG_DP_SHARED_EXPERT_LOCAL): shared expert is computed on the LOCAL + # hidden in the decoder layer (before the dp gather) and added after the + # reduce_scatterv. When set, never compute/add it here (on the global buffer). + shared_output = None if hidden_states.shape[0] > 0: - if not defer_shared and not self._fuse_shared_experts_inside_sbo: + if ( + not defer_shared + and not self._fuse_shared_experts_inside_sbo + and not skip_shared_experts + ): shared_output = self._forward_shared_experts( hidden_states, gemm_output_zero_allocator ) @@ -1005,7 +1016,7 @@ class DeepseekV2MoE(nn.Module): hidden_states.device, layer_id=self.layer_id ) - if self._fuse_shared_experts_inside_sbo: + if self._fuse_shared_experts_inside_sbo and not skip_shared_experts: shared_output = None def _pre_combine_hook( @@ -1053,6 +1064,7 @@ class DeepseekV2MoE(nn.Module): defer_shared and hidden_states.shape[0] > 0 and not self._fuse_shared_experts_inside_sbo + and not skip_shared_experts ): shared_output = self._forward_shared_experts( hidden_states, gemm_output_zero_allocator diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index a5520d91c..e73895f8e 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -157,6 +157,10 @@ def _is_fused_mhc_post_pre_enabled() -> bool: _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +# PoC: compute the (replicated TP1) shared expert on LOCAL hidden before the dp +# gather instead of on the gathered global buffer. Requires +# SGLANG_SHARED_EXPERT_TP1=1 (replicated shared expert). Default OFF. +_SHARED_EXPERT_LOCAL = get_bool_env_var("SGLANG_DP_SHARED_EXPERT_LOCAL") _is_gfx95_supported = is_gfx95_supported() if _use_aiter: @@ -1580,6 +1584,22 @@ class DeepseekV4DecoderLayer(nn.Module): and forward_batch.dp_padding_mode is not None and not forward_batch.dp_padding_mode.is_max_len() ) + # PoC (SGLANG_DP_SHARED_EXPERT_LOCAL): compute the replicated shared expert + # on LOCAL hidden before the gather and add it back after the combine + # (reduce_scatterv OR dp_scatter), instead of on the gathered global buffer. + # Applies to BOTH prefill and decode: the shared expert is a per-token MLP, + # so computing it on this rank's local tokens (M_local rows) is identical to + # computing it on the gathered global buffer (M_global rows) and keeping the + # local slice -- but costs 1/dp_size the rows. With a replicated (TP1) shared + # expert this cancels the TP1 "full-dim" cost in decode (M_local * dim == + # M_global * dim/tp), so decode no longer pays the ~dp_size x penalty. + _shared_local = None + _do_shared_local = ( + _SHARED_EXPERT_LOCAL + and _use_tp_moe_gather + and getattr(self.mlp, "shared_experts", None) is not None + and getattr(self.mlp, "_shared_expert_tp1", False) + ) if _use_cp: if get_moe_a2a_backend().is_none(): hidden_states = dsa_cp_gather_hidden_states(hidden_states) @@ -1593,6 +1613,8 @@ class DeepseekV4DecoderLayer(nn.Module): get_global_dp_buffer(get_tp_group()), hidden_states, ) + if _do_shared_local and local_hidden_states.shape[0] > 0: + _shared_local = self.mlp._forward_shared_experts(local_hidden_states) dp_gather_partial(hidden_states, local_hidden_states, forward_batch) _a2a_scatter_chunks: Optional[List[torch.Tensor]] = None if _use_tp_attn_a2a_scatter: @@ -1609,6 +1631,7 @@ class DeepseekV4DecoderLayer(nn.Module): # Skip the MoE-internal post-experts all_reduce when we will do the # reduce via reduce_scatterv at the combine below (else double-reduce). use_reduce_scatter=_use_cp or _use_gatherv_pair, + skip_shared_experts=_do_shared_local, ) if _use_cp and get_moe_a2a_backend().is_none(): hidden_states = dsa_cp_reduce_scatter_hidden_states(hidden_states) @@ -1629,6 +1652,12 @@ class DeepseekV4DecoderLayer(nn.Module): ) else: dp_scatter(hidden_states, global_hidden_states, forward_batch) + # PoC: add the locally-computed shared-expert output to this rank's + # reduce-scattered / dp-scattered local slice (skipped inside self.mlp + # above). Covers both prefill (gatherv) and decode (dp_scatter). + if _shared_local is not None: + n = hidden_states.shape[0] + hidden_states = hidden_states + _shared_local[:n] if _use_tp_attn_a2a_scatter: assert _a2a_scatter_chunks is not None gathered = [torch.empty_like(t) for t in _a2a_scatter_chunks]