diff --git a/python/sglang/kernels/ops/moe/ep_moe_kernels.py b/python/sglang/kernels/ops/moe/ep_moe_kernels.py index 20263ae3ab0c..a9911fead638 100644 --- a/python/sglang/kernels/ops/moe/ep_moe_kernels.py +++ b/python/sglang/kernels/ops/moe/ep_moe_kernels.py @@ -1267,6 +1267,7 @@ def ep_scatter_from_psum( m_indices: torch.Tensor, output_index: torch.Tensor, scale_ue8m0: bool = False, + expert_start: int = 0, ): BLOCK_E = 128 BLOCK_D = 128 @@ -1315,6 +1316,8 @@ def ep_scatter_from_psum( output_index, output_index.stride(0), output_index.stride(1), + expert_start, + num_experts, topk_num=recv_topk.shape[1], num_warps=num_warps, HIDDEN_SIZE=hidden_size, diff --git a/test/registered/kernels/ops/moe/test_ep_scatter.py b/test/registered/kernels/ops/moe/test_ep_scatter.py new file mode 100644 index 000000000000..0fb927f3d471 --- /dev/null +++ b/test/registered/kernels/ops/moe/test_ep_scatter.py @@ -0,0 +1,175 @@ +import pytest +import torch + +from sglang.kernels.ops.moe.ep_moe_kernels import ep_scatter_from_psum +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +dev = "cuda" +BLOCK_D = 128 +BLOCK_E = 128 + + +def _build(num_tokens, num_experts, top_k, pad_frac, expert_start, seed): + """Route tokens to local experts, then describe the routing the way the + DeepEPv2 permute path does: an INCLUSIVE prefix sum over local experts, and + `recv_topk` holding expert ids offset by `expert_start` (-1 = padding).""" + g = torch.Generator(device="cpu").manual_seed(seed) + local_ids = torch.full((num_tokens, top_k), -1, dtype=torch.int32) + counts = [0] * num_experts + for t in range(num_tokens): + experts = torch.randperm(num_experts, generator=g)[:top_k].tolist() + for slot, e in enumerate(experts): + if pad_frac > 0 and torch.rand(1, generator=g).item() < pad_frac: + continue + local_ids[t, slot] = e + counts[e] += 1 + + psum = torch.tensor(counts, dtype=torch.int32).cumsum(0).to(torch.int32) + recv_topk = torch.where(local_ids >= 0, local_ids + expert_start, local_ids) + return local_ids.to(dev), recv_topk.to(dev), psum.to(dev) + + +def _run(recv_topk, psum, hidden, dtype, expert_start): + num_tokens = recv_topk.shape[0] + total = int(psum[-1].item()) + scale_hidden = hidden // BLOCK_D + is_fp8 = dtype != torch.bfloat16 + + recv_x = torch.randn(num_tokens, hidden, dtype=torch.bfloat16, device=dev).to(dtype) + if is_fp8: + recv_x_scale = torch.rand( + num_tokens, scale_hidden, dtype=torch.float32, device=dev + ) + output_tensor_scale = torch.zeros( + total, scale_hidden, dtype=torch.float32, device=dev + ) + else: + recv_x_scale = None + output_tensor_scale = None + + output_tensor = torch.zeros(total, hidden, dtype=dtype, device=dev) + # m_indices is padded up to BLOCK_E; -1 marks slots no expert owns. + m_indices = torch.full( + ((total + BLOCK_E - 1) // BLOCK_E * BLOCK_E,), -1, dtype=torch.int32, device=dev + ) + expert_start_loc = torch.empty_like(psum) + output_index = torch.empty_like(recv_topk) + + # Leave expert_start at its default in the local-id case, so this exercises + # the call shape the DeepEPv2 permute path actually uses. + kwargs = {"expert_start": expert_start} if expert_start else {} + ep_scatter_from_psum( + recv_x, + recv_x_scale, + recv_topk, + psum, + expert_start_loc, + output_tensor, + output_tensor_scale, + m_indices, + output_index, + **kwargs, + ) + return dict( + recv_x=recv_x, + recv_x_scale=recv_x_scale, + output_tensor=output_tensor, + output_tensor_scale=output_tensor_scale, + m_indices=m_indices, + expert_start_loc=expert_start_loc, + output_index=output_index, + total=total, + ) + + +def _check(out, local_ids, psum): + """The scatter's destination slots are chosen by atomic_add, so the order + within one expert is not deterministic. Assert the invariants instead — they + pin the result down completely apart from that order.""" + output_index = out["output_index"] + total = out["total"] + valid = local_ids >= 0 + + # 1. -1 sentinel exactly on the padding lanes. The post-permute kernel gates + # on this, so a stale value here silently reads a garbage row. + assert torch.equal(output_index < 0, ~valid) + + # 2. Every valid lane lands inside its own expert's slab. + starts = torch.cat([torch.zeros(1, dtype=psum.dtype, device=dev), psum[:-1]]) + e = local_ids[valid].long() + dst = output_index[valid].long() + assert torch.all(dst >= starts[e]) and torch.all(dst < psum[e]) + + # 3. ... and the slabs are filled exactly, with no slot used twice. + assert dst.numel() == total + assert torch.equal(torch.sort(dst).values, torch.arange(total, device=dev)) + + # 4. m_indices agrees with the routing, and the BLOCK_E tail stays -1. + assert torch.equal(out["m_indices"][dst], local_ids[valid].to(torch.int32)) + assert torch.all(out["m_indices"][total:] == -1) + + # 5. expert_start_loc has been advanced to each expert's end. + assert torch.equal(out["expert_start_loc"], psum) + + # 6. The rows themselves arrived, scales included. + src = torch.nonzero(valid)[:, 0] + assert torch.equal( + out["output_tensor"][dst].view(torch.uint8), + out["recv_x"][src].view(torch.uint8), + ) + if out["recv_x_scale"] is not None: + assert torch.equal(out["output_tensor_scale"][dst], out["recv_x_scale"][src]) + + +@pytest.mark.parametrize("num_tokens", [1, 8, 256, 4096]) +@pytest.mark.parametrize("pad_frac", [0.0, 0.3]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_ep_scatter_from_psum(num_tokens, pad_frac, dtype): + num_experts, top_k, hidden = 8, 8, 512 + local_ids, recv_topk, psum = _build( + num_tokens, num_experts, top_k, pad_frac, 0, seed=num_tokens + ) + out = _run(recv_topk, psum, hidden, dtype, 0) + _check(out, local_ids, psum) + + +@pytest.mark.parametrize("expert_start", [0, 8, 120]) +def test_ep_scatter_from_psum_expert_start(expert_start): + """`expert_start` shifts which global ids count as local. Offsetting both the + ids and `expert_start` must not change the result, and out-of-range ids must + be dropped like padding.""" + num_tokens, num_experts, top_k, hidden = 512, 8, 8, 512 + local_ids, recv_topk, psum = _build( + num_tokens, num_experts, top_k, 0.2, expert_start, seed=7 + ) + out = _run(recv_topk, psum, hidden, torch.float8_e4m3fn, expert_start) + _check(out, local_ids, psum) + + +def test_ep_scatter_from_psum_drops_non_local_experts(): + num_tokens, num_experts, top_k, hidden = 256, 8, 8, 512 + expert_start = 8 + local_ids, recv_topk, psum = _build( + num_tokens, num_experts, top_k, 0.0, expert_start, seed=11 + ) + # Send half of slot 0 to an expert this rank does not own; it must be treated + # exactly like padding, not scattered somewhere out of bounds. + foreign = torch.arange(num_tokens, device=dev) % 2 == 0 + recv_topk[foreign, 0] = expert_start + num_experts + 3 + dropped = local_ids[foreign, 0].clone() + local_ids[foreign, 0] = -1 + counts = torch.bincount(dropped.long(), minlength=num_experts).to(psum.dtype) + psum = ( + (torch.cat([psum[:1], psum[1:] - psum[:-1]]) - counts).cumsum(0).to(psum.dtype) + ) + + out = _run(recv_topk, psum, hidden, torch.float8_e4m3fn, expert_start) + _check(out, local_ids, psum) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v", "-s"]))