455 lines
18 KiB
Python
455 lines
18 KiB
Python
"""Real-CUDA multi-rank distributed test for LPLBSolver.
|
|
|
|
Supplements the CPU-only `test_lplb.py` checks: this test spawns 2
|
|
CUDA-bound ranks on a single node via `torch.multiprocessing.spawn` and
|
|
exercises the actual code paths the production LP server uses, including
|
|
the EP all-reduce and the DP-attention "empty rank" case (rank with zero
|
|
tokens for a given forward).
|
|
|
|
Critical invariants this test catches that the CPU tests cannot:
|
|
1. The all-reduce inside `LPLBSolver.solve()` actually completes when
|
|
one rank has empty `topk_ids` (i.e., the empty-rank deadlock fix is
|
|
in effect — without it, the empty rank would skip `solver.solve()`
|
|
and hang the collective).
|
|
2. Both ranks receive the **same** counts post-all-reduce and produce
|
|
identical outputs.
|
|
3. The output equals the torch-IPM oracle computed from the ORACLE
|
|
summed counts (this is the discriminator that tells "all-reduce
|
|
happened" apart from "rank silently used local-only counts and
|
|
coincidentally produced a finite tensor").
|
|
4. The fused-only contract holds end-to-end: a fused-kernel failure
|
|
must surface as an exception (no silent torch fallback exists).
|
|
5. Post-rebalance reinit produces solvers tied to the new metadata —
|
|
outputs change when `phy2log` changes.
|
|
|
|
Pattern follows `test/registered/layers/mamba/test_mamba2_mixer.py` —
|
|
`torch.multiprocessing.spawn` is required because `torch.distributed`
|
|
plays poorly with arbitrary subprocess launchers.
|
|
"""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
|
|
update_environment_variables,
|
|
)
|
|
from sglang.srt.distributed.parallel_state import (
|
|
init_distributed_environment,
|
|
initialize_model_parallel,
|
|
)
|
|
from sglang.test.ci.ci_register import register_cuda_ci
|
|
from sglang.test.test_utils import publish_build_topology
|
|
|
|
register_cuda_ci(est_time=18, stage="base-b", runner_config="2-gpu-large")
|
|
|
|
NUM_GPUS = 2
|
|
TOPK = 2
|
|
|
|
# Smallest replicated-metadata config that still exercises a non-trivial LP:
|
|
# num_gpus = 2
|
|
# num_logical = 4 (3 single-copy + 1 replicated)
|
|
# num_phy = 6 (4 + 2 redundant copies of expert 0)
|
|
# → NC = 1 + 2 = 3, NV = 2 + 2 + 2 = 6 (well within H20 shmem budget)
|
|
NUM_LOGICAL = 4
|
|
NUM_PHY = 6
|
|
|
|
|
|
def _make_metadata(rebalanced: bool = False):
|
|
"""Build (phy2log, log2phy, num_valid) for the test config.
|
|
|
|
When `rebalanced=True`, swap the redundant logical from expert 0 to
|
|
expert 1 — same shapes, different mapping, so a stale solver would
|
|
produce wrong outputs.
|
|
"""
|
|
replicated_logical = 1 if rebalanced else 0
|
|
phy2log = torch.tensor(
|
|
[0, 1, 2, 3, replicated_logical, replicated_logical], dtype=torch.int64
|
|
)
|
|
log2phy = torch.full((NUM_LOGICAL, 3), -1, dtype=torch.int64)
|
|
for i in range(NUM_LOGICAL):
|
|
log2phy[i, 0] = i
|
|
log2phy[replicated_logical, 1] = 4
|
|
log2phy[replicated_logical, 2] = 5
|
|
num_valid = torch.ones(NUM_LOGICAL, dtype=torch.int64)
|
|
num_valid[replicated_logical] = 3
|
|
return phy2log, log2phy, num_valid
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="This test requires CUDA",
|
|
)
|
|
def test_dispatch_probability_matches_torch_reference():
|
|
"""The fused CUDA `dispatch_probability` and the pure-torch reference
|
|
must produce identical outputs for the same ``random_vals``. Single-rank,
|
|
runs on any CUDA GPU."""
|
|
from sglang.kernels.ops.lplb.cuda_solver import (
|
|
dispatch_probability,
|
|
dispatch_probability_torch_reference,
|
|
)
|
|
|
|
torch.manual_seed(0)
|
|
device = torch.device("cuda:0")
|
|
num_logical, max_copies, num_tokens, topk = 8, 3, 64, 4
|
|
|
|
# Random LP-style probabilities (some rows zeroed to exercise the fallback).
|
|
log2phy_prob = torch.rand(num_logical, max_copies, device=device)
|
|
log2phy_prob[0] = 0 # force fallback row
|
|
log2phy_map = torch.full(
|
|
(num_logical, max_copies), -1, dtype=torch.int32, device=device
|
|
)
|
|
for i in range(num_logical):
|
|
log2phy_map[i, 0] = i # primary replica
|
|
log2phy_map[i, 1] = num_logical + (i % 4) # one redundant
|
|
# leave last slot at -1 to exercise the masked-replica path
|
|
|
|
topk_ids = torch.randint(
|
|
0, num_logical, (num_tokens, topk), dtype=torch.int32, device=device
|
|
)
|
|
random_vals = torch.rand(num_tokens * topk, device=device, dtype=torch.float32)
|
|
|
|
cuda_out = dispatch_probability(topk_ids, log2phy_prob, log2phy_map, random_vals)
|
|
torch_out = dispatch_probability_torch_reference(
|
|
topk_ids, log2phy_prob, log2phy_map, random_vals
|
|
)
|
|
|
|
assert torch.equal(cuda_out, torch_out), (
|
|
f"dispatch_probability disagrees with torch reference: "
|
|
f"{(cuda_out != torch_out).sum().item()}/{cuda_out.numel()} mismatches"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available(),
|
|
reason="This test requires CUDA",
|
|
)
|
|
def test_solve_ipm_matches_torch_reference():
|
|
"""The fused CUDA IPM kernel and the pure-torch reference should agree to
|
|
a small tolerance on a *real* LPLB LP. They are NOT bit-equivalent — the
|
|
kernel factors the KKT system with a hand-written block Cholesky while the
|
|
reference uses torch.linalg.solve (LU) — so we compare with allclose and
|
|
print the max abs difference (the numerical-difference number the review
|
|
asked about).
|
|
|
|
The LP is built the way LPLBSolver does (Big-M column + normalized
|
|
counts), because an arbitrary random Ax=b never satisfies the solver's
|
|
convergence test (the Big-M slack must reach ~0) and both backends would
|
|
just return the 0.5 non-convergence sentinel — agreeing trivially without
|
|
exercising the solve. Single-rank, any CUDA GPU."""
|
|
from sglang.kernels.ops.lplb.cuda_solver import solve_ipm as cuda_solve_ipm
|
|
from sglang.kernels.ops.lplb.torch_solver import solve_ipm_torch_reference
|
|
from sglang.srt.eplb.lplb_solver import LPLBSolver
|
|
|
|
torch.manual_seed(0)
|
|
device = torch.device("cuda:0")
|
|
phy2log, log2phy, num_valid = _make_metadata()
|
|
solver = LPLBSolver(
|
|
phy2log=phy2log.to(device),
|
|
log2phy=log2phy.to(device),
|
|
num_gpus=NUM_GPUS,
|
|
ep_group=None,
|
|
logical_to_all_physical_map_num_valid=num_valid.to(device),
|
|
)
|
|
|
|
# Realistic counts: the replicated expert (logical 0) carries heavy load
|
|
# to distribute across its 3 physical copies; the rest are moderate. This
|
|
# gives the barrier method a well-posed instance that converges in 5 iters.
|
|
global_counts = torch.tensor(
|
|
[120.0, 30.0, 25.0, 20.0], dtype=torch.float32, device=device
|
|
)
|
|
|
|
# Reconstruct the LP exactly as LPLBSolver._solve builds it.
|
|
counts_norm = global_counts / global_counts.sum().clamp(min=1.0)
|
|
t1 = counts_norm[solver.log_single]
|
|
b1 = counts_norm[solver.log_replicated]
|
|
b2 = -(solver.B1 @ t1).flatten()
|
|
b = torch.cat([b1, b2])
|
|
big_M_col = b - solver._A_base_row_sum
|
|
A_full = torch.hstack([solver.A_base, big_M_col.unsqueeze(1)])
|
|
|
|
cuda_x = cuda_solve_ipm(A_full, b, solver.c_vec)
|
|
torch_x = solve_ipm_torch_reference(A_full, b, solver.c_vec)
|
|
|
|
converged = not torch.allclose(cuda_x, torch.full_like(cuda_x, 0.5))
|
|
max_diff = (cuda_x - torch_x).abs().max().item()
|
|
print(
|
|
f"\n[ipm-compare] converged={converged} max|cuda-torch|={max_diff:.3e} "
|
|
f"cuda={[round(v, 4) for v in cuda_x.tolist()]} "
|
|
f"torch={[round(v, 4) for v in torch_x.tolist()]}"
|
|
)
|
|
assert converged, (
|
|
"IPM returned the 0.5 non-convergence sentinel — the comparison would "
|
|
"be trivial. Adjust the LP instance so it converges."
|
|
)
|
|
assert torch.allclose(cuda_x, torch_x, atol=1e-2, rtol=1e-2), (
|
|
f"fused IPM diverges from torch reference: max abs diff {max_diff:.3e}"
|
|
)
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not torch.cuda.is_available() or torch.cuda.device_count() < NUM_GPUS,
|
|
reason=f"This test requires at least {NUM_GPUS} CUDA devices",
|
|
)
|
|
def test_lplb_distributed_two_rank():
|
|
"""Driver: spawn `NUM_GPUS` worker processes that each run the
|
|
distributed-LPLB invariant suite."""
|
|
torch.multiprocessing.spawn(
|
|
_worker_main,
|
|
args=(NUM_GPUS,),
|
|
nprocs=NUM_GPUS,
|
|
)
|
|
|
|
|
|
def _worker_main(local_rank: int, world_size: int):
|
|
"""Per-rank entry point under torch.multiprocessing.spawn."""
|
|
# Inject minimal ServerArgs before any LPLB module reads the global state.
|
|
# Silent fallbacks no longer exist — the fused CUDA path is the only LP
|
|
# path, so this test relies on hard failures, not gating flags.
|
|
from sglang.srt.server_args import (
|
|
ServerArgs,
|
|
set_global_server_args_for_scheduler,
|
|
)
|
|
|
|
set_global_server_args_for_scheduler(
|
|
ServerArgs(
|
|
model_path="dummy",
|
|
# Match the tp/ep width initialize_model_parallel is about to
|
|
# build below -- get_parallel()'s derived widths (attn_tp_size,
|
|
# moe_ep_size, ...) are projected from this at publish time, and
|
|
# nothing here should leave that projection reflecting a width
|
|
# this process never actually runs at.
|
|
tp_size=world_size,
|
|
ep_size=world_size,
|
|
)
|
|
)
|
|
|
|
device = torch.device(f"cuda:{local_rank}")
|
|
torch.cuda.set_device(device)
|
|
torch.set_default_device(device)
|
|
|
|
update_environment_variables(
|
|
{
|
|
"RANK": str(local_rank),
|
|
"LOCAL_RANK": str(local_rank),
|
|
"WORLD_SIZE": str(world_size),
|
|
"MASTER_ADDR": "localhost",
|
|
"MASTER_PORT": "12347", # Distinct from other tests' ports.
|
|
}
|
|
)
|
|
init_distributed_environment(
|
|
world_size=world_size, rank=local_rank, local_rank=local_rank
|
|
)
|
|
publish_build_topology(
|
|
tp_size=world_size, ep_size=world_size, world_rank=local_rank
|
|
)
|
|
initialize_model_parallel()
|
|
|
|
from sglang.srt.eplb.lplb_solver import clear_global_lplb_solvers
|
|
|
|
try:
|
|
# Clear the global solver registry between subtests so a stale solver
|
|
# can't shadow a freshly-built one.
|
|
clear_global_lplb_solvers()
|
|
_check_solver_with_empty_rank(local_rank, world_size, device)
|
|
clear_global_lplb_solvers()
|
|
_check_all_ranks_empty(local_rank, world_size, device)
|
|
clear_global_lplb_solvers()
|
|
_check_solver_determinism(local_rank, world_size, device)
|
|
clear_global_lplb_solvers()
|
|
_check_post_rebalance_reinit(local_rank, world_size, device)
|
|
finally:
|
|
# Use parallel_state's destroy helpers rather than raw
|
|
# `dist.destroy_process_group()` so we don't leave _WORLD/_MOE_EP/_TP
|
|
# globals stale for later tests.
|
|
from sglang.srt.distributed.parallel_state import (
|
|
destroy_distributed_environment,
|
|
destroy_model_parallel,
|
|
)
|
|
|
|
clear_global_lplb_solvers()
|
|
destroy_model_parallel()
|
|
destroy_distributed_environment()
|
|
|
|
|
|
def _build_solver():
|
|
"""Construct an LPLBSolver bound to the live moe_ep_group."""
|
|
from sglang.srt.distributed.parallel_state import get_moe_ep_group
|
|
from sglang.srt.eplb.lplb_solver import LPLBSolver
|
|
|
|
phy2log, log2phy, num_valid = _make_metadata(rebalanced=False)
|
|
return (
|
|
LPLBSolver(
|
|
phy2log=phy2log.cuda(),
|
|
log2phy=log2phy.cuda(),
|
|
num_gpus=NUM_GPUS,
|
|
ep_group=get_moe_ep_group(),
|
|
logical_to_all_physical_map_num_valid=num_valid.cuda(),
|
|
),
|
|
phy2log,
|
|
log2phy,
|
|
num_valid,
|
|
)
|
|
|
|
|
|
def _expected_output(solver, expected_global_counts: torch.Tensor) -> torch.Tensor:
|
|
"""Oracle: what `solver.solve(*)` should produce given known global
|
|
counts. Uses the solver's own `_solve` so we match implementation
|
|
quirks (regularization, on-device clamp normalization, etc.).
|
|
"""
|
|
return solver._solve(expected_global_counts.float().cuda())
|
|
|
|
|
|
def _check_solver_with_empty_rank(rank: int, world_size: int, device: torch.device):
|
|
"""Rank 0 has real `topk_ids`; rank 1 has `(0, topk)` empty.
|
|
|
|
The all-reduce inside `solve()` must complete on both ranks, and both
|
|
must produce output equal to the oracle solver-output for the SUMMED
|
|
global counts (rank 0's bincount + rank 1's zero contribution).
|
|
"""
|
|
solver, _, _, _ = _build_solver()
|
|
|
|
# Rank 0 owns these tokens; rank 1 is idle (DP-attention empty case).
|
|
rank0_topk = torch.tensor(
|
|
[[0, 1], [0, 2], [3, 0], [0, 1]], dtype=torch.int32, device=device
|
|
)
|
|
if rank == 0:
|
|
topk_ids = rank0_topk
|
|
else:
|
|
topk_ids = torch.empty((0, TOPK), dtype=torch.int32, device=device)
|
|
|
|
# Compute the oracle: the global counts are exactly the bincount of
|
|
# rank 0's tokens (rank 1 contributes zero).
|
|
flat = rank0_topk.flatten().long()
|
|
expected_counts = torch.bincount(flat, minlength=NUM_LOGICAL)
|
|
expected = _expected_output(solver, expected_counts)
|
|
|
|
# Run the actual collective.
|
|
actual = solver.solve(topk_ids)
|
|
|
|
# Shape + finiteness sanity.
|
|
assert actual.shape == (
|
|
NUM_LOGICAL,
|
|
solver.max_copies,
|
|
), f"rank {rank}: bad output shape {actual.shape}"
|
|
assert torch.isfinite(actual).all(), f"rank {rank}: non-finite values in output"
|
|
assert (actual >= 0).all(), f"rank {rank}: negative probabilities"
|
|
|
|
# The strong invariant: output matches the summed-count oracle.
|
|
# If the all-reduce silently failed, rank 1 would feed all-zero local
|
|
# counts to its solver, the oracle wouldn't match, and this assertion
|
|
# would catch it.
|
|
assert torch.allclose(actual, expected, atol=1e-4, rtol=1e-3), (
|
|
f"rank {rank}: solve() output disagrees with summed-count oracle "
|
|
f"(max abs diff {(actual - expected).abs().max().item():.2e}). "
|
|
"This usually indicates the EP all-reduce inside solve() was "
|
|
"skipped on the empty rank, or rank 1 used local counts only."
|
|
)
|
|
|
|
# 0.5-uniform sentinel detection: torch IPM's failure path returns a
|
|
# vector of all-0.5. If we see that, the IPM diverged silently.
|
|
is_uniform_half = torch.allclose(
|
|
actual,
|
|
torch.full_like(actual, 0.5),
|
|
atol=1e-6,
|
|
)
|
|
assert not is_uniform_half, (
|
|
f"rank {rank}: output is the 0.5-uniform fallback, indicating "
|
|
"torch IPM convergence failure"
|
|
)
|
|
|
|
|
|
def _check_all_ranks_empty(rank: int, world_size: int, device: torch.device):
|
|
"""Both ranks have empty topk_ids (idle batch). The collective must
|
|
still complete; the output should be the all-zero-counts oracle."""
|
|
solver, _, _, _ = _build_solver()
|
|
|
|
empty = torch.empty((0, TOPK), dtype=torch.int32, device=device)
|
|
expected_counts = torch.zeros(NUM_LOGICAL, dtype=torch.int64)
|
|
expected = _expected_output(solver, expected_counts)
|
|
|
|
actual = solver.solve(empty)
|
|
|
|
assert torch.isfinite(actual).all(), f"rank {rank}: non-finite for empty-batch"
|
|
assert (actual >= 0).all(), f"rank {rank}: negative for empty-batch"
|
|
assert torch.allclose(actual, expected, atol=1e-4, rtol=1e-3), (
|
|
f"rank {rank}: empty-batch output disagrees with all-zero oracle"
|
|
)
|
|
|
|
|
|
def _check_solver_determinism(rank: int, world_size: int, device: torch.device):
|
|
"""Same input, two `solve()` calls in a row → identical output."""
|
|
solver, _, _, _ = _build_solver()
|
|
if rank == 0:
|
|
topk_ids = torch.tensor([[0, 1], [3, 0]], dtype=torch.int32, device=device)
|
|
else:
|
|
topk_ids = torch.empty((0, TOPK), dtype=torch.int32, device=device)
|
|
|
|
out1 = solver.solve(topk_ids)
|
|
out2 = solver.solve(topk_ids)
|
|
assert torch.equal(out1, out2), (
|
|
f"rank {rank}: solve() not deterministic across calls "
|
|
f"(max diff {(out1 - out2).abs().max().item():.2e})"
|
|
)
|
|
|
|
|
|
def _check_post_rebalance_reinit(rank: int, world_size: int, device: torch.device):
|
|
"""Build a NEW solver with different `phy2log` (the redundant logical
|
|
moves from expert 0 to expert 1). Same input → DIFFERENT output, since
|
|
the LP problem changed."""
|
|
from sglang.srt.distributed.parallel_state import get_moe_ep_group
|
|
from sglang.srt.eplb.lplb_solver import LPLBSolver
|
|
|
|
solver_old, _, _, _ = _build_solver()
|
|
phy2log_new, log2phy_new, num_valid_new = _make_metadata(rebalanced=True)
|
|
solver_new = LPLBSolver(
|
|
phy2log=phy2log_new.cuda(),
|
|
log2phy=log2phy_new.cuda(),
|
|
num_gpus=NUM_GPUS,
|
|
ep_group=get_moe_ep_group(),
|
|
logical_to_all_physical_map_num_valid=num_valid_new.cuda(),
|
|
)
|
|
|
|
# An input that hits the replicated logical in BOTH configs (so the
|
|
# LP solution is sensitive to which logical is the replicated one).
|
|
rank0_topk = torch.tensor(
|
|
[[0, 1], [0, 1], [0, 1], [2, 3]], dtype=torch.int32, device=device
|
|
)
|
|
if rank == 0:
|
|
topk_ids = rank0_topk
|
|
else:
|
|
topk_ids = torch.empty((0, TOPK), dtype=torch.int32, device=device)
|
|
|
|
out_old = solver_old.solve(topk_ids)
|
|
out_new = solver_new.solve(topk_ids)
|
|
|
|
# The two outputs should differ — same `topk_ids`, different metadata
|
|
# ⇒ different LP solution. If they're identical, one of the solvers
|
|
# used stale state.
|
|
assert not torch.equal(out_old, out_new), (
|
|
f"rank {rank}: post-rebalance solver produced same output as "
|
|
"pre-rebalance solver despite different phy2log mappings — "
|
|
"stale solver state likely"
|
|
)
|
|
|
|
# Positive oracle: each output should match its OWN solver's
|
|
# `_solve` of the SUMMED global counts (rank 0's bincount; rank 1
|
|
# contributes zero). Both ranks compare against the same global oracle
|
|
# since solve() all-reduces the counts.
|
|
expected_counts = torch.bincount(rank0_topk.flatten().long(), minlength=NUM_LOGICAL)
|
|
expected_old = _expected_output(solver_old, expected_counts)
|
|
expected_new = _expected_output(solver_new, expected_counts)
|
|
assert torch.allclose(out_old, expected_old, atol=1e-4, rtol=1e-3), (
|
|
f"rank {rank}: pre-rebalance output disagrees with own oracle "
|
|
f"(diff {(out_old - expected_old).abs().max().item():.2e})"
|
|
)
|
|
assert torch.allclose(out_new, expected_new, atol=1e-4, rtol=1e-3), (
|
|
f"rank {rank}: post-rebalance output disagrees with own oracle "
|
|
f"(diff {(out_new - expected_new).abs().max().item():.2e})"
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main([__file__, "-v"]))
|