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

(perf): Shard Kimi-K2.5 Eagle3 draft fc + symm-mem AG (#29223)

Browse Source 
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
Khoa Pham
2026-06-27 17:40:32 -07:00
committed by GitHub
co-authored by Cursor
parent 073de15053
commit da802ddcaf
5 changed files with 885 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/triton_symm_mem_ag.py
+525
View File
@@ -0,0 +1,525 @@
# SPDX-License-Identifier: Apache-2.0
"""Symmetric-memory ``multimem.st`` all-gather along the hidden (last) dim.
Each rank stores its ``[T, H/TP]`` shard into a multicast buffer in one NVLink
pass instead of an NCCL ring; ``create_state`` rendezvous once so launches are
CUDA-graph capturable.
"""
import logging
from dataclasses import dataclass
from typing import Any
import torch
import torch.distributed as dist
import torch.distributed._symmetric_memory as symm_mem
import triton
import triton.language as tl
logger = logging.getLogger(__name__)
# Each thread moves _NUMEL_PER_THREAD bf16 via one 128-bit multimem op; the
# grid-strided block count is tunable in [_MIN_BLOCKS, _MAX_BLOCKS].
_BLOCK_THREADS = 1024
_NUMEL_PER_THREAD = 8
_MIN_BLOCKS = 4
_MAX_BLOCKS = 32
# ------------------------------------------------------------------------------
# Low-level PTX helpers
# ------------------------------------------------------------------------------
@triton.jit
def _multimem_st_128(multicast_ptrs, x, y, z, w, mask):
return tl.inline_asm_elementwise(
"""
{
.reg .pred %p0;
setp.eq.s32 %p0, $6, 1;
@!%p0 bra end;
multimem.st.relaxed.sys.global.v4.f32 [$1], {$2, $3, $4, $5};
end:
}
""",
"=r,l,r,r,r,r,r",
args=[multicast_ptrs, x, y, z, w, mask.to(tl.int32)],
dtype=(tl.uint32),
is_pure=False,
pack=1,
)
@triton.jit
def _local_ld_128(in_ptr, mask):
return tl.inline_asm_elementwise(
"""
{
.reg .pred %p0;
setp.eq.s32 %p0, $5, 1;
@!%p0 bra end;
ld.relaxed.sys.global.v4.b32 {$0, $1, $2, $3}, [$4];
end:
}
""",
"=r,=r,=r,=r,l,r",
args=[in_ptr, mask.to(tl.int32)],
dtype=(tl.uint32, tl.uint32, tl.uint32, tl.uint32),
is_pure=True,
pack=1,
)
@triton.jit
def _get_tid():
return tl.inline_asm_elementwise(
"""
mov.u32 $0, %tid.x;
mov.u32 $1, %tid.y;
mov.u32 $2, %tid.z;
""",
"=r,=r,=r",
[],
dtype=(tl.uint32, tl.uint32, tl.uint32),
is_pure=True,
pack=1,
)
@triton.jit
def _get_ntid():
return tl.inline_asm_elementwise(
"""
mov.u32 $0, %ntid.x;
mov.u32 $1, %ntid.y;
mov.u32 $2, %ntid.z;
""",
"=r,=r,=r",
[],
dtype=(tl.uint32, tl.uint32, tl.uint32),
is_pure=True,
pack=1,
)
@triton.jit
def _get_flat_tid():
tid_x, tid_y, tid_z = _get_tid()
ntid_x, ntid_y, _ = _get_ntid()
return tid_z * ntid_y * ntid_x + tid_y * ntid_x + tid_x
@triton.jit
def _sync_threads():
tl.inline_asm_elementwise(
"bar.sync 0;", "=r", [], dtype=tl.int32, is_pure=False, pack=1
)
@triton.jit
def _send_signal(addrs):
tl.inline_asm_elementwise(
"""
{
.reg .u32 %tmp32_<1>;
.reg .pred %p<1>;
send_signal:
atom.global.relaxed.sys.cas.b32 %tmp32_0, [$1], 0, 1;
setp.eq.u32 %p0, %tmp32_0, 0;
@!%p0 bra send_signal;
}
""",
"=r, l",
[addrs],
dtype=tl.int32,
is_pure=False,
pack=1,
)
@triton.jit
def _send_signal_release(addrs):
tl.inline_asm_elementwise(
"""
{
.reg .u32 %tmp32_<1>;
.reg .pred %p<1>;
send_signal:
atom.global.release.sys.cas.b32 %tmp32_0, [$1], 0, 1;
setp.eq.u32 %p0, %tmp32_0, 0;
@!%p0 bra send_signal;
}
""",
"=r, l",
[addrs],
dtype=tl.int32,
is_pure=False,
pack=1,
)
@triton.jit
def _wait_signal(addrs):
tl.inline_asm_elementwise(
"""
{
.reg .u32 %tmp32_<1>;
.reg .pred %p<1>;
wait_signal:
atom.global.sys.relaxed.cas.b32 %tmp32_0, [$1], 1, 0;
setp.eq.u32 %p0, %tmp32_0, 1;
@!%p0 bra wait_signal;
}
""",
"=r, l",
[addrs],
dtype=tl.int32,
is_pure=False,
pack=1,
)
@triton.jit
def _wait_signal_acquire(addrs):
tl.inline_asm_elementwise(
"""
{
.reg .u32 %tmp32_<1>;
.reg .pred %p<1>;
wait_signal:
atom.global.sys.acquire.cas.b32 %tmp32_0, [$1], 1, 0;
setp.eq.u32 %p0, %tmp32_0, 1;
@!%p0 bra wait_signal;
}
""",
"=r, l",
[addrs],
dtype=tl.int32,
is_pure=False,
pack=1,
)
@triton.jit
def _blockwise_barrier(
signal_pad_ptrs,
rank: tl.constexpr,
world_size: tl.constexpr,
sem: tl.constexpr,
):
block_id = (
tl.program_id(2) * tl.num_programs(1) * tl.num_programs(0)
+ tl.program_id(1) * tl.num_programs(0)
+ tl.program_id(0)
)
flat_tid = _get_flat_tid()
remote_ranks = tl.arange(0, world_size)
signal_pad_ptrs = signal_pad_ptrs.to(tl.pointer_type(tl.uint64))
remote_signal_pad_addrs = tl.load(signal_pad_ptrs + remote_ranks).to(
tl.pointer_type(tl.uint32)
)
send_addrs = remote_signal_pad_addrs + block_id * world_size + rank
local_signal_pad_addr = tl.load(signal_pad_ptrs + rank).to(
tl.pointer_type(tl.uint32)
)
wait_addrs = local_signal_pad_addr + block_id * world_size + remote_ranks
if flat_tid < world_size:
if sem == "relaxed":
_send_signal(send_addrs)
_wait_signal(wait_addrs)
else:
_send_signal_release(send_addrs)
_wait_signal_acquire(wait_addrs)
@triton.jit
def _all_gather_kernel_inner(
input_ptr,
multicast_ptr,
signal_pad_ptr,
total_tokens,
hidden_offset,
LOCAL_HIDDEN: tl.constexpr,
TOTAL_HIDDEN: tl.constexpr,
BLOCK_SIZE: tl.constexpr,
NUMEL_PER_THREAD: tl.constexpr,
RANK: tl.constexpr,
WORLD_SIZE: tl.constexpr,
SKIP_ENTRY_SYNC: tl.constexpr,
) -> None:
if SKIP_ENTRY_SYNC == 0:
_blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem="relaxed")
_sync_threads()
chunks_per_row: tl.constexpr = LOCAL_HIDDEN // NUMEL_PER_THREAD
total_hidden_chunks: tl.constexpr = TOTAL_HIDDEN // NUMEL_PER_THREAD
hidden_offset_chunks = hidden_offset // NUMEL_PER_THREAD
total_chunks = total_tokens * chunks_per_row
pid = tl.program_id(axis=0)
tid = _get_flat_tid()
block_start = pid * BLOCK_SIZE
while block_start < total_chunks:
chunk = block_start + tid
mask = chunk < total_chunks
row = chunk // chunks_per_row
col_chunk = chunk % chunks_per_row
in_ptr = input_ptr.to(tl.pointer_type(tl.uint64)) + chunk * 2
out_chunk = row * total_hidden_chunks + hidden_offset_chunks + col_chunk
out_ptr = (
multicast_ptr.to(tl.int64).to(tl.pointer_type(tl.uint64)) + out_chunk * 2
)
x, y, z, w = _local_ld_128(in_ptr, mask)
_multimem_st_128(out_ptr, x, y, z, w, mask)
block_start += tl.num_programs(axis=0) * BLOCK_SIZE
_sync_threads()
_blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem="acq_rel")
# ------------------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------------------
@dataclass
class MultimemAllGatherState:
group: dist.ProcessGroup
rank_in_group: int
world_size: int
device: torch.device
max_token_num: int
hidden_dim: int
comm_buff: torch.Tensor
# Rendezvous handle; stable for the buffer's lifetime, resolved once.
symm_mem_hdl: Any
def create_state(
group: dist.ProcessGroup,
rank_in_group: int,
max_tokens: int,
hidden_size: int,
device: torch.device | None = None,
) -> MultimemAllGatherState:
"""Allocate and rendezvous the symmetric-memory buffer. Collective: call
once outside CUDA-graph capture with identical args on every rank."""
assert type(group) is dist.ProcessGroup, f"Expected ProcessGroup, got {type(group)}"
assert hidden_size % _NUMEL_PER_THREAD == 0, (
f"hidden_size={hidden_size} must be a multiple of {_NUMEL_PER_THREAD} "
f"bf16 for 16-byte multimem.st row alignment"
)
device = device or torch.device(f"cuda:{torch.cuda.current_device()}")
# Pad holds _MAX_BLOCKS * world_size uint32 slots; max() never shrinks it.
pad_bytes = _MAX_BLOCKS * group.size() * 4
symm_mem.set_signal_pad_size(max(symm_mem.get_signal_pad_size(), pad_bytes))
with torch.inference_mode(False), torch.no_grad():
comm_buff = symm_mem.empty(
(max_tokens, hidden_size), dtype=torch.bfloat16, device=device
)
hdl = symm_mem.rendezvous(comm_buff, group=group)
assert hdl.rank == rank_in_group, (
f"symm_mem handle rank {hdl.rank} != rank_in_group {rank_in_group}; the "
f"hidden-shard offset would be wrong"
)
return MultimemAllGatherState(
group=group,
rank_in_group=rank_in_group,
world_size=group.size(),
device=device,
max_token_num=max_tokens,
hidden_dim=hidden_size,
comm_buff=comm_buff,
symm_mem_hdl=hdl,
)
def _launch_config(local_numel: int):
assert local_numel % _NUMEL_PER_THREAD == 0
return _MIN_BLOCKS, _BLOCK_THREADS, _BLOCK_THREADS // 32, _NUMEL_PER_THREAD
def all_gather_inner(
state: MultimemAllGatherState,
hidden_states: torch.Tensor,
tp_hidden_dim: int,
skip_entry_sync: bool = False,
safe: bool = True,
) -> torch.Tensor:
"""Gather ``[T, H/TP]`` shards into ``[T, H]`` along the hidden dim.
``tp_hidden_dim`` is the gathered width ``H``. Returns a clone when ``safe``,
else a view into the symmetric buffer (valid until the next collective)."""
world_size = state.world_size
assert hidden_states.dtype == torch.bfloat16, "Only bfloat16 is supported"
assert hidden_states.is_contiguous(), "hidden_states must be contiguous"
assert hidden_states.data_ptr() % 16 == 0, (
f"hidden_states.data_ptr()={hex(hidden_states.data_ptr())} must be "
f"16-byte aligned for 128-bit multimem.st"
)
assert (
tp_hidden_dim % world_size == 0
), f"tp_hidden_dim={tp_hidden_dim} must be divisible by world_size={world_size}"
local_hidden = tp_hidden_dim // world_size
assert local_hidden % _NUMEL_PER_THREAD == 0, (
f"per-rank hidden shard ({local_hidden}) must be a multiple of "
f"{_NUMEL_PER_THREAD} bf16"
)
assert tp_hidden_dim <= state.hidden_dim, (
f"comm buffer too narrow: tp_hidden_dim={tp_hidden_dim} > "
f"state.hidden_dim={state.hidden_dim}"
)
total_tokens, in_hidden = hidden_states.shape
assert (
in_hidden == local_hidden
), f"input hidden ({in_hidden}) != this rank's shard ({local_hidden})"
assert (
total_tokens <= state.max_token_num
), f"total_tokens={total_tokens} exceeds max_token_num={state.max_token_num}"
hidden_offset = local_hidden * state.rank_in_group
symm_mem_hdl = state.symm_mem_hdl
num_blocks, block_size, num_warps, numel_per_thread = _launch_config(
total_tokens * local_hidden
)
grid = (num_blocks, 1, 1)
_all_gather_kernel_inner[grid](
input_ptr=hidden_states,
multicast_ptr=symm_mem_hdl.multicast_ptr,
signal_pad_ptr=symm_mem_hdl.signal_pad_ptrs_dev,
total_tokens=total_tokens,
hidden_offset=hidden_offset,
LOCAL_HIDDEN=local_hidden,
TOTAL_HIDDEN=state.hidden_dim,
BLOCK_SIZE=block_size,
NUMEL_PER_THREAD=numel_per_thread,
RANK=symm_mem_hdl.rank,
WORLD_SIZE=symm_mem_hdl.world_size,
SKIP_ENTRY_SYNC=1 if skip_entry_sync else 0,
num_warps=num_warps,
)
output = state.comm_buff[:total_tokens, :tp_hidden_dim]
return output.clone() if safe else output
# ------------------------------------------------------------------------------
# Guarded wrapper
# ------------------------------------------------------------------------------
def recommended_max_tokens(include_prefill: bool, floor: int = 0) -> int:
"""Largest batch (tokens) to keep on the fast path; bigger falls back to
NCCL. Covers the spec-decode batch plus, if ``include_prefill``, a prefill
chunk. Returns ``floor`` if server args are unavailable."""
try:
from sglang.srt.server_args import get_global_server_args
sa = get_global_server_args()
def g(name: str) -> int:
v = getattr(sa, name, 0)
return v if isinstance(v, int) and v > 0 else 0
tokens = g("max_running_requests") * max(
g("speculative_num_draft_tokens"), g("speculative_eagle_topk"), 1
)
if include_prefill:
tokens = max(tokens, g("chunked_prefill_size"), g("max_prefill_tokens"))
return max(tokens, floor)
except Exception:
return floor
class MultimemAllGatherer:
"""Guarded multimem all-gather (last dim) with NCCL fallback; the single
entry point for every caller. Owns one symmetric buffer built lazily on the
first eager call, and uses the kernel only when the input fits its
dtype/shape/alignment contract. Guards use TP-replicated quantities so all
ranks pick the same path. ``skip_entry_sync=True`` drops the entry barrier;
only safe when a cross-rank sync sits between consecutive calls."""
_UNINIT = object()
def __init__(
self,
max_tokens: int,
*,
enabled: bool = True,
skip_entry_sync: bool = False,
):
self._max_tokens = int(max_tokens)
self._skip_entry_sync = skip_entry_sync
# None => always NCCL; _UNINIT => build on first eager call.
self._state = self._UNINIT if enabled else None
def __call__(self, x: torch.Tensor) -> torch.Tensor:
state = self._state
if state is self._UNINIT:
state = self._build(x)
if state is not self._UNINIT:
self._state = state
if (
state is not None
and state is not self._UNINIT
and x.dtype == torch.bfloat16
and x.dim() == 2
and x.is_contiguous()
and 0 < x.shape[0] <= state.max_token_num
and x.data_ptr() % 16 == 0
and x.shape[-1] * state.world_size <= state.hidden_dim
):
return all_gather_inner(
state,
x,
tp_hidden_dim=x.shape[-1] * state.world_size,
skip_entry_sync=self._skip_entry_sync,
safe=False,
)
# Lazy import avoids a module-load dependency on the distributed facade.
from sglang.srt.distributed import tensor_model_parallel_all_gather
return tensor_model_parallel_all_gather(x, dim=-1)
def _build(self, x: torch.Tensor):
if x.dim() != 2 or x.dtype != torch.bfloat16:
return None
if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing():
# Can't allocate under capture; retry later.
return self._UNINIT
if x.shape[-1] % _NUMEL_PER_THREAD != 0:
return None
try:
from sglang.srt.distributed import get_tp_group
tp_group = get_tp_group()
if tp_group.world_size <= 1:
return None
state = create_state(
group=tp_group.device_group,
rank_in_group=tp_group.rank_in_group,
max_tokens=self._max_tokens,
hidden_size=x.shape[-1] * tp_group.world_size,
)
if state.symm_mem_hdl.multicast_ptr == 0:
# No multicast for this world size / arch; multimem.st would
# write nowhere. Fall back to NCCL.
logger.warning(
"multimem all-gather disabled (no multicast for world_size=%d)",
tp_group.world_size,
)
return None
return state
except Exception as e:
logger.warning("multimem all-gather disabled (%s)", e)
return None
python/sglang/srt/layers/logits_processor.py
+10 -4
View File
@@ -21,9 +21,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
import torch
from torch import nn
from sglang.srt.distributed import (
tensor_model_parallel_all_gather,
)
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
from sglang.srt.environ import envs
from sglang.srt.layers.dp_attention import (
DpPaddingMode,
@@ -296,6 +294,14 @@ class LogitsProcessor(nn.Module):
self.return_full_logits = return_full_logits
self.enable_mis = get_global_server_args().enable_mis
self._logits_gatherer = triton_symm_mem_ag.MultimemAllGatherer(
max_tokens=triton_symm_mem_ag.recommended_max_tokens(
include_prefill=False, floor=128
),
enabled=self.do_tensor_parallel_all_gather and not self.use_attn_tp_group,
skip_entry_sync=True,
)
# enable chunked logprobs processing
self.enable_logprobs_chunk = envs.SGLANG_ENABLE_LOGITS_PROCESSER_CHUNK.get()
# chunk size for logprobs processing
@@ -853,7 +859,7 @@ class LogitsProcessor(nn.Module):
if self.use_attn_tp_group:
logits = self._gather_attn_tp_logits(logits)
else:
logits = tensor_model_parallel_all_gather(logits)
logits = self._logits_gatherer(logits)
logits = self._scatter_dp_attn_logits(
logits, local_hidden_states, logits_metadata
python/sglang/srt/models/kimi_k25_eagle3.py
+15 -3
View File
@@ -25,9 +25,10 @@ from torch import nn
from transformers import PretrainedConfig
from sglang.srt.distributed import get_pp_group
from sglang.srt.distributed.device_communicators import triton_symm_mem_ag
from sglang.srt.layers.communicator import AttentionInputs, get_attn_tp_context
from sglang.srt.layers.layernorm import RMSNorm
from sglang.srt.layers.linear import ReplicatedLinear
from sglang.srt.layers.linear import ColumnParallelLinear, ReplicatedLinear
from sglang.srt.layers.logits_processor import LogitsProcessor
from sglang.srt.layers.quantization.base_config import QuantizationConfig
from sglang.srt.layers.vocab_parallel_embedding import (
@@ -207,10 +208,20 @@ class Eagle3MLAModel(nn.Module):
getattr(config, "target_hidden_size", None) or config.hidden_size
)
self.num_aux_hidden_states = _get_eagle_aux_layer_count(config)
self.fc = nn.Linear(
self.fc = ColumnParallelLinear(
target_hidden_size * self.num_aux_hidden_states,
config.hidden_size,
bias=getattr(config, "bias", False),
gather_output=False,
quant_config=quant_config,
prefix=add_prefix("fc", prefix),
)
# Guarded multimem all-gather for the fc output; buffer covers prefill.
self._fc_gatherer = triton_symm_mem_ag.MultimemAllGatherer(
max_tokens=triton_symm_mem_ag.recommended_max_tokens(
include_prefill=True, floor=512
),
skip_entry_sync=False,
)
# Per-aux RMSNorm before fc; enabled via `fc_norm` or legacy
@@ -278,7 +289,8 @@ class Eagle3MLAModel(nn.Module):
[norm(chunk) for norm, chunk in zip(self.fc_norm, chunks)],
dim=-1,
)
hidden_states = self.fc(hidden_states)
hidden_states, _ = self.fc(hidden_states)
hidden_states = self._fc_gatherer(hidden_states)
if hidden_states.shape[0] == 0:
return hidden_states, [hidden_states]
test/registered/jit/benchmark/bench_symm_mem_all_gather.py
+169
View File
@@ -0,0 +1,169 @@
"""Benchmark the symmetric-memory multimem all-gather vs NCCL.
Providers:
- ``nccl`` : ``all_gather_into_tensor`` + concat-along-hidden reshape
(what ``tensor_model_parallel_all_gather(dim=-1)`` does)
- ``mm_safe`` : multimem kernel, ``safe=True`` (clones the buffer view)
- ``mm`` : multimem kernel, ``safe=False`` (fc gather config)
- ``mm_skipsync`` : multimem kernel, ``safe=False, skip_entry_sync=True``
(logits gather config)
Usage::
# Benchmark on the default world sizes (2, 4, 8 GPUs):
python test/registered/jit/benchmark/bench_symm_mem_all_gather.py
# Pick a specific world size (or comma-separated list):
python test/registered/jit/benchmark/bench_symm_mem_all_gather.py --num-gpu 8
"""
from __future__ import annotations
import atexit
import logging
import os
import torch
import torch.distributed as dist
import sglang.srt.distributed.parallel_state as ps
from sglang.jit_kernel.benchmark import marker
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, multigpu_bench_main
from sglang.jit_kernel.utils import cache_once
from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
all_gather_inner,
create_state,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(
est_time=120,
suite="base-b-kernel-benchmark-1-gpu-large",
disabled="requires multi-GPU, self-skips in CI",
)
# ---------------------------------------------------------------------------
# Sweep parameters
# ---------------------------------------------------------------------------
DTYPE = torch.bfloat16
PROVIDERS = ["nccl", "mm_safe", "mm", "mm_skipsync"]
# Full gathered hidden width H (per-rank shard is H / world_size).
HIDDENS = [7168, 16384, 163840]
NUM_TOKENS = [1, 8, 16, 32, 64, 128]
WORLD_SIZES = list(range(2, 9))
HIDDENS = get_benchmark_range(HIDDENS, [7168, 163840])
NUM_TOKENS = get_benchmark_range(NUM_TOKENS, [16, 64])
WORLD_SIZES = get_benchmark_range(WORLD_SIZES, [2, 4, 8])
MAX_HIDDEN = max(HIDDENS)
MAX_TOKENS = max(NUM_TOKENS)
# ---------------------------------------------------------------------------
# Per-rank distributed init (run once per torchrun worker)
# ---------------------------------------------------------------------------
@cache_once
def _init_cpu_group() -> dist.ProcessGroup:
local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="gloo")
ps._WORLD = ps.init_world_group(
ranks=list(range(world_size)),
local_rank=local_rank,
backend="nccl",
)
atexit.register(dist.destroy_process_group)
logging.disable(logging.INFO)
torch.cuda.set_stream(torch.cuda.Stream())
return ps._WORLD.cpu_group
@cache_once
def _init_nccl_group() -> dist.ProcessGroup:
_init_cpu_group()
coord = ps._WORLD
assert coord is not None and coord.device_group is not None
return coord.device_group
@cache_once
def _init_state():
_init_cpu_group()
coord = ps._WORLD
return create_state(
group=coord.device_group,
rank_in_group=coord.rank_in_group,
max_tokens=MAX_TOKENS,
hidden_size=MAX_HIDDEN,
)
# ---------------------------------------------------------------------------
# Benchmark
# ---------------------------------------------------------------------------
@marker.parametrize("hidden", HIDDENS)
@marker.parametrize("num_tokens", NUM_TOKENS)
@marker.benchmark("provider", PROVIDERS)
def benchmark(num_tokens: int, hidden: int, provider: str):
gpu_group = _init_nccl_group()
state = _init_state()
world_size = state.world_size
local_hidden = hidden // world_size
if hidden % world_size != 0 or local_hidden % 8 != 0:
marker.skip(f"hidden={hidden} incompatible with world_size={world_size}")
if provider != "nccl" and state.symm_mem_hdl.multicast_ptr == 0:
marker.skip(f"multimem multicast unavailable for world_size={world_size}")
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
x = torch.randn(num_tokens, local_hidden, dtype=DTYPE, device=device)
if provider == "nccl":
out_buf = torch.empty(
world_size * num_tokens, local_hidden, dtype=DTYPE, device=device
)
def fn(inp: torch.Tensor) -> torch.Tensor:
dist.all_gather_into_tensor(out_buf, inp, group=gpu_group)
return (
out_buf.reshape(world_size, num_tokens, local_hidden)
.movedim(0, 1)
.reshape(num_tokens, hidden)
)
else:
safe = provider == "mm_safe"
skip_entry_sync = provider == "mm_skipsync"
def fn(inp: torch.Tensor) -> torch.Tensor:
return all_gather_inner(
state,
inp,
tp_hidden_dim=hidden,
skip_entry_sync=skip_entry_sync,
safe=safe,
)
return marker.do_bench(
fn,
input_args=(x,),
graph_clone_args=(0,),
sync_multigpu_fn=lambda: dist.barrier(gpu_group),
# Footprint = the gathered output every rank ends up with.
memory_args=None,
memory_output=None,
extra_memory_footprint=num_tokens * hidden * x.element_size(),
)
if __name__ == "__main__":
multigpu_bench_main(
name=__name__,
file=__file__,
num_gpus=WORLD_SIZES,
main_fn=benchmark.run,
)
test/registered/jit/test_symm_mem_all_gather.py
+166
View File
@@ -0,0 +1,166 @@
"""Correctness test for the symmetric-memory multimem all-gather kernel.
Compares ``all_gather_inner`` (concat-along-hidden multimem.st gather) against
NCCL all-gather for a sweep of token counts, hidden widths, and the
``safe`` / ``skip_entry_sync`` knobs, in both eager and CUDA-graph modes.
Usage::
# Run on the default world sizes (2, 4, 8 GPUs):
python test/registered/jit/test_symm_mem_all_gather.py
# Pick a specific world size (or comma-separated list):
python test/registered/jit/test_symm_mem_all_gather.py --num-gpu 4
python test/registered/jit/test_symm_mem_all_gather.py --num-gpu 2,4,8
# Extra pytest args (forwarded to each torchrun worker):
python test/registered/jit/test_symm_mem_all_gather.py -k 16384
"""
from __future__ import annotations
import atexit
import logging
import os
import pytest
import torch
import torch.distributed as dist
import sglang.srt.distributed.parallel_state as ps
from sglang.jit_kernel.tests.utils import multigpu_pytest_main
from sglang.jit_kernel.utils import cache_once, get_ci_test_range
from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import (
all_gather_inner,
create_state,
)
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=240, suite="base-b-kernel-unit-8-gpu-h200")
register_cuda_ci(est_time=240, suite="nightly-kernel-8-gpu-h200", nightly=True)
# ---------------------------------------------------------------------------
# Test parameters
# ---------------------------------------------------------------------------
# Full gathered hidden width H (per-rank shard is H / world_size). Each value
# is a multiple of 8 * 8 so it stays valid for world sizes 2 / 4 / 8.
TEST_HIDDEN = [2048, 7168, 16384]
TEST_NUM_TOKENS = [1, 8, 16, 128]
TEST_LOOP = 8
TEST_HIDDEN = get_ci_test_range(TEST_HIDDEN, [7168])
TEST_NUM_TOKENS = get_ci_test_range(TEST_NUM_TOKENS, [16])
MAX_HIDDEN = max(TEST_HIDDEN)
MAX_TOKENS = max(TEST_NUM_TOKENS)
# ---------------------------------------------------------------------------
# Per-rank distributed setup (run once per torchrun worker)
# ---------------------------------------------------------------------------
@cache_once
def _init_cpu_group_once() -> dist.ProcessGroup:
local_rank = int(os.environ["LOCAL_RANK"])
world_size = int(os.environ["WORLD_SIZE"])
torch.cuda.set_device(local_rank)
dist.init_process_group(backend="gloo")
ps._WORLD = ps.init_world_group(
ranks=list(range(world_size)),
local_rank=local_rank,
backend="nccl",
)
atexit.register(dist.destroy_process_group)
logging.disable(logging.INFO)
torch.cuda.set_stream(torch.cuda.Stream())
cpu_group = ps._WORLD.cpu_group
assert isinstance(cpu_group, dist.ProcessGroup)
return cpu_group
@cache_once
def _init_nccl_group_once() -> dist.ProcessGroup:
_init_cpu_group_once()
coord = ps._WORLD
assert coord is not None and coord.device_group is not None
return coord.device_group
@cache_once
def _init_state_once():
_init_cpu_group_once()
coord = ps._WORLD
return create_state(
group=coord.device_group,
rank_in_group=coord.rank_in_group,
max_tokens=MAX_TOKENS,
hidden_size=MAX_HIDDEN,
)
def _nccl_all_gather(x: torch.Tensor, group: dist.ProcessGroup, world_size: int):
"""Reference gather matching ``tensor_model_parallel_all_gather(dim=-1)``:
concat per-rank ``[T, H/W]`` shards in rank order into ``[T, H]``."""
num_tokens, local_hidden = x.shape
gathered = torch.empty(
world_size * num_tokens, local_hidden, dtype=x.dtype, device=x.device
)
dist.all_gather_into_tensor(gathered, x.contiguous(), group=group)
return (
gathered.reshape(world_size, num_tokens, local_hidden)
.movedim(0, 1)
.reshape(num_tokens, world_size * local_hidden)
)
# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("skip_entry_sync", [False, True])
@pytest.mark.parametrize("safe", [False, True])
@pytest.mark.parametrize("hidden", TEST_HIDDEN)
@pytest.mark.parametrize("num_tokens", TEST_NUM_TOKENS)
@torch.inference_mode()
def test_symm_mem_all_gather(
num_tokens: int,
hidden: int,
safe: bool,
skip_entry_sync: bool,
) -> None:
nccl_group = _init_nccl_group_once()
state = _init_state_once()
world_size = state.world_size
device = torch.device(f"cuda:{int(os.environ['LOCAL_RANK'])}")
if state.symm_mem_hdl.multicast_ptr == 0:
pytest.skip(f"multimem multicast unavailable for world_size={world_size}")
local_hidden = hidden // world_size
if hidden % world_size != 0 or local_hidden % 8 != 0:
pytest.skip(f"hidden={hidden} incompatible with world_size={world_size}")
def gather(x: torch.Tensor) -> torch.Tensor:
return all_gather_inner(
state,
x,
tp_hidden_dim=hidden,
skip_entry_sync=skip_entry_sync,
safe=safe,
).clone()
for _ in range(TEST_LOOP):
# Entry barrier may be skipped on the kernel side; make sure every rank's
# input is ready and the buffer is free before the next gather.
dist.barrier(nccl_group)
x = torch.randn(num_tokens, local_hidden, dtype=torch.bfloat16, device=device)
ref = _nccl_all_gather(x, nccl_group, world_size)
out = gather(x)
# Pure copy gather: exact bitwise equality.
torch.testing.assert_close(out, ref, atol=0, rtol=0)
if __name__ == "__main__":
# multimem multicast needs world_size in {4, 6, 8} (cc9) or {6, 8} (cc10);
# unsupported sizes self-skip via the multicast_ptr guard above.
multigpu_pytest_main(__name__, __file__, num_gpus=(4, 8))