From 3f7c95d6cccb59413089505b4bcd1956e8611f52 Mon Sep 17 00:00:00 2001 From: Qi Yuhang <45795032+HydraQYH@users.noreply.github.com> Date: Wed, 29 Apr 2026 22:50:09 +0800 Subject: [PATCH] [JIT Kernel][1/2]Migrate MXFP8 Group GEMM & Quant into JIT (#23833) Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- .../jit_kernel/benchmark/bench_mxfp8_moe.py | 290 +++++++++++ ...es_sm100_mxfp8_blockscaled_group_quant.cuh | 461 ++++++++++++++++++ ...sm100_mxfp8_blockscaled_moe_group_gemm.cuh | 217 +++++++++ ...fp8_blockscaled_moe_group_gemm_functor.cuh | 64 +++ ...xfp8_blockscaled_moe_group_gemm_traits.cuh | 123 +++++ python/sglang/jit_kernel/mxfp8.py | 136 ++++++ .../sglang/jit_kernel/tests/test_mxfp8_moe.py | 153 ++++++ 7 files changed, 1444 insertions(+) create mode 100644 python/sglang/jit_kernel/benchmark/bench_mxfp8_moe.py create mode 100644 python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh create mode 100644 python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh create mode 100644 python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh create mode 100644 python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh create mode 100644 python/sglang/jit_kernel/mxfp8.py create mode 100644 python/sglang/jit_kernel/tests/test_mxfp8_moe.py diff --git a/python/sglang/jit_kernel/benchmark/bench_mxfp8_moe.py b/python/sglang/jit_kernel/benchmark/bench_mxfp8_moe.py new file mode 100644 index 000000000..11ae67904 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_mxfp8_moe.py @@ -0,0 +1,290 @@ +from __future__ import annotations + +import sys +from typing import Any + +import torch +import triton + +from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark +from sglang.jit_kernel.mxfp8 import ( + es_sm100_mxfp8_blockscaled_grouped_quant, + es_sm100_mxfp8_blockscaled_moe_grouped_gemm, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-kernel-benchmark-1-gpu-large") + + +def is_sm100_supported(device=None) -> bool: + if not torch.cuda.is_available(): + return False + return (torch.cuda.get_device_capability(device)[0] == 10) and ( + torch.version.cuda >= "12.8" + ) + + +_SM100_SUPPORTED = is_sm100_supported() + + +def _probe_sgl_kernel_group_mm() -> tuple[bool, str]: + if not _SM100_SUPPORTED: + return False, "MXFP8 MoE benchmark requires sm100+ with CUDA 12.8+." + try: + import sgl_kernel # noqa: F401 + except Exception as e: + return False, f"import sgl_kernel failed: {e}" + if not hasattr(sgl_kernel, "es_sm100_mxfp8_blockscaled_grouped_mm"): + return False, "sgl_kernel.es_sm100_mxfp8_blockscaled_grouped_mm is missing." + try: + pass + + # We assume if it's imported, it works + except Exception as e: + return False, f"calling sgl-kernel grouped_mm op failed: {e}" + return True, "" + + +_SGL_KERNEL_AVAILABLE, _SGL_KERNEL_REASON = _probe_sgl_kernel_group_mm() + + +def align(val: int, alignment: int = 128) -> int: + return int((val + alignment - 1) // alignment * alignment) + + +def _prepare_case( + total_tokens: int, n_g: int, k_g: int, num_experts: int, dtype: torch.dtype +) -> dict[str, Any]: + device = torch.device("cuda") + base = total_tokens // num_experts + rem = total_tokens % num_experts + m_per_expert = [base + (1 if i < rem else 0) for i in range(num_experts)] + + expert_offset = 0 + expert_offsets = [] + aux_expert_offset = 0 + aux_expert_offsets = [] + a_blockscale_offset = 0 + a_blockscale_offsets = [] + b_blockscale_offset = 0 + b_blockscale_offsets = [] + tokens_per_expert_list = [] + expert_ranges = [] + problem_sizes = [] + + a_list = [] + b_list = [] + for g in range(num_experts): + m_g = m_per_expert[g] + tokens_per_expert_list.append(m_g) + expert_ranges.append((expert_offset, expert_offset + m_g)) + expert_offsets.append(expert_offset) + expert_offset += m_g + + aux_expert_offsets.append(aux_expert_offset) + aux_expert_offset += n_g + + a_blockscale_offsets.append(a_blockscale_offset) + a_blockscale_offset += align(m_g, 128) + + b_blockscale_offsets.append(b_blockscale_offset) + b_blockscale_offset += n_g # n_g already align to 128 in practice + + problem_sizes.append([m_g, n_g, k_g]) + + a = torch.randn((m_g, k_g), device=device, dtype=dtype) * 0.1 + b = torch.randn((n_g, k_g), device=device, dtype=dtype) * 0.1 + a_list.append(a) + b_list.append(b) + + a = torch.concat(a_list, dim=0) + b = torch.concat(b_list, dim=0) + + _expert_offsets = torch.tensor(expert_offsets).to(device=device, dtype=torch.int32) + _aux_expert_offsets = torch.tensor(aux_expert_offsets).to( + device=device, dtype=torch.int32 + ) + _a_blockscale_offsets = torch.tensor(a_blockscale_offsets).to( + device=device, dtype=torch.int32 + ) + _b_blockscale_offsets = torch.tensor(b_blockscale_offsets).to( + device=device, dtype=torch.int32 + ) + _tokens_per_expert = torch.tensor(tokens_per_expert_list).to( + device=device, dtype=torch.int32 + ) + _problem_sizes = torch.tensor(problem_sizes).to(device=device, dtype=torch.int32) + + a_quant = torch.zeros_like(a, dtype=torch.float8_e4m3fn, device=device) + a_scale_factor = torch.zeros( + (a_blockscale_offset, k_g // 32), dtype=torch.uint8, device=device + ) + + b_quant = torch.zeros_like(b, dtype=torch.float8_e4m3fn, device=device) + b_scale_factor = torch.zeros( + (num_experts * n_g, k_g // 32), dtype=torch.uint8, device=device + ) + + # Use a global workspace to avoid allocating 1GB every time + workspace = torch.empty((1024, 1024, 1024), dtype=torch.uint8, device=device) + + es_sm100_mxfp8_blockscaled_grouped_quant( + a, + _tokens_per_expert, + _expert_offsets, + _a_blockscale_offsets, + a_quant, + a_scale_factor, + ) + + es_sm100_mxfp8_blockscaled_grouped_quant( + b, + torch.ones_like(_tokens_per_expert) * n_g, + _aux_expert_offsets, + _b_blockscale_offsets, + b_quant, + b_scale_factor, + ) + + b_quant = b_quant.view(num_experts, n_g, k_g) + b_scale_factor = b_scale_factor.view(num_experts, n_g, k_g // 32) + + sgl_b_quant = b_quant.transpose(1, 2) + sgl_b_scale_factor = b_scale_factor.transpose(1, 2) + + return { + "a": a, + "b": b.view(num_experts, n_g, k_g), + "b_quant": b_quant, + "a_quant": a_quant, + "b_scale_factor": b_scale_factor, + "a_scale_factor": a_scale_factor, + "expert_offsets": _expert_offsets, + "a_blockscale_offsets": _a_blockscale_offsets, + "tokens_per_expert": _tokens_per_expert, + "problem_sizes": _problem_sizes, + "sgl_b_quant": sgl_b_quant, + "sgl_b_scale_factor": sgl_b_scale_factor, + "workspace": workspace, + "expert_ranges": expert_ranges, + "dtype": dtype, + } + + +def _sgl_kernel_group_mm(case: dict[str, Any]) -> torch.Tensor: + from sgl_kernel import es_sm100_mxfp8_blockscaled_grouped_mm + + a_quant = case["a_quant"] + sgl_b_quant = case["sgl_b_quant"] + a_scale_factor = case["a_scale_factor"] + sgl_b_scale_factor = case["sgl_b_scale_factor"] + problem_sizes = case["problem_sizes"] + expert_offsets = case["expert_offsets"] + a_blockscale_offsets = case["a_blockscale_offsets"] + dtype = case["dtype"] + + total_tokens = a_quant.shape[0] + n_g = sgl_b_quant.shape[2] + + # sgl-kernel takes output pre-allocated + d = torch.empty((total_tokens, n_g), device=a_quant.device, dtype=dtype) + es_sm100_mxfp8_blockscaled_grouped_mm( + d, + a_quant, + sgl_b_quant, + a_scale_factor, + sgl_b_scale_factor, + problem_sizes, + expert_offsets, + a_blockscale_offsets, + ) + return d + + +shape_range = get_benchmark_range( + full_range=[ + # (total_tokens, n_g, k_g, num_experts) + (1024, 4096, 4096, 64), + (2048, 4096, 4096, 64), + (4096, 4096, 4096, 64), + ] + + [ + (total_tokens, n_g, k_g, num_experts) + for total_tokens in [32 * (2**i) for i in range(9)] # 32 to 8192 + for n_g, k_g, num_experts in [ + # DeepSeek-V3/R1, gateup, TP = 1, EP = 8 + (4096, 7168, 32), + # DeepSeek-V3/R1, down, TP = 1, EP = 8 + (7168, 2048, 32), + ] + ], + ci_range=[(1024, 2048, 2048, 8)], +) + +line_vals = ["jit"] +line_names = ["JIT MXFP8 MoE GroupMM"] +styles = [("green", "-")] + +if _SGL_KERNEL_AVAILABLE: + line_vals.append("sgl_kernel") + line_names.append("sgl-kernel MXFP8 MoE GroupMM") + styles.append(("orange", "-")) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["total_tokens", "n_g", "k_g", "num_experts"], + x_vals=shape_range, + x_log=False, + line_arg="provider", + line_vals=line_vals, + line_names=line_names, + styles=styles, + ylabel="us", + plot_name="mxfp8-moe-groupmm-performance", + args={}, + ) +) +def benchmark(total_tokens, n_g, k_g, num_experts, provider): + case = _prepare_case(total_tokens, n_g, k_g, num_experts, torch.bfloat16) + + if provider == "jit": + fn = lambda: es_sm100_mxfp8_blockscaled_moe_grouped_gemm( + case["b_quant"], + case["a_quant"], + case["b_scale_factor"], + case["a_scale_factor"], + case["expert_offsets"], + case["a_blockscale_offsets"], + case["tokens_per_expert"], + case["workspace"], + case["dtype"], + ) + elif provider == "sgl_kernel": + fn = lambda: _sgl_kernel_group_mm(case) + else: + raise ValueError(f"Unknown provider: {provider}") + + # Warm up + fn() + + # Profile + if provider == "jit": + torch.cuda.nvtx.range_push("jit") + fn() + torch.cuda.nvtx.range_pop() + elif provider == "sgl_kernel": + torch.cuda.nvtx.range_push("sgl_kernel") + fn() + torch.cuda.nvtx.range_pop() + + return run_benchmark(fn) + + +if __name__ == "__main__": + if not _SM100_SUPPORTED: + print("[skip] MXFP8 MoE GroupMM benchmark requires sm100+ with CUDA 12.8+.") + sys.exit(0) + if not _SGL_KERNEL_AVAILABLE: + print(f"[info] sgl-kernel baseline unavailable: {_SGL_KERNEL_REASON}") + benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh b/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh new file mode 100644 index 000000000..ced6e5f0a --- /dev/null +++ b/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh @@ -0,0 +1,461 @@ +#pragma once +#include +#include + +#include +#include + +#include +#include + +#include "cute/tensor.hpp" +#include +#include +#include + +namespace expert_specialization { + +using namespace cute; + +constexpr uint32_t THREAD_BLOCK_SIZE = 128; +constexpr uint32_t WARP_SIZE = 32; +constexpr int BLOCK_M = 128; +constexpr int BLOCK_K = 128; +using ThrLayout = Layout, Stride<_8, _1>>; +using ValLayout = Layout>; +using SfR2SThrLayout = Layout, Stride<_4, _1>>; +using SfR2SValLayout = Layout>; +using ScaleFactorTileLayout = Layout, _4>, Stride, _1>>; + +// Fast reciprocal. +inline __device__ float reciprocal_approximate_ftz(float a) { + float b; + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a)); + return b; +} + +// Some code references TRT-LLM: +// https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/quantization.cuh +template +__inline__ __device__ uint8_t cvt_warp_fp16_to_mxfp8(FragmentS& fragment_s, FragmentD& fragment_d) { + using FragmentSLayout = typename FragmentS::layout_type; + using FragmentDLayout = typename FragmentD::layout_type; + FragmentSLayout fragment_s_layout; + FragmentDLayout fragment_d_layout; + static_assert(is_static::value && size(fragment_s_layout) == 16); + static_assert(is_static::value && size(fragment_d_layout) == 16); + + constexpr int eles_per_thr = 16; + using ValType = typename FragmentS::element_type; + using VecType = std::conditional_t, __nv_bfloat162, __half2>; + VecType vec[8]; + // Assign vals + vec[0].x = fragment_s(Int<0>{}); + vec[0].y = fragment_s(Int<1>{}); + vec[1].x = fragment_s(Int<2>{}); + vec[1].y = fragment_s(Int<3>{}); + vec[2].x = fragment_s(Int<4>{}); + vec[2].y = fragment_s(Int<5>{}); + vec[3].x = fragment_s(Int<6>{}); + vec[3].y = fragment_s(Int<7>{}); + vec[4].x = fragment_s(Int<8>{}); + vec[4].y = fragment_s(Int<9>{}); + vec[5].x = fragment_s(Int<10>{}); + vec[5].y = fragment_s(Int<11>{}); + vec[6].x = fragment_s(Int<12>{}); + vec[6].y = fragment_s(Int<13>{}); + vec[7].x = fragment_s(Int<14>{}); + vec[7].y = fragment_s(Int<15>{}); + + auto local_max = __habs2(vec[0]); + for (int i = 1; i < eles_per_thr / 2; i++) { + local_max = __hmax2(__habs2(vec[i]), local_max); + } + local_max = __hmax2(__shfl_xor_sync(uint32_t(-1), local_max, 1), local_max); + + // Get the final absolute maximum values. + float block_max(0.0f); + if constexpr (std::is_same_v) { + block_max = __bfloat162float(__hmax(local_max.x, local_max.y)); + } else { + block_max = __half2float(__hmax(local_max.x, local_max.y)); + } + // Get the SF (max value of the vector / max value of mxfp8). + float sf_val = block_max * reciprocal_approximate_ftz(448.0f); + // 8 bits representation of the SF. + uint8_t fp8_sf_val; + + __nv_fp8_e8m0 tmp_sf_val; + tmp_sf_val.__x = __nv_cvt_float_to_e8m0(sf_val, __NV_SATFINITE, cudaRoundPosInf); + sf_val = static_cast(tmp_sf_val); + fp8_sf_val = tmp_sf_val.__x; + // Get the output scale (reciprocal of the SFValue). + float output_scale = block_max != 0.f ? reciprocal_approximate_ftz(sf_val) : 0.0f; + + // Convert the input to float. + float2 fp2_vals[eles_per_thr / 2]; + +#pragma unroll + for (int i = 0; i < eles_per_thr / 2; i++) { + if constexpr (std::is_same_v) { + fp2_vals[i] = __half22float2(vec[i]); + } else { + fp2_vals[i] = __bfloat1622float2(vec[i]); + } + fp2_vals[i].x *= output_scale; + fp2_vals[i].y *= output_scale; + } + union { + uint8_t bytes[16]; + __nv_fp8x2_e4m3 elts[8]; + } u; + u.elts[0] = __nv_fp8x2_e4m3(fp2_vals[0]); + u.elts[1] = __nv_fp8x2_e4m3(fp2_vals[1]); + u.elts[2] = __nv_fp8x2_e4m3(fp2_vals[2]); + u.elts[3] = __nv_fp8x2_e4m3(fp2_vals[3]); + u.elts[4] = __nv_fp8x2_e4m3(fp2_vals[4]); + u.elts[5] = __nv_fp8x2_e4m3(fp2_vals[5]); + u.elts[6] = __nv_fp8x2_e4m3(fp2_vals[6]); + u.elts[7] = __nv_fp8x2_e4m3(fp2_vals[7]); + fragment_d(Int<0>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[0]); + fragment_d(Int<1>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[1]); + fragment_d(Int<2>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[2]); + fragment_d(Int<3>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[3]); + fragment_d(Int<4>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[4]); + fragment_d(Int<5>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[5]); + fragment_d(Int<6>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[6]); + fragment_d(Int<7>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[7]); + fragment_d(Int<8>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[8]); + fragment_d(Int<9>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[9]); + fragment_d(Int<10>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[10]); + fragment_d(Int<11>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[11]); + fragment_d(Int<12>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[12]); + fragment_d(Int<13>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[13]); + fragment_d(Int<14>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[14]); + fragment_d(Int<15>{}) = cutlass::float_e4m3_t::bitcast(u.bytes[15]); + return fp8_sf_val; +} + +template < + typename TensorS, + typename TensorP, + typename TensorD, + typename TensorSharedSF, + typename TensorSF, + typename TiledCopyG2R, + typename TiledCopyR2G, + typename TiledCopyR2S> +__inline__ __device__ void mxfp8_group_quant_tile( + TensorS& tensor_s, + TensorP& tensor_p, + TensorD& tensor_d, + TensorSharedSF& tensor_shared_sf, + TensorSF& tensor_sf, + int m, + TiledCopyG2R& tiled_copy_g2r, + TiledCopyR2G& tiled_copy_r2g, + TiledCopyR2S& tiled_copy_r2s) { + static_assert( + size(get<0>(typename TensorS::layout_type{})) == 128 && size(get<1>(typename TensorS::layout_type{})) == 128 && + stride(get<1>(typename TensorS::layout_type{})) == 1); + static_assert( + size(get<0>(typename TensorD::layout_type{})) == 128 && size(get<1>(typename TensorD::layout_type{})) == 128 && + stride(get<1>(typename TensorD::layout_type{})) == 1); + static_assert( + size(get<0>(typename TensorP::layout_type{})) == 128 && size(get<1>(typename TensorP::layout_type{})) == 128); + static_assert( + size(get<0>(typename TensorSharedSF::layout_type{})) == 128 && + size(get<1>(typename TensorSharedSF::layout_type{})) == 4); + static_assert( + size(get<0>(typename TensorSF::layout_type{})) == 128 && size(get<1>(typename TensorSF::layout_type{})) == 4); + + using Tiler_MN = typename TiledCopyG2R::Tiler_MN; + auto tiler_mn = Tiler_MN{}; + static_assert(size<0>(tiler_mn) == 16 && size<1>(tiler_mn) == 128); + + auto tiled_tensor_s = tiled_divide(tensor_s, tiler_mn); + auto tiled_tensor_p = tiled_divide(tensor_p, tiler_mn); + auto tiled_tensor_d = tiled_divide(tensor_d, tiler_mn); + static_assert(size<2>(tiled_tensor_s) == 1); + static_assert(size<2>(tiled_tensor_p) == 1); + static_assert(size<2>(tiled_tensor_d) == 1); + auto squeeze_tiled_tensor_s = take<0, 2>(tiled_tensor_s); + auto squeeze_tiled_tensor_p = take<0, 2>(tiled_tensor_p); + auto squeeze_tiled_tensor_d = take<0, 2>(tiled_tensor_d); + + using SF_Tiler_MN = typename TiledCopyR2S::Tiler_MN; + auto sf_tiler_mn = SF_Tiler_MN{}; + static_assert(size<0>(sf_tiler_mn) == 16 && size<1>(sf_tiler_mn) == 4); + + auto tiled_tensor_sf = tiled_divide(tensor_sf, sf_tiler_mn); + auto tiled_tensor_shared_sf = tiled_divide(tensor_shared_sf, sf_tiler_mn); + auto squeeze_tiled_tensor_sf = take<0, 2>(tiled_tensor_sf); + auto squeeze_tiled_tensor_shared_sf = take<0, 2>(tiled_tensor_shared_sf); + + constexpr int tile_loop_count = size<1>(tiled_tensor_s); + constexpr int rows_in_tile = 16; + // We don't need to clear shared memory + // clear(squeeze_tiled_tensor_shared_sf); +#pragma unroll 4 + for (int t = 0; t < tile_loop_count; t++) { + if (t * rows_in_tile >= m) { + break; + } + auto current_copy_tile_s = tensor<0>(squeeze_tiled_tensor_s(_, t)); + auto current_copy_tile_p = tensor<0>(squeeze_tiled_tensor_p(_, t)); + auto current_copy_tile_d = tensor<0>(squeeze_tiled_tensor_d(_, t)); + auto current_copy_tile_sf = tensor<0>(squeeze_tiled_tensor_sf(_, t)); + auto current_copy_tile_shared_sf = tensor<0>(squeeze_tiled_tensor_shared_sf(_, t)); + + // Global to Register copy + auto thr_copy_g2r = tiled_copy_g2r.get_thread_slice(threadIdx.x); + auto thr_tile_g2r_s = thr_copy_g2r.partition_S(current_copy_tile_s); + auto thr_tile_g2r_p = thr_copy_g2r.partition_S(current_copy_tile_p); + auto input_fragment = make_fragment_like(thr_tile_g2r_s); + + // Register to Global copy + auto thr_copy_r2g = tiled_copy_r2g.get_thread_slice(threadIdx.x); + auto thr_tile_r2g_d = thr_copy_r2g.partition_D(current_copy_tile_d); + auto thr_tile_r2g_p = thr_copy_r2g.partition_D(current_copy_tile_p); + auto output_fragment = make_fragment_like(thr_tile_r2g_d); + + // Register to Shared copy + auto thr_copy_r2s = tiled_copy_r2s.get_thread_slice(threadIdx.x / 2); + auto thr_tile_r2s_shared_sf = thr_copy_r2s.partition_D(current_copy_tile_shared_sf); + auto shared_sf_fragment = make_fragment_like(thr_tile_r2s_shared_sf); + + // CopyG2R & convert & CopyR2G + copy_if(tiled_copy_g2r, thr_tile_g2r_p, thr_tile_g2r_s, input_fragment); + uint8_t fp8_sf_val = cvt_warp_fp16_to_mxfp8(input_fragment, output_fragment); + copy_if(tiled_copy_r2g, thr_tile_r2g_p, output_fragment, thr_tile_r2g_d); + shared_sf_fragment[0] = fp8_sf_val; + + // Before first copy r2s, clear shared memory and wait previous group + if (t == 0 && threadIdx.x == 0) { + // Wait for the group to have completed reading from shared memory. + cuda::ptx::cp_async_bulk_wait_group_read(cuda::ptx::n32_t<0>()); + } + __syncthreads(); + + if (threadIdx.x % 2 == 0) { + copy(tiled_copy_r2s, shared_sf_fragment, thr_tile_r2s_shared_sf); + } + __syncthreads(); + } + + // Wait for shared memory writes to be visible to TMA engine. + cuda::ptx::fence_proxy_async(cuda::ptx::space_shared); // b) + __syncthreads(); + + if (threadIdx.x == 0) { + cuda::ptx::cp_async_bulk( + cuda::ptx::space_global, + cuda::ptx::space_shared, + squeeze_tiled_tensor_sf.data().get(), + squeeze_tiled_tensor_shared_sf.data().get(), + 512); + // Wait for TMA transfer to have finished reading shared memory. + // Create a "bulk async-group" out of the previous bulk copy operation. + cuda::ptx::cp_async_bulk_commit_group(); + } + __syncthreads(); +} + +template +__global__ void mxfp8_group_quant( + const T_IN* input, + const int* tokens_per_expert, + const int* expert_offsets, + const int* blockscale_offsets, + cutlass::float_e4m3_t* quant_output, + uint8_t* scale_factor, + int groups, + int k, + TiledCopyG2R tiled_copy_g2r, + TiledCopyR2G tiled_copy_r2g, + TiledCopyR2S tiled_copy_r2s) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000 + __shared__ __align__(512) uint8_t shared_memory[512]; + ScaleFactorTileLayout scale_factor_tile_layout{}; + auto scale_factor_shared = make_tensor( + make_smem_ptr(shared_memory), + scale_factor_tile_layout); // ((_32,_4), _4):((_16,_4), _1) + // Transform Groupwise Schedule into Flatten Schedule + uint group_total_tiles = 0; + uint head_cta_id = 0; + for (int g = 0; g < groups; g++) { + int m = tokens_per_expert[g]; + int64_t expert_offset = static_cast(expert_offsets[g]); + int64_t blockscale_offset = static_cast(blockscale_offsets[g]); + + auto input_tensor = make_tensor( + make_gmem_ptr(input + expert_offset * k), + make_layout(make_shape(m, k), LayoutRight{})); // (M, K):(K, 1) half_t/bfloat16_t + + auto quant_output_tensor = make_tensor( + make_gmem_ptr(quant_output + expert_offset * k), + make_layout(make_shape(m, k), LayoutRight{})); // (M, K):(K, 1) cutlass::float_e4m3_t + + auto scale_factor_shape = make_shape(ceil_div(m, 128) * 128, k / 32); + auto scale_factor_layout = tile_to_shape(scale_factor_tile_layout, scale_factor_shape, LayoutRight{}); + // layout<0>(layout<0>(scale_factor_layout)) (_32,_4):(_16,_4) -- static + // layout<1>(layout<0>(scale_factor_layout)) M_align_128 / 128 -- dynamic shape dynamic stride + // layout<0>(layout<1>(scale_factor_layout)) _4:_1 -- static + // layout<1>(layout<1>(scale_factor_layout)) (K / 32) / 4 : _512 -- dynamic shape static stride + + // Reshape to zipped layout for 1D indexing + auto zipped_scale_factor_layout = make_layout( + make_layout(layout<0>(layout<0>(scale_factor_layout)), layout<0>(layout<1>(scale_factor_layout))), + make_layout( + layout<1>(layout<0>(scale_factor_layout)), + layout<1>(layout<1>( + scale_factor_layout)))); // (((_32,_4),_4),(M_align_128 / 128,(K / 32) / 4)):(((_16,_4),_1),(?,_512)) + + auto scale_factor_tensor = + make_tensor(make_gmem_ptr(scale_factor + blockscale_offset * (k / 32)), zipped_scale_factor_layout); + + // Used for cases where M is not divisible by 128 (most scenarios). + auto input_shape = shape(input_tensor); // (M, K):(K, 1) + auto identity_tensor = make_identity_tensor(input_shape); + auto predict_tensor = cute::lazy::transform(identity_tensor, [&](auto c) { return elem_less(c, input_shape); }); + + // (_128, _128) + auto tiler = make_shape(Int{}, Int{}); + + auto tiled_input_tensor = zipped_divide(input_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128))) + auto tiled_quant_output_tensor = + zipped_divide(quant_output_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128))) + auto tiled_predict_tensor = zipped_divide(predict_tensor, tiler); // ((128, 128), (cdiv(M, 128), cdiv(K, 128))) + + auto total_tiles = size<1>(tiled_input_tensor); // cdiv(M, 128) * cdiv(K, 128) + group_total_tiles += total_tiles; + auto blk_offset = (blockIdx.x + (gridDim.x - head_cta_id)) % gridDim.x; + head_cta_id = group_total_tiles % gridDim.x; + while (blk_offset < total_tiles) { + auto current_input_tile = tensor<0>(tiled_input_tensor(_, blk_offset)); + auto current_quant_output_tile = tensor<0>(tiled_quant_output_tensor(_, blk_offset)); + auto current_predict_tile = tensor<0>(tiled_predict_tensor(_, blk_offset)); + auto current_scale_factor_tile = tensor<0>(scale_factor_tensor(_, blk_offset)); + + mxfp8_group_quant_tile< + decltype(current_input_tile), + decltype(current_predict_tile), + decltype(current_quant_output_tile), + decltype(scale_factor_shared), + decltype(current_scale_factor_tile), + TiledCopyG2R, + TiledCopyR2G, + TiledCopyR2S>( + current_input_tile, + current_predict_tile, + current_quant_output_tile, + scale_factor_shared, + current_scale_factor_tile, + m, + tiled_copy_g2r, + tiled_copy_r2g, + tiled_copy_r2s); + blk_offset += gridDim.x; + } + } +#endif +} + +template +void launch_es_sm100_mxfp8_blockscaled_grouped_quant( + const T_IN* input, + const int* tokens_per_expert, + const int* expert_offsets, + const int* blockscale_offsets, + cutlass::float_e4m3_t* quant_output, + uint8_t* scale_factor, + int num_experts, + int k, + int sm_count, + cudaStream_t stream) { + ThrLayout thr_layout{}; + ValLayout val_layout{}; + SfR2SThrLayout r2s_thr_layout{}; + SfR2SValLayout r2s_val_layout{}; + + using CopyOpG2R = UniversalCopy>; + using CopyAtomG2R = cute::Copy_Atom; + auto tiled_copy_g2r = cute::make_tiled_copy(CopyAtomG2R{}, thr_layout, val_layout); // Tiler_MN: (16, 128) + + using CopyOpR2G = UniversalCopy>; + using CopyAtomR2G = cute::Copy_Atom; + auto tiled_copy_r2g = cute::make_tiled_copy(CopyAtomR2G{}, thr_layout, val_layout); // Tiler_MN: (16, 128) + + using CopyOpR2S = UniversalCopy>; + using CopyAtomR2S = cute::Copy_Atom; + auto tiled_copy_r2s = cute::make_tiled_copy(CopyAtomR2S{}, r2s_thr_layout, r2s_val_layout); // Tiler_MN: (16, 4) + + int max_active_blocks_per_sm = -1; + auto error_code = cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks_per_sm, + mxfp8_group_quant, + THREAD_BLOCK_SIZE, + 0); + host::RuntimeCheck(error_code == cudaSuccess, "cudaOccupancyMaxActiveBlocksPerMultiprocessor failed"); + + dim3 grid(sm_count * max_active_blocks_per_sm, 1, 1); + dim3 block(THREAD_BLOCK_SIZE, 1, 1); + mxfp8_group_quant + <<>>( + input, + tokens_per_expert, + expert_offsets, + blockscale_offsets, + quant_output, + scale_factor, + num_experts, + k, + tiled_copy_g2r, + tiled_copy_r2g, + tiled_copy_r2s); +} + +} // namespace expert_specialization + +template +struct EsSm100MXFP8BlockscaledGroupQuant { + static void + run(const tvm::ffi::TensorView input, + const tvm::ffi::TensorView tokens_per_expert, + const tvm::ffi::TensorView expert_offsets, + const tvm::ffi::TensorView blockscale_offsets, + tvm::ffi::TensorView quant_output, + tvm::ffi::TensorView scale_factor) { + using namespace host; + auto N = SymbolicSize{"num_tokens"}; + auto D = SymbolicSize{"hidden_size"}; + auto G = SymbolicSize{"num_experts"}; + auto N_SF_Alinged = SymbolicSize{"num_tokens_sf_aligned"}; + auto D_SF = SymbolicSize{"hidden_size_sf"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, D}).with_strides({D, 1}).with_dtype().with_device(device).verify(input); + TensorMatcher({G}).with_dtype().with_device(device).verify(tokens_per_expert); + TensorMatcher({G}).with_dtype().with_device(device).verify(expert_offsets); + TensorMatcher({G}).with_dtype().with_device(device).verify(blockscale_offsets); + RuntimeCheck(D.unwrap() % 128 == 0, "k must align to 128"); + + TensorMatcher({N, D}).with_strides({D, 1}).with_dtype().with_device(device).verify(quant_output); + TensorMatcher({N_SF_Alinged, D_SF}).with_dtype().with_device(device).verify(scale_factor); + RuntimeCheck(D.unwrap() / 32 == D_SF.unwrap(), "Scale factor K should be hidden_size / 32"); + + cudaStream_t stream = LaunchKernel::resolve_device(device.unwrap()); + expert_specialization::launch_es_sm100_mxfp8_blockscaled_grouped_quant( + reinterpret_cast(input.data_ptr()), + reinterpret_cast(tokens_per_expert.data_ptr()), + reinterpret_cast(expert_offsets.data_ptr()), + reinterpret_cast(blockscale_offsets.data_ptr()), + reinterpret_cast(quant_output.data_ptr()), + reinterpret_cast(scale_factor.data_ptr()), + static_cast(G.unwrap()), + static_cast(D.unwrap()), + runtime::get_sm_count(device.unwrap().device_id), + stream); + } +}; diff --git a/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh b/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh new file mode 100644 index 000000000..8af85af96 --- /dev/null +++ b/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh @@ -0,0 +1,217 @@ +#pragma once +#include +#include + +#include +#include + +#include "cute/tensor.hpp" +#include "es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh" +#include "es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh" + +namespace expert_specialization { + +using namespace host; + +template +void es_sm100_mxfp8_blockscaled_moe_group_gemm_pre_compute( + tvm::ffi::TensorView b, + tvm::ffi::TensorView sfb, + tvm::ffi::TensorView expert_offsets, + tvm::ffi::TensorView blockscale_offsets, + tvm::ffi::TensorView b_ptrs, + tvm::ffi::TensorView sfb_ptrs, + tvm::ffi::TensorView d, + tvm::ffi::TensorView d_ptrs, + int num_experts, + int m, + int k, + cudaStream_t stream) { + using OffsetFunctor = Sm100Mxfp8BlockScaledMoeGroupGemmOffsetFunctor; + using ElementB = typename OffsetFunctor::ElementB; + using ElementSF = typename OffsetFunctor::ElementSF; + using ElementD = typename OffsetFunctor::ElementD; + + host::RuntimeCheck(num_experts <= 1024, "num_experts more than 1024"); + OffsetFunctor offset_functor( + reinterpret_cast(expert_offsets.data_ptr()), + reinterpret_cast(blockscale_offsets.data_ptr()), + reinterpret_cast(b.data_ptr()), + reinterpret_cast(sfb.data_ptr()), + reinterpret_cast(d.data_ptr()), + reinterpret_cast(b_ptrs.data_ptr()), + reinterpret_cast(sfb_ptrs.data_ptr()), + reinterpret_cast(d_ptrs.data_ptr())); + + sm100_mxfp8_blockscaled_moe_group_gemm_pre_compute_kernel<<<1, num_experts, 0, stream>>>(offset_functor, m, k); +} + +template +void es_sm100_mxfp8_blockscaled_moe_group_gemm( + tvm::ffi::TensorView a, + tvm::ffi::TensorView sfa, + tvm::ffi::TensorView tokens_per_expert, + tvm::ffi::TensorView b_ptrs, + tvm::ffi::TensorView sfb_ptrs, + tvm::ffi::TensorView d_ptrs, + tvm::ffi::TensorView workspace, + int num_experts, + int m, + int n, + int k, + int device_id, + int sm_count, + cudaStream_t stream) { + using Gemm = typename GemmTraits::Gemm; + using ElementA = typename Gemm::ElementA; + using ElementB = typename Gemm::ElementB; + using ElementSF = typename GemmTraits::ElementSF; + using ElementD = typename GemmTraits::ElementD; + + cutlass::KernelHardwareInfo hw_info; + hw_info.device_id = device_id; + hw_info.sm_count = sm_count; + hw_info.cluster_shape = GemmTraits::MMAConfig::preferred_cluster; + hw_info.cluster_shape_fallback = GemmTraits::MMAConfig::fallback_cluster; + + typename Gemm::Arguments arguments = { + cutlass::gemm::GemmUniversalMode::kGrouped, + {m, n, k, num_experts, reinterpret_cast(tokens_per_expert.data_ptr())}, + {reinterpret_cast(a.data_ptr()), + reinterpret_cast(b_ptrs.data_ptr()), + reinterpret_cast(sfa.data_ptr()), + reinterpret_cast(sfb_ptrs.data_ptr())}, + {{}, nullptr, nullptr, reinterpret_cast(d_ptrs.data_ptr()), nullptr}, + hw_info, + {} // Scheduler + }; + + Gemm gemm; + + auto can_implement_status = gemm.can_implement(arguments); + host::RuntimeCheck(can_implement_status == cutlass::Status::kSuccess, "Can not implement MoE Group GEMM"); + + auto status = gemm.initialize(arguments, reinterpret_cast(workspace.data_ptr()), stream); + host::RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to initialize MoE Group GEMM"); + + status = gemm.run(stream, nullptr); + host::RuntimeCheck(status == cutlass::Status::kSuccess, "Failed to run MoE Group GEMM"); +} + +template // CUTLASS dtype +void es_sm100_mxfp8_blockscaled_moe_group_gemm_dispatch_dtype( + tvm::ffi::TensorView a, + tvm::ffi::TensorView b, + tvm::ffi::TensorView sfa, + tvm::ffi::TensorView sfb, + tvm::ffi::TensorView expert_offsets, + tvm::ffi::TensorView blockscale_offsets, + tvm::ffi::TensorView tokens_per_expert, + tvm::ffi::TensorView b_ptrs, + tvm::ffi::TensorView sfb_ptrs, + tvm::ffi::TensorView d, + tvm::ffi::TensorView d_ptrs, + tvm::ffi::TensorView workspace, + int num_experts, + int m, + int n, + int k, + int device_id, + int sm_count, + cudaStream_t stream) { + using GemmTraits = ExpertSpecializationSm100MXFP8BlockscaledMoeGroupGemmTraits; + + es_sm100_mxfp8_blockscaled_moe_group_gemm_pre_compute( + b, sfb, expert_offsets, blockscale_offsets, b_ptrs, sfb_ptrs, d, d_ptrs, num_experts, m, k, stream); + es_sm100_mxfp8_blockscaled_moe_group_gemm( + a, + sfa, + tokens_per_expert, + b_ptrs, + sfb_ptrs, + d_ptrs, + workspace, + num_experts, + m, + n, + k, + device_id, + sm_count, + stream); +} + +} // namespace expert_specialization + +template +struct EsSm100MXFP8BlockscaledMoeGroupGemm { + static void + run(tvm::ffi::TensorView a, + tvm::ffi::TensorView b, + tvm::ffi::TensorView sfa, + tvm::ffi::TensorView sfb, + tvm::ffi::TensorView expert_offsets, + tvm::ffi::TensorView blockscale_offsets, + tvm::ffi::TensorView tokens_per_expert, + tvm::ffi::TensorView b_ptrs, + tvm::ffi::TensorView sfb_ptrs, + tvm::ffi::TensorView d, + tvm::ffi::TensorView d_ptrs, + tvm::ffi::TensorView workspace) { + using namespace host; + auto num_tokens = SymbolicSize{"num_tokens"}; + auto num_sf_tokens = SymbolicSize{"num_sf_tokens"}; + auto hidden_size = SymbolicSize{"hidden_size"}; + auto num_experts = SymbolicSize{"num_experts"}; + auto M = SymbolicSize{"M"}; + auto K = SymbolicSize{"K"}; + auto M_SF = SymbolicSize{"M_SF"}; + auto K_SF = SymbolicSize{"K_SF"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({num_experts, M, K}).with_dtype().with_device(device).verify(a); + TensorMatcher({num_tokens, K}).with_dtype().with_device(device).verify(b); + TensorMatcher({num_experts, M_SF, K_SF}).with_dtype().with_device(device).verify(sfa); + TensorMatcher({num_sf_tokens, K_SF}).with_dtype().with_device(device).verify(sfb); + RuntimeCheck(K.unwrap() % 128 == 0, "K should align 128"); + RuntimeCheck(K.unwrap() / 32 == K_SF.unwrap(), "K dimension mismatch"); + + TensorMatcher({num_experts}).with_dtype().with_device(device).verify(expert_offsets); + TensorMatcher({num_experts}).with_dtype().with_device(device).verify(blockscale_offsets); + TensorMatcher({num_experts}).with_dtype().with_device(device).verify(tokens_per_expert); + TensorMatcher({num_experts}).with_dtype().with_device(device).verify(b_ptrs); + TensorMatcher({num_experts}).with_dtype().with_device(device).verify(sfb_ptrs); + TensorMatcher({num_experts}).with_dtype().with_device(device).verify(d_ptrs); + // Check output + TensorMatcher({num_tokens, M}).with_strides({M, 1}).with_dtype().with_device(device).verify(d); + + cudaStream_t stream = LaunchKernel::resolve_device(device.unwrap()); + int device_id = device.unwrap().device_id; + + if constexpr (std::is_same_v || std::is_same_v) { + using CUTLASS_DTYPE = std::conditional_t, cutlass::bfloat16_t, cutlass::half_t>; + expert_specialization::es_sm100_mxfp8_blockscaled_moe_group_gemm_dispatch_dtype( + a, + b, + sfa, + sfb, + expert_offsets, + blockscale_offsets, + tokens_per_expert, + b_ptrs, + sfb_ptrs, + d, + d_ptrs, + workspace, + static_cast(num_experts.unwrap()), + static_cast(M.unwrap()), + static_cast(num_tokens.unwrap()), + static_cast(K.unwrap()), + device_id, + static_cast(runtime::get_sm_count(device_id)), + stream); + } else { + Panic("Unsupported dtype"); + } + } +}; diff --git a/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh b/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh new file mode 100644 index 000000000..7c4cb86ff --- /dev/null +++ b/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_functor.cuh @@ -0,0 +1,64 @@ +#pragma once +#include "cute/tensor.hpp" +#include "es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh" +#include + +namespace expert_specialization { + +using namespace cute; + +template +struct Sm100Mxfp8BlockScaledMoeGroupGemmOffsetFunctor { + using ElementB = typename GemmTraits::Gemm::ElementB; + using ElementSF = typename GemmTraits::ElementSF; + using ElementD = typename GemmTraits::ElementD; + // Input + int* expert_offsets{nullptr}; + int* blockscale_offsets{nullptr}; + // Output + ElementB* b_base{nullptr}; + ElementSF* sfb_base{nullptr}; + ElementD* d_base{nullptr}; + ElementB** b_offsets{nullptr}; + ElementSF** sfb_offsets{nullptr}; + ElementD** d_offsets{nullptr}; + + Sm100Mxfp8BlockScaledMoeGroupGemmOffsetFunctor() = default; + Sm100Mxfp8BlockScaledMoeGroupGemmOffsetFunctor( + int* _expert_offsets, + int* _blockscale_offsets, + ElementB* _b_base, + ElementSF* _sfb_base, + ElementD* _d_base, + ElementB** _b_offsets, + ElementSF** _sfb_offsets, + ElementD** _d_offsets) + : expert_offsets{_expert_offsets}, + blockscale_offsets{_blockscale_offsets}, + b_base(_b_base), + sfb_base(_sfb_base), + d_base(_d_base), + b_offsets(_b_offsets), + sfb_offsets(_sfb_offsets), + d_offsets(_d_offsets) {} + + void CUTE_DEVICE operator()(int expert_id, int m, int k) { + int64_t expert_offset = static_cast(expert_offsets[expert_id]); + int64_t blockscale_offset = static_cast(blockscale_offsets[expert_id]); + int64_t b_stride = expert_offset * k; + int64_t sfb_stride = blockscale_offset * (k / 32); + int64_t d_stride = expert_offset * m; + + b_offsets[expert_id] = b_base + b_stride; + sfb_offsets[expert_id] = sfb_base + sfb_stride; + d_offsets[expert_id] = d_base + d_stride; + } +}; + +template +__global__ void sm100_mxfp8_blockscaled_moe_group_gemm_pre_compute_kernel(OffsetFunctor offset_functor, int m, int k) { + int expert_id = static_cast(threadIdx.x); + offset_functor(expert_id, m, k); +} + +} // namespace expert_specialization diff --git a/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh b/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh new file mode 100644 index 000000000..4e0c250ff --- /dev/null +++ b/python/sglang/jit_kernel/csrc/moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm_traits.cuh @@ -0,0 +1,123 @@ +#pragma once + +// Misc +#include "cute/tensor.hpp" +#include "cutlass/arch/arch.h" +#include "cutlass/arch/mma.h" +#include "cutlass/cutlass.h" +#include "cutlass/detail/sm100_blockscaled_layout.hpp" +#include "cutlass/epilogue/dispatch_policy.hpp" +#include "cutlass/gemm/dispatch_policy.hpp" +#include "cutlass/gemm/group_array_problem_shape.hpp" +#include "cutlass/layout/layout.h" +#include "cutlass/numeric_conversion.h" +#include "cutlass/numeric_size.h" + +// Collective Builder +#include "cutlass/epilogue/collective/collective_builder.hpp" +#include "cutlass/epilogue/fusion/sm90_callbacks_tma_warpspecialized.hpp" +#include "cutlass/epilogue/thread/activation.h" +#include "cutlass/gemm/collective/collective_builder.hpp" + +// Integration +#include "cutlass/gemm/device/gemm_universal_adapter.h" +#include "cutlass/gemm/kernel/gemm_universal.hpp" + +namespace expert_specialization { + +using namespace cute; + +// Different configs for 1SM and 2SM MMA kernel +struct MMA2SMConfig { + using MmaTileShape = Shape<_256, _128, _128>; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecialized2SmMxf8f6f4Sm100; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized2Sm; + const static dim3 preferred_cluster; + const static dim3 fallback_cluster; +}; +const dim3 MMA2SMConfig::preferred_cluster(4, 1, 1); +const dim3 MMA2SMConfig::fallback_cluster(2, 1, 1); + +template +struct ExpertSpecializationSm100MXFP8BlockscaledMoeGroupGemmTraits { + using MMAConfig = _MMAConfig; + using ElementInput = cutlass::float_e4m3_t; + using ElementOutput = OutputDtype; + using ProblemShape = cutlass::gemm::MoEProblemShape>; // per group + + // A matrix configuration + using ElementA = cutlass::mx_float8_t; + using LayoutA = cutlass::layout::RowMajor; + constexpr static int AlignmentA = 16; + + // B matrix configuration + using ElementB = cutlass::mx_float8_t; + using LayoutB = cutlass::layout::ColumnMajor; + constexpr static int AlignmentB = 16; + + // C/D matrix configuration + using ElementC = void; + using ElementD = ElementOutput; + using LayoutC = cutlass::layout::ColumnMajor; + using LayoutD = cutlass::layout::ColumnMajor; + constexpr static int AlignmentC = 128 / cutlass::sizeof_bits::value; + constexpr static int AlignmentD = 128 / cutlass::sizeof_bits::value; + using ElementAccumulator = float; + + static constexpr auto RoundStyle = cutlass::FloatRoundStyle::round_to_nearest; + using CustomEVTIdentity = // acc + cutlass::epilogue::fusion::Sm90EVT< + cutlass::epilogue::fusion:: + Sm90Compute, + cutlass::epilogue::fusion::Sm90AccFetch>; + + // Core kernel configurations + using ArchTag = cutlass::arch::Sm100; + using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; + using StageCountType = cutlass::gemm::collective::StageCountAuto; + + // Runtime Cluster Shape + using ClusterShape = Shape; + + // Define Epilogue + using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + typename MMAConfig::MmaTileShape, + ClusterShape, + cutlass::epilogue::collective::EpilogueTileAuto, + ElementAccumulator, + ElementAccumulator, + ElementC, + LayoutC*, + AlignmentC, + ElementD, + LayoutD*, + AlignmentD, + typename MMAConfig::EpilogueSchedule, + CustomEVTIdentity>::CollectiveOp; + + // Define Mainloop + using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder< + ArchTag, + OperatorClass, + ElementA, + LayoutA, + AlignmentA, + ElementB, + LayoutB*, + AlignmentB, + ElementAccumulator, + typename MMAConfig::MmaTileShape, + ClusterShape, + cutlass::gemm::collective::StageCountAutoCarveout( + sizeof(typename CollectiveEpilogue::SharedStorage))>, + typename MMAConfig::KernelSchedule>::CollectiveOp; + + // Define GemmKernel + using GemmKernel = cutlass::gemm::kernel::GemmUniversal; + using Gemm = cutlass::gemm::device::GemmUniversalAdapter; + using ElementSF = typename GemmKernel::ElementSF; +}; + +} // namespace expert_specialization diff --git a/python/sglang/jit_kernel/mxfp8.py b/python/sglang/jit_kernel/mxfp8.py new file mode 100644 index 000000000..2f0a91f9f --- /dev/null +++ b/python/sglang/jit_kernel/mxfp8.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from sglang.jit_kernel.utils import ( + cache_once, + load_jit, + make_cpp_args, + override_jit_cuda_arch, +) + +if TYPE_CHECKING: + from tvm_ffi.module import Module + + +def _mxfp8_cuda_flags() -> list[str]: + return [ + "-DNDEBUG", + "-DCUTLASS_ENABLE_TENSOR_CORE_MMA=1", + "-DCUTLASS_VERSIONS_GENERATED", + "-DCUTLASS_DEBUG_TRACE_LEVEL=0", + "--expt-extended-lambda", + ] + + +def _mxfp8_arch_env(): + if not torch.cuda.is_available(): + raise RuntimeError("MXFP8 JIT kernels require CUDA.") + major, minor = torch.cuda.get_device_capability() + if major < 10: + raise RuntimeError( + f"MXFP8 JIT kernels require compute capability >= 10.0, got {major}.{minor}." + ) + # MXFP8 kernels use architecture-family-specific instructions and must be + # compiled for `sm_*a` targets (e.g. sm_100a), not plain sm_100. + # JIT compilation targets only the current device, unlike AOT fat-binaries; + # adding extra architectures here would clash with the single SGL_CUDA_ARCH + # value injected by load_jit(). + return override_jit_cuda_arch(major, minor, suffix="a") + + +@cache_once +def _jit_es_sm100_mxfp8_blockscaled_group_quant(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + with _mxfp8_arch_env(): + return load_jit( + "es_sm100_mxfp8_blockscaled_group_quant", + *args, + cuda_files=[ + "moe/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cuh" + ], + cuda_wrappers=[ + ( + "es_sm100_mxfp8_blockscaled_group_quant", + f"EsSm100MXFP8BlockscaledGroupQuant<{args}>::run", + ) + ], + extra_dependencies=["cutlass"], + extra_cuda_cflags=_mxfp8_cuda_flags(), + ) + + +@cache_once +def _jit_es_sm100_mxfp8_blockscaled_moe_group_gemm(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + with _mxfp8_arch_env(): + return load_jit( + "es_sm100_mxfp8_blockscaled_moe_group_gemm", + *args, + cuda_files=[ + "moe/expert_specialization/es_sm100_mxfp8_blockscaled_moe_group_gemm.cuh" + ], + cuda_wrappers=[ + ( + "es_sm100_mxfp8_blockscaled_moe_group_gemm", + f"EsSm100MXFP8BlockscaledMoeGroupGemm<{args}>::run", + ) + ], + extra_dependencies=["cutlass"], + extra_cuda_cflags=_mxfp8_cuda_flags(), + ) + + +def es_sm100_mxfp8_blockscaled_grouped_quant( + input: torch.Tensor, + tokens_per_expert: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + quant_output: torch.Tensor, + scale_factor: torch.Tensor, +) -> None: + module = _jit_es_sm100_mxfp8_blockscaled_group_quant(input.dtype) + module.es_sm100_mxfp8_blockscaled_group_quant( + input, + tokens_per_expert, + expert_offsets, + blockscale_offsets, + quant_output, + scale_factor, + ) + + +def es_sm100_mxfp8_blockscaled_moe_grouped_gemm( + a: torch.Tensor, + b: torch.Tensor, + sfa: torch.Tensor, + sfb: torch.Tensor, + expert_offsets: torch.Tensor, + blockscale_offsets: torch.Tensor, + tokens_per_expert: torch.Tensor, + workspace: torch.Tensor, + dtype: torch.dtype, +) -> torch.Tensor: + num_experts, m, tokens = a.shape[0], a.shape[1], b.shape[0] + d = torch.empty((tokens, m), device=a.device, dtype=dtype) + d_ptrs = torch.empty((num_experts,), device=a.device, dtype=torch.int64) + b_ptrs = torch.empty((num_experts,), device=a.device, dtype=torch.int64) + sfb_ptrs = torch.empty((num_experts,), device=a.device, dtype=torch.int64) + module = _jit_es_sm100_mxfp8_blockscaled_moe_group_gemm(dtype) + module.es_sm100_mxfp8_blockscaled_moe_group_gemm( + a, + b, + sfa, + sfb, + expert_offsets, + blockscale_offsets, + tokens_per_expert, + b_ptrs, + sfb_ptrs, + d, + d_ptrs, + workspace, + ) + return d diff --git a/python/sglang/jit_kernel/tests/test_mxfp8_moe.py b/python/sglang/jit_kernel/tests/test_mxfp8_moe.py new file mode 100644 index 000000000..12157b7aa --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_mxfp8_moe.py @@ -0,0 +1,153 @@ +import random +import sys + +import pytest +import torch + +from sglang.jit_kernel.mxfp8 import ( + es_sm100_mxfp8_blockscaled_grouped_quant, + es_sm100_mxfp8_blockscaled_moe_grouped_gemm, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=5, suite="stage-b-kernel-unit-1-gpu-large") +register_cuda_ci(est_time=120, suite="nightly-kernel-1-gpu", nightly=True) + + +def align(val: int, alignment: int = 128) -> int: + return int((val + alignment - 1) // alignment * alignment) + + +# Copy from: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/utils.py +def calc_diff(x, y): + x, y = x.double(), y.double() + denominator = (x * x + y * y).sum() + sim = 2 * (x * y).sum() / denominator + return 1 - sim + + +def is_sm100_supported(device=None) -> bool: + return (torch.cuda.get_device_capability(device)[0] == 10) and ( + torch.version.cuda >= "12.8" + ) + + +@pytest.mark.skipif( + not is_sm100_supported(), + reason="test_mxfp8_moe at jit kernen is only supported on sm100", +) +@pytest.mark.parametrize("num_experts", [8, 16, 32, 64]) +@pytest.mark.parametrize("out_dtype", [torch.half, torch.bfloat16]) +def test_es_sm100_mxfp8_blockscaled_grouped_mm(num_experts, out_dtype): + device = "cuda" + alignment = 128 + n_g = random.randint(1, 64) * alignment + k_g = random.randint(1, 64) * alignment + + expert_offset = 0 + expert_offsets = [] + aux_expert_offset = 0 + aux_expert_offsets = [] + a_blockscale_offset = 0 + a_blockscale_offsets = [] + b_blockscale_offset = 0 + b_blockscale_offsets = [] + a_list = [] + b_list = [] + ref_d_list = [] + tokens_per_expert = [] + + for g in range(num_experts): + m_g = random.randint(1, 512) + tokens_per_expert.append(m_g) + expert_offsets.append(expert_offset) + expert_offset += m_g + aux_expert_offsets.append(aux_expert_offset) + aux_expert_offset += n_g + a_blockscale_offsets.append(a_blockscale_offset) + a_blockscale_offset += align(m_g, 128) + b_blockscale_offsets.append(b_blockscale_offset) + b_blockscale_offset += n_g # n_g already align to 128 + + a = torch.normal( + 0.0, std=1.0, size=(m_g, k_g), device=device, dtype=out_dtype + ) # (M, K):(K, 1) + b = torch.normal( + 0.0, std=1.0, size=(n_g, k_g), device=device, dtype=out_dtype + ) # (N, K):(K, 1) + + a_list.append(a) + b_list.append(b) + ref_d = a @ b.T + ref_d_list.append(ref_d) + a = torch.concat(a_list, dim=0) + b = torch.concat(b_list, dim=0) + + _expert_offsets = torch.tensor(expert_offsets).to(device=device, dtype=torch.int32) + _aux_expert_offsets = torch.tensor(aux_expert_offsets).to( + device=device, dtype=torch.int32 + ) + _a_blockscale_offsets = torch.tensor(a_blockscale_offsets).to( + device=device, dtype=torch.int32 + ) + _b_blockscale_offsets = torch.tensor(b_blockscale_offsets).to( + device=device, dtype=torch.int32 + ) + + a_quant = torch.zeros_like(a, dtype=torch.float8_e4m3fn, device=device) + a_scale_factor = torch.zeros( + (a_blockscale_offset, k_g // 32), dtype=torch.uint8, device=device + ) + + b_quant = torch.zeros_like(b, dtype=torch.float8_e4m3fn, device=device) + b_scale_factor = torch.zeros( + (num_experts * n_g, k_g // 32), dtype=torch.uint8, device=device + ) + tokens_per_expert = torch.tensor(tokens_per_expert).to( + device=device, dtype=torch.int32 + ) + workspace = torch.empty((1024, 1024, 1024), dtype=torch.uint8, device=device) + + es_sm100_mxfp8_blockscaled_grouped_quant( + a, + tokens_per_expert, + _expert_offsets, + _a_blockscale_offsets, + a_quant, + a_scale_factor, + ) + es_sm100_mxfp8_blockscaled_grouped_quant( + b, + torch.ones_like(tokens_per_expert) * n_g, + _aux_expert_offsets, + _b_blockscale_offsets, + b_quant, + b_scale_factor, + ) + + b_quant = b_quant.view(num_experts, n_g, k_g) + b_scale_factor = b_scale_factor.view(num_experts, n_g, k_g // 32) + d = es_sm100_mxfp8_blockscaled_moe_grouped_gemm( + b_quant, + a_quant, + b_scale_factor, + a_scale_factor, + _expert_offsets, + _a_blockscale_offsets, + tokens_per_expert, + workspace, + a.dtype, + ) + + for g in range(num_experts): + baseline = ref_d_list[g] + actual = d[expert_offsets[g] : (expert_offsets[g] + tokens_per_expert[g])] + diff = calc_diff(actual, baseline) + assert diff < 0.001 + print( + f"m_g={baseline.shape[0]} n_g={n_g} k_g={k_g} num_experts={num_experts}, out_dtype={out_dtype}, diff={diff:.5f}: OK" + ) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__]))