microsoft/mscclpp
1
0
Fork 0
mirror of https://github.com/microsoft/mscclpp.git synced 2026-07-14 19:27:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

merge main

Browse Source
This commit is contained in:
Binyang Li
2026-07-13 22:01:36 +00:00
parent 359e681e8e 9b42ad3062
commit bc4374e650
102 changed files with 4323 additions and 576 deletions
Download Patch File Download Diff File Expand all files Collapse all files

105
python/test/executor_test.py
View File

@@ -14,6 +14,7 @@ from mscclpp import CommGroup, GpuBuffer
from mscclpp.utils import KernelBuilder, pack
import os
import struct
from typing import Callable
import cupy as cp
from mpi4py import MPI
@@ -34,13 +35,13 @@ def parse_dtype(dtype_str):
raise ValueError(f"Unknown data type: {dtype_str}")
def bench_time(n_iters: int, n_graph_iters: int, func):
# capture cuda graph for n_iters of the kernel launch
def bench_time(n_iters: int, n_graph_iters: int, funcs: list[Callable]):
"""Benchmark execution time. `funcs` is a list of callables; iteration i runs funcs[i % len(funcs)]."""
stream = cp.cuda.Stream(non_blocking=True)
with stream:
stream.begin_capture()
for i in range(n_iters):
func(stream)
funcs[i % len(funcs)](stream)
graph = stream.end_capture()
# now run a warm up round
@@ -61,15 +62,17 @@ def bench_time(n_iters: int, n_graph_iters: int, func):
def bench_correctness(
collective: str,
input_buf: cp.ndarray,
result_buf: cp.ndarray,
test_buf: cp.ndarray,
input_bufs: list[cp.ndarray],
result_bufs: list[cp.ndarray],
test_bufs: list[cp.ndarray],
dtype_str: str,
rank: int,
num_ranks: int,
n_iters: int,
func,
funcs: list[Callable],
split_mask: int = 0,
):
"""Validate correctness. Buffers and funcs are parallel lists; iteration i uses index i % len(funcs)."""
type_size = cp.dtype(parse_dtype(dtype_str)).itemsize
fill_data_kernel_name = "fill_data_%s" % dtype_str
@@ -79,8 +82,12 @@ def bench_correctness(
coll = "reduce_scatter"
elif "allreduce" in collective:
coll = "all_reduce"
else:
elif "alltoall" in collective:
coll = "all_to_all"
elif "sendrecv" in collective:
coll = "send_recv"
else:
raise ValueError(f"Unknown collective: {collective}")
test_data_kernel_name = "test_data_%s_%s" % (coll, dtype_str)
file_dir = os.path.dirname(os.path.abspath(__file__))
@@ -97,11 +104,20 @@ def bench_correctness(
with stream:
stream.begin_capture()
for i in range(n_iters):
fill_data_params = pack(input_buf) + struct.pack("Q", input_buf.nbytes // type_size) + pack(rank, i)
idx = i % len(funcs)
cur_input = input_bufs[idx]
cur_result = result_bufs[idx]
cur_test = test_bufs[idx]
fill_data_params = (
pack(cur_input) + struct.pack("Q", cur_input.nbytes // type_size) + pack(rank, i, split_mask)
)
fill_data_kernel.launch_kernel(fill_data_params, nblocks, nthreads, 0, stream)
func(stream)
funcs[idx](stream)
test_data_params = (
pack(result_buf, test_buf) + struct.pack("Q", input_buf.nbytes // type_size) + pack(num_ranks, rank, i)
pack(cur_result, cur_test)
+ struct.pack("Q", cur_input.nbytes // type_size)
+ pack(num_ranks, rank, i, split_mask)
)
test_data_kernel.launch_kernel(test_data_params, nblocks, nthreads, 0, stream)
graph = stream.end_capture()
@@ -143,10 +159,20 @@ def build_bufs(
rank: int,
num_ranks: int,
):
"""Allocate input/result/test buffers. Returns parallel lists (length 2 for sendrecv double-buffering,
length 1 otherwise) so callers can iterate uniformly."""
type_size = cp.dtype(dtype).itemsize
assert (size % type_size) == 0, "size %d not multiple of type size %d" % (size, type_size)
nelems = size // type_size
# Sendrecv uses double buffering: build two parallel buffer slots.
if "sendrecv" in collective:
n_slots = 2
input_bufs = [GpuBuffer(nelems, dtype=dtype) for _ in range(n_slots)]
result_bufs = [GpuBuffer(nelems, dtype=dtype) for _ in range(n_slots)]
test_bufs = [cp.zeros(nelems, dtype=dtype) for _ in range(n_slots)]
return input_bufs, result_bufs, test_bufs, nelems
if "allgather" in collective:
assert (nelems % num_ranks) == 0, "nelems %d not multiple of num_ranks %d" % (nelems, num_ranks)
nelems_input = nelems if in_place else nelems // num_ranks
@@ -173,7 +199,7 @@ def build_bufs(
test_buf = cp.zeros(nelems, dtype=dtype)
return input_buf, result_buf, test_buf
return [input_buf], [result_buf], [test_buf], nelems
def main(
@@ -184,8 +210,14 @@ def main(
packet_type: PacketType = PacketType.LL16,
n_iters: int = 10,
n_graph_iters: int = 10,
split_mask: int = 0,
):
mscclpp_group = CommGroup(MPI.COMM_WORLD)
if split_mask < 0 or (split_mask & (split_mask + 1)) != 0 or mscclpp_group.nranks % (split_mask + 1) != 0:
raise ValueError(
f"split_mask must be of the form 2^k - 1 and nranks ({mscclpp_group.nranks}) must be divisible "
f"by group_size ({split_mask + 1}), got split_mask={hex(split_mask)}"
)
cp.cuda.Device(mscclpp_group.my_rank % mscclpp_group.nranks_per_node).use()
executor = Executor(mscclpp_group.communicator)
npkit_dump_dir = env().npkit_dump_dir
@@ -195,7 +227,7 @@ def main(
collective = execution_plan.collective
dtype = parse_dtype(dtype_str)
input_buf, result_buf, test_buf = build_bufs(
input_bufs, result_bufs, test_bufs, nelem = build_bufs(
collective,
size,
in_place,
@@ -204,39 +236,48 @@ def main(
mscclpp_group.nranks,
)
executor_func = lambda stream: executor.execute(
mscclpp_group.my_rank,
input_buf.data.ptr,
result_buf.data.ptr,
input_buf.nbytes,
result_buf.nbytes,
dtype_to_mscclpp_dtype(dtype_str),
execution_plan,
stream.ptr,
packet_type,
)
executor_funcs = [
(
lambda stream, inp=inp, res=res: executor.execute(
mscclpp_group.my_rank,
inp.data.ptr,
res.data.ptr,
inp.nbytes,
res.nbytes,
dtype_to_mscclpp_dtype(dtype_str),
execution_plan,
stream.ptr,
packet_type,
)
)
for inp, res in zip(input_bufs, result_bufs)
]
mscclpp_group.barrier()
bench_correctness(
collective,
input_buf,
result_buf,
test_buf,
input_bufs,
result_bufs,
test_bufs,
dtype_str,
mscclpp_group.my_rank,
mscclpp_group.nranks,
n_iters,
executor_func,
executor_funcs,
split_mask=split_mask,
)
mscclpp_group.barrier()
execution_time = bench_time(n_iters, n_graph_iters, executor_func)
execution_time = bench_time(n_iters, n_graph_iters, executor_funcs)
if npkit_dump_dir is not None:
npkit.dump(npkit_dump_dir)
npkit.shutdown()
result_nbytes = result_bufs[0].nbytes
print(
f"Rank: {mscclpp_group.my_rank} Execution time: {execution_time} us, "
f"data size: {result_buf.nbytes} bytes data type: {dtype_str} "
f"data size: {result_nbytes} bytes data type: {dtype_str} "
f"bandwidth: {result_nbytes / (execution_time * 1e-6) / (1024**3):.2f} GB/s, "
f"packet type: {packet_type}"
)
executor = None
@@ -252,6 +293,9 @@ if __name__ == "__main__":
parser.add_argument("--packet_type", type=str, default="LL16", help="Choose from LL8, LL16")
parser.add_argument("--n_iters", type=int, default=10)
parser.add_argument("--n_graph_iters", type=int, default=10)
parser.add_argument(
"--split_mask", type=lambda x: int(x, 0), default=0x0, help="split mask for sendrecv (e.g. 0x3)"
)
args = parser.parse_args()
packet_type = PacketType.LL16
@@ -267,4 +311,5 @@ if __name__ == "__main__":
packet_type,
args.n_iters,
args.n_graph_iters,
args.split_mask,
)

155
python/test/executor_test_verifier.cu
View File

@@ -22,14 +22,19 @@ static __device__ unsigned int ranqd1(unsigned int seed) {
// fill/test kernel pairs must have the same thread block size to
// match their random number series.
#define FILL_DATA(FuncNameType, DataType) \
extern "C" __global__ void __launch_bounds__(1024, 1) \
fill_data_##FuncNameType(DataType* input_buf, size_t num_elems, int rank, int seq) { \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
seed = ranqd1(seed); \
input_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \
} \
// `split_mask` groups ranks together: group_size = split_mask + 1, group_id = rank / group_size.
// Data is seeded by group_id so that all ranks within a group produce the same fill, and ranks
// in different groups produce different fills. With split_mask == 0 this reduces to per-rank
// seeding (group_id == rank).
#define FILL_DATA(FuncNameType, DataType) \
extern "C" __global__ void __launch_bounds__(1024, 1) \
fill_data_##FuncNameType(DataType* input_buf, size_t num_elems, int rank, int seq, int split_mask) { \
int seed_rank = rank / (split_mask + 1); \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
seed = ranqd1(seed); \
input_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \
} \
}
FILL_DATA(bfloat16, __nv_bfloat16)
@@ -37,18 +42,20 @@ FILL_DATA(float16, __half)
FILL_DATA(float32, float)
FILL_DATA(int32, int)
#define TEST_DATA_ALL_GATHER(FuncNameType, DataType) \
extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_gather_##FuncNameType( \
DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) { \
for (int rank = 0; rank < num_ranks; rank++) { \
size_t rank_offset = rank * num_elems; \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
seed = ranqd1(seed); \
test_buf[rank_offset + i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \
assert(result_buf[rank_offset + i] == test_buf[rank_offset + i]); \
} \
} \
#define TEST_DATA_ALL_GATHER(FuncNameType, DataType) \
extern "C" __global__ void __launch_bounds__(1024, 1) \
test_data_all_gather_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \
int my_rank, int seq, int split_mask) { \
for (int rank = 0; rank < num_ranks; rank++) { \
size_t rank_offset = rank * num_elems; \
int seed_rank = rank / (split_mask + 1); \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
seed = ranqd1(seed); \
test_buf[rank_offset + i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \
assert(result_buf[rank_offset + i] == test_buf[rank_offset + i]); \
} \
} \
}
TEST_DATA_ALL_GATHER(bfloat16, __nv_bfloat16)
@@ -56,25 +63,27 @@ TEST_DATA_ALL_GATHER(float16, __half)
TEST_DATA_ALL_GATHER(float32, float)
TEST_DATA_ALL_GATHER(int32, int)
#define TEST_DATA_ALL_REDUCE(FuncNameType, DataType, Eps) \
extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_reduce_##FuncNameType( \
DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) { \
for (int rank = 0; rank < num_ranks; rank++) { \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
if (rank == 0) { \
test_buf[i] = 0; \
} \
seed = ranqd1(seed); \
test_buf[i] += DataType(seed % blockDim.x) / DataType(blockDim.x); \
} \
} \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
float expected = float(test_buf[i]); \
float result = float(result_buf[i]); \
float tol = Eps * num_ranks * (1.0f + abs(expected)); \
assert(abs(result - expected) <= tol); \
} \
#define TEST_DATA_ALL_REDUCE(FuncNameType, DataType, Eps) \
extern "C" __global__ void __launch_bounds__(1024, 1) \
test_data_all_reduce_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \
int my_rank, int seq, int split_mask) { \
for (int rank = 0; rank < num_ranks; rank++) { \
int seed_rank = rank / (split_mask + 1); \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
if (rank == 0) { \
test_buf[i] = 0; \
} \
seed = ranqd1(seed); \
test_buf[i] += DataType(seed % blockDim.x) / DataType(blockDim.x); \
} \
} \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
float expected = float(test_buf[i]); \
float result = float(result_buf[i]); \
float tol = Eps * num_ranks * (1.0f + abs(expected)); \
assert(abs(result - expected) <= tol); \
} \
}
TEST_DATA_ALL_REDUCE(bfloat16, __nv_bfloat16, 7.8125e-3f)
@@ -83,12 +92,14 @@ TEST_DATA_ALL_REDUCE(float32, float, 1.1920929e-7f)
TEST_DATA_ALL_REDUCE(int32, int, 0.0f)
#define TEST_DATA_REDUCE_SCATTER(FuncNameType, DataType, Eps) \
extern "C" __global__ void __launch_bounds__(1024, 1) test_data_reduce_scatter_##FuncNameType( \
DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) { \
extern "C" __global__ void __launch_bounds__(1024, 1) \
test_data_reduce_scatter_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, \
int num_ranks, int my_rank, int seq, int split_mask) { \
int nem_elems_per_rank = num_elems / num_ranks; \
int offset = nem_elems_per_rank * my_rank; \
for (int rank = 0; rank < num_ranks; rank++) { \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \
int seed_rank = rank / (split_mask + 1); \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
if (rank == 0) { \
test_buf[i] = 0; \
@@ -112,25 +123,51 @@ TEST_DATA_REDUCE_SCATTER(float16, __half, 9.765625e-4f)
TEST_DATA_REDUCE_SCATTER(float32, float, 1.1920929e-7f)
TEST_DATA_REDUCE_SCATTER(int32, int, 0.0f)
#define TEST_DATA_ALL_TO_ALL(FuncNameType, DataType) \
extern "C" __global__ void __launch_bounds__(1024, 1) test_data_all_to_all_##FuncNameType( \
DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, int my_rank, int seq) { \
int nem_elems_per_rank = num_elems / num_ranks; \
int offset = nem_elems_per_rank * my_rank; \
for (int rank = 0; rank < num_ranks; rank++) { \
size_t rank_offset = rank * nem_elems_per_rank; \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
seed = ranqd1(seed); \
if (i >= my_rank * nem_elems_per_rank && i < (my_rank + 1) * nem_elems_per_rank) { \
test_buf[rank_offset + i - offset] = DataType(seed % blockDim.x) / DataType(blockDim.x); \
assert(result_buf[rank_offset + i - offset] == test_buf[rank_offset + i - offset]); \
} \
} \
} \
#define TEST_DATA_ALL_TO_ALL(FuncNameType, DataType) \
extern "C" __global__ void __launch_bounds__(1024, 1) \
test_data_all_to_all_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \
int my_rank, int seq, int split_mask) { \
int nem_elems_per_rank = num_elems / num_ranks; \
int offset = nem_elems_per_rank * my_rank; \
for (int rank = 0; rank < num_ranks; rank++) { \
size_t rank_offset = rank * nem_elems_per_rank; \
int seed_rank = rank / (split_mask + 1); \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + seed_rank + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
seed = ranqd1(seed); \
if (i >= my_rank * nem_elems_per_rank && i < (my_rank + 1) * nem_elems_per_rank) { \
test_buf[rank_offset + i - offset] = DataType(seed % blockDim.x) / DataType(blockDim.x); \
assert(result_buf[rank_offset + i - offset] == test_buf[rank_offset + i - offset]); \
} \
} \
} \
}
TEST_DATA_ALL_TO_ALL(bfloat16, __nv_bfloat16)
TEST_DATA_ALL_TO_ALL(float16, __half)
TEST_DATA_ALL_TO_ALL(float32, float)
TEST_DATA_ALL_TO_ALL(int32, int)
TEST_DATA_ALL_TO_ALL(int32, int)
// Sendrecv verification: receive from the prev group in the ring.
// fill_data seeds by group_id (rank / (split_mask + 1)); the receiver in group g expects the
// data produced by group (g - 1 + num_groups) % num_groups, so we recompute that seed here.
#define TEST_DATA_SEND_RECV(FuncNameType, DataType) \
extern "C" __global__ void __launch_bounds__(1024, 1) \
test_data_send_recv_##FuncNameType(DataType* result_buf, DataType* test_buf, size_t num_elems, int num_ranks, \
int my_rank, int seq, int split_mask) { \
int group_size = split_mask + 1; \
int num_groups = num_ranks / group_size; \
int my_group_id = my_rank / group_size; \
int prev_group_id = (my_group_id - 1 + num_groups) % num_groups; \
unsigned int seed = (unsigned int)(blockIdx.x * blockDim.x + threadIdx.x + prev_group_id + seq); \
for (size_t i = blockIdx.x * blockDim.x + threadIdx.x; i < num_elems; i += blockDim.x * gridDim.x) { \
seed = ranqd1(seed); \
test_buf[i] = DataType(seed % blockDim.x) / DataType(blockDim.x); \
assert(result_buf[i] == test_buf[i]); \
} \
}
TEST_DATA_SEND_RECV(bfloat16, __nv_bfloat16)
TEST_DATA_SEND_RECV(float16, __half)
TEST_DATA_SEND_RECV(float32, float)
TEST_DATA_SEND_RECV(int32, int)

