[MoE] Single-launch moe_align for tiny batches with many experts (#32395)

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.",
+    )
+)

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,
+    )

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:

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__]))