From 5fdf6cd18f978b84fb02b7ac82c740a4d2d7d5e3 Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Sat, 8 Aug 2026 16:08:25 +0800 Subject: [PATCH] [MoE] Single-launch moe_align for tiny batches with many experts (#32395) --- python/sglang/kernels/ops/moe/__init__.py | 17 ++ .../kernels/ops/moe/moe_align_small_numel.py | 147 ++++++++++++ .../triton_utils/moe_align_block_size.py | 32 +++ .../ops/moe/test_moe_align_small_numel.py | 214 ++++++++++++++++++ 4 files changed, 410 insertions(+) create mode 100644 python/sglang/kernels/ops/moe/moe_align_small_numel.py create mode 100644 test/registered/kernels/ops/moe/test_moe_align_small_numel.py diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py index 492e171fd..00359e8fd 100644 --- a/python/sglang/kernels/ops/moe/__init__.py +++ b/python/sglang/kernels/ops/moe/__init__.py @@ -169,3 +169,20 @@ register_kernel( target="sglang.kernels.ops.moe.pack_topk_ids:PackTopkIds.triton", ) ) + +# Single-CTA align for tiny batches: covers the corner the AOT/JIT +# moe_align_block_size small-batch path leaves out (num_experts > 64), and is +# selected by the moe_runner call site on numel <= SMALL_NUMEL_LIMIT. +register_kernel( + KernelSpec( + op="moe.moe_align_small_numel", + backend=KernelBackend.TRITON, + target="sglang.kernels.ops.moe.moe_align_small_numel:moe_align_small_numel", + capabilities=_CUDA, + format_signature=FormatSignature( + in_place=True, + description="align/sort expert token ids into block-padded buffers", + ), + description="MoE align-block-size, single-launch triton variant.", + ) +) diff --git a/python/sglang/kernels/ops/moe/moe_align_small_numel.py b/python/sglang/kernels/ops/moe/moe_align_small_numel.py new file mode 100644 index 000000000..add436ede --- /dev/null +++ b/python/sglang/kernels/ops/moe/moe_align_small_numel.py @@ -0,0 +1,147 @@ +"""Single-launch moe_align for tiny batches with many experts. + +The CUDA small-batch align kernel is gated to ``num_experts <= 64`` (its shared +memory grows as O(threads x experts)), so bs=1 decode on a MoE with a wider +expert dimension always paid the generic two-kernel (align + count_and_sort) +path. This kernel covers that corner in a single launch, at any expert count, +for ``numel <= SMALL_NUMEL_LIMIT``. +""" + +import torch +import triton +import triton.language as tl + +from sglang.kernels.jit.utils import is_arch_support_pdl + +# Largest numel routed to this kernel. Its [NP, NP] pairwise tensors fit in +# registers at NP=64 (~4 us, on par with the two CUDA launches it replaces) but +# spill to local memory at NP=256 (~230 us measured). +SMALL_NUMEL_LIMIT = 64 + + +@triton.jit +def _moe_align_small_numel_kernel( + topk_ids_ptr, # [numel] int, flattened (token, slot) expert ids, -1 = filtered + sorted_token_ids_ptr, # [max_num_tokens_padded] int32 + expert_ids_ptr, # [max_num_m_blocks] int32 + num_tokens_post_pad_ptr, # [1] int32 + num_experts, # E + 1 (the "+1 offset" convention's bucket count) + block_size, + numel, + NP: tl.constexpr, # power-of-2 >= numel + NB: tl.constexpr, # power-of-2 >= max blocks used + USE_GDC: tl.constexpr = False, +): + """Single-CTA moe_align for tiny batches with MANY experts. + + Everything works on the PAIR axis ([NP, NP] pairwise comparisons plus a + rank-0 representative per bucket) -- an expert-axis formulation (histogram + / cumsum over ~1k buckets) is ~3x more single-SM work and measured slower + than the two-kernel path it replaces. + + Reference semantics reproduced: + - "+1 offset" convention: expert -1 (EP-filtered) maps to bucket 0 and its + blocks get expert_ids = -1 (skipped by fused_moe's filter_expert); + - every bucket is padded to a block_size multiple, offsets in bucket order; + - pad slots inside [0, num_tokens_post_pad) hold `numel`. + + Intended deviations, both invisible to fused_moe: + - intra-bucket order is stable in pair index (the reference's atomicAdd + order is scheduling-dependent; every pair writes its own output row); + - sorted_token_ids beyond num_tokens_post_pad is left unwritten (the + reference pre-fills the whole buffer; consumers only read below the + published total). + """ + if USE_GDC: + # Consumer side of the router top-k that produced topk_ids. + tl.extra.cuda.gdc_wait() + + offs_p = tl.arange(0, NP) + mask_p = offs_p < numel + ids = tl.load(topk_ids_ptr + offs_p, mask=mask_p, other=-2) + # Padded lanes get an out-of-range bucket and are masked out everywhere. + bucket = tl.where(mask_p, (ids + 1).to(tl.int32), num_experts) + + # Pairwise stats: stable rank within the bucket and bucket population. + same = (bucket[None, :] == bucket[:, None]) & mask_p[None, :] & mask_p[:, None] + earlier = offs_p[None, :] < offs_p[:, None] + rank = tl.sum((same & earlier).to(tl.int32), axis=1) # [NP] + cnt = tl.sum(same.to(tl.int32), axis=1) # [NP], own-bucket population + padded_cnt = ((cnt + block_size - 1) // block_size) * block_size + is_rep = (rank == 0) & mask_p # one representative pair per bucket + + # Bucket-ordered exclusive offsets: sum the padded counts of every + # representative with a strictly smaller bucket id. + smaller_rep = (bucket[None, :] < bucket[:, None]) & is_rep[None, :] + excl = tl.sum(smaller_rep.to(tl.int32) * padded_cnt[None, :], axis=1) # [NP] + + total = tl.sum(tl.where(is_rep, padded_cnt, 0), axis=0) + tl.store(num_tokens_post_pad_ptr, total.to(tl.int32)) + + # expert_ids per used block: representative r owns blocks + # [excl[r], excl[r] + padded_cnt[r]); the written id is bucket - 1 + # (bucket 0 = filtered -> -1). + offs_b = tl.arange(0, NB) + block_start = offs_b * block_size + in_range = ( + (block_start[:, None] >= excl[None, :]) + & (block_start[:, None] < (excl + padded_cnt)[None, :]) + & is_rep[None, :] + ) + eid = tl.sum(in_range.to(tl.int32) * (bucket[None, :] - 1), axis=1) + tl.store(expert_ids_ptr + offs_b, eid.to(tl.int32), mask=block_start < total) + + # Fill the used region's pad slots with `numel`, then scatter the real + # pair indices over them. The barrier is required: fill and scatter run on + # different warps of this CTA, and a scatter store must not be overtaken + # by a later-warp fill store to the same address. + n_fill = (total + NP - 1) // NP + for it in range(n_fill): + f_offs = it * NP + offs_p + tl.store( + sorted_token_ids_ptr + f_offs, + tl.full([NP], 0, tl.int32) + numel, + mask=f_offs < total, + ) + tl.debug_barrier() + pos = excl + rank + tl.store(sorted_token_ids_ptr + pos, offs_p.to(tl.int32), mask=mask_p) + + if USE_GDC: + tl.extra.cuda.gdc_launch_dependents() + + +def moe_align_small_numel( + topk_ids: torch.Tensor, + num_experts: int, + block_size: int, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_pad: torch.Tensor, +) -> None: + """Align and sort expert token ids into block-padded buffers, in one launch. + + Buffer contract matches ``sglang.kernels.ops.moe.moe_align_block_size`` + (minus its ``cumsum_buffer``, which a single CTA does not need): + ``num_experts`` is the bucket count ``E + 1`` under the "+1 offset" + convention, and the three output buffers are written in place. + + Callers gate on ``topk_ids.numel() <= SMALL_NUMEL_LIMIT``; the kernel stays + correct above it, but its pairwise tensors spill and it stops being faster + than the two-kernel path. + """ + numel = topk_ids.numel() + pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {} + _moe_align_small_numel_kernel[(1,)]( + topk_ids, + sorted_token_ids, + expert_ids, + num_tokens_post_pad, + num_experts, + block_size, + numel, + NP=triton.next_power_of_2(max(numel, 2)), + NB=triton.next_power_of_2(max(expert_ids.numel(), 2)), + num_warps=4, + **pdl_kwargs, + ) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py index 8879e351d..9c611a468 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/moe_align_block_size.py @@ -18,6 +18,16 @@ _is_musa = is_musa() if _is_cuda or _is_hip or _is_xpu or _is_musa: from sglang.kernels.ops.moe import moe_align_block_size as sgl_moe_align_block_size +if _is_cuda: + from sglang.kernels.ops.moe.moe_align_small_numel import ( + SMALL_NUMEL_LIMIT, + moe_align_small_numel, + ) + +# Where the CUDA kernel's own small-batch single-block path stops: its +# per-thread histogram costs 4 * (buckets + 1) ** 2 bytes of shared memory. +_CUDA_SMALL_BATCH_MAX_BUCKETS = 64 + def moe_align_block_size( topk_ids: torch.Tensor, @@ -93,6 +103,28 @@ def moe_align_block_size( (num_experts + 2,), dtype=torch.int32, device=topk_ids.device ) + # Tiny-batch fast path (bs=1 decode): one single-CTA triton launch replaces + # the generic align + count_and_sort pair, covering the corner the CUDA + # small-batch kernel cannot reach. Below that bucket limit the CUDA kernel + # is already a single launch and does O(numel) work where this one does + # O(numel ** 2) pairwise, so leave that side to it. ignore_invalid_expert is + # a different contract from the "+1 offset" convention this kernel implements. + if ( + _is_cuda + and topk_ids.numel() <= SMALL_NUMEL_LIMIT + and num_experts + 1 > _CUDA_SMALL_BATCH_MAX_BUCKETS + and not ignore_invalid_expert + ): + moe_align_small_numel( + topk_ids, + num_experts + 1, + block_size, + sorted_ids, + expert_ids, + num_tokens_post_pad, + ) + return sorted_ids, expert_ids, num_tokens_post_pad + # ===== TO BE REFACTORED ==== use_jit_align = False if _SGLANG_EXPERIMENTAL_LORA_OPTI: diff --git a/test/registered/kernels/ops/moe/test_moe_align_small_numel.py b/test/registered/kernels/ops/moe/test_moe_align_small_numel.py new file mode 100644 index 000000000..fd67fecab --- /dev/null +++ b/test/registered/kernels/ops/moe/test_moe_align_small_numel.py @@ -0,0 +1,214 @@ +"""Correctness of the single-launch tiny-numel moe_align triton kernel. + +The oracle is a plain-torch implementation of the documented contract, so it +does not depend on any other kernel's shape support; the AOT `sgl_kernel` path +is cross-checked on top of it to back the drop-in-replacement claim. +""" + +import itertools +import sys + +import pytest +import torch +import triton + +from sglang.kernels.jit.utils import get_ci_test_range +from sglang.kernels.ops.moe import moe_align_block_size as cuda_moe_align_block_size +from sglang.kernels.ops.moe.moe_align_small_numel import ( + SMALL_NUMEL_LIMIT, + moe_align_small_numel, +) +from sglang.srt.layers.moe.moe_runner.triton_utils.moe_align_block_size import ( + moe_align_block_size as runner_moe_align_block_size, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +# Bucket counts above this are not uniformly supported by the AOT sgl_kernel +# path across wheel versions, so the cross-check against it stops here; the +# kernel under test has no expert limit and the oracle covers it past this bound. +CUDA_XCHECK_MAX_EXPERTS = 1023 + + +def _reference(topk_ids, block_size, num_experts): + """The contract, in plain torch, on CPU: bucket = expert + 1 (so EP-filtered + -1 lands in bucket 0), each bucket padded to a block_size multiple, blocks in + bucket order with expert_ids = bucket - 1, pad slots holding numel, pairs + placed in ascending pair index within their bucket.""" + numel = topk_ids.numel() + bucket = (topk_ids.flatten().to(torch.int64) + 1).cpu() + counts = torch.bincount(bucket, minlength=num_experts + 1) + padded = ((counts + block_size - 1) // block_size) * block_size + offsets = torch.cumsum(padded, 0) - padded + total = int(padded.sum()) + + non_empty = torch.nonzero(padded, as_tuple=True)[0] + expert_ids = torch.repeat_interleave( + non_empty - 1, padded[non_empty] // block_size + ).to(torch.int32) + + sorted_ids = torch.full((total,), numel, dtype=torch.int32) + cursor = offsets.clone() + for pair in range(numel): + b = int(bucket[pair]) + sorted_ids[cursor[b]] = pair + cursor[b] += 1 + return sorted_ids, expert_ids, total + + +def _alloc(numel, block_size, num_experts): + """Output buffers sized exactly as the moe_runner call site sizes them.""" + if numel < num_experts + 1: + max_num_tokens_padded = numel * block_size + else: + max_num_tokens_padded = numel + (num_experts + 1) * (block_size - 1) + max_num_m_blocks = triton.cdiv(max_num_tokens_padded, block_size) + return ( + torch.empty((max_num_tokens_padded,), dtype=torch.int32, device="cuda"), + torch.empty((max_num_m_blocks,), dtype=torch.int32, device="cuda"), + torch.empty((1,), dtype=torch.int32, device="cuda"), + ) + + +def _run_triton(topk_ids, block_size, num_experts): + sorted_ids, expert_ids, num_post_pad = _alloc( + topk_ids.numel(), block_size, num_experts + ) + moe_align_small_numel( + topk_ids, num_experts + 1, block_size, sorted_ids, expert_ids, num_post_pad + ) + return sorted_ids, expert_ids, num_post_pad + + +def _run_cuda(topk_ids, block_size, num_experts, ignore_invalid_expert=False): + sorted_ids, expert_ids, num_post_pad = _alloc( + topk_ids.numel(), block_size, num_experts + ) + cumsum_buffer = torch.empty((num_experts + 2,), dtype=torch.int32, device="cuda") + cuda_moe_align_block_size( + topk_ids, + num_experts + 1, + block_size, + sorted_ids, + expert_ids, + num_post_pad, + cumsum_buffer, + True, + ignore_invalid_expert, + ) + return sorted_ids, expert_ids, num_post_pad + + +def _assert_exact(got, ref, block_size): + """Full equality, valid against the oracle because its intra-bucket order + matches the kernel's (stable in pair index). The tail past the published + total is left unwritten by design, so nothing is asserted there.""" + got_sorted, got_expert, got_total = got + ref_sorted, ref_expert, ref_total = ref + assert got_total.item() == ref_total, "num_tokens_post_pad" + num_blocks = ref_total // block_size + assert torch.equal(got_expert[:num_blocks].cpu(), ref_expert), "expert_ids" + assert torch.equal(got_sorted[:ref_total].cpu(), ref_sorted), "sorted_token_ids" + + +def _assert_blockwise(got, ref, block_size): + """Per-block multiset equality -- the comparison that also holds against the + CUDA kernel, whose intra-bucket order is atomicAdd scheduling order.""" + got_sorted, got_expert, got_total = got + ref_sorted, ref_expert, ref_total = ref + assert got_total.item() == ref_total.item(), "num_tokens_post_pad" + total = ref_total.item() + num_blocks = total // block_size + assert torch.equal(got_expert[:num_blocks], ref_expert[:num_blocks]), "expert_ids" + got_blocks = got_sorted[:total].view(num_blocks, block_size).sort(dim=1).values + ref_blocks = ref_sorted[:total].view(num_blocks, block_size).sort(dim=1).values + assert torch.equal(got_blocks, ref_blocks), "sorted_token_ids block contents" + + +# num_experts straddles the CUDA small-batch kernel's 64-bucket limit (the corner +# this kernel exists to cover) and goes past what the AOT path handles at all. +ALIGN_CASES = get_ci_test_range( + [ + (block_size, num_experts, topk, num_tokens) + for block_size, num_experts, topk, num_tokens in itertools.product( + [16, 32, 64, 128], [8, 64, 65, 129, 1024], [1, 2, 4, 8], [1, 4, 8] + ) + if topk * num_tokens <= SMALL_NUMEL_LIMIT + ], + [ + (16, 65, 1, 1), + (32, 65, 8, 8), + (64, 129, 4, 4), + (128, 1024, 8, 8), + (128, 8, 2, 4), + ], +) + + +@pytest.mark.parametrize("block_size,num_experts,topk,num_tokens", ALIGN_CASES) +def test_matches_reference(block_size, num_experts, topk, num_tokens): + """Exact against the oracle, plus drop-in equivalence with the CUDA path + this replaces wherever that path supports the bucket count.""" + torch.manual_seed(0) + topk_ids = torch.randint( + 0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda" + ) + got = _run_triton(topk_ids, block_size, num_experts) + _assert_exact(got, _reference(topk_ids, block_size, num_experts), block_size) + if num_experts <= CUDA_XCHECK_MAX_EXPERTS: + _assert_blockwise(got, _run_cuda(topk_ids, block_size, num_experts), block_size) + + +def test_ep_filtered_ids_map_to_expert_minus_one(): + """EP-filtered pairs (-1) collect in bucket 0, whose blocks carry -1 so + fused_moe's filter_expert skips them.""" + torch.manual_seed(1) + num_experts, block_size = 1024, 64 + topk_ids = torch.randint(0, num_experts, (8, 4), dtype=torch.int32, device="cuda") + topk_ids[0] = -1 + topk_ids[3][2] = -1 + + ref = _reference(topk_ids, block_size, num_experts) + _assert_exact(_run_triton(topk_ids, block_size, num_experts), ref, block_size) + assert -1 in ref[1].tolist(), "filtered pairs must produce an expert_id == -1 block" + + +@pytest.mark.parametrize( + "numel", [SMALL_NUMEL_LIMIT - 1, SMALL_NUMEL_LIMIT, SMALL_NUMEL_LIMIT + 1] +) +def test_runner_dispatch_boundary(numel): + """Both sides of the moe_runner gate must agree with the CUDA reference, so + a future change to the limit cannot silently ship an unvalidated path.""" + torch.manual_seed(2) + num_experts, block_size = 129, 32 + topk_ids = torch.randint( + 0, num_experts, (numel, 1), dtype=torch.int32, device="cuda" + ) + _assert_blockwise( + runner_moe_align_block_size(topk_ids, block_size, num_experts), + _run_cuda(topk_ids, block_size, num_experts), + block_size, + ) + + +def test_runner_defers_for_ignore_invalid_expert(): + """ignore_invalid_expert is a different contract than the '+1 offset' + convention the triton kernel implements, so the runner must keep producing + what the CUDA kernel produces under that flag.""" + torch.manual_seed(3) + num_experts, block_size = 128, 32 + topk_ids = torch.randint(0, num_experts, (8, 4), dtype=torch.int32, device="cuda") + topk_ids[1] = -1 + + _assert_blockwise( + runner_moe_align_block_size( + topk_ids, block_size, num_experts, ignore_invalid_expert=True + ), + _run_cuda(topk_ids, block_size, num_experts, ignore_invalid_expert=True), + block_size, + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__]))