minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Improve performance of dsv4 in high concurrency (#28938)

Browse Source 
Co-authored-by: wunhuang <wunhuang@amd.com>
This commit is contained in:
kk
2026-06-23 02:20:48 -07:00
committed by GitHub
co-authored by wunhuang
parent a65c68c1fb
commit af9027f6c9
5 changed files with 111 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/parallel_state.py
+20 -3
View File
@@ -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)
python/sglang/srt/layers/deepseek_v4_rope.py
+43 -32
View File
@@ -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)
python/sglang/srt/layers/dp_attention.py
+5 -7
View File
@@ -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(
python/sglang/srt/models/deepseek_v2.py
+14 -2
View File
@@ -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
python/sglang/srt/models/deepseek_v4.py
+29
View File
@@ -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]