10
python/test/test_fp8_accum.py
View File

@@ -167,7 +167,7 @@ else:
# ---------------------------------------------------------------------------
# FP8 E4M3B15 helpers (bias=15, encode saturates to ±1.75, no NaN)
# FP8 E4M3B15 helpers (bias=15, float source saturates to ±1.875, no NaN)
# Matches Triton's fp8e4b15: all 256 bit patterns are finite.
# ---------------------------------------------------------------------------
@@ -193,7 +193,7 @@ def float_to_e4m3b15(f32_array, chunk_size=65536):
"""Encode a cupy float32 array to uint8 E4M3B15 bit patterns.
Same lookup-table approach as float_to_e4m3fn.
Saturates to ±1.75 (0x7e/0xfe), matching Triton's fp8e4b15.
Saturates to ±1.875 (0x7f/0xff), matching the device float32 → e4m3b15 path.
"""
# Build lookup table of all 128 positive E4M3B15 values (0x00..0x7F)
all_bytes = cp.arange(128, dtype=cp.uint8)
@@ -203,7 +203,7 @@ def float_to_e4m3b15(f32_array, chunk_size=65536):
values = f32_array.astype(cp.float32)
signs = cp.signbit(values).astype(cp.uint8)
absval = cp.abs(values)
absval = cp.clip(absval, cp.float32(0.0), cp.float32(1.75))
absval = cp.clip(absval, cp.float32(0.0), cp.float32(1.875))
result = cp.zeros(absval.shape, dtype=cp.uint8)
n = absval.size
@@ -442,8 +442,8 @@ def test_fp8_e4m3b15_accum(mpi_group: MpiGroup, algo_name: str, size: int):
bits_r = cp.asarray(rng_r.randint(0, 256, (size,)).astype(np.uint8))
ref_f32 += e4m3b15_to_float(bits_r)
# Clamp reference to e4m3b15 representable range
ref_f32 = cp.clip(ref_f32, -1.75, 1.75)
# Clamp reference to e4m3b15 representable range (float source saturates at ±1.875)
ref_f32 = cp.clip(ref_f32, -1.875, 1.875)
# Compute errors
abs_err = cp.abs(result_f32 - ref_f32)