From bac001112b895fdff8cb0f167abca012198c14b9 Mon Sep 17 00:00:00 2001 From: Johnsonms Date: Fri, 20 Feb 2026 08:31:10 +0000 Subject: [PATCH 1/4] [jit_kernel] Add prepare_moe_input JIT kernels (port from sgl-kernel AOT) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Port sgl-kernel/csrc/moe/prepare_moe_input.cu to the lightweight JIT kernel system, matching the sgl_kernel call signatures. Replaces CUTLASS (cutlass::Array) and FlashInfer (flashinfer::vec_t) dependencies with self-contained lightweight equivalents: - cutlass::Array → uint4 for 128-bit dtype-agnostic row copies - flashinfer::vec_t → FloatVec (float accumulator with cast_load/cast_store using CUDA intrinsics) Three exported functions: - prepare_moe_input: compute expert offsets, GEMM problem sizes, and token permutations (input_permutation, output_permutation); supports optional blockscale_offsets for FP8 MX block-scale workflows - shuffle_rows: dtype-agnostic 128-bit vectorized row gather (returns new tensor matching sgl_kernel.shuffle_rows API) - apply_shuffle_mul_sum: gather rows by permutation, scale by topk weights, and reduce over the topk dimension using float32 accumulators Files: - csrc/moe/prepare_moe_input.cuh: all kernels + tvm-ffi host launchers - prepare_moe_input.py: cache_once/load_jit wrappers, register_custom_op - tests/test_prepare_moe_input.py: 75 correctness tests vs PyTorch reference and optional cross-validation vs sgl_kernel AOT - benchmark/bench_prepare_moe_input.py: JIT vs AOT throughput comparison --- .../benchmark/bench_prepare_moe_input.py | 244 +++++++++++ .../jit_kernel/csrc/moe/prepare_moe_input.cuh | 385 ++++++++++++++++++ python/sglang/jit_kernel/prepare_moe_input.py | 209 ++++++++++ .../tests/test_prepare_moe_input.py | 303 ++++++++++++++ 4 files changed, 1141 insertions(+) create mode 100644 python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py create mode 100644 python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh create mode 100644 python/sglang/jit_kernel/prepare_moe_input.py create mode 100644 python/sglang/jit_kernel/tests/test_prepare_moe_input.py diff --git a/python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py b/python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py new file mode 100644 index 000000000000..fefe4c78a658 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py @@ -0,0 +1,244 @@ +""" +Benchmark: prepare_moe_input JIT vs AOT (sgl_kernel) + +Measures throughput (µs) for prepare_moe_input, shuffle_rows, +and apply_shuffle_mul_sum across typical MoE configurations. + +Run: + python python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py +""" + +import itertools + +import torch +import triton +import triton.testing + +from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.jit_kernel.prepare_moe_input import ( + apply_shuffle_mul_sum as apply_shuffle_mul_sum_jit, +) +from sglang.jit_kernel.prepare_moe_input import ( + prepare_moe_input as prepare_moe_input_jit, +) +from sglang.jit_kernel.prepare_moe_input import shuffle_rows as shuffle_rows_jit + +try: + from sgl_kernel import apply_shuffle_mul_sum as apply_shuffle_mul_sum_aot + from sgl_kernel import prepare_moe_input as prepare_moe_input_aot + from sgl_kernel import shuffle_rows as shuffle_rows_aot + + AOT_AVAILABLE = True +except ImportError: + apply_shuffle_mul_sum_aot = None + prepare_moe_input_aot = None + shuffle_rows_aot = None + AOT_AVAILABLE = False + +# --------------------------------------------------------------------------- +# Benchmark configuration +# --------------------------------------------------------------------------- + +NUM_TOKENS_RANGE = get_benchmark_range( + full_range=[1, 64, 256, 1024, 4096], + ci_range=[64, 512], +) + +# (topk, num_experts) typical MoE configs +MOE_CONFIGS = get_benchmark_range( + full_range=[(2, 8), (4, 16), (8, 64)], + ci_range=[(4, 16)], +) + +LINE_VALS = ["jit", "aot"] if AOT_AVAILABLE else ["jit"] +LINE_NAMES = ["JIT (new)", "AOT sgl_kernel"] if AOT_AVAILABLE else ["JIT (new)"] +STYLES = [(("blue", "--")), (("orange", "-"))] if AOT_AVAILABLE else [("blue", "--")] + + +# --------------------------------------------------------------------------- +# Benchmark: prepare_moe_input +# --------------------------------------------------------------------------- + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens", "topk", "num_experts"], + x_vals=[(nt, tk, ne) for nt, (tk, ne) in itertools.product(NUM_TOKENS_RANGE, MOE_CONFIGS)], + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="prepare-moe-input-performance", + args={}, + ) +) +def bench_prepare_moe_input(num_tokens: int, topk: int, num_experts: int, provider: str): + device = "cuda" + n, k = 4096, 4096 + topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32, device=device) + + def alloc(): + return ( + torch.empty(num_experts + 1, dtype=torch.int32, device=device), + torch.empty((num_experts, 3), dtype=torch.int32, device=device), + torch.empty((num_experts, 3), dtype=torch.int32, device=device), + torch.empty(num_tokens * topk, dtype=torch.int32, device=device), + torch.empty(num_tokens * topk, dtype=torch.int32, device=device), + ) + + eo, ps1, ps2, ip, op = alloc() + + if provider == "jit": + fn = lambda: prepare_moe_input_jit(topk_ids, eo, ps1, ps2, ip, op, num_experts, n, k) + elif provider == "aot": + fn = lambda: prepare_moe_input_aot(topk_ids, eo, ps1, ps2, ip, op, num_experts, n, k) + else: + raise ValueError(f"Unknown provider: {provider}") + + return run_benchmark(fn) + + +# --------------------------------------------------------------------------- +# Benchmark: shuffle_rows +# --------------------------------------------------------------------------- + +SHUFFLE_CONFIGS = get_benchmark_range( + full_range=[(4096, 256), (4096, 4096), (32768, 512)], + ci_range=[(4096, 4096)], +) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_rows", "num_cols"], + x_vals=SHUFFLE_CONFIGS, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="shuffle-rows-performance", + args={}, + ) +) +def bench_shuffle_rows(num_rows: int, num_cols: int, provider: str): + device = "cuda" + dtype = torch.bfloat16 + input_t = torch.randn((num_rows, num_cols), dtype=dtype, device=device) + dst2src = torch.randperm(num_rows, device=device).to(torch.int32) + + if provider == "jit": + fn = lambda: shuffle_rows_jit(input_t, dst2src, (num_rows, num_cols)) + elif provider == "aot": + fn = lambda: shuffle_rows_aot(input_t, dst2src, (num_rows, num_cols)) + else: + raise ValueError(f"Unknown provider: {provider}") + + return run_benchmark(fn) + + +# --------------------------------------------------------------------------- +# Benchmark: apply_shuffle_mul_sum +# --------------------------------------------------------------------------- + +APPLY_CONFIGS = get_benchmark_range( + full_range=[(512, 2, 4096), (1024, 4, 4096), (2048, 8, 4096)], + ci_range=[(512, 4, 4096)], +) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["m", "topk", "k"], + x_vals=APPLY_CONFIGS, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="apply-shuffle-mul-sum-performance", + args={}, + ) +) +def bench_apply_shuffle_mul_sum(m: int, topk: int, k: int, provider: str): + device = "cuda" + dtype = torch.bfloat16 + perm = torch.randperm(m * topk, device=device).to(torch.int32) + input_t = torch.randn((m * topk, k), dtype=dtype, device=device) + factors = torch.rand((m * topk,), dtype=dtype, device=device) + output = torch.empty((m, k), dtype=dtype, device=device) + + if provider == "jit": + fn = lambda: apply_shuffle_mul_sum_jit(input_t, output, perm, factors) + elif provider == "aot": + fn = lambda: apply_shuffle_mul_sum_aot(input_t, output, perm, factors) + else: + raise ValueError(f"Unknown provider: {provider}") + + return run_benchmark(fn) + + +# --------------------------------------------------------------------------- +# Quick correctness diff +# --------------------------------------------------------------------------- + + +def calculate_diff(): + if not AOT_AVAILABLE: + print("sgl_kernel not available — skipping AOT diff check") + return + + device = "cuda" + print("Correctness diff (JIT vs AOT):") + + # prepare_moe_input + for num_tokens, topk, num_experts in [(64, 4, 8), (256, 2, 16)]: + topk_ids = torch.randint( + 0, num_experts, (num_tokens, topk), dtype=torch.int32, device=device + ) + n, k = 256, 512 + + def alloc(): + return ( + torch.empty(num_experts + 1, dtype=torch.int32, device=device), + torch.empty((num_experts, 3), dtype=torch.int32, device=device), + torch.empty((num_experts, 3), dtype=torch.int32, device=device), + torch.empty(num_tokens * topk, dtype=torch.int32, device=device), + torch.empty(num_tokens * topk, dtype=torch.int32, device=device), + ) + + eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit = alloc() + eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot = alloc() + prepare_moe_input_jit(topk_ids, eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit, num_experts, n, k) + prepare_moe_input_aot(topk_ids, eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot, num_experts, n, k) + + match = torch.equal(eo_jit, eo_aot) and torch.equal(ps1_jit, ps1_aot) + status = "OK" if match else "MISMATCH" + print( + f" prepare_moe_input tokens={num_tokens:4d} topk={topk} experts={num_experts:3d}" + f" [{status}]" + ) + + # apply_shuffle_mul_sum + for dtype, m, topk, k in [(torch.bfloat16, 64, 4, 1024), (torch.float16, 128, 2, 512)]: + perm = torch.randperm(m * topk, device=device).to(torch.int32) + input_t = torch.randn((m * topk, k), dtype=dtype, device=device) + factors = torch.rand((m * topk,), dtype=dtype, device=device) + out_jit = torch.empty((m, k), dtype=dtype, device=device) + out_aot = torch.empty((m, k), dtype=dtype, device=device) + apply_shuffle_mul_sum_jit(input_t, out_jit, perm, factors) + apply_shuffle_mul_sum_aot(input_t, out_aot, perm, factors) + match = torch.allclose(out_jit, out_aot, atol=1e-2, rtol=1e-3) + status = "OK" if match else "MISMATCH" + print( + f" apply_shuffle_mul_sum dtype={dtype} m={m} topk={topk} k={k} [{status}]" + ) + + +if __name__ == "__main__": + calculate_diff() + print() + bench_prepare_moe_input.run(print_data=True) + bench_shuffle_rows.run(print_data=True) + bench_apply_shuffle_mul_sum.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh b/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh new file mode 100644 index 000000000000..ea814c19362c --- /dev/null +++ b/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh @@ -0,0 +1,385 @@ +// Adapt from sgl-kernel/csrc/moe/prepare_moe_input.cu +#pragma once + +#include +#include +#include + +#include +#include + +#include +#include + +using tvm::ffi::TensorView; + +namespace { + +constexpr uint64_t kThreadsPerExpert = 512; + +// --------------------------------------------------------------------------- +// prepare_moe_input kernels — compute problem sizes and sorted permutations +// --------------------------------------------------------------------------- + +__global__ void compute_problem_sizes_kernel( + const int32_t* __restrict__ topk_ids, + int32_t* __restrict__ problem_sizes1, + int32_t* __restrict__ problem_sizes2, + int32_t* __restrict__ atomic_buffer, + const int64_t topk_length, + const int64_t n, + const int64_t k) { + int expert_id = blockIdx.x; + int occurrences = 0; + for (int64_t i = threadIdx.x; i < topk_length; i += kThreadsPerExpert) { + occurrences += (topk_ids[i] == expert_id); + } + atomicAdd(&atomic_buffer[expert_id], occurrences); + __syncthreads(); + + if (threadIdx.x == 0) { + int final_occurrences = atomic_buffer[expert_id]; + problem_sizes1[expert_id * 3] = final_occurrences; + problem_sizes1[expert_id * 3 + 1] = static_cast(2 * n); + problem_sizes1[expert_id * 3 + 2] = static_cast(k); + problem_sizes2[expert_id * 3] = final_occurrences; + problem_sizes2[expert_id * 3 + 1] = static_cast(k); + problem_sizes2[expert_id * 3 + 2] = static_cast(n); + } +} + +__global__ void compute_expert_offsets_kernel( + const int32_t* __restrict__ problem_sizes1, + int32_t* __restrict__ expert_offsets, + int32_t* __restrict__ atomic_buffer, + const int64_t num_experts) { + int32_t tot_offset = 0; + expert_offsets[0] = 0; + for (int64_t i = 0; i < num_experts; ++i) { + atomic_buffer[i] = tot_offset; + tot_offset += problem_sizes1[i * 3]; + expert_offsets[i + 1] = tot_offset; + } +} + +__global__ void compute_expert_blockscale_offsets_kernel( + const int32_t* __restrict__ problem_sizes1, + int32_t* __restrict__ expert_offsets, + int32_t* __restrict__ blockscale_offsets, + int32_t* __restrict__ atomic_buffer, + const int64_t num_experts) { + int32_t tot_offset = 0; + int32_t tot_rounded_offset = 0; + expert_offsets[0] = 0; + blockscale_offsets[0] = 0; + for (int64_t i = 0; i < num_experts; ++i) { + atomic_buffer[i] = tot_offset; + int num_tokens = problem_sizes1[i * 3]; + int rounded_num_tokens = (num_tokens + 127) / 128 * 128; + tot_offset += num_tokens; + tot_rounded_offset += rounded_num_tokens; + expert_offsets[i + 1] = tot_offset; + blockscale_offsets[i + 1] = tot_rounded_offset; + } +} + +__global__ void compute_arg_sorts_kernel( + const int32_t* __restrict__ topk_ids, + int32_t* __restrict__ input_permutation, + int32_t* __restrict__ output_permutation, + int32_t* __restrict__ atomic_buffer, + const int64_t topk_length, + const int64_t topk) { + int expert_id = blockIdx.x; + for (int64_t i = threadIdx.x; i < topk_length; i += kThreadsPerExpert) { + if (topk_ids[i] == expert_id) { + int start = atomicAdd(&atomic_buffer[expert_id], 1); + input_permutation[start] = i / topk; + output_permutation[i] = start; + } + } +} + +// --------------------------------------------------------------------------- +// shuffle_rows kernel — dtype-agnostic 128-bit vectorized row copy +// --------------------------------------------------------------------------- + +__global__ void shuffle_rows_kernel( + const void* __restrict__ input, + const int32_t* __restrict__ dst2src_map, + void* __restrict__ output, + int64_t num_dst_rows, + int64_t row_bytes) { + int64_t dst_row = static_cast(blockIdx.x); + if (dst_row >= num_dst_rows) return; + + int64_t src_row = dst2src_map[dst_row]; + const char* src_ptr = static_cast(input) + src_row * row_bytes; + char* dst_ptr = static_cast(output) + dst_row * row_bytes; + + // 128-bit (16-byte) vectorized copy + int64_t num_vecs = row_bytes / 16; + const uint4* src_vec = reinterpret_cast(src_ptr); + uint4* dst_vec = reinterpret_cast(dst_ptr); + for (int64_t i = threadIdx.x; i < num_vecs; i += blockDim.x) { + dst_vec[i] = src_vec[i]; + } + // scalar remainder + int64_t rem_start = num_vecs * 16; + for (int64_t b = rem_start + threadIdx.x; b < row_bytes; b += blockDim.x) { + dst_ptr[b] = src_ptr[b]; + } +} + +// --------------------------------------------------------------------------- +// apply_shuffle_mul_sum kernel — gather, scale, and reduce over topk +// Accumulates in float32 regardless of input dtype. +// --------------------------------------------------------------------------- + +template +struct FloatVec { + static constexpr uint32_t N = 16u / sizeof(scalar_t); + float data[N]; + + __device__ void fill(float v) { +#pragma unroll + for (uint32_t i = 0; i < N; ++i) data[i] = v; + } + + __device__ void cast_load(const scalar_t* ptr) { +#pragma unroll + for (uint32_t i = 0; i < N; ++i) { + if constexpr (std::is_same_v) + data[i] = __half2float(ptr[i]); + else if constexpr (std::is_same_v) + data[i] = __bfloat162float(ptr[i]); + else + data[i] = static_cast(ptr[i]); + } + } + + __device__ void cast_store(scalar_t* ptr) const { +#pragma unroll + for (uint32_t i = 0; i < N; ++i) { + if constexpr (std::is_same_v) + ptr[i] = __float2half(data[i]); + else if constexpr (std::is_same_v) + ptr[i] = __float2bfloat16(data[i]); + else + ptr[i] = static_cast(data[i]); + } + } + + __device__ float& operator[](uint32_t i) { return data[i]; } + __device__ const float& operator[](uint32_t i) const { return data[i]; } +}; + +template +__global__ void apply_shuffle_mul_sum_kernel( + const scalar_t* __restrict__ input, // [m * topk, row_stride] + scalar_t* __restrict__ output, // [m, row_stride] + const int32_t* __restrict__ permutation, // [m * topk] + int m, + int topk, + int row_stride, + const scalar_t* __restrict__ factors) // [m * topk], or nullptr +{ + int i = blockIdx.x; + if (i >= m) return; + + using vec_t = FloatVec; + constexpr uint32_t vec_size = vec_t::N; + + int thread_idx = threadIdx.x; + int stride = blockDim.x; + + // Vectorized part + for (int d_vec = thread_idx; d_vec < row_stride / (int)vec_size; d_vec += stride) { + int d = d_vec * (int)vec_size; + vec_t sum_vec; + sum_vec.fill(0.0f); + + for (int j = 0; j < topk; ++j) { + int token_major_idx = i * topk + j; + int src_row = permutation[token_major_idx]; + + vec_t val_vec; + val_vec.cast_load(input + src_row * row_stride + d); + + float factor = (factors != nullptr) ? static_cast(factors[token_major_idx]) : 1.0f; +#pragma unroll + for (uint32_t k = 0; k < vec_size; ++k) { + sum_vec[k] += factor * val_vec[k]; + } + } + sum_vec.cast_store(output + i * row_stride + d); + } + + // Scalar remainder + int rem_start = (row_stride / (int)vec_size) * (int)vec_size; + for (int d = rem_start + thread_idx; d < row_stride; d += stride) { + float sum_val = 0.0f; + for (int j = 0; j < topk; ++j) { + int token_major_idx = i * topk + j; + int src_row = permutation[token_major_idx]; + float val = static_cast(input[src_row * row_stride + d]); + float factor = (factors != nullptr) ? static_cast(factors[token_major_idx]) : 1.0f; + sum_val += factor * val; + } + if constexpr (std::is_same_v) + output[i * row_stride + d] = __float2half(sum_val); + else if constexpr (std::is_same_v) + output[i * row_stride + d] = __float2bfloat16(sum_val); + else + output[i * row_stride + d] = static_cast(sum_val); + } +} + +} // namespace + +// --------------------------------------------------------------------------- +// Host launchers (tvm-ffi interface) +// --------------------------------------------------------------------------- + +void prepare_moe_input( + TensorView topk_ids, + TensorView expert_offsets, + tvm::ffi::Optional blockscale_offsets, + TensorView problem_sizes1, + TensorView problem_sizes2, + TensorView input_permutation, + TensorView output_permutation, + TensorView atomic_buffer, + int64_t num_experts, + int64_t n, + int64_t k) { + using namespace host; + + RuntimeCheck(topk_ids.dtype().code == kDLInt && topk_ids.dtype().bits == 32, "topk_ids must be int32"); + RuntimeCheck(num_experts > 0, "num_experts must be positive"); + + const int64_t topk_length = [&] { + int64_t total = 1; + for (int d = 0; d < topk_ids.dim(); ++d) total *= topk_ids.size(d); + return total; + }(); + + const int64_t topk = (topk_ids.dim() >= 2) ? topk_ids.size(1) : 1; + + const int32_t* topk_ids_ptr = static_cast(topk_ids.data_ptr()); + int32_t* expert_offsets_ptr = static_cast(expert_offsets.data_ptr()); + int32_t* problem_sizes1_ptr = static_cast(problem_sizes1.data_ptr()); + int32_t* problem_sizes2_ptr = static_cast(problem_sizes2.data_ptr()); + int32_t* input_perm_ptr = static_cast(input_permutation.data_ptr()); + int32_t* output_perm_ptr = static_cast(output_permutation.data_ptr()); + int32_t* atomic_ptr = static_cast(atomic_buffer.data_ptr()); + + cudaStream_t stream = LaunchKernel::resolve_device(topk_ids.device()); + + uint32_t num_threads = static_cast(std::min((int64_t)kThreadsPerExpert, topk_length)); + uint32_t num_blocks = static_cast(num_experts); + + compute_problem_sizes_kernel<<>>( + topk_ids_ptr, problem_sizes1_ptr, problem_sizes2_ptr, atomic_ptr, topk_length, n, k); + + if (blockscale_offsets.has_value()) { + int32_t* bs_offsets_ptr = static_cast(blockscale_offsets.value().data_ptr()); + compute_expert_blockscale_offsets_kernel<<<1, 1, 0, stream>>>( + problem_sizes1_ptr, expert_offsets_ptr, bs_offsets_ptr, atomic_ptr, num_experts); + } else { + compute_expert_offsets_kernel<<<1, 1, 0, stream>>>( + problem_sizes1_ptr, expert_offsets_ptr, atomic_ptr, num_experts); + } + + compute_arg_sorts_kernel<<>>( + topk_ids_ptr, input_perm_ptr, output_perm_ptr, atomic_ptr, topk_length, topk); +} + +void shuffle_rows(TensorView input, TensorView dst2src_map, TensorView output) { + using namespace host; + + RuntimeCheck(input.dim() == 2, "input must be 2-D [num_src_rows, num_cols]"); + RuntimeCheck(output.dim() == 2, "output must be 2-D [num_dst_rows, num_cols]"); + RuntimeCheck(dst2src_map.dim() == 1, "dst2src_map must be 1-D"); + RuntimeCheck(input.size(1) == output.size(1), "num_cols mismatch between input and output"); + RuntimeCheck(output.size(0) == dst2src_map.size(0), "num_dst_rows mismatch"); + + const int64_t num_dst_rows = output.size(0); + const int64_t num_cols = input.size(1); + + // Compute element byte size from dtype + const int64_t elem_bytes = static_cast(input.dtype().bits) / 8; + const int64_t row_bytes = num_cols * elem_bytes; + + RuntimeCheck(row_bytes % 16 == 0, "row_bytes must be a multiple of 16 for vectorized copy"); + + cudaStream_t stream = LaunchKernel::resolve_device(input.device()); + + dim3 grid(static_cast(num_dst_rows)); + dim3 block(256); + + shuffle_rows_kernel<<>>( + input.data_ptr(), + static_cast(dst2src_map.data_ptr()), + output.data_ptr(), + num_dst_rows, + row_bytes); +} + +template +static void launch_apply_shuffle_mul_sum( + TensorView input, + TensorView output, + TensorView permutation, + tvm::ffi::Optional factors) { + using namespace host; + const int m = static_cast(output.size(0)); + const int topk_times_m = static_cast(input.size(0)); + const int topk = (m > 0) ? topk_times_m / m : 1; + const int row_stride = static_cast(output.size(1)); + + constexpr uint32_t vec_size = 16u / sizeof(scalar_t); + const uint32_t block_threads = static_cast( + std::min(static_cast(row_stride / (int)vec_size), 1024)); + + const scalar_t* factors_ptr = factors.has_value() + ? static_cast(factors.value().data_ptr()) + : nullptr; + + cudaStream_t stream = LaunchKernel::resolve_device(input.device()); + + dim3 grid(static_cast(m)); + dim3 block(block_threads > 0 ? block_threads : 1u); + + apply_shuffle_mul_sum_kernel<<>>( + static_cast(input.data_ptr()), + static_cast(output.data_ptr()), + static_cast(permutation.data_ptr()), + m, + topk, + row_stride, + factors_ptr); +} + +void apply_shuffle_mul_sum( + TensorView input, + TensorView output, + TensorView permutation, + tvm::ffi::Optional factors) { + using namespace host; + + RuntimeCheck(input.dim() == 2, "input must be 2-D [m * topk, row_stride]"); + RuntimeCheck(output.dim() == 2, "output must be 2-D [m, row_stride]"); + RuntimeCheck(permutation.dim() == 1, "permutation must be 1-D [m * topk]"); + RuntimeCheck(input.size(1) == output.size(1), "row_stride mismatch"); + + auto dtype = output.dtype(); + if (host::is_type(dtype)) + launch_apply_shuffle_mul_sum(input, output, permutation, factors); + else if (host::is_type(dtype)) + launch_apply_shuffle_mul_sum(input, output, permutation, factors); + else if (host::is_type(dtype)) + launch_apply_shuffle_mul_sum(input, output, permutation, factors); + else + RuntimeCheck(false, "apply_shuffle_mul_sum: unsupported dtype (expected fp16, bf16, or fp32)"); +} diff --git a/python/sglang/jit_kernel/prepare_moe_input.py b/python/sglang/jit_kernel/prepare_moe_input.py new file mode 100644 index 000000000000..686ab7324add --- /dev/null +++ b/python/sglang/jit_kernel/prepare_moe_input.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Optional + +import torch + +from sglang.jit_kernel.utils import cache_once, load_jit +from sglang.srt.utils.custom_op import register_custom_op + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +@cache_once +def _jit_prepare_moe_input_module() -> Module: + return load_jit( + "prepare_moe_input", + cuda_files=["moe/prepare_moe_input.cuh"], + cuda_wrappers=[ + ("prepare_moe_input", "prepare_moe_input"), + ("shuffle_rows", "shuffle_rows"), + ("apply_shuffle_mul_sum", "apply_shuffle_mul_sum"), + ], + ) + + +@register_custom_op( + op_name="prepare_moe_input_out", + mutates_args=[ + "expert_offsets", + "problem_sizes1", + "problem_sizes2", + "input_permutation", + "output_permutation", + ], +) +def prepare_moe_input_out( + topk_ids: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: Optional[torch.Tensor], + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + input_permutation: torch.Tensor, + output_permutation: torch.Tensor, + num_experts: int, + n: int, + k: int, +) -> None: + """ + Compute MoE routing metadata and sorted token permutations. + + Args: + topk_ids: [num_tokens, topk] int32 — expert indices per token + expert_offsets: [num_experts + 1] int32 — output cumulative token offsets + blockscale_offsets: [num_experts + 1] int32 (optional) — rounded offsets for FP8 block scales + problem_sizes1: [num_experts, 3] int32 — GEMM problem sizes (gate/up projection) + problem_sizes2: [num_experts, 3] int32 — GEMM problem sizes (down projection) + input_permutation: [num_tokens * topk] int32 — token index per expert slot + output_permutation: [num_tokens * topk] int32 — expert slot index per token + num_experts: total number of experts + n: hidden size / 2 (intermediate size per partition) + k: input hidden size + """ + atomic_buffer = torch.zeros(num_experts, dtype=torch.int32, device=topk_ids.device) + module = _jit_prepare_moe_input_module() + module.prepare_moe_input( + topk_ids, + expert_offsets, + blockscale_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + atomic_buffer, + num_experts, + n, + k, + ) + + +def prepare_moe_input( + topk_ids: torch.Tensor, + expert_offsets: torch.Tensor, + problem_sizes1: torch.Tensor, + problem_sizes2: torch.Tensor, + input_permutation: torch.Tensor, + output_permutation: torch.Tensor, + num_experts: int, + n: int, + k: int, + blockscale_offsets: Optional[torch.Tensor] = None, +) -> None: + """ + Compute MoE routing metadata and sorted token permutations. + + Matches the call signature of ``sgl_kernel.prepare_moe_input``. + + Args: + topk_ids: [num_tokens, topk] int32 + expert_offsets: [num_experts + 1] int32 — written in-place + problem_sizes1: [num_experts, 3] int32 — written in-place + problem_sizes2: [num_experts, 3] int32 — written in-place + input_permutation: [num_tokens * topk] int32 — written in-place + output_permutation: [num_tokens * topk] int32 — written in-place + num_experts: total number of experts + n: intermediate size (half of gate/up projection output) + k: input hidden size + blockscale_offsets: optional [num_experts + 1] int32 for FP8 block-scale offsets + """ + prepare_moe_input_out( + topk_ids, + expert_offsets, + blockscale_offsets, + problem_sizes1, + problem_sizes2, + input_permutation, + output_permutation, + num_experts, + n, + k, + ) + + +@register_custom_op( + op_name="shuffle_rows_out", + mutates_args=["output"], +) +def shuffle_rows_out( + input: torch.Tensor, + dst2src_map: torch.Tensor, + output: torch.Tensor, +) -> None: + """ + Shuffle rows of ``input`` according to ``dst2src_map``, writing to ``output``. + + Args: + input: [num_src_rows, num_cols] + dst2src_map: [num_dst_rows] int32 — dst_row[i] = input[dst2src_map[i]] + output: [num_dst_rows, num_cols] — pre-allocated output + """ + module = _jit_prepare_moe_input_module() + module.shuffle_rows(input, dst2src_map, output) + + +def shuffle_rows( + input: torch.Tensor, + dst2src_map: torch.Tensor, + output_shape: tuple, +) -> torch.Tensor: + """ + Shuffle rows of ``input`` according to ``dst2src_map``. + + Matches the call signature of ``sgl_kernel.shuffle_rows``. + + Args: + input: [num_src_rows, num_cols] + dst2src_map: [num_dst_rows] int32 + output_shape: (num_dst_rows, num_cols) — shape of the returned tensor + + Returns: + output tensor [num_dst_rows, num_cols], same dtype as input + """ + output = torch.empty(output_shape, dtype=input.dtype, device=input.device) + shuffle_rows_out(input, dst2src_map, output) + return output + + +@register_custom_op( + op_name="apply_shuffle_mul_sum_out", + mutates_args=["output"], +) +def apply_shuffle_mul_sum_out( + input: torch.Tensor, + output: torch.Tensor, + permutation: torch.Tensor, + factors: Optional[torch.Tensor], +) -> None: + """ + Gather rows of ``input`` by ``permutation``, scale by ``factors``, and sum over topk. + + Equivalent to: + (input[permutation].view(m, topk, k) * factors.view(m, topk, 1)).sum(dim=1) + + Args: + input: [m * topk, k] fp32/fp16/bf16 + output: [m, k] — written in-place + permutation: [m * topk] int32 — maps output positions to input rows + factors: [m * topk] optional scaling factors (topk weights) + """ + module = _jit_prepare_moe_input_module() + module.apply_shuffle_mul_sum(input, output, permutation, factors) + + +def apply_shuffle_mul_sum( + input: torch.Tensor, + output: torch.Tensor, + permutation: torch.Tensor, + factors: Optional[torch.Tensor] = None, +) -> None: + """ + Gather-scale-reduce: matches the call signature of ``sgl_kernel.apply_shuffle_mul_sum``. + + Args: + input: [m * topk, k] fp32/fp16/bf16 + output: [m, k] — written in-place + permutation: [m * topk] int32 + factors: optional [m * topk] scaling factors + """ + apply_shuffle_mul_sum_out(input, output, permutation, factors) diff --git a/python/sglang/jit_kernel/tests/test_prepare_moe_input.py b/python/sglang/jit_kernel/tests/test_prepare_moe_input.py new file mode 100644 index 000000000000..b56a745f7507 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_prepare_moe_input.py @@ -0,0 +1,303 @@ +""" +Correctness tests for the prepare_moe_input JIT kernels. + +Validates prepare_moe_input, shuffle_rows, and apply_shuffle_mul_sum +against pure-PyTorch references and (when available) sgl_kernel AOT. +""" + +import itertools +import os + +import pytest +import torch + +from sglang.jit_kernel.prepare_moe_input import ( + apply_shuffle_mul_sum, + prepare_moe_input, + shuffle_rows, +) + +try: + from sgl_kernel import apply_shuffle_mul_sum as apply_shuffle_mul_sum_aot + from sgl_kernel import prepare_moe_input as prepare_moe_input_aot + from sgl_kernel import shuffle_rows as shuffle_rows_aot + + AOT_AVAILABLE = True +except ImportError: + AOT_AVAILABLE = False + +# --------------------------------------------------------------------------- +# CI / full-range helpers +# --------------------------------------------------------------------------- + +_is_ci = ( + os.getenv("CI", "false").lower() == "true" + or os.getenv("GITHUB_ACTIONS", "false").lower() == "true" +) + +NUM_TOKENS_FULL = [1, 16, 128, 512] +NUM_TOKENS_CI = [1, 64, 256] + +TOPK_FULL = [1, 2, 4, 8] +TOPK_CI = [2, 4] + +NUM_EXPERTS_FULL = [8, 16, 64] +NUM_EXPERTS_CI = [8, 32] + +NUM_TOKENS = NUM_TOKENS_CI if _is_ci else NUM_TOKENS_FULL +TOPK_LIST = TOPK_CI if _is_ci else TOPK_FULL +NUM_EXPERTS_LIST = NUM_EXPERTS_CI if _is_ci else NUM_EXPERTS_FULL + + +# --------------------------------------------------------------------------- +# Pure-PyTorch reference for prepare_moe_input +# --------------------------------------------------------------------------- + + +def prepare_moe_input_ref(topk_ids, num_experts, n, k): + """Compute expert offsets, problem sizes, and permutations using PyTorch.""" + num_tokens, topk = topk_ids.shape + topk_flat = topk_ids.flatten() # [num_tokens * topk] + + # Count tokens per expert + counts = torch.zeros(num_experts, dtype=torch.int32, device=topk_ids.device) + for eid in range(num_experts): + counts[eid] = (topk_flat == eid).sum() + + # Expert offsets (exclusive prefix sum) + expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=topk_ids.device) + for i in range(num_experts): + expert_offsets[i + 1] = expert_offsets[i] + counts[i] + + # Problem sizes + problem_sizes1 = torch.zeros((num_experts, 3), dtype=torch.int32, device=topk_ids.device) + problem_sizes2 = torch.zeros((num_experts, 3), dtype=torch.int32, device=topk_ids.device) + for eid in range(num_experts): + c = counts[eid].item() + problem_sizes1[eid] = torch.tensor([c, 2 * n, k], dtype=torch.int32) + problem_sizes2[eid] = torch.tensor([c, k, n], dtype=torch.int32) + + # Permutations: sort topk_flat by expert + input_perm = torch.empty(num_tokens * topk, dtype=torch.int32, device=topk_ids.device) + output_perm = torch.empty(num_tokens * topk, dtype=torch.int32, device=topk_ids.device) + expert_cursor = expert_offsets[:-1].clone() # running offset per expert + for i in range(num_tokens * topk): + eid = topk_flat[i].item() + slot = expert_cursor[eid].item() + input_perm[slot] = i // topk + output_perm[i] = slot + expert_cursor[eid] += 1 + + return expert_offsets, problem_sizes1, problem_sizes2, input_perm, output_perm + + +# --------------------------------------------------------------------------- +# Tests: prepare_moe_input +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "num_tokens, topk, num_experts", + list(itertools.product(NUM_TOKENS, TOPK_LIST, NUM_EXPERTS_LIST)), +) +def test_prepare_moe_input_vs_ref(num_tokens, topk, num_experts): + torch.manual_seed(num_tokens * topk * num_experts) + topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda") + n, k = 128, 256 + + expert_offsets_jit = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda") + problem_sizes1_jit = torch.empty((num_experts, 3), dtype=torch.int32, device="cuda") + problem_sizes2_jit = torch.empty((num_experts, 3), dtype=torch.int32, device="cuda") + input_perm_jit = torch.empty(num_tokens * topk, dtype=torch.int32, device="cuda") + output_perm_jit = torch.empty(num_tokens * topk, dtype=torch.int32, device="cuda") + + prepare_moe_input( + topk_ids, + expert_offsets_jit, + problem_sizes1_jit, + problem_sizes2_jit, + input_perm_jit, + output_perm_jit, + num_experts, + n, + k, + ) + + ref = prepare_moe_input_ref(topk_ids, num_experts, n, k) + expert_offsets_ref, ps1_ref, ps2_ref, ip_ref, op_ref = ref + + assert torch.equal(expert_offsets_jit, expert_offsets_ref), "expert_offsets mismatch" + assert torch.equal(problem_sizes1_jit, ps1_ref), "problem_sizes1 mismatch" + assert torch.equal(problem_sizes2_jit, ps2_ref), "problem_sizes2 mismatch" + # Permutations may differ in tie-breaking order within an expert; + # verify they produce equivalent sorted token counts per expert. + counts_jit = torch.bincount( + topk_ids.flatten()[input_perm_jit.long()], minlength=num_experts + ) + counts_ref = torch.bincount( + topk_ids.flatten()[ip_ref.long()], minlength=num_experts + ) + assert torch.equal(counts_jit, counts_ref), "input_permutation expert counts mismatch" + + +def test_prepare_moe_input_with_blockscale_offsets(): + num_tokens, topk, num_experts = 64, 4, 8 + n, k = 128, 256 + topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda") + + expert_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda") + blockscale_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda") + problem_sizes1 = torch.empty((num_experts, 3), dtype=torch.int32, device="cuda") + problem_sizes2 = torch.empty((num_experts, 3), dtype=torch.int32, device="cuda") + input_perm = torch.empty(num_tokens * topk, dtype=torch.int32, device="cuda") + output_perm = torch.empty(num_tokens * topk, dtype=torch.int32, device="cuda") + + prepare_moe_input( + topk_ids, + expert_offsets, + problem_sizes1, + problem_sizes2, + input_perm, + output_perm, + num_experts, + n, + k, + blockscale_offsets=blockscale_offsets, + ) + + # blockscale_offsets[i+1] - blockscale_offsets[i] must be a multiple of 128 + diffs = blockscale_offsets[1:] - blockscale_offsets[:-1] + assert (diffs % 128 == 0).all(), "blockscale_offsets spacing must be multiple of 128" + assert blockscale_offsets[0].item() == 0, "blockscale_offsets must start at 0" + + +# --------------------------------------------------------------------------- +# Tests: shuffle_rows +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("num_dst_rows, num_cols", [(64, 256), (128, 4096), (1, 512)]) +def test_shuffle_rows_vs_ref(dtype, num_dst_rows, num_cols): + torch.manual_seed(num_dst_rows * num_cols) + num_src_rows = num_dst_rows + 32 + input_t = torch.randn((num_src_rows, num_cols), dtype=dtype, device="cuda") + dst2src = torch.randperm(num_src_rows, device="cuda")[:num_dst_rows].to(torch.int32) + + out_jit = shuffle_rows(input_t, dst2src, (num_dst_rows, num_cols)) + out_ref = input_t[dst2src.long()] + + assert torch.equal(out_jit, out_ref), f"shuffle_rows mismatch (dtype={dtype})" + assert out_jit.shape == (num_dst_rows, num_cols) + assert out_jit.dtype == dtype + + +# --------------------------------------------------------------------------- +# Tests: apply_shuffle_mul_sum +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("m, topk, k", [(32, 2, 256), (64, 4, 1024), (128, 8, 512)]) +def test_apply_shuffle_mul_sum_vs_ref(dtype, m, topk, k): + torch.manual_seed(m * topk * k) + device = "cuda" + + # Generate a random permutation mapping [m*topk] -> row indices in input + perm = torch.randperm(m * topk, device=device).to(torch.int32) + input_t = torch.randn((m * topk, k), dtype=dtype, device=device) + factors = torch.rand((m * topk,), dtype=dtype, device=device) + output_jit = torch.empty((m, k), dtype=dtype, device=device) + + apply_shuffle_mul_sum(input_t, output_jit, perm, factors) + + # Reference: gather, reshape, weight, sum + gathered = input_t[perm.long()].view(m, topk, k).to(torch.float32) + w = factors.view(m, topk, 1).to(torch.float32) + ref = (gathered * w).sum(dim=1).to(dtype) + + atol = 0.05 if dtype != torch.float32 else 1e-4 + rtol = 1e-2 if dtype != torch.float32 else 1e-4 + assert torch.allclose(output_jit, ref, atol=atol, rtol=rtol), ( + f"apply_shuffle_mul_sum mismatch (dtype={dtype}, m={m}, topk={topk}, k={k})" + ) + + +def test_apply_shuffle_mul_sum_no_factors(): + """Without factors, should be equivalent to a simple sum (factor=1).""" + m, topk, k = 32, 4, 256 + device = "cuda" + dtype = torch.float32 + perm = torch.randperm(m * topk, device=device).to(torch.int32) + input_t = torch.randn((m * topk, k), dtype=dtype, device=device) + output = torch.empty((m, k), dtype=dtype, device=device) + + apply_shuffle_mul_sum(input_t, output, perm, None) + ref = input_t[perm.long()].view(m, topk, k).sum(dim=1) + assert torch.allclose(output, ref, atol=1e-5), "apply_shuffle_mul_sum (no factors) mismatch" + + +# --------------------------------------------------------------------------- +# Cross-validation against AOT sgl_kernel +# --------------------------------------------------------------------------- + + +@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel not available") +@pytest.mark.parametrize("num_tokens, topk, num_experts", [(64, 4, 8), (256, 2, 16)]) +def test_prepare_moe_input_vs_aot(num_tokens, topk, num_experts): + torch.manual_seed(42) + topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda") + n, k = 256, 512 + + def alloc(): + return ( + torch.empty(num_experts + 1, dtype=torch.int32, device="cuda"), + torch.empty((num_experts, 3), dtype=torch.int32, device="cuda"), + torch.empty((num_experts, 3), dtype=torch.int32, device="cuda"), + torch.empty(num_tokens * topk, dtype=torch.int32, device="cuda"), + torch.empty(num_tokens * topk, dtype=torch.int32, device="cuda"), + ) + + eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit = alloc() + eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot = alloc() + + prepare_moe_input(topk_ids, eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit, num_experts, n, k) + prepare_moe_input_aot(topk_ids, eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot, num_experts, n, k) + + assert torch.equal(eo_jit, eo_aot), "expert_offsets mismatch vs AOT" + assert torch.equal(ps1_jit, ps1_aot), "problem_sizes1 mismatch vs AOT" + assert torch.equal(ps2_jit, ps2_aot), "problem_sizes2 mismatch vs AOT" + + +@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel not available") +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) +def test_shuffle_rows_vs_aot(dtype): + torch.manual_seed(0) + input_t = torch.randn((128, 1024), dtype=dtype, device="cuda") + dst2src = torch.randperm(128, device="cuda").to(torch.int32) + + out_jit = shuffle_rows(input_t, dst2src, (128, 1024)) + out_aot = shuffle_rows_aot(input_t, dst2src, (128, 1024)) + assert torch.equal(out_jit, out_aot), f"shuffle_rows vs AOT mismatch (dtype={dtype})" + + +@pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel not available") +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_apply_shuffle_mul_sum_vs_aot(dtype): + torch.manual_seed(0) + m, topk, k = 64, 4, 1024 + perm = torch.randperm(m * topk, device="cuda").to(torch.int32) + input_t = torch.randn((m * topk, k), dtype=dtype, device="cuda") + factors = torch.rand((m * topk,), dtype=dtype, device="cuda") + + out_jit = torch.empty((m, k), dtype=dtype, device="cuda") + out_aot = torch.empty((m, k), dtype=dtype, device="cuda") + apply_shuffle_mul_sum(input_t, out_jit, perm, factors) + apply_shuffle_mul_sum_aot(input_t, out_aot, perm, factors) + assert torch.allclose(out_jit, out_aot, atol=1e-3, rtol=1e-3), ( + f"apply_shuffle_mul_sum vs AOT mismatch (dtype={dtype})" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 7245f04f240ceb00eba9fc6e219982b23d14665a Mon Sep 17 00:00:00 2001 From: Johnsonms Date: Fri, 20 Feb 2026 08:32:41 +0000 Subject: [PATCH 2/4] [jit_kernel] Fix formatting in prepare_moe_input Apply clang-format and black fixes from pre-commit hooks. --- .../benchmark/bench_prepare_moe_input.py | 34 ++++++--- .../jit_kernel/csrc/moe/prepare_moe_input.cuh | 40 +++++------ .../tests/test_prepare_moe_input.py | 72 +++++++++++++------ 3 files changed, 96 insertions(+), 50 deletions(-) diff --git a/python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py b/python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py index fefe4c78a658..0fc1c8bb0899 100644 --- a/python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py +++ b/python/sglang/jit_kernel/benchmark/bench_prepare_moe_input.py @@ -63,7 +63,10 @@ @triton.testing.perf_report( triton.testing.Benchmark( x_names=["num_tokens", "topk", "num_experts"], - x_vals=[(nt, tk, ne) for nt, (tk, ne) in itertools.product(NUM_TOKENS_RANGE, MOE_CONFIGS)], + x_vals=[ + (nt, tk, ne) + for nt, (tk, ne) in itertools.product(NUM_TOKENS_RANGE, MOE_CONFIGS) + ], line_arg="provider", line_vals=LINE_VALS, line_names=LINE_NAMES, @@ -73,10 +76,14 @@ args={}, ) ) -def bench_prepare_moe_input(num_tokens: int, topk: int, num_experts: int, provider: str): +def bench_prepare_moe_input( + num_tokens: int, topk: int, num_experts: int, provider: str +): device = "cuda" n, k = 4096, 4096 - topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32, device=device) + topk_ids = torch.randint( + 0, num_experts, (num_tokens, topk), dtype=torch.int32, device=device + ) def alloc(): return ( @@ -90,9 +97,13 @@ def alloc(): eo, ps1, ps2, ip, op = alloc() if provider == "jit": - fn = lambda: prepare_moe_input_jit(topk_ids, eo, ps1, ps2, ip, op, num_experts, n, k) + fn = lambda: prepare_moe_input_jit( + topk_ids, eo, ps1, ps2, ip, op, num_experts, n, k + ) elif provider == "aot": - fn = lambda: prepare_moe_input_aot(topk_ids, eo, ps1, ps2, ip, op, num_experts, n, k) + fn = lambda: prepare_moe_input_aot( + topk_ids, eo, ps1, ps2, ip, op, num_experts, n, k + ) else: raise ValueError(f"Unknown provider: {provider}") @@ -210,8 +221,12 @@ def alloc(): eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit = alloc() eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot = alloc() - prepare_moe_input_jit(topk_ids, eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit, num_experts, n, k) - prepare_moe_input_aot(topk_ids, eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot, num_experts, n, k) + prepare_moe_input_jit( + topk_ids, eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit, num_experts, n, k + ) + prepare_moe_input_aot( + topk_ids, eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot, num_experts, n, k + ) match = torch.equal(eo_jit, eo_aot) and torch.equal(ps1_jit, ps1_aot) status = "OK" if match else "MISMATCH" @@ -221,7 +236,10 @@ def alloc(): ) # apply_shuffle_mul_sum - for dtype, m, topk, k in [(torch.bfloat16, 64, 4, 1024), (torch.float16, 128, 2, 512)]: + for dtype, m, topk, k in [ + (torch.bfloat16, 64, 4, 1024), + (torch.float16, 128, 2, 512), + ]: perm = torch.randperm(m * topk, device=device).to(torch.int32) input_t = torch.randn((m * topk, k), dtype=dtype, device=device) factors = torch.rand((m * topk,), dtype=dtype, device=device) diff --git a/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh b/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh index ea814c19362c..a9cb144b1c8a 100644 --- a/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh +++ b/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh @@ -3,6 +3,7 @@ #include #include + #include #include @@ -143,7 +144,8 @@ struct FloatVec { __device__ void fill(float v) { #pragma unroll - for (uint32_t i = 0; i < N; ++i) data[i] = v; + for (uint32_t i = 0; i < N; ++i) + data[i] = v; } __device__ void cast_load(const scalar_t* ptr) { @@ -170,14 +172,18 @@ struct FloatVec { } } - __device__ float& operator[](uint32_t i) { return data[i]; } - __device__ const float& operator[](uint32_t i) const { return data[i]; } + __device__ float& operator[](uint32_t i) { + return data[i]; + } + __device__ const float& operator[](uint32_t i) const { + return data[i]; + } }; template __global__ void apply_shuffle_mul_sum_kernel( - const scalar_t* __restrict__ input, // [m * topk, row_stride] - scalar_t* __restrict__ output, // [m, row_stride] + const scalar_t* __restrict__ input, // [m * topk, row_stride] + scalar_t* __restrict__ output, // [m, row_stride] const int32_t* __restrict__ permutation, // [m * topk] int m, int topk, @@ -260,7 +266,8 @@ void prepare_moe_input( const int64_t topk_length = [&] { int64_t total = 1; - for (int d = 0; d < topk_ids.dim(); ++d) total *= topk_ids.size(d); + for (int d = 0; d < topk_ids.dim(); ++d) + total *= topk_ids.size(d); return total; }(); @@ -287,8 +294,7 @@ void prepare_moe_input( compute_expert_blockscale_offsets_kernel<<<1, 1, 0, stream>>>( problem_sizes1_ptr, expert_offsets_ptr, bs_offsets_ptr, atomic_ptr, num_experts); } else { - compute_expert_offsets_kernel<<<1, 1, 0, stream>>>( - problem_sizes1_ptr, expert_offsets_ptr, atomic_ptr, num_experts); + compute_expert_offsets_kernel<<<1, 1, 0, stream>>>(problem_sizes1_ptr, expert_offsets_ptr, atomic_ptr, num_experts); } compute_arg_sorts_kernel<<>>( @@ -328,10 +334,7 @@ void shuffle_rows(TensorView input, TensorView dst2src_map, TensorView output) { template static void launch_apply_shuffle_mul_sum( - TensorView input, - TensorView output, - TensorView permutation, - tvm::ffi::Optional factors) { + TensorView input, TensorView output, TensorView permutation, tvm::ffi::Optional factors) { using namespace host; const int m = static_cast(output.size(0)); const int topk_times_m = static_cast(input.size(0)); @@ -339,12 +342,10 @@ static void launch_apply_shuffle_mul_sum( const int row_stride = static_cast(output.size(1)); constexpr uint32_t vec_size = 16u / sizeof(scalar_t); - const uint32_t block_threads = static_cast( - std::min(static_cast(row_stride / (int)vec_size), 1024)); + const uint32_t block_threads = static_cast(std::min(static_cast(row_stride / (int)vec_size), 1024)); - const scalar_t* factors_ptr = factors.has_value() - ? static_cast(factors.value().data_ptr()) - : nullptr; + const scalar_t* factors_ptr = + factors.has_value() ? static_cast(factors.value().data_ptr()) : nullptr; cudaStream_t stream = LaunchKernel::resolve_device(input.device()); @@ -362,10 +363,7 @@ static void launch_apply_shuffle_mul_sum( } void apply_shuffle_mul_sum( - TensorView input, - TensorView output, - TensorView permutation, - tvm::ffi::Optional factors) { + TensorView input, TensorView output, TensorView permutation, tvm::ffi::Optional factors) { using namespace host; RuntimeCheck(input.dim() == 2, "input must be 2-D [m * topk, row_stride]"); diff --git a/python/sglang/jit_kernel/tests/test_prepare_moe_input.py b/python/sglang/jit_kernel/tests/test_prepare_moe_input.py index b56a745f7507..983dd13a61bb 100644 --- a/python/sglang/jit_kernel/tests/test_prepare_moe_input.py +++ b/python/sglang/jit_kernel/tests/test_prepare_moe_input.py @@ -65,21 +65,31 @@ def prepare_moe_input_ref(topk_ids, num_experts, n, k): counts[eid] = (topk_flat == eid).sum() # Expert offsets (exclusive prefix sum) - expert_offsets = torch.zeros(num_experts + 1, dtype=torch.int32, device=topk_ids.device) + expert_offsets = torch.zeros( + num_experts + 1, dtype=torch.int32, device=topk_ids.device + ) for i in range(num_experts): expert_offsets[i + 1] = expert_offsets[i] + counts[i] # Problem sizes - problem_sizes1 = torch.zeros((num_experts, 3), dtype=torch.int32, device=topk_ids.device) - problem_sizes2 = torch.zeros((num_experts, 3), dtype=torch.int32, device=topk_ids.device) + problem_sizes1 = torch.zeros( + (num_experts, 3), dtype=torch.int32, device=topk_ids.device + ) + problem_sizes2 = torch.zeros( + (num_experts, 3), dtype=torch.int32, device=topk_ids.device + ) for eid in range(num_experts): c = counts[eid].item() problem_sizes1[eid] = torch.tensor([c, 2 * n, k], dtype=torch.int32) problem_sizes2[eid] = torch.tensor([c, k, n], dtype=torch.int32) # Permutations: sort topk_flat by expert - input_perm = torch.empty(num_tokens * topk, dtype=torch.int32, device=topk_ids.device) - output_perm = torch.empty(num_tokens * topk, dtype=torch.int32, device=topk_ids.device) + input_perm = torch.empty( + num_tokens * topk, dtype=torch.int32, device=topk_ids.device + ) + output_perm = torch.empty( + num_tokens * topk, dtype=torch.int32, device=topk_ids.device + ) expert_cursor = expert_offsets[:-1].clone() # running offset per expert for i in range(num_tokens * topk): eid = topk_flat[i].item() @@ -102,7 +112,9 @@ def prepare_moe_input_ref(topk_ids, num_experts, n, k): ) def test_prepare_moe_input_vs_ref(num_tokens, topk, num_experts): torch.manual_seed(num_tokens * topk * num_experts) - topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda") + topk_ids = torch.randint( + 0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda" + ) n, k = 128, 256 expert_offsets_jit = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda") @@ -126,7 +138,9 @@ def test_prepare_moe_input_vs_ref(num_tokens, topk, num_experts): ref = prepare_moe_input_ref(topk_ids, num_experts, n, k) expert_offsets_ref, ps1_ref, ps2_ref, ip_ref, op_ref = ref - assert torch.equal(expert_offsets_jit, expert_offsets_ref), "expert_offsets mismatch" + assert torch.equal( + expert_offsets_jit, expert_offsets_ref + ), "expert_offsets mismatch" assert torch.equal(problem_sizes1_jit, ps1_ref), "problem_sizes1 mismatch" assert torch.equal(problem_sizes2_jit, ps2_ref), "problem_sizes2 mismatch" # Permutations may differ in tie-breaking order within an expert; @@ -137,13 +151,17 @@ def test_prepare_moe_input_vs_ref(num_tokens, topk, num_experts): counts_ref = torch.bincount( topk_ids.flatten()[ip_ref.long()], minlength=num_experts ) - assert torch.equal(counts_jit, counts_ref), "input_permutation expert counts mismatch" + assert torch.equal( + counts_jit, counts_ref + ), "input_permutation expert counts mismatch" def test_prepare_moe_input_with_blockscale_offsets(): num_tokens, topk, num_experts = 64, 4, 8 n, k = 128, 256 - topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda") + topk_ids = torch.randint( + 0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda" + ) expert_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda") blockscale_offsets = torch.empty(num_experts + 1, dtype=torch.int32, device="cuda") @@ -167,7 +185,9 @@ def test_prepare_moe_input_with_blockscale_offsets(): # blockscale_offsets[i+1] - blockscale_offsets[i] must be a multiple of 128 diffs = blockscale_offsets[1:] - blockscale_offsets[:-1] - assert (diffs % 128 == 0).all(), "blockscale_offsets spacing must be multiple of 128" + assert ( + diffs % 128 == 0 + ).all(), "blockscale_offsets spacing must be multiple of 128" assert blockscale_offsets[0].item() == 0, "blockscale_offsets must start at 0" @@ -218,9 +238,9 @@ def test_apply_shuffle_mul_sum_vs_ref(dtype, m, topk, k): atol = 0.05 if dtype != torch.float32 else 1e-4 rtol = 1e-2 if dtype != torch.float32 else 1e-4 - assert torch.allclose(output_jit, ref, atol=atol, rtol=rtol), ( - f"apply_shuffle_mul_sum mismatch (dtype={dtype}, m={m}, topk={topk}, k={k})" - ) + assert torch.allclose( + output_jit, ref, atol=atol, rtol=rtol + ), f"apply_shuffle_mul_sum mismatch (dtype={dtype}, m={m}, topk={topk}, k={k})" def test_apply_shuffle_mul_sum_no_factors(): @@ -234,7 +254,9 @@ def test_apply_shuffle_mul_sum_no_factors(): apply_shuffle_mul_sum(input_t, output, perm, None) ref = input_t[perm.long()].view(m, topk, k).sum(dim=1) - assert torch.allclose(output, ref, atol=1e-5), "apply_shuffle_mul_sum (no factors) mismatch" + assert torch.allclose( + output, ref, atol=1e-5 + ), "apply_shuffle_mul_sum (no factors) mismatch" # --------------------------------------------------------------------------- @@ -246,7 +268,9 @@ def test_apply_shuffle_mul_sum_no_factors(): @pytest.mark.parametrize("num_tokens, topk, num_experts", [(64, 4, 8), (256, 2, 16)]) def test_prepare_moe_input_vs_aot(num_tokens, topk, num_experts): torch.manual_seed(42) - topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda") + topk_ids = torch.randint( + 0, num_experts, (num_tokens, topk), dtype=torch.int32, device="cuda" + ) n, k = 256, 512 def alloc(): @@ -261,8 +285,12 @@ def alloc(): eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit = alloc() eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot = alloc() - prepare_moe_input(topk_ids, eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit, num_experts, n, k) - prepare_moe_input_aot(topk_ids, eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot, num_experts, n, k) + prepare_moe_input( + topk_ids, eo_jit, ps1_jit, ps2_jit, ip_jit, op_jit, num_experts, n, k + ) + prepare_moe_input_aot( + topk_ids, eo_aot, ps1_aot, ps2_aot, ip_aot, op_aot, num_experts, n, k + ) assert torch.equal(eo_jit, eo_aot), "expert_offsets mismatch vs AOT" assert torch.equal(ps1_jit, ps1_aot), "problem_sizes1 mismatch vs AOT" @@ -278,7 +306,9 @@ def test_shuffle_rows_vs_aot(dtype): out_jit = shuffle_rows(input_t, dst2src, (128, 1024)) out_aot = shuffle_rows_aot(input_t, dst2src, (128, 1024)) - assert torch.equal(out_jit, out_aot), f"shuffle_rows vs AOT mismatch (dtype={dtype})" + assert torch.equal( + out_jit, out_aot + ), f"shuffle_rows vs AOT mismatch (dtype={dtype})" @pytest.mark.skipif(not AOT_AVAILABLE, reason="sgl_kernel not available") @@ -294,9 +324,9 @@ def test_apply_shuffle_mul_sum_vs_aot(dtype): out_aot = torch.empty((m, k), dtype=dtype, device="cuda") apply_shuffle_mul_sum(input_t, out_jit, perm, factors) apply_shuffle_mul_sum_aot(input_t, out_aot, perm, factors) - assert torch.allclose(out_jit, out_aot, atol=1e-3, rtol=1e-3), ( - f"apply_shuffle_mul_sum vs AOT mismatch (dtype={dtype})" - ) + assert torch.allclose( + out_jit, out_aot, atol=1e-3, rtol=1e-3 + ), f"apply_shuffle_mul_sum vs AOT mismatch (dtype={dtype})" if __name__ == "__main__": From 0964eccfde0db2cf01dfce4602a7a2c442b44e21 Mon Sep 17 00:00:00 2001 From: Johnsonms Date: Fri, 20 Feb 2026 18:14:35 +0000 Subject: [PATCH 3/4] [jit_kernel] Use blockDim.x stride in prepare_moe_input kernels Replace hardcoded kThreadsPerExpert stride with blockDim.x in compute_problem_sizes_kernel and compute_arg_sorts_kernel loops, making the kernels robust to future changes in launch configuration. --- python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh b/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh index a9cb144b1c8a..03a7f1ae07f8 100644 --- a/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh +++ b/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh @@ -32,7 +32,7 @@ __global__ void compute_problem_sizes_kernel( const int64_t k) { int expert_id = blockIdx.x; int occurrences = 0; - for (int64_t i = threadIdx.x; i < topk_length; i += kThreadsPerExpert) { + for (int64_t i = threadIdx.x; i < topk_length; i += blockDim.x) { occurrences += (topk_ids[i] == expert_id); } atomicAdd(&atomic_buffer[expert_id], occurrences); @@ -92,7 +92,7 @@ __global__ void compute_arg_sorts_kernel( const int64_t topk_length, const int64_t topk) { int expert_id = blockIdx.x; - for (int64_t i = threadIdx.x; i < topk_length; i += kThreadsPerExpert) { + for (int64_t i = threadIdx.x; i < topk_length; i += blockDim.x) { if (topk_ids[i] == expert_id) { int start = atomicAdd(&atomic_buffer[expert_id], 1); input_permutation[start] = i / topk; From 8ec04b2adce80257d82c2abedb8c5f5e48d947fe Mon Sep 17 00:00:00 2001 From: Johnsonms Date: Fri, 20 Feb 2026 20:34:36 +0000 Subject: [PATCH 4/4] [jit_kernel] Use LaunchKernel for apply_shuffle_mul_sum kernel launch Replace raw <<<>>> launch with LaunchKernel for apply_shuffle_mul_sum_kernel to get post-launch CUDA error checking via RuntimeDeviceCheck. --- python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh b/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh index 03a7f1ae07f8..d5224f359827 100644 --- a/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh +++ b/python/sglang/jit_kernel/csrc/moe/prepare_moe_input.cuh @@ -352,7 +352,8 @@ static void launch_apply_shuffle_mul_sum( dim3 grid(static_cast(m)); dim3 block(block_threads > 0 ? block_threads : 1u); - apply_shuffle_mul_sum_kernel<<>>( + LaunchKernel(grid, block, input.device())( + apply_shuffle_mul_sum_kernel, static_cast(input.data_ptr()), static_cast(output.data_ptr()), static_cast(permutation.data_ptr()),