[3/N] CI refactor: move some manually triggered tests. (#13448)
This commit is contained in:
Liangsheng Yin
@@ -0,0 +1,445 @@
|
||||
# Copy from deepseek-ai/DeepEP/tests/test_internode.py
|
||||
|
||||
import os
|
||||
import time
|
||||
|
||||
# noinspection PyUnresolvedReferences
|
||||
import deep_ep
|
||||
|
||||
# Test compatibility with low latency functions
|
||||
import test_deepep_low_latency
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.test.test_deepep_utils import (
|
||||
bench,
|
||||
calc_diff,
|
||||
create_grouped_scores,
|
||||
init_dist,
|
||||
inplace_unique,
|
||||
per_token_cast_back,
|
||||
per_token_cast_to_fp8,
|
||||
)
|
||||
|
||||
|
||||
def test_main(
|
||||
num_sms: int,
|
||||
local_rank: int,
|
||||
num_local_ranks: int,
|
||||
num_ranks: int,
|
||||
num_nodes: int,
|
||||
rank: int,
|
||||
buffer: deep_ep.Buffer,
|
||||
group: dist.ProcessGroup,
|
||||
):
|
||||
# Settings
|
||||
num_tokens, hidden, num_topk_groups, num_topk, num_experts = (
|
||||
4096,
|
||||
7168,
|
||||
min(num_nodes, 4),
|
||||
8,
|
||||
(256 // num_ranks) * num_ranks,
|
||||
)
|
||||
assert num_experts % num_ranks == 0 and num_local_ranks == 8
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk_groups={num_topk_groups}, num_topk={num_topk}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# Random data
|
||||
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank
|
||||
x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
|
||||
x_e4m3 = per_token_cast_to_fp8(x)
|
||||
scores = (
|
||||
torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
|
||||
+ 1
|
||||
)
|
||||
group_scores = scores.view(num_tokens, num_nodes, -1).amax(dim=-1)
|
||||
group_idx = torch.topk(
|
||||
group_scores, k=num_topk_groups, dim=-1, sorted=False
|
||||
).indices
|
||||
masked_scores = create_grouped_scores(scores, group_idx, num_nodes)
|
||||
topk_idx = torch.topk(masked_scores, num_topk, dim=-1, largest=True, sorted=False)[
|
||||
1
|
||||
]
|
||||
topk_weights = (
|
||||
torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank
|
||||
)
|
||||
topk_weights_pure_rand = torch.randn(
|
||||
(num_tokens, num_topk), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
rank_idx = topk_idx // (num_experts // num_ranks)
|
||||
rank_idx.masked_fill_(topk_idx == -1, -1)
|
||||
inplace_unique(rank_idx, num_ranks)
|
||||
rdma_rank_idx = rank_idx // num_local_ranks
|
||||
rdma_rank_idx.masked_fill_(rank_idx == -1, -1)
|
||||
inplace_unique(rdma_rank_idx, num_nodes)
|
||||
|
||||
# RDMA dispatch counts
|
||||
rdma_idx = topk_idx // (num_experts // num_nodes)
|
||||
rdma_idx.masked_fill_(topk_idx == -1, -1)
|
||||
inplace_unique(rdma_idx, num_nodes)
|
||||
num_rdma_token_sent = rdma_idx.ne(-1).sum().item()
|
||||
|
||||
# Expert meta
|
||||
num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda")
|
||||
for i in range(num_experts):
|
||||
num_tokens_per_expert[i] = (topk_idx == i).sum()
|
||||
gbl_num_tokens_per_expert = num_tokens_per_expert.clone()
|
||||
dist.all_reduce(gbl_num_tokens_per_expert, group=group)
|
||||
|
||||
# Rank layout meta
|
||||
num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda")
|
||||
num_tokens_per_rdma_rank = torch.empty((num_nodes,), dtype=torch.int, device="cuda")
|
||||
token_idx_in_rank = torch.full(
|
||||
(num_ranks, num_tokens), -1, dtype=torch.long, device="cuda"
|
||||
)
|
||||
for i in range(num_ranks):
|
||||
num_tokens_per_rank[i] = (rank_idx == i).sum()
|
||||
token_sel = (rank_idx == i).max(dim=-1)[0]
|
||||
count = token_sel.sum().item()
|
||||
tokens = torch.sort(token_sel.to(torch.int), descending=True)[1]
|
||||
tokens[:count] = torch.sort(tokens[:count])[0]
|
||||
token_idx_in_rank[i][tokens[:count]] = torch.arange(
|
||||
count, dtype=torch.long, device="cuda"
|
||||
)
|
||||
for i in range(num_nodes):
|
||||
num_tokens_per_rdma_rank[i] = (rdma_rank_idx == i).sum()
|
||||
token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int)
|
||||
is_token_in_rank = token_idx_in_rank >= 0
|
||||
gbl_num_tokens_per_rank = num_tokens_per_rank.clone()
|
||||
dist.all_reduce(gbl_num_tokens_per_rank, group=group)
|
||||
|
||||
(
|
||||
ref_num_tokens_per_rank,
|
||||
ref_num_tokens_per_rdma_rank,
|
||||
ref_num_tokens_per_expert,
|
||||
ref_is_token_in_rank,
|
||||
_,
|
||||
) = buffer.get_dispatch_layout(topk_idx, num_experts)
|
||||
assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank)
|
||||
assert torch.allclose(ref_num_tokens_per_rdma_rank, num_tokens_per_rdma_rank)
|
||||
assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert)
|
||||
assert torch.allclose(ref_is_token_in_rank, is_token_in_rank)
|
||||
t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0]
|
||||
if local_rank == 0:
|
||||
print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True)
|
||||
print("", flush=True)
|
||||
group.barrier()
|
||||
time.sleep(1)
|
||||
|
||||
# Config
|
||||
rdma_buffer_size, nvl_buffer_size = 128, (720 if num_ranks in (144, 160) else 512)
|
||||
config = deep_ep.Config(num_sms, 8, nvl_buffer_size, 16, rdma_buffer_size)
|
||||
|
||||
# Test dispatch
|
||||
# noinspection PyShadowingNames
|
||||
def check_data(check_x, recv_gbl_rank_prefix_sum):
|
||||
assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1))
|
||||
check_start = 0
|
||||
for i in range(num_ranks):
|
||||
check_end = recv_gbl_rank_prefix_sum[i].item()
|
||||
assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0
|
||||
check_start = check_end
|
||||
|
||||
for previous_mode in (False, True):
|
||||
for async_mode in (False, True):
|
||||
for current_x in (x_pure_rand, x, x_e4m3):
|
||||
for with_topk in (False, True):
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
|
||||
flush=True,
|
||||
end="",
|
||||
)
|
||||
dispatch_args = {
|
||||
"x": current_x,
|
||||
"num_tokens_per_rank": num_tokens_per_rank,
|
||||
"num_tokens_per_rdma_rank": num_tokens_per_rdma_rank,
|
||||
"is_token_in_rank": is_token_in_rank,
|
||||
"num_tokens_per_expert": num_tokens_per_expert,
|
||||
"config": config,
|
||||
"async_finish": async_mode,
|
||||
}
|
||||
if with_topk:
|
||||
dispatch_args.update(
|
||||
{
|
||||
"topk_idx": topk_idx,
|
||||
"topk_weights": (
|
||||
topk_weights_pure_rand
|
||||
if current_x is x_pure_rand
|
||||
else topk_weights
|
||||
),
|
||||
}
|
||||
)
|
||||
if previous_mode:
|
||||
dispatch_args.update({"previous_event": buffer.capture()})
|
||||
(
|
||||
recv_x,
|
||||
recv_topk_idx,
|
||||
recv_topk_weights,
|
||||
recv_num_tokens_per_expert_list,
|
||||
handle,
|
||||
event,
|
||||
) = buffer.dispatch(**dispatch_args)
|
||||
event.current_stream_wait() if async_mode else ()
|
||||
recv_x = (
|
||||
per_token_cast_back(*recv_x)
|
||||
if isinstance(recv_x, tuple)
|
||||
else recv_x
|
||||
)
|
||||
|
||||
# Checks
|
||||
recv_gbl_rank_prefix_sum = handle[-4]
|
||||
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
|
||||
0
|
||||
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
|
||||
assert (
|
||||
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
|
||||
== recv_num_tokens_per_expert_list
|
||||
)
|
||||
if current_x is not x_pure_rand:
|
||||
check_data(recv_x, recv_gbl_rank_prefix_sum)
|
||||
if with_topk:
|
||||
# Check `topk_idx`
|
||||
assert (
|
||||
recv_topk_idx.eq(-1)
|
||||
| (
|
||||
(recv_topk_idx >= 0)
|
||||
& (recv_topk_idx < (num_experts // num_ranks))
|
||||
)
|
||||
).sum().item() == recv_topk_idx.numel()
|
||||
for i, count in enumerate(recv_num_tokens_per_expert_list):
|
||||
assert recv_topk_idx.eq(i).sum().item() == count
|
||||
|
||||
# Check `topk_weights`
|
||||
if current_x is not x_pure_rand:
|
||||
recv_topk_weights[recv_topk_idx.eq(-1)] = (
|
||||
recv_topk_weights.amax(dim=1, keepdim=True).expand_as(
|
||||
recv_topk_weights
|
||||
)[recv_topk_idx.eq(-1)]
|
||||
)
|
||||
check_data(recv_topk_weights, recv_gbl_rank_prefix_sum)
|
||||
|
||||
# Test cached dispatch (must without top-k staffs)
|
||||
if not with_topk:
|
||||
dispatch_args = {
|
||||
"x": current_x,
|
||||
"handle": handle,
|
||||
"config": config,
|
||||
"async_finish": async_mode,
|
||||
}
|
||||
if previous_mode:
|
||||
dispatch_args.update({"previous_event": buffer.capture()})
|
||||
recv_x, _, _, _, _, event = buffer.dispatch(**dispatch_args)
|
||||
event.current_stream_wait() if async_mode else ()
|
||||
recv_x = (
|
||||
per_token_cast_back(*recv_x)
|
||||
if isinstance(recv_x, tuple)
|
||||
else recv_x
|
||||
)
|
||||
if current_x is not x_pure_rand:
|
||||
check_data(recv_x, recv_gbl_rank_prefix_sum)
|
||||
|
||||
# Test combine
|
||||
combine_args = {
|
||||
"x": recv_x,
|
||||
"handle": handle,
|
||||
"config": config,
|
||||
"async_finish": async_mode,
|
||||
}
|
||||
if with_topk:
|
||||
combine_args.update({"topk_weights": recv_topk_weights})
|
||||
if previous_mode:
|
||||
dispatch_args.update({"previous_event": buffer.capture()})
|
||||
combined_x, combined_topk_weights, event = buffer.combine(
|
||||
**combine_args
|
||||
)
|
||||
event.current_stream_wait() if async_mode else ()
|
||||
check_x = combined_x.float() / is_token_in_rank.sum(
|
||||
dim=1
|
||||
).unsqueeze(1)
|
||||
ref_x = x_pure_rand if current_x is x_pure_rand else x
|
||||
assert calc_diff(check_x, ref_x) < 5e-6
|
||||
if with_topk:
|
||||
check_topk_weights = (
|
||||
combined_topk_weights
|
||||
if (current_x is x_pure_rand)
|
||||
else (
|
||||
combined_topk_weights
|
||||
/ is_token_in_rank.sum(dim=1).unsqueeze(1)
|
||||
)
|
||||
)
|
||||
ref_topk_weights = (
|
||||
topk_weights_pure_rand
|
||||
if current_x is x_pure_rand
|
||||
else topk_weights
|
||||
)
|
||||
assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9
|
||||
|
||||
# For later tuning
|
||||
dispatch_bf16_rdma_send_bytes = num_rdma_token_sent * hidden * 2
|
||||
dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2
|
||||
combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes
|
||||
combine_bf16_rdma_recv_bytes = dispatch_bf16_rdma_send_bytes
|
||||
|
||||
if local_rank == 0:
|
||||
print(" passed", flush=True)
|
||||
if local_rank == 0:
|
||||
print("", flush=True)
|
||||
|
||||
# Tune dispatch performance
|
||||
best_dispatch_results = None
|
||||
fp8_factor = (1 + 4 / 128) / 2
|
||||
for current_x in (x_e4m3, x):
|
||||
best_time, best_results = 1e10, None
|
||||
rdma_send_bytes = (
|
||||
(dispatch_bf16_rdma_send_bytes * fp8_factor)
|
||||
if isinstance(current_x, tuple)
|
||||
else dispatch_bf16_rdma_send_bytes
|
||||
)
|
||||
nvl_recv_bytes = (
|
||||
(dispatch_bf16_nvl_recv_bytes * fp8_factor)
|
||||
if isinstance(current_x, tuple)
|
||||
else dispatch_bf16_nvl_recv_bytes
|
||||
)
|
||||
for nvl_chunk_size in range(4, 33, 4):
|
||||
for rdma_chunk_size in range(4, 33, 4):
|
||||
config = deep_ep.Config(
|
||||
num_sms,
|
||||
nvl_chunk_size,
|
||||
nvl_buffer_size,
|
||||
rdma_chunk_size,
|
||||
rdma_buffer_size,
|
||||
)
|
||||
tune_args = {"x": current_x, "handle": handle, "config": config}
|
||||
t = bench(lambda: buffer.dispatch(**tune_args))[0]
|
||||
if t < best_time:
|
||||
best_time, best_results = t, (
|
||||
num_sms,
|
||||
nvl_chunk_size,
|
||||
rdma_chunk_size,
|
||||
)
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {rdma_send_bytes / 1e9 / t:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ",
|
||||
flush=True,
|
||||
)
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {rdma_send_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
|
||||
flush=True,
|
||||
)
|
||||
print("", flush=True)
|
||||
|
||||
if isinstance(current_x, tuple):
|
||||
# Gather FP8 the best config from rank 0
|
||||
best_dispatch_results = torch.tensor(
|
||||
[best_results[0], best_results[1], best_results[2]],
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
all_best_fp8_results_list = [
|
||||
torch.zeros_like(best_dispatch_results)
|
||||
for _ in range(torch.distributed.get_world_size())
|
||||
]
|
||||
dist.all_gather(
|
||||
all_best_fp8_results_list, best_dispatch_results, group=group
|
||||
)
|
||||
best_dispatch_results = all_best_fp8_results_list[0].tolist()
|
||||
dispatch_config = deep_ep.Config(
|
||||
best_dispatch_results[0],
|
||||
best_dispatch_results[1],
|
||||
nvl_buffer_size,
|
||||
best_dispatch_results[2],
|
||||
rdma_buffer_size,
|
||||
)
|
||||
|
||||
dispatch_args = {
|
||||
"x": x,
|
||||
"num_tokens_per_rank": num_tokens_per_rank,
|
||||
"num_tokens_per_rdma_rank": num_tokens_per_rdma_rank,
|
||||
"is_token_in_rank": is_token_in_rank,
|
||||
"num_tokens_per_expert": num_tokens_per_expert,
|
||||
"config": dispatch_config if dispatch_config is not None else config,
|
||||
}
|
||||
recv_x, _, _, _, handle, _ = buffer.dispatch(**dispatch_args)
|
||||
|
||||
# Tune combine performance
|
||||
best_time, best_results = 1e10, None
|
||||
for nvl_chunk_size in range(1, 5, 1):
|
||||
for rdma_chunk_size in range(8, 33, 4):
|
||||
config = deep_ep.Config(
|
||||
num_sms,
|
||||
nvl_chunk_size,
|
||||
nvl_buffer_size,
|
||||
rdma_chunk_size,
|
||||
rdma_buffer_size,
|
||||
)
|
||||
tune_args = {"x": recv_x, "handle": handle, "config": config}
|
||||
t = bench(lambda: buffer.combine(**tune_args))[0]
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}, RDMA chunk {rdma_chunk_size}: {combine_bf16_rdma_recv_bytes / 1e9 / t:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ",
|
||||
flush=True,
|
||||
)
|
||||
if t < best_time:
|
||||
best_time, best_results = t, (
|
||||
num_sms,
|
||||
nvl_chunk_size,
|
||||
rdma_chunk_size,
|
||||
)
|
||||
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}, RDMA chunk {best_results[2]}: {combine_bf16_rdma_recv_bytes / 1e9 / best_time:.2f} GB/s (RDMA), {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
|
||||
flush=True,
|
||||
)
|
||||
print("", flush=True)
|
||||
|
||||
|
||||
# noinspection PyUnboundLocalVariable
|
||||
def test_loop(local_rank: int, num_local_ranks: int):
|
||||
num_nodes = int(os.getenv("WORLD_SIZE", 1))
|
||||
rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
|
||||
test_ll_compatibility = False
|
||||
if test_ll_compatibility:
|
||||
ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9
|
||||
|
||||
buffer = deep_ep.Buffer(
|
||||
group,
|
||||
int(1e9),
|
||||
int(1e9),
|
||||
low_latency_mode=test_ll_compatibility,
|
||||
num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1),
|
||||
)
|
||||
assert num_local_ranks == 8 and num_ranks > 8
|
||||
torch.manual_seed(rank)
|
||||
|
||||
for i in (24,):
|
||||
test_main(
|
||||
i, local_rank, num_local_ranks, num_ranks, num_nodes, rank, buffer, group
|
||||
)
|
||||
if local_rank == 0:
|
||||
print("", flush=True)
|
||||
|
||||
# Test compatibility with low latency functions
|
||||
if test_ll_compatibility:
|
||||
buffer.clean_low_latency_buffer(ll_num_tokens, ll_hidden, ll_num_experts)
|
||||
test_deepep_low_latency.test_main(
|
||||
ll_num_tokens,
|
||||
ll_hidden,
|
||||
ll_num_experts,
|
||||
ll_num_topk,
|
||||
rank,
|
||||
num_ranks,
|
||||
group,
|
||||
buffer,
|
||||
seed=1,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
num_processes = 8
|
||||
torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)
|
||||
@@ -0,0 +1,378 @@
|
||||
# Copy from deepseek-ai/DeepEP/tests/test_intranode.py
|
||||
|
||||
import time
|
||||
|
||||
# noinspection PyUnresolvedReferences
|
||||
import deep_ep
|
||||
|
||||
# Test compatibility with low latency functions
|
||||
import test_deepep_low_latency
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.test.test_deepep_utils import (
|
||||
bench,
|
||||
calc_diff,
|
||||
init_dist,
|
||||
inplace_unique,
|
||||
per_token_cast_back,
|
||||
per_token_cast_to_fp8,
|
||||
)
|
||||
|
||||
|
||||
def test_main(
|
||||
num_sms: int,
|
||||
local_rank: int,
|
||||
num_ranks: int,
|
||||
rank: int,
|
||||
buffer: deep_ep.Buffer,
|
||||
group: dist.ProcessGroup,
|
||||
):
|
||||
# Settings
|
||||
num_tokens, hidden, num_topk, num_experts = (
|
||||
4096,
|
||||
7168,
|
||||
8,
|
||||
(256 // num_ranks) * num_ranks,
|
||||
)
|
||||
assert num_experts % num_ranks == 0
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f"[config] num_tokens={num_tokens}, hidden={hidden}, num_topk={num_topk}",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# Random data
|
||||
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * rank
|
||||
x_pure_rand = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
|
||||
x_e4m3 = per_token_cast_to_fp8(x)
|
||||
scores = (
|
||||
torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
|
||||
+ 1
|
||||
)
|
||||
topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)[1]
|
||||
topk_weights = (
|
||||
torch.ones((num_tokens, num_topk), dtype=torch.float32, device="cuda") * rank
|
||||
)
|
||||
topk_weights_pure_rand = torch.randn(
|
||||
(num_tokens, num_topk), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
rank_idx = topk_idx // (num_experts // num_ranks)
|
||||
rank_idx.masked_fill_(topk_idx == -1, -1)
|
||||
inplace_unique(rank_idx, num_ranks)
|
||||
|
||||
# Expert meta
|
||||
num_tokens_per_expert = torch.zeros((num_experts,), dtype=torch.int, device="cuda")
|
||||
for i in range(num_experts):
|
||||
num_tokens_per_expert[i] = (topk_idx == i).sum()
|
||||
gbl_num_tokens_per_expert = num_tokens_per_expert.clone()
|
||||
dist.all_reduce(gbl_num_tokens_per_expert, group=group)
|
||||
|
||||
# Rank layout meta
|
||||
num_tokens_per_rank = torch.empty((num_ranks,), dtype=torch.int, device="cuda")
|
||||
token_idx_in_rank = torch.full(
|
||||
(num_ranks, num_tokens), -1, dtype=torch.long, device="cuda"
|
||||
)
|
||||
for i in range(num_ranks):
|
||||
num_tokens_per_rank[i] = (rank_idx == i).sum()
|
||||
token_sel = (rank_idx == i).max(dim=-1)[0]
|
||||
count = token_sel.sum().item()
|
||||
tokens = torch.sort(token_sel.to(torch.int), descending=True)[1]
|
||||
tokens[:count] = torch.sort(tokens[:count])[0]
|
||||
token_idx_in_rank[i][tokens[:count]] = torch.arange(
|
||||
count, dtype=torch.long, device="cuda"
|
||||
)
|
||||
token_idx_in_rank = token_idx_in_rank.T.contiguous().to(torch.int)
|
||||
is_token_in_rank = token_idx_in_rank >= 0
|
||||
gbl_num_tokens_per_rank = num_tokens_per_rank.clone()
|
||||
dist.all_reduce(gbl_num_tokens_per_rank, group=group)
|
||||
|
||||
ref_num_tokens_per_rank, _, ref_num_tokens_per_expert, ref_is_token_in_rank, _ = (
|
||||
buffer.get_dispatch_layout(topk_idx, num_experts)
|
||||
)
|
||||
assert torch.allclose(ref_num_tokens_per_rank, num_tokens_per_rank)
|
||||
assert torch.allclose(ref_num_tokens_per_expert, num_tokens_per_expert)
|
||||
assert torch.allclose(ref_is_token_in_rank, is_token_in_rank)
|
||||
t = bench(lambda: buffer.get_dispatch_layout(topk_idx, num_experts))[0]
|
||||
if local_rank == 0:
|
||||
print(f"[layout] Kernel performance: {t * 1000:.3f} ms", flush=True)
|
||||
print("", flush=True)
|
||||
group.barrier()
|
||||
time.sleep(1)
|
||||
|
||||
# Config
|
||||
nvl_buffer_size = 256
|
||||
config = deep_ep.Config(num_sms, 8, nvl_buffer_size)
|
||||
|
||||
# Test dispatch
|
||||
# noinspection PyShadowingNames
|
||||
def check_data(check_x, rank_prefix_matrix):
|
||||
assert torch.allclose(check_x.amin(dim=1), check_x.amax(dim=1))
|
||||
check_start = 0
|
||||
for i in range(num_ranks):
|
||||
check_end = rank_prefix_matrix[i][rank].item()
|
||||
assert (check_x[check_start:check_end, :].int() - i).sum().item() == 0
|
||||
check_start = check_end
|
||||
|
||||
for previous_mode in (False, True):
|
||||
for async_mode in (False, True):
|
||||
for current_x in (x_pure_rand, x, x_e4m3):
|
||||
for with_topk in (False, True):
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f'[testing] Running with {"FP8" if isinstance(current_x, tuple) else "BF16"}, {"with" if with_topk else "without"} top-k (async={async_mode}, previous={previous_mode}) ...',
|
||||
flush=True,
|
||||
end="",
|
||||
)
|
||||
dispatch_args = {
|
||||
"x": current_x,
|
||||
"num_tokens_per_rank": num_tokens_per_rank,
|
||||
"is_token_in_rank": is_token_in_rank,
|
||||
"num_tokens_per_expert": num_tokens_per_expert,
|
||||
"config": config,
|
||||
"async_finish": async_mode,
|
||||
}
|
||||
if with_topk:
|
||||
dispatch_args.update(
|
||||
{
|
||||
"topk_idx": topk_idx,
|
||||
"topk_weights": (
|
||||
topk_weights_pure_rand
|
||||
if current_x is x_pure_rand
|
||||
else topk_weights
|
||||
),
|
||||
}
|
||||
)
|
||||
if previous_mode:
|
||||
dispatch_args.update({"previous_event": buffer.capture()})
|
||||
(
|
||||
recv_x,
|
||||
recv_topk_idx,
|
||||
recv_topk_weights,
|
||||
recv_num_tokens_per_expert_list,
|
||||
handle,
|
||||
event,
|
||||
) = buffer.dispatch(**dispatch_args)
|
||||
event.current_stream_wait() if async_mode else ()
|
||||
recv_x = (
|
||||
per_token_cast_back(*recv_x)
|
||||
if isinstance(recv_x, tuple)
|
||||
else recv_x
|
||||
)
|
||||
|
||||
# Checks
|
||||
rank_prefix_matrix = handle[0]
|
||||
assert gbl_num_tokens_per_rank[rank].item() == recv_x.size(
|
||||
0
|
||||
), f"{gbl_num_tokens_per_rank[rank].item()} != {recv_x.size(0)}"
|
||||
assert (
|
||||
gbl_num_tokens_per_expert.view(num_ranks, -1)[rank].tolist()
|
||||
== recv_num_tokens_per_expert_list
|
||||
)
|
||||
if current_x is not x_pure_rand:
|
||||
check_data(recv_x, rank_prefix_matrix)
|
||||
if with_topk:
|
||||
# Check `topk_idx`
|
||||
assert (
|
||||
recv_topk_idx.eq(-1)
|
||||
| (
|
||||
(recv_topk_idx >= 0)
|
||||
& (recv_topk_idx < (num_experts // num_ranks))
|
||||
)
|
||||
).sum().item() == recv_topk_idx.numel()
|
||||
for i, count in enumerate(recv_num_tokens_per_expert_list):
|
||||
assert recv_topk_idx.eq(i).sum().item() == count
|
||||
|
||||
# Check `topk_weights`
|
||||
if current_x is not x_pure_rand:
|
||||
recv_topk_weights[recv_topk_idx.eq(-1)] = (
|
||||
recv_topk_weights.amax(dim=1, keepdim=True).expand_as(
|
||||
recv_topk_weights
|
||||
)[recv_topk_idx.eq(-1)]
|
||||
)
|
||||
check_data(recv_topk_weights, rank_prefix_matrix)
|
||||
|
||||
# Test cached dispatch (must without top-k staffs)
|
||||
if not with_topk:
|
||||
dispatch_args = {
|
||||
"x": current_x,
|
||||
"handle": handle,
|
||||
"config": config,
|
||||
"async_finish": async_mode,
|
||||
}
|
||||
if previous_mode:
|
||||
dispatch_args.update({"previous_event": buffer.capture()})
|
||||
recv_x, _, _, _, _, event = buffer.dispatch(**dispatch_args)
|
||||
event.current_stream_wait() if async_mode else ()
|
||||
recv_x = (
|
||||
per_token_cast_back(*recv_x)
|
||||
if isinstance(recv_x, tuple)
|
||||
else recv_x
|
||||
)
|
||||
if current_x is not x_pure_rand:
|
||||
check_data(recv_x, rank_prefix_matrix)
|
||||
|
||||
# Test combine
|
||||
combine_args = {
|
||||
"x": recv_x,
|
||||
"handle": handle,
|
||||
"config": config,
|
||||
"async_finish": async_mode,
|
||||
}
|
||||
if with_topk:
|
||||
combine_args.update({"topk_weights": recv_topk_weights})
|
||||
if previous_mode:
|
||||
dispatch_args.update({"previous_event": buffer.capture()})
|
||||
combined_x, combined_topk_weights, event = buffer.combine(
|
||||
**combine_args
|
||||
)
|
||||
event.current_stream_wait() if async_mode else ()
|
||||
check_x = combined_x.float() / is_token_in_rank.sum(
|
||||
dim=1
|
||||
).unsqueeze(1)
|
||||
ref_x = x_pure_rand if current_x is x_pure_rand else x
|
||||
assert calc_diff(check_x, ref_x) < 5e-6
|
||||
if with_topk:
|
||||
check_topk_weights = (
|
||||
combined_topk_weights
|
||||
if (current_x is x_pure_rand)
|
||||
else (
|
||||
combined_topk_weights
|
||||
/ is_token_in_rank.sum(dim=1).unsqueeze(1)
|
||||
)
|
||||
)
|
||||
ref_topk_weights = (
|
||||
topk_weights_pure_rand
|
||||
if current_x is x_pure_rand
|
||||
else topk_weights
|
||||
)
|
||||
assert calc_diff(check_topk_weights, ref_topk_weights) < 1e-9
|
||||
|
||||
# For later tuning
|
||||
dispatch_bf16_nvl_recv_bytes = recv_x.numel() * 2
|
||||
combine_bf16_nvl_send_bytes = dispatch_bf16_nvl_recv_bytes
|
||||
|
||||
if local_rank == 0:
|
||||
print(" passed", flush=True)
|
||||
if local_rank == 0:
|
||||
print("", flush=True)
|
||||
|
||||
# Tune dispatch performance
|
||||
best_dispatch_results = None
|
||||
fp8_factor = (1 + 4 / 128) / 2
|
||||
for current_x in (x_e4m3, x):
|
||||
best_time, best_results = 1e10, None
|
||||
nvl_recv_bytes = (
|
||||
(dispatch_bf16_nvl_recv_bytes * fp8_factor)
|
||||
if isinstance(current_x, tuple)
|
||||
else dispatch_bf16_nvl_recv_bytes
|
||||
)
|
||||
for nvl_chunk_size in range(4, 33, 4):
|
||||
config = deep_ep.Config(num_sms, nvl_chunk_size, nvl_buffer_size)
|
||||
tune_args = {"x": current_x, "handle": handle, "config": config}
|
||||
t = bench(lambda: buffer.dispatch(**tune_args))[0]
|
||||
if t < best_time:
|
||||
best_time, best_results = t, (num_sms, nvl_chunk_size)
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}: {nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ",
|
||||
flush=True,
|
||||
)
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)',
|
||||
flush=True,
|
||||
)
|
||||
print("", flush=True)
|
||||
|
||||
if isinstance(current_x, tuple):
|
||||
# Gather FP8 the best config from rank 0
|
||||
best_dispatch_results = torch.tensor(
|
||||
[best_results[0], best_results[1]], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
all_best_fp8_results_list = [
|
||||
torch.zeros_like(best_dispatch_results)
|
||||
for _ in range(torch.distributed.get_world_size())
|
||||
]
|
||||
dist.all_gather(
|
||||
all_best_fp8_results_list, best_dispatch_results, group=group
|
||||
)
|
||||
best_dispatch_results = all_best_fp8_results_list[0].tolist()
|
||||
dispatch_config = deep_ep.Config(
|
||||
best_dispatch_results[0], best_dispatch_results[1], nvl_buffer_size
|
||||
)
|
||||
|
||||
dispatch_args = {
|
||||
"x": x,
|
||||
"num_tokens_per_rank": num_tokens_per_rank,
|
||||
"is_token_in_rank": is_token_in_rank,
|
||||
"num_tokens_per_expert": num_tokens_per_expert,
|
||||
"config": dispatch_config if dispatch_config is not None else config,
|
||||
}
|
||||
recv_x, _, _, _, handle, _ = buffer.dispatch(**dispatch_args)
|
||||
|
||||
# Tune combine performance
|
||||
best_time, best_results = 1e10, None
|
||||
for nvl_chunk_size in range(1, 7, 1):
|
||||
config = deep_ep.Config(num_sms, nvl_chunk_size, nvl_buffer_size)
|
||||
tune_args = {"x": recv_x, "handle": handle, "config": config}
|
||||
t = bench(lambda: buffer.combine(**tune_args))[0]
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f"[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size}: {combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ",
|
||||
flush=True,
|
||||
)
|
||||
if t < best_time:
|
||||
best_time, best_results = t, (num_sms, nvl_chunk_size)
|
||||
|
||||
if local_rank == 0:
|
||||
print(
|
||||
f"[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}: {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)",
|
||||
flush=True,
|
||||
)
|
||||
print("", flush=True)
|
||||
|
||||
|
||||
# noinspection PyUnboundLocalVariable
|
||||
def test_loop(local_rank: int, num_local_ranks: int):
|
||||
rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
|
||||
test_ll_compatibility, num_rdma_bytes = False, 0
|
||||
if test_ll_compatibility:
|
||||
ll_num_tokens, ll_hidden, ll_num_experts, ll_num_topk = 16, 5120, 256, 9
|
||||
num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
|
||||
ll_num_tokens, ll_hidden, num_ranks, ll_num_experts
|
||||
)
|
||||
|
||||
buffer = deep_ep.Buffer(
|
||||
group,
|
||||
int(1e9),
|
||||
num_rdma_bytes,
|
||||
low_latency_mode=test_ll_compatibility,
|
||||
num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1),
|
||||
)
|
||||
torch.manual_seed(rank)
|
||||
|
||||
for i in (24,):
|
||||
test_main(i, local_rank, num_ranks, rank, buffer, group)
|
||||
if local_rank == 0:
|
||||
print("", flush=True)
|
||||
|
||||
# Test compatibility with low latency functions
|
||||
if test_ll_compatibility:
|
||||
buffer.clean_low_latency_buffer(ll_num_tokens, ll_hidden, ll_num_experts)
|
||||
test_deepep_low_latency.test_main(
|
||||
ll_num_tokens,
|
||||
ll_hidden,
|
||||
ll_num_experts,
|
||||
ll_num_topk,
|
||||
rank,
|
||||
num_ranks,
|
||||
group,
|
||||
buffer,
|
||||
seed=1,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
num_processes = 8
|
||||
torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)
|
||||
@@ -0,0 +1,325 @@
|
||||
# Copy from deepseek-ai/DeepEP/tests/test_low_latency.py
|
||||
|
||||
import random
|
||||
from functools import partial
|
||||
|
||||
import deep_ep
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
|
||||
from sglang.test.test_deepep_utils import (
|
||||
bench,
|
||||
bench_kineto,
|
||||
calc_diff,
|
||||
hash_tensor,
|
||||
init_dist,
|
||||
per_token_cast_back,
|
||||
)
|
||||
|
||||
|
||||
def test_main(
|
||||
num_tokens: int,
|
||||
hidden: int,
|
||||
num_experts: int,
|
||||
num_topk: int,
|
||||
rank: int,
|
||||
num_ranks: int,
|
||||
group: dist.ProcessGroup,
|
||||
buffer: deep_ep.Buffer,
|
||||
seed: int = 0,
|
||||
):
|
||||
torch.manual_seed(seed + rank)
|
||||
random.seed(seed + rank)
|
||||
|
||||
assert num_experts % num_ranks == 0
|
||||
num_local_experts = num_experts // num_ranks
|
||||
|
||||
# NOTES: the integers greater than 256 exceeds the BF16 precision limit
|
||||
rank_offset = 128
|
||||
assert (
|
||||
num_ranks - rank_offset < 257
|
||||
), "Too many ranks (exceeding test precision limit)"
|
||||
|
||||
x = torch.ones((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") * (
|
||||
rank - rank_offset
|
||||
)
|
||||
x[:, -128:] = torch.arange(num_tokens, device="cuda").to(torch.bfloat16).view(-1, 1)
|
||||
scores = (
|
||||
torch.randn((num_tokens, num_experts), dtype=torch.float32, device="cuda").abs()
|
||||
+ 1
|
||||
)
|
||||
topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=True)[1]
|
||||
topk_weights = torch.randn(
|
||||
(num_tokens, num_topk), dtype=torch.float32, device="cuda"
|
||||
).abs()
|
||||
|
||||
# Randomly mask some positions
|
||||
for i in range(10):
|
||||
topk_idx[random.randint(0, num_tokens - 1), random.randint(0, num_topk - 1)] = (
|
||||
-1
|
||||
)
|
||||
|
||||
# Check dispatch correctness
|
||||
do_check = True
|
||||
hash_value, num_times = 0, 0
|
||||
for return_recv_hook in (False, True):
|
||||
for dispatch_use_fp8 in (False, True):
|
||||
num_times += 1
|
||||
for i in range((num_times % 2) + 1):
|
||||
packed_recv_x, packed_recv_count, handle, event, hook = (
|
||||
buffer.low_latency_dispatch(
|
||||
x,
|
||||
topk_idx,
|
||||
num_tokens,
|
||||
num_experts,
|
||||
use_fp8=dispatch_use_fp8,
|
||||
async_finish=not return_recv_hook,
|
||||
return_recv_hook=return_recv_hook,
|
||||
)
|
||||
)
|
||||
hook() if return_recv_hook else event.current_stream_wait()
|
||||
packed_recv_x = (
|
||||
(packed_recv_x[0], packed_recv_x[1].contiguous())
|
||||
if dispatch_use_fp8
|
||||
else packed_recv_x
|
||||
)
|
||||
simulated_gemm_x = (
|
||||
per_token_cast_back(
|
||||
packed_recv_x[0].view(-1, hidden),
|
||||
packed_recv_x[1].view(-1, hidden // 128),
|
||||
).view(packed_recv_x[0].shape)
|
||||
if dispatch_use_fp8
|
||||
else packed_recv_x.clone()
|
||||
)
|
||||
all_topk_idx = torch.empty(
|
||||
(num_ranks, num_tokens, num_topk), dtype=topk_idx.dtype, device="cuda"
|
||||
)
|
||||
dist.all_gather_into_tensor(all_topk_idx, topk_idx, group=group)
|
||||
for i in range(num_local_experts if do_check else 0):
|
||||
expert_id = rank * num_local_experts + i
|
||||
recv_x = (
|
||||
per_token_cast_back(packed_recv_x[0][i], packed_recv_x[1][i])
|
||||
if dispatch_use_fp8
|
||||
else packed_recv_x[i]
|
||||
)
|
||||
recv_count, recv_src_info, recv_layout_range = (
|
||||
packed_recv_count[i],
|
||||
handle[0][i],
|
||||
handle[1][i],
|
||||
)
|
||||
|
||||
# Check expert indices
|
||||
int_mask = (2**32) - 1
|
||||
num_valid_tokens = recv_count.item()
|
||||
assert (
|
||||
num_valid_tokens == (recv_layout_range & int_mask).sum().item()
|
||||
), f"{num_valid_tokens} != {recv_layout_range & int_mask}.sum().item()"
|
||||
assert (
|
||||
num_valid_tokens == (all_topk_idx == expert_id).sum().item()
|
||||
), f"{num_valid_tokens} != {(all_topk_idx == expert_id).sum().item()}"
|
||||
|
||||
# Check received data
|
||||
recv_x = recv_x[:num_valid_tokens]
|
||||
recv_x_amin = recv_x[:, :-128].amin(dim=-1)
|
||||
recv_src_info = recv_src_info[:num_valid_tokens]
|
||||
assert torch.equal(recv_x_amin, recv_x[:, :-128].amax(dim=-1))
|
||||
assert (
|
||||
recv_x[:, -128:] - recv_src_info.view(-1, 1) % num_tokens
|
||||
).sum().item() == 0
|
||||
for j in range(num_ranks):
|
||||
begin_idx, count = (recv_layout_range[j] >> 32).item(), (
|
||||
recv_layout_range[j] & int_mask
|
||||
).item()
|
||||
assert (recv_x_amin == j - rank_offset).sum().item() == (
|
||||
all_topk_idx[j] == expert_id
|
||||
).sum().item()
|
||||
assert (
|
||||
recv_x[begin_idx : begin_idx + count][:-128] - j
|
||||
).sum().item() == 0
|
||||
if dispatch_use_fp8:
|
||||
hash_value ^= hash_tensor(packed_recv_x[0][i, :num_valid_tokens])
|
||||
hash_value ^= hash_tensor(packed_recv_x[1][i, :num_valid_tokens])
|
||||
else:
|
||||
hash_value ^= hash_tensor(packed_recv_x[i, :num_valid_tokens])
|
||||
|
||||
# Check combine correctness
|
||||
for zero_copy in (False, True):
|
||||
if zero_copy:
|
||||
buffer.get_next_low_latency_combine_buffer(handle)[
|
||||
:, :, :
|
||||
] = simulated_gemm_x
|
||||
out = torch.empty(
|
||||
(num_tokens, hidden), dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
combined_x, event, hook = buffer.low_latency_combine(
|
||||
simulated_gemm_x,
|
||||
topk_idx,
|
||||
topk_weights,
|
||||
handle,
|
||||
async_finish=not return_recv_hook,
|
||||
zero_copy=zero_copy,
|
||||
return_recv_hook=return_recv_hook,
|
||||
out=out,
|
||||
)
|
||||
hook() if return_recv_hook else event.current_stream_wait()
|
||||
if do_check:
|
||||
diff = calc_diff(
|
||||
x
|
||||
* topk_weights.masked_fill(topk_idx == -1, 0)
|
||||
.sum(dim=1)
|
||||
.view(-1, 1),
|
||||
combined_x,
|
||||
)
|
||||
assert torch.isnan(combined_x).sum().item() == 0
|
||||
assert diff < 1e-5, f"Error: {diff=}, {zero_copy=}"
|
||||
hash_value ^= hash_tensor(combined_x)
|
||||
|
||||
def create_test_cast_with_outliers(num_outliers):
|
||||
tmp = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda")
|
||||
tmp /= tmp.abs().amax(dim=1).view(-1, 1)
|
||||
assert tmp.abs().amax().item() <= 1
|
||||
|
||||
# Create some amax outliers
|
||||
for i in range(num_outliers):
|
||||
tmp[random.randint(0, num_tokens - 1)] *= 1e3
|
||||
return tmp
|
||||
|
||||
# noinspection PyShadowingNames
|
||||
def large_gemm_with_hook(hook):
|
||||
mat_0 = torch.randn((8192, 8192), dtype=torch.float)
|
||||
mat_1 = torch.randn((8192, 8192), dtype=torch.float)
|
||||
mat_0 @ mat_1
|
||||
hook()
|
||||
|
||||
# noinspection PyShadowingNames
|
||||
def test_func(zero_copy: bool, return_recv_hook: bool):
|
||||
recv_x, recv_count, handle, event, hook = buffer.low_latency_dispatch(
|
||||
x,
|
||||
topk_idx,
|
||||
num_tokens,
|
||||
num_experts,
|
||||
async_finish=False,
|
||||
return_recv_hook=return_recv_hook,
|
||||
)
|
||||
large_gemm_with_hook(hook) if return_recv_hook else None
|
||||
if zero_copy:
|
||||
buffer.get_next_low_latency_combine_buffer(handle)[
|
||||
:, :, :
|
||||
] = simulated_gemm_x
|
||||
combined_x, event, hook = buffer.low_latency_combine(
|
||||
simulated_gemm_x,
|
||||
topk_idx,
|
||||
topk_weights,
|
||||
handle,
|
||||
zero_copy=zero_copy,
|
||||
return_recv_hook=return_recv_hook,
|
||||
)
|
||||
large_gemm_with_hook(hook) if return_recv_hook else None
|
||||
|
||||
# Calculate bandwidth
|
||||
num_fp8_bytes, num_bf16_bytes = (hidden + hidden / 128 * 4 + 16), hidden * 2
|
||||
num_dispatch_comm_bytes, num_combine_comm_bytes = 0, 0
|
||||
for i in range(num_tokens):
|
||||
num_selections = (topk_idx[i] != -1).sum().item()
|
||||
num_dispatch_comm_bytes += num_fp8_bytes * num_selections
|
||||
num_combine_comm_bytes += num_bf16_bytes * num_selections
|
||||
|
||||
# Dispatch + combine testing
|
||||
avg_t, min_t, max_t = bench(
|
||||
partial(test_func, zero_copy=False, return_recv_hook=False)
|
||||
)
|
||||
print(
|
||||
f"[rank {rank}] Dispatch + combine bandwidth: {(num_dispatch_comm_bytes + num_combine_comm_bytes) / 1e9 / avg_t:.2f} GB/s, "
|
||||
f"avg_t={avg_t * 1e6:.2f} us, min_t={min_t * 1e6:.2f} us, max_t={max_t * 1e6:.2f} us",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
# Separate profiling
|
||||
for return_recv_hook in (False, True):
|
||||
group.barrier()
|
||||
dispatch_t, combine_t = bench_kineto(
|
||||
partial(test_func, zero_copy=True, return_recv_hook=return_recv_hook),
|
||||
kernel_names=("dispatch", "combine"),
|
||||
barrier_comm_profiling=True,
|
||||
suppress_kineto_output=True,
|
||||
)
|
||||
if not return_recv_hook:
|
||||
print(
|
||||
f"[rank {rank}] Dispatch bandwidth: {num_dispatch_comm_bytes / 1e9 / dispatch_t:.2f} GB/s, avg_t={dispatch_t * 1e6:.2f} us | "
|
||||
f"Combine bandwidth: {num_combine_comm_bytes / 1e9 / combine_t:.2f} GB/s, avg_t={combine_t * 1e6:.2f} us",
|
||||
flush=True,
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"[rank {rank}] Dispatch send/recv time: {dispatch_t * 2 * 1e6:.2f} us | "
|
||||
f"Combine send/recv time: {combine_t * 2 * 1e6:.2f} us",
|
||||
flush=True,
|
||||
)
|
||||
|
||||
return hash_value
|
||||
|
||||
|
||||
# noinspection PyUnboundLocalVariable
|
||||
def test_loop(local_rank: int, num_local_ranks: int):
|
||||
rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
|
||||
num_tokens, hidden, num_topk, num_experts = 128, 7168, 8, 288
|
||||
|
||||
num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint(
|
||||
num_tokens, hidden, num_ranks, num_experts
|
||||
)
|
||||
if local_rank == 0:
|
||||
print(f"Allocating buffer size: {num_rdma_bytes / 1e6} MB ...", flush=True)
|
||||
buffer = deep_ep.Buffer(
|
||||
group,
|
||||
num_rdma_bytes=num_rdma_bytes,
|
||||
low_latency_mode=True,
|
||||
num_qps_per_rank=num_experts // num_ranks,
|
||||
)
|
||||
test_main(
|
||||
num_tokens,
|
||||
hidden,
|
||||
num_experts,
|
||||
num_topk,
|
||||
rank,
|
||||
num_ranks,
|
||||
group,
|
||||
buffer,
|
||||
seed=1,
|
||||
)
|
||||
|
||||
do_pressure_test = False
|
||||
for seed in range(int(1e9) if do_pressure_test else 0):
|
||||
if local_rank == 0:
|
||||
print(f"Testing with seed {seed} ...", flush=True)
|
||||
ref_hash = test_main(
|
||||
num_tokens,
|
||||
hidden,
|
||||
num_experts,
|
||||
num_topk,
|
||||
rank,
|
||||
num_ranks,
|
||||
group,
|
||||
buffer,
|
||||
seed=seed,
|
||||
)
|
||||
for i in range(20):
|
||||
assert (
|
||||
test_main(
|
||||
num_tokens,
|
||||
hidden,
|
||||
num_experts,
|
||||
num_topk,
|
||||
rank,
|
||||
num_ranks,
|
||||
group,
|
||||
buffer,
|
||||
seed=seed,
|
||||
)
|
||||
== ref_hash
|
||||
), f"Error: seed={seed}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# TODO: you may modify NUMA binding for less CPU overhead
|
||||
num_processes = 8
|
||||
torch.multiprocessing.spawn(test_loop, args=(num_processes,), nprocs=num_processes)
|
||||
Executable
+154
@@ -0,0 +1,154 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class _BaseTestDynamicEPLB(CustomTestCase):
|
||||
extra_args = []
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
with (
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.override(False),
|
||||
envs.SGLANG_EXPERT_LOCATION_UPDATER_CANARY.override(True),
|
||||
):
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"2",
|
||||
"--dp",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
"--deepep-mode",
|
||||
"normal",
|
||||
"--disable-cuda-graph",
|
||||
"--enable-eplb",
|
||||
"--ep-num-redundant-experts",
|
||||
"4",
|
||||
"--eplb-rebalance-num-iterations",
|
||||
"50",
|
||||
"--expert-distribution-recorder-buffer-size",
|
||||
"50",
|
||||
# TODO pr-chain: enable later
|
||||
# "--enable-expert-distribution-metrics",
|
||||
# TODO auto determine these flags
|
||||
"--expert-distribution-recorder-mode",
|
||||
"stat",
|
||||
"--ep-dispatch-algorithm",
|
||||
"static",
|
||||
*cls.extra_args,
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreater(metrics["score"], 0.5)
|
||||
|
||||
|
||||
class TestDynamicEPLBSimple(_BaseTestDynamicEPLB):
|
||||
pass
|
||||
|
||||
|
||||
class TestDynamicEPLBMultiChunk(_BaseTestDynamicEPLB):
|
||||
extra_args = ["--eplb-rebalance-layers-per-chunk", "1"]
|
||||
|
||||
|
||||
class TestStaticEPLB(CustomTestCase):
|
||||
def test_save_expert_distribution_and_init_expert_location(self):
|
||||
envs.SGLANG_ENABLE_JIT_DEEPGEMM.set(False)
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir:
|
||||
engine_kwargs = dict(
|
||||
model_path=DEFAULT_MLA_MODEL_NAME_FOR_TEST,
|
||||
trust_remote_code=True,
|
||||
ep_num_redundant_experts=4,
|
||||
enable_dp_attention=True,
|
||||
moe_a2a_backend="deepep",
|
||||
disable_cuda_graph=True,
|
||||
expert_distribution_recorder_mode="stat",
|
||||
tp_size=2,
|
||||
dp_size=2,
|
||||
log_level="info",
|
||||
# TODO pr-chain: enable later
|
||||
# enable_expert_distribution_metrics=True,
|
||||
)
|
||||
|
||||
print(f"Action: start engine")
|
||||
envs.SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR.set(tmp_dir)
|
||||
engine = sgl.Engine(
|
||||
**engine_kwargs,
|
||||
disable_overlap_schedule=True,
|
||||
)
|
||||
engine.start_expert_distribution_record()
|
||||
self._assert_engine_generate_correct(engine)
|
||||
|
||||
print(f"Action: dump_expert_distribution_record")
|
||||
engine.dump_expert_distribution_record()
|
||||
snapshot_path = list(Path(tmp_dir).glob("*.pt"))[0]
|
||||
assert snapshot_path is not None
|
||||
print(f"{snapshot_path=}")
|
||||
|
||||
print(f"Action: shutdown engine")
|
||||
engine.shutdown()
|
||||
del engine
|
||||
|
||||
print(f"Action: start engine with init_expert_location")
|
||||
engine = sgl.Engine(
|
||||
**engine_kwargs,
|
||||
init_expert_location=str(snapshot_path),
|
||||
port=21000,
|
||||
# TODO auto determine these flags
|
||||
ep_dispatch_algorithm="static",
|
||||
)
|
||||
self._assert_engine_generate_correct(engine)
|
||||
print(f"Action: shutdown engine")
|
||||
engine.shutdown()
|
||||
del engine
|
||||
|
||||
def _assert_engine_generate_correct(self, engine: sgl.Engine):
|
||||
output = engine.generate(
|
||||
prompt=["1+1=2, 2+2=4", "One plus one is two, two plus two is four"],
|
||||
sampling_params=dict(max_new_tokens=8, temperature=0.0),
|
||||
)
|
||||
print(f"engine.generate {output=}")
|
||||
self.assertEqual(
|
||||
[x["text"] for x in output],
|
||||
[", 4+4=8,", ", four plus four is eight, eight"],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,116 @@
|
||||
import json
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_MLA_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestPureTP(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"2",
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
"--disable-cuda-graph",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreater(metrics["score"], 0.5)
|
||||
|
||||
|
||||
class TestDPAttn(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_MLA_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
with envs.SGLANG_ENABLE_JIT_DEEPGEMM.override(False):
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"2",
|
||||
"--dp",
|
||||
"2",
|
||||
"--enable-dp-attention",
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
"--deepep-mode",
|
||||
"normal",
|
||||
"--disable-cuda-graph",
|
||||
# Test custom config
|
||||
"--deepep-config",
|
||||
json.dumps(
|
||||
{
|
||||
"normal_dispatch": {
|
||||
"num_sms": 20,
|
||||
"num_max_nvl_chunked_send_tokens": 16,
|
||||
"num_max_nvl_chunked_recv_tokens": 256,
|
||||
"num_max_rdma_chunked_send_tokens": 6,
|
||||
"num_max_rdma_chunked_recv_tokens": 128,
|
||||
},
|
||||
"normal_combine": {
|
||||
"num_sms": 20,
|
||||
"num_max_nvl_chunked_send_tokens": 6,
|
||||
"num_max_nvl_chunked_recv_tokens": 256,
|
||||
"num_max_rdma_chunked_send_tokens": 6,
|
||||
"num_max_rdma_chunked_recv_tokens": 128,
|
||||
},
|
||||
}
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
self.assertGreater(metrics["score"], 0.5)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,75 @@
|
||||
"""
|
||||
Usage:
|
||||
python -m unittest test_moe_deepep_eval_accuracy_large.TestMoEDeepEPEvalAccuracyLarge.test_mmlu
|
||||
"""
|
||||
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
|
||||
from sglang.test.run_eval import run_eval
|
||||
from sglang.test.test_utils import (
|
||||
DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST,
|
||||
DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
DEFAULT_URL_FOR_TEST,
|
||||
CustomTestCase,
|
||||
popen_launch_server,
|
||||
)
|
||||
|
||||
|
||||
class TestMoEDeepEPEvalAccuracyLarge(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.model = DEFAULT_DEEPEP_MODEL_NAME_FOR_TEST
|
||||
cls.base_url = DEFAULT_URL_FOR_TEST
|
||||
cls.process = popen_launch_server(
|
||||
cls.model,
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
"--trust-remote-code",
|
||||
"--tp",
|
||||
"8",
|
||||
"--moe-a2a-backend",
|
||||
"deepep",
|
||||
"--cuda-graph-max-bs",
|
||||
"128",
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
kill_process_tree(cls.process.pid)
|
||||
|
||||
def test_gsm8k(self):
|
||||
args = SimpleNamespace(
|
||||
num_shots=8,
|
||||
data_path=None,
|
||||
num_questions=200,
|
||||
parallel=64,
|
||||
max_new_tokens=512,
|
||||
host="http://127.0.0.1",
|
||||
port=int(self.base_url.split(":")[-1]),
|
||||
)
|
||||
metrics = run_eval_few_shot_gsm8k(args)
|
||||
print(f"Eval accuracy of GSM8K: {metrics=}")
|
||||
|
||||
self.assertGreater(metrics["accuracy"], 0.93)
|
||||
|
||||
def test_mmlu(self):
|
||||
args = SimpleNamespace(
|
||||
base_url=self.base_url,
|
||||
model=self.model,
|
||||
eval_name="mmlu",
|
||||
num_examples=64,
|
||||
num_threads=32,
|
||||
)
|
||||
|
||||
metrics = run_eval(args)
|
||||
print(f"Eval accuracy of MMLU: {metrics=}")
|
||||
self.assertGreater(metrics["score"], 0.87)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,51 @@
|
||||
import unittest
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST, CustomTestCase
|
||||
|
||||
|
||||
class TestBind(CustomTestCase):
|
||||
backend = None
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.backend = sgl.Runtime(model_path=DEFAULT_MODEL_NAME_FOR_TEST)
|
||||
sgl.set_default_backend(cls.backend)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls.backend.shutdown()
|
||||
|
||||
def test_bind(self):
|
||||
@sgl.function
|
||||
def few_shot_qa(s, prompt, question):
|
||||
s += prompt
|
||||
s += "Q: What is the capital of France?\n"
|
||||
s += "A: Paris\n"
|
||||
s += "Q: " + question + "\n"
|
||||
s += "A:" + sgl.gen("answer", stop="\n")
|
||||
|
||||
few_shot_qa_2 = few_shot_qa.bind(
|
||||
prompt="The following are questions with answers.\n\n"
|
||||
)
|
||||
|
||||
tracer = few_shot_qa_2.trace()
|
||||
print(tracer.last_node.print_graph_dfs() + "\n")
|
||||
|
||||
def test_cache(self):
|
||||
@sgl.function
|
||||
def few_shot_qa(s, prompt, question):
|
||||
s += prompt
|
||||
s += "Q: What is the capital of France?\n"
|
||||
s += "A: Paris\n"
|
||||
s += "Q: " + question + "\n"
|
||||
s += "A:" + sgl.gen("answer", stop="\n")
|
||||
|
||||
few_shot_qa_2 = few_shot_qa.bind(
|
||||
prompt="Answer the following questions as if you were a 5-year-old kid.\n\n"
|
||||
)
|
||||
few_shot_qa_2.cache(self.backend)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,91 @@
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from sglang.lang.choices import (
|
||||
greedy_token_selection,
|
||||
token_length_normalized,
|
||||
unconditional_likelihood_normalized,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
MOCK_CHOICES_INPUT_DATA = {
|
||||
"choices": [
|
||||
"organ", # ["organ"]
|
||||
"organism", # ["organ", "ism"]
|
||||
"antidisestablishmentarianism", # ["ant", "id", "is", "est", "ablish", "ment", "arian", "ism"]
|
||||
],
|
||||
"normalized_prompt_logprobs": [-0.1, -0.2, -0.05],
|
||||
"input_token_logprobs": [
|
||||
[[-0.1, 1, None]],
|
||||
[[-0.1, 1, None], [-0.3, 2, None]],
|
||||
[
|
||||
[-0.4, 3, None],
|
||||
[-0.25, 4, None],
|
||||
[-0.1, 5, None],
|
||||
[-0.01, 6, None],
|
||||
[-0.01, 7, None],
|
||||
[-0.01, 8, None],
|
||||
[-0.01, 9, None],
|
||||
[-0.01, 2, None],
|
||||
],
|
||||
],
|
||||
"output_token_logprobs": [
|
||||
[[-0.1, 10, None]],
|
||||
[[-0.1, 10, None]],
|
||||
[[-0.1, 10, None]],
|
||||
],
|
||||
"unconditional_token_logprobs": [
|
||||
[[None, 1, None]],
|
||||
[[None, 1, None], [-1.4, 2, None]],
|
||||
[
|
||||
[None, 3, None],
|
||||
[-0.25, 4, None],
|
||||
[-0.1, 5, None],
|
||||
[-0.01, 6, None],
|
||||
[-0.01, 7, None],
|
||||
[-0.01, 8, None],
|
||||
[-0.01, 9, None],
|
||||
[-0.01, 2, None],
|
||||
],
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class TestChoices(CustomTestCase):
|
||||
|
||||
def test_token_length_normalized(self):
|
||||
"""Confirm 'antidisestablishmentarianism' is selected due to high confidences for
|
||||
its later tokens resulting in highest token length normalized prompt logprob."""
|
||||
decision = token_length_normalized(**MOCK_CHOICES_INPUT_DATA)
|
||||
assert decision.decision == "antidisestablishmentarianism"
|
||||
|
||||
def test_greedy_token_selection(self):
|
||||
"""Confirm 'organ' is selected due it having the joint highest initial token
|
||||
logprob, and a higher average logprob than organism's second token."""
|
||||
decision = greedy_token_selection(**MOCK_CHOICES_INPUT_DATA)
|
||||
assert decision.decision == "organ"
|
||||
assert np.allclose(
|
||||
decision.meta_info["greedy_logprob_matrix"],
|
||||
[
|
||||
[-0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1, -0.1],
|
||||
[-0.1, -0.3, -0.2, -0.2, -0.2, -0.2, -0.2, -0.2],
|
||||
[-0.4, -0.25, -0.1, -0.01, -0.01, -0.01, -0.01, -0.01],
|
||||
],
|
||||
atol=0.01,
|
||||
)
|
||||
|
||||
def test_unconditional_likelihood_normalized(self):
|
||||
"""Confirm 'organism' is selected due to it having the highest average token logprob
|
||||
once normalized by the unconditional token logprobs."""
|
||||
decision = unconditional_likelihood_normalized(**MOCK_CHOICES_INPUT_DATA)
|
||||
assert decision.decision == "organism"
|
||||
assert np.allclose(
|
||||
decision.meta_info["normalized_unconditional_prompt_logprobs"],
|
||||
[-0.1, 0.5, -0.05],
|
||||
atol=0.01,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,138 @@
|
||||
import argparse
|
||||
from enum import Enum
|
||||
|
||||
from pydantic import BaseModel, constr
|
||||
|
||||
import sglang as sgl
|
||||
from sglang.srt.constrained.outlines_backend import build_regex_from_object
|
||||
from sglang.test.test_utils import (
|
||||
add_common_sglang_args_and_parse,
|
||||
select_sglang_backend,
|
||||
)
|
||||
|
||||
IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)"
|
||||
|
||||
ip_jump_forward = (
|
||||
r"The google's DNS sever address is "
|
||||
+ IP_REGEX
|
||||
+ r" and "
|
||||
+ IP_REGEX
|
||||
+ r". "
|
||||
+ r"The google's website domain name is "
|
||||
+ r"www\.(\w)+\.(\w)+"
|
||||
+ r"."
|
||||
)
|
||||
|
||||
|
||||
# fmt: off
|
||||
@sgl.function
|
||||
def regex_gen(s):
|
||||
s += "Q: What is the IP address of the Google DNS servers?\n"
|
||||
s += "A: " + sgl.gen(
|
||||
"answer",
|
||||
max_tokens=128,
|
||||
temperature=0,
|
||||
regex=ip_jump_forward,
|
||||
)
|
||||
# fmt: on
|
||||
|
||||
json_jump_forward = (
|
||||
r"""The information about Hogwarts is in the following JSON format\.\n"""
|
||||
+ r"""\n\{\n"""
|
||||
+ r""" "name": "[\w\d\s]*",\n"""
|
||||
+ r""" "country": "[\w\d\s]*",\n"""
|
||||
+ r""" "latitude": [-+]?[0-9]*\.?[0-9]+,\n"""
|
||||
+ r""" "population": [-+]?[0-9]+,\n"""
|
||||
+ r""" "top 3 landmarks": \["[\w\d\s]*", "[\w\d\s]*", "[\w\d\s]*"\],\n"""
|
||||
+ r"""\}\n"""
|
||||
)
|
||||
|
||||
# fmt: off
|
||||
@sgl.function
|
||||
def json_gen(s):
|
||||
s += sgl.gen(
|
||||
"json",
|
||||
max_tokens=128,
|
||||
temperature=0,
|
||||
regex=json_jump_forward,
|
||||
)
|
||||
# fmt: on
|
||||
|
||||
|
||||
class Weapon(str, Enum):
|
||||
sword = "sword"
|
||||
axe = "axe"
|
||||
mace = "mace"
|
||||
spear = "spear"
|
||||
bow = "bow"
|
||||
crossbow = "crossbow"
|
||||
|
||||
|
||||
class Armor(str, Enum):
|
||||
leather = "leather"
|
||||
chainmail = "chainmail"
|
||||
plate = "plate"
|
||||
|
||||
|
||||
class Character(BaseModel):
|
||||
name: constr(max_length=10)
|
||||
age: int
|
||||
armor: Armor
|
||||
weapon: Weapon
|
||||
strength: int
|
||||
|
||||
|
||||
@sgl.function
|
||||
def character_gen(s):
|
||||
s += "Give me a character description who is a wizard.\n"
|
||||
s += sgl.gen(
|
||||
"character",
|
||||
max_tokens=128,
|
||||
temperature=0,
|
||||
regex=build_regex_from_object(Character),
|
||||
)
|
||||
|
||||
|
||||
def main(args):
|
||||
# Select backend
|
||||
backend = select_sglang_backend(args)
|
||||
sgl.set_default_backend(backend)
|
||||
|
||||
state = regex_gen.run(temperature=0)
|
||||
|
||||
print("=" * 20, "IP TEST", "=" * 20)
|
||||
print(state.text())
|
||||
|
||||
state = json_gen.run(temperature=0)
|
||||
|
||||
print("=" * 20, "JSON TEST", "=" * 20)
|
||||
print(state.text())
|
||||
|
||||
state = character_gen.run(temperature=0)
|
||||
|
||||
print("=" * 20, "CHARACTER TEST", "=" * 20)
|
||||
print(state.text())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
args = add_common_sglang_args_and_parse(parser)
|
||||
main(args)
|
||||
|
||||
# ==================== IP TEST ====================
|
||||
# Q: What is the IP address of the Google DNS servers?
|
||||
# A: The google's DNS sever address is 8.8.8.8 and 8.8.4.4. The google's website domain name is www.google.com.
|
||||
# ==================== JSON TEST ====================
|
||||
# The information about Hogwarts is in the following JSON format.
|
||||
|
||||
# {
|
||||
# "name": "Hogwarts School of Witchcraft and Wizardry",
|
||||
# "country": "Scotland",
|
||||
# "latitude": 55.566667,
|
||||
# "population": 1000,
|
||||
# "top 3 landmarks": ["Hogwarts Castle", "The Great Hall", "The Forbidden Forest"],
|
||||
# }
|
||||
|
||||
# ==================== CHARACTER TEST ====================
|
||||
# Give me a character description who is a wizard.
|
||||
# { "name" : "Merlin", "age" : 500, "armor" : "chainmail" , "weapon" : "sword" , "strength" : 10 }
|
||||
@@ -0,0 +1,92 @@
|
||||
import unittest
|
||||
|
||||
from sglang import OpenAI, set_default_backend
|
||||
from sglang.test.test_programs import (
|
||||
test_chat_completion_speculative,
|
||||
test_completion_speculative,
|
||||
test_decode_int,
|
||||
test_decode_json,
|
||||
test_expert_answer,
|
||||
test_few_shot_qa,
|
||||
test_image_qa,
|
||||
test_mt_bench,
|
||||
test_parallel_decoding,
|
||||
test_parallel_encoding,
|
||||
test_react,
|
||||
test_select,
|
||||
test_stream,
|
||||
test_tool_use,
|
||||
)
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestOpenAIBackend(CustomTestCase):
|
||||
instruct_backend = None
|
||||
chat_backend = None
|
||||
chat_vision_backend = None
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.instruct_backend = OpenAI("gpt-3.5-turbo-instruct")
|
||||
cls.chat_backend = OpenAI("gpt-3.5-turbo")
|
||||
cls.chat_vision_backend = OpenAI("gpt-4-turbo")
|
||||
|
||||
def test_few_shot_qa(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_few_shot_qa()
|
||||
|
||||
def test_mt_bench(self):
|
||||
set_default_backend(self.chat_backend)
|
||||
test_mt_bench()
|
||||
|
||||
def test_select(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_select(check_answer=True)
|
||||
|
||||
def test_decode_int(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_decode_int()
|
||||
|
||||
def test_decode_json(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_decode_json()
|
||||
|
||||
def test_expert_answer(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_expert_answer()
|
||||
|
||||
def test_tool_use(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_tool_use()
|
||||
|
||||
def test_react(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_react()
|
||||
|
||||
def test_parallel_decoding(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_parallel_decoding()
|
||||
|
||||
def test_parallel_encoding(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_parallel_encoding()
|
||||
|
||||
def test_image_qa(self):
|
||||
set_default_backend(self.chat_vision_backend)
|
||||
test_image_qa()
|
||||
|
||||
def test_stream(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_stream()
|
||||
|
||||
def test_completion_speculative(self):
|
||||
set_default_backend(self.instruct_backend)
|
||||
test_completion_speculative()
|
||||
|
||||
def test_chat_completion_speculative(self):
|
||||
set_default_backend(self.chat_backend)
|
||||
test_chat_completion_speculative()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,68 @@
|
||||
"""
|
||||
Tests for the separate_reasoning functionality in sglang.
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test/lang/test_separate_reasoning.py
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
from sglang import gen, separate_reasoning
|
||||
from sglang.lang.ir import SglExprList, SglSeparateReasoning
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class TestSeparateReasoning(CustomTestCase):
|
||||
def test_separate_reasoning_creation(self):
|
||||
"""Test that SglSeparateReasoning objects are created correctly."""
|
||||
# Test with valid model type and gen expression
|
||||
test_gen = gen("test")
|
||||
expr = separate_reasoning(test_gen, model_type="deepseek-r1")
|
||||
self.assertIsInstance(expr, SglExprList)
|
||||
self.assertEqual(len(expr.expr_list), 2)
|
||||
self.assertEqual(expr.expr_list[0], test_gen)
|
||||
reasoning_expr = expr.expr_list[1]
|
||||
self.assertIsInstance(reasoning_expr, SglSeparateReasoning)
|
||||
self.assertEqual(reasoning_expr.model_type, "deepseek-r1")
|
||||
self.assertEqual(reasoning_expr.name, "test_reasoning_content")
|
||||
|
||||
# Test with another valid model type
|
||||
expr = separate_reasoning(test_gen, model_type="qwen3")
|
||||
self.assertIsInstance(expr, SglExprList)
|
||||
self.assertEqual(expr.expr_list[1].model_type, "qwen3")
|
||||
|
||||
def test_separate_reasoning_name_processing(self):
|
||||
"""Test that separate_reasoning correctly processes names."""
|
||||
test_gen = gen("test_var")
|
||||
expr = separate_reasoning(test_gen, model_type="deepseek-r1")
|
||||
reasoning_expr = expr.expr_list[1]
|
||||
self.assertEqual(reasoning_expr.name, "test_var_reasoning_content")
|
||||
|
||||
# Test the process_name_for_reasoning method
|
||||
self.assertEqual(
|
||||
reasoning_expr.process_name_for_reasoning("another_var"),
|
||||
"another_var_reasoning_content",
|
||||
)
|
||||
|
||||
def test_separate_reasoning_repr(self):
|
||||
"""Test the string representation of SglSeparateReasoning."""
|
||||
test_gen = gen("test_var")
|
||||
expr = separate_reasoning(test_gen, model_type="deepseek-r1")
|
||||
reasoning_expr = expr.expr_list[1]
|
||||
self.assertEqual(
|
||||
repr(reasoning_expr),
|
||||
"SeparateReasoning(model_type=deepseek-r1, name=test_var_reasoning_content)",
|
||||
)
|
||||
|
||||
def test_separate_reasoning_with_invalid_model_type(self):
|
||||
"""Test that separate_reasoning accepts any model type during creation."""
|
||||
# Create with invalid model type
|
||||
test_gen = gen("test")
|
||||
expr = separate_reasoning(test_gen, model_type="invalid-model")
|
||||
self.assertIsInstance(expr, SglExprList)
|
||||
self.assertEqual(expr.expr_list[1].model_type, "invalid-model")
|
||||
# The actual validation happens in the ReasoningParser constructor
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,193 @@
|
||||
"""
|
||||
Tests for the execution of separate_reasoning functionality in sglang.
|
||||
|
||||
Usage:
|
||||
python3 -m unittest test/lang/test_separate_reasoning_execution.py
|
||||
"""
|
||||
|
||||
import threading
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from sglang.lang.interpreter import StreamExecutor
|
||||
from sglang.lang.ir import SglGen, SglSeparateReasoning
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
# Helper function to create events that won't block program exit
|
||||
def create_daemon_event():
|
||||
event = threading.Event()
|
||||
return event
|
||||
|
||||
|
||||
class MockReasoningParser:
|
||||
def __init__(self, model_type):
|
||||
self.model_type = model_type
|
||||
self.parse_non_stream_called = False
|
||||
self.parse_stream_chunk_called = False
|
||||
|
||||
def parse_non_stream(self, full_text):
|
||||
self.parse_non_stream_called = True
|
||||
# Simulate parsing by adding a prefix to indicate reasoning
|
||||
reasoning = f"[REASONING from {self.model_type}]: {full_text}"
|
||||
normal_text = f"[NORMAL from {self.model_type}]: {full_text}"
|
||||
return reasoning, normal_text
|
||||
|
||||
def parse_stream_chunk(self, chunk_text):
|
||||
self.parse_stream_chunk_called = True
|
||||
# Simulate parsing by adding a prefix to indicate reasoning
|
||||
reasoning = f"[REASONING from {self.model_type}]: {chunk_text}"
|
||||
normal_text = f"[NORMAL from {self.model_type}]: {chunk_text}"
|
||||
return reasoning, normal_text
|
||||
|
||||
|
||||
class TestSeparateReasoningExecution(CustomTestCase):
|
||||
def setUp(self):
|
||||
"""Set up for the test."""
|
||||
super().setUp()
|
||||
# Store any events created during the test
|
||||
self.events = []
|
||||
|
||||
def tearDown(self):
|
||||
"""Clean up any threads that might have been created during the test."""
|
||||
super().tearDown()
|
||||
|
||||
# Set all events to ensure any waiting threads are released
|
||||
for event in self.events:
|
||||
event.set()
|
||||
|
||||
def tearDown(self):
|
||||
super().tearDown()
|
||||
# wake up all threads
|
||||
for ev in self.events:
|
||||
ev.set()
|
||||
|
||||
@patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
|
||||
def test_execute_separate_reasoning(self, mock_parser_class):
|
||||
"""Test that _execute_separate_reasoning correctly calls the ReasoningParser."""
|
||||
# Setup mock parser
|
||||
mock_parser = MockReasoningParser("deepseek-r1")
|
||||
mock_parser_class.return_value = mock_parser
|
||||
|
||||
# Create a mock backend to avoid AttributeError in __del__
|
||||
mock_backend = MagicMock()
|
||||
|
||||
# Create a StreamExecutor with necessary setup
|
||||
executor = StreamExecutor(
|
||||
backend=mock_backend,
|
||||
arguments={},
|
||||
default_sampling_para={},
|
||||
chat_template={
|
||||
"role_map": {"user": "user", "assistant": "assistant"}
|
||||
}, # Simple chat template
|
||||
stream=False,
|
||||
use_thread=False,
|
||||
)
|
||||
|
||||
# Set up the executor with a variable and its value
|
||||
var_name = "test_var"
|
||||
reasoning_name = f"{var_name}_reasoning_content"
|
||||
var_value = "Test content"
|
||||
executor.variables = {var_name: var_value}
|
||||
|
||||
# Create events and track them for cleanup
|
||||
var_event = create_daemon_event()
|
||||
reasoning_event = create_daemon_event()
|
||||
self.events.extend([var_event, reasoning_event])
|
||||
|
||||
executor.variable_event = {var_name: var_event, reasoning_name: reasoning_event}
|
||||
executor.variable_event[var_name].set() # Mark as ready
|
||||
|
||||
# Set up the current role
|
||||
executor.cur_role = "assistant"
|
||||
executor.cur_role_begin_pos = 0
|
||||
executor.text_ = var_value
|
||||
|
||||
# Create a gen expression and a separate_reasoning expression
|
||||
gen_expr = SglGen(var_name)
|
||||
expr = SglSeparateReasoning("deepseek-r1", expr=gen_expr)
|
||||
|
||||
# Execute separate_reasoning
|
||||
executor._execute_separate_reasoning(expr)
|
||||
|
||||
# Verify that the parser was created with the correct model type
|
||||
mock_parser_class.assert_called_once_with("deepseek-r1")
|
||||
|
||||
# Verify that parse_non_stream was called
|
||||
self.assertTrue(mock_parser.parse_non_stream_called)
|
||||
|
||||
# Verify that the variables were updated correctly
|
||||
reasoning_name = f"{var_name}_reasoning_content"
|
||||
self.assertIn(reasoning_name, executor.variables)
|
||||
self.assertEqual(
|
||||
executor.variables[reasoning_name],
|
||||
f"[REASONING from deepseek-r1]: {var_value}",
|
||||
)
|
||||
self.assertEqual(
|
||||
executor.variables[var_name], f"[NORMAL from deepseek-r1]: {var_value}"
|
||||
)
|
||||
|
||||
# Verify that the variable event was set
|
||||
self.assertIn(reasoning_name, executor.variable_event)
|
||||
self.assertTrue(executor.variable_event[reasoning_name].is_set())
|
||||
|
||||
# Verify that the text was updated
|
||||
self.assertEqual(executor.text_, f"[NORMAL from deepseek-r1]: {var_value}")
|
||||
|
||||
@patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
|
||||
def test_reasoning_parser_integration(self, mock_parser_class):
|
||||
"""Test the integration between separate_reasoning and ReasoningParser."""
|
||||
# Setup mock parsers for different model types
|
||||
deepseek_parser = MockReasoningParser("deepseek-r1")
|
||||
qwen_parser = MockReasoningParser("qwen3")
|
||||
|
||||
# Configure the mock to return different parsers based on model type
|
||||
def get_parser(model_type):
|
||||
if model_type == "deepseek-r1":
|
||||
return deepseek_parser
|
||||
elif model_type == "qwen3":
|
||||
return qwen_parser
|
||||
else:
|
||||
raise ValueError(f"Unsupported model type: {model_type}")
|
||||
|
||||
mock_parser_class.side_effect = get_parser
|
||||
|
||||
# Test with DeepSeek-R1 model
|
||||
test_text = "This is a test"
|
||||
reasoning, normal_text = deepseek_parser.parse_non_stream(test_text)
|
||||
|
||||
self.assertEqual(reasoning, f"[REASONING from deepseek-r1]: {test_text}")
|
||||
self.assertEqual(normal_text, f"[NORMAL from deepseek-r1]: {test_text}")
|
||||
|
||||
# Test with Qwen3 model
|
||||
reasoning, normal_text = qwen_parser.parse_non_stream(test_text)
|
||||
|
||||
self.assertEqual(reasoning, f"[REASONING from qwen3]: {test_text}")
|
||||
self.assertEqual(normal_text, f"[NORMAL from qwen3]: {test_text}")
|
||||
|
||||
@patch("sglang.srt.parser.reasoning_parser.ReasoningParser")
|
||||
def test_reasoning_parser_invalid_model(self, mock_parser_class):
|
||||
"""Test that ReasoningParser raises an error for invalid model types."""
|
||||
|
||||
# Configure the mock to raise an error for invalid model types
|
||||
def get_parser(model_type):
|
||||
if model_type in ["deepseek-r1", "qwen3"]:
|
||||
return MockReasoningParser(model_type)
|
||||
elif model_type is None:
|
||||
raise ValueError("Model type must be specified")
|
||||
else:
|
||||
raise ValueError(f"Unsupported model type: {model_type}")
|
||||
|
||||
mock_parser_class.side_effect = get_parser
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
mock_parser_class("invalid-model")
|
||||
self.assertIn("Unsupported model type", str(context.exception))
|
||||
|
||||
with self.assertRaises(ValueError) as context:
|
||||
mock_parser_class(None)
|
||||
self.assertIn("Model type must be specified", str(context.exception))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user