diff --git a/sgl-kernel/benchmark/bench_es_fp8_blockwise_grouped_gemm.py b/sgl-kernel/benchmark/bench_es_fp8_blockwise_grouped_gemm.py index 7591c5dd1..a725f7e71 100644 --- a/sgl-kernel/benchmark/bench_es_fp8_blockwise_grouped_gemm.py +++ b/sgl-kernel/benchmark/bench_es_fp8_blockwise_grouped_gemm.py @@ -43,31 +43,23 @@ def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: ) -def create_unbalanced_expert_token_distribution(max_num_experts): - ratios = [random.random() for _ in range(max_num_experts)] - - def convert_to_tokens(ratio: float): - if ratio <= 0.7: - return random.randint(1, 32) - elif ratio > 0.7 and ratio <= 0.85: - return random.randint(32, 64) - elif ratio > 0.85 and ratio <= 0.95: - return random.randint(64, 128) - elif ratio > 0.95: - return random.randint(128, 1024) - else: - return 128 - - group_ms = [convert_to_tokens(ratio) for ratio in ratios] +def create_unbalanced_expert_token_distribution( + batch_size: int, topk: int, num_experts: int +): + expert_ids = np.random.randint(0, num_experts, size=(batch_size * topk,)).tolist() + expert_to_count = dict() + for expert_id in range(num_experts): + expert_to_count[expert_id] = 0 + for expert_id in expert_ids: + expert_to_count[expert_id] += 1 + group_ms = [] + for expert_id in range(num_experts): + group_ms.append(expert_to_count[expert_id]) return group_ms -group_ms = create_unbalanced_expert_token_distribution(8192) -# group_ms = [128 for _ in range(8192)] -# group_ms = [128 if i % 2 == 0 else 64 for i in range(8192)] - - def bench_es( + group_ms: List[int], n: int, k: int, num_groups: int, @@ -94,12 +86,13 @@ def bench_es( m_g = group_ms[g] expert_offsets[g + 1] = expert_offsets[g] + m_g problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device) + if m_g != 0: + a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device)) + a_tensors.append(a_g) + a_scales_tensors.append(a_scale) - a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device)) b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t()) - a_tensors.append(a_g) b_tensors.append(b_g) - a_scales_tensors.append(a_scale) b_scales_tensors.append(b_scale) a_stack = torch.empty( @@ -109,8 +102,11 @@ def bench_es( (num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn ) + _aux_idx = 0 for g in range(num_groups): - a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[g] + if group_ms[g] != 0: + a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[_aux_idx] + _aux_idx += 1 b_stack[g] = b_tensors[g].t() b_stack = b_stack.transpose(1, 2) @@ -121,11 +117,17 @@ def bench_es( (num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32 ) + _aux_idx = 0 for g in range(num_groups): - a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[g] + if group_ms[g] != 0: + a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[ + _aux_idx + ] + _aux_idx += 1 b_scale_stack[g] = b_scales_tensors[g].t() b_scale_stack = b_scale_stack.transpose(1, 2) + workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8) c_out = torch.empty((expert_offsets[-1], n_g), device=device, dtype=out_dtype) a_strides = torch.full( (num_groups,), a_stack.stride(0), device=device, dtype=torch.int64 @@ -133,7 +135,6 @@ def bench_es( d_strides = torch.full( (num_groups,), c_out.stride(0), device=device, dtype=torch.int64 ) - workspace = torch.empty((1024 * 1024 * 1024), device=device, dtype=torch.uint8) def run_cutlass(): es_fp8_blockwise_scaled_grouped_mm( @@ -171,6 +172,7 @@ def bench_es( def bench_sgl( + group_ms: List[int], n: int, k: int, num_groups: int, @@ -197,12 +199,13 @@ def bench_sgl( m_g = group_ms[g] expert_offsets[g + 1] = expert_offsets[g] + m_g problem_sizes[g][:] = torch.tensor([m_g, n_g, k_g], device=device) + if m_g != 0: + a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device)) + a_tensors.append(a_g) + a_scales_tensors.append(a_scale) - a_g, a_scale = per_token_cast_to_fp8(torch.randn((m_g, k_g), device=device)) b_g, b_scale = per_block_cast_to_fp8(torch.randn((n_g, k_g), device=device).t()) - a_tensors.append(a_g) b_tensors.append(b_g) - a_scales_tensors.append(a_scale) b_scales_tensors.append(b_scale) a_stack = torch.empty( @@ -212,8 +215,11 @@ def bench_sgl( (num_groups, n_g, k_g), device=device, dtype=torch.float8_e4m3fn ) + _aux_idx = 0 for g in range(num_groups): - a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[g] + if group_ms[g] != 0: + a_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[_aux_idx] + _aux_idx += 1 b_stack[g] = b_tensors[g].t() b_stack = b_stack.transpose(1, 2) @@ -224,8 +230,13 @@ def bench_sgl( (num_groups, n_g // 128, k_g // 128), device=device, dtype=torch.float32 ) + _aux_idx = 0 for g in range(num_groups): - a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[g] + if group_ms[g] != 0: + a_scale_stack[expert_offsets[g] : expert_offsets[g + 1]] = a_scales_tensors[ + _aux_idx + ] + _aux_idx += 1 b_scale_stack[g] = b_scales_tensors[g].t() b_scale_stack = b_scale_stack.transpose(1, 2) @@ -300,16 +311,36 @@ def benchmark_one_shape( num_run: int, ): for shape in shape_args: - print(f"\nBenchmark: n={shape.n}, k={shape.k}, num_groups={shape.num_groups}") - for kernel_name, kernel_func in benchmark_kernels.items(): - average_time, m = kernel_func( - shape.n, - shape.k, - shape.num_groups, - num_warmup, - num_run, + for batch_size in [ + 128, + 256, + 384, + 512, + 640, + 768, + 896, + 1024, + 1280, + 1536, + 2048, + 3072, + ]: + group_ms = create_unbalanced_expert_token_distribution( + batch_size, 8, shape.num_groups ) - print(f"{kernel_name}: {average_time} us") + print( + f"\nBenchmark: batch_size={batch_size}, n={shape.n}, k={shape.k}, num_groups={shape.num_groups}" + ) + for kernel_name, kernel_func in benchmark_kernels.items(): + average_time, m = kernel_func( + group_ms, + shape.n, + shape.k, + shape.num_groups, + num_warmup, + num_run, + ) + print(f"{kernel_name}: {average_time} us") def main(): @@ -317,18 +348,22 @@ def main(): parser.add_argument("--num-warmup", type=int, default=3) parser.add_argument("--num-run", type=int, default=20) shape_args = [ - # Prefill, DeepSeek-R1, gateup, chunk_size = 4096, TP = 8 + # DeepSeek-R1, gateup, TP = 8 ShapeArg(n=512, k=7168, num_groups=256), - # Prefill, DeepSeek-R1, down, chunk_size = 4096, TP = 8 + # DeepSeek-R1, down, TP = 8 ShapeArg(n=7168, k=256, num_groups=256), - # Prefill, Qwen3-235B-A22B-FP8, gateup, TP = 4 + # DeepSeek-R1, gateup, TP = 4 + ShapeArg(n=1024, k=7168, num_groups=256), + # DeepSeek-R1, down, TP = 4 + ShapeArg(n=7168, k=512, num_groups=256), + # Qwen3-235B-A22B-FP8, gateup, TP = 4 ShapeArg(n=768, k=4096, num_groups=128), - # Prefill, Qwen3-235B-A22B-FP8, down, TP = 4 + # Qwen3-235B-A22B-FP8, down, TP = 4 ShapeArg(n=4096, k=384, num_groups=128), - # Decode, DeepSeek-R1, gateup, bs = 128, EP = 8 - ShapeArg(n=4096, k=7168, num_groups=32), - # Decode, DeepSeek-R1, gateup, bs = 256, EP = 16 - ShapeArg(n=4096, k=7168, num_groups=16), + # Qwen3-235B-A22B-FP8, gateup, TP = 2 + ShapeArg(n=1536, k=4096, num_groups=128), + # Qwen3-235B-A22B-FP8, down, TP = 2 + ShapeArg(n=4096, k=768, num_groups=128), ] args = parser.parse_args() benchmark_one_shape(shape_args, args.num_warmup, args.num_run) diff --git a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise.cu b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise.cu index f05209b37..f985226d7 100644 --- a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise.cu +++ b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise.cu @@ -1,3 +1,4 @@ +#include #include #include @@ -70,10 +71,18 @@ void es_fp8_blockwise_scaled_grouped_mm( torch::Tensor mm_problem_sizes = torch::empty({num_experts, 3}, options_int32); torch::Tensor hm_problem_sizes = torch::empty({num_experts, 3}, options_int32); + torch::Tensor backup_workspace_0 = torch::empty_like(workspace); + torch::Tensor backup_workspace_1 = torch::empty_like(workspace); + const std::string H20_device_type_str("NVIDIA H20"); bool is_h20_device = std::string(at::cuda::getCurrentDeviceProperties()->name) == H20_device_type_str; - at::cuda::CUDAGuard device_guard{(char)a.get_device()}; - cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device()); + + auto stream = at::cuda::getCurrentCUDAStream(); + static auto backup_stream_0 = at::cuda::getStreamFromPool(); + static auto backup_stream_1 = at::cuda::getStreamFromPool(); + at::cuda::CUDAEvent start_event; + at::cuda::CUDAEvent end_event_0; + at::cuda::CUDAEvent end_event_1; if (output.dtype() == torch::kBFloat16) { expert_specialization::es_sm90_fp8_blockwise_scaled_group_mm_pre_compute( @@ -95,7 +104,7 @@ void es_fp8_blockwise_scaled_grouped_mm( problem_sizes, expert_offsets, is_h20_device, - stream); + stream.stream()); } else if (output.dtype() == torch::kFloat16) { expert_specialization::es_sm90_fp8_blockwise_scaled_group_mm_pre_compute( out_ptrs, @@ -116,11 +125,15 @@ void es_fp8_blockwise_scaled_grouped_mm( problem_sizes, expert_offsets, is_h20_device, - stream); + stream.stream()); } else { TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)"); } + start_event.recordOnce(stream); + start_event.block(backup_stream_0); + start_event.block(backup_stream_1); + if (output.dtype() == torch::kBFloat16) { expert_specialization::es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype( out_ptrs, @@ -137,8 +150,12 @@ void es_fp8_blockwise_scaled_grouped_mm( mm_problem_sizes, hm_problem_sizes, workspace, + backup_workspace_0, + backup_workspace_1, is_h20_device, - stream); + stream.stream(), + backup_stream_0.stream(), + backup_stream_1.stream()); } else if (output.dtype() == torch::kFloat16) { expert_specialization::es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype( out_ptrs, @@ -155,11 +172,20 @@ void es_fp8_blockwise_scaled_grouped_mm( mm_problem_sizes, hm_problem_sizes, workspace, + backup_workspace_0, + backup_workspace_1, is_h20_device, - stream); + stream.stream(), + backup_stream_0.stream(), + backup_stream_1.stream()); } else { TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)"); } + + end_event_0.recordOnce(backup_stream_0); + end_event_1.recordOnce(backup_stream_1); + end_event_0.block(stream); + end_event_1.block(stream); #else TORCH_CHECK_NOT_IMPLEMENTED( can_implement, "No implemented fp8_blockwise_scaled_grouped_mm for current compute capability: ", sm_version); diff --git a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh index db7f430f2..4c1f5b37e 100644 --- a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh +++ b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_functor.cuh @@ -126,7 +126,12 @@ struct Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor { Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor(int* _problem_sizes) : problem_sizes(_problem_sizes) {} void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) { - if (m < 64) { + float m_f = __int2float_rn(m); + float n_f = __int2float_rn(n); + float k_f = __int2float_rn(k); + float arithmetic_intensity = 2.0f * m_f * n_f * k_f / (m_f * k_f + k_f * n_f + 2.0f * m_f * n_f); + + if (m <= 32 || arithmetic_intensity < 70.0f) { // Swap A/B problem_sizes[expert_id * 3 + 0] = n; problem_sizes[expert_id * 3 + 1] = m; @@ -168,7 +173,12 @@ struct Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor { Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor(int* _problem_sizes) : problem_sizes(_problem_sizes) {} void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) { - if (m >= 64 && m < 128) { + float m_f = __int2float_rn(m); + float n_f = __int2float_rn(n); + float k_f = __int2float_rn(k); + float arithmetic_intensity = 2.0f * m_f * n_f * k_f / (m_f * k_f + k_f * n_f + 2.0f * m_f * n_f); + + if ((!(m <= 32 || arithmetic_intensity < 70.0f)) && m <= 64) { problem_sizes[expert_id * 3 + 0] = m; problem_sizes[expert_id * 3 + 1] = n; problem_sizes[expert_id * 3 + 2] = k; @@ -208,7 +218,12 @@ struct Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor { Fp8BlockwiseGroupedGemmProblemSizeFilterFunctor(int* _problem_sizes) : problem_sizes(_problem_sizes) {} void CUTE_DEVICE operator()(int64_t expert_id, int m, int n, int k) { - if (m >= 128) { + float m_f = __int2float_rn(m); + float n_f = __int2float_rn(n); + float k_f = __int2float_rn(k); + float arithmetic_intensity = 2.0f * m_f * n_f * k_f / (m_f * k_f + k_f * n_f + 2.0f * m_f * n_f); + + if ((!(m <= 32 || arithmetic_intensity < 70.0f)) && m > 64) { problem_sizes[expert_id * 3 + 0] = m; problem_sizes[expert_id * 3 + 1] = n; problem_sizes[expert_id * 3 + 2] = k; diff --git a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh index 3ed98821d..d816eb109 100644 --- a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh +++ b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_launcher.cuh @@ -99,7 +99,8 @@ void launch_sm90_fp8_blockwise_scaled_group_mm( const torch::Tensor& layout_sfb, const torch::Tensor& problem_sizes, const torch::Tensor& workspace, - cudaStream_t stream) { + cudaStream_t stream, + int sm_count) { using ElementA = typename GemmTraits::ElementA; using StrideA = typename GemmTraits::StrideA; using ElementB = typename GemmTraits::ElementB; @@ -128,7 +129,7 @@ void launch_sm90_fp8_blockwise_scaled_group_mm( cutlass::KernelHardwareInfo hw_info; hw_info.device_id = c10::cuda::current_device(); - hw_info.sm_count = at::cuda::getCurrentDeviceProperties()->multiProcessorCount; + hw_info.sm_count = sm_count; typename GemmKernel::EpilogueArguments epilogue_args{ {}, nullptr, nullptr, static_cast(out_ptrs.data_ptr()), static_cast(stride_d.data_ptr())}; @@ -147,7 +148,7 @@ void launch_sm90_fp8_blockwise_scaled_group_mm( auto status = gemm_op.initialize(args, workspace.data_ptr(), stream); TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM"); - status = gemm_op.run(stream, nullptr, true); // Enable PDL + status = gemm_op.run(stream, nullptr); TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM"); } @@ -167,8 +168,12 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype( const torch::Tensor& mm_problem_sizes, const torch::Tensor& hm_problem_sizes, const torch::Tensor& workspace, + const torch::Tensor& backup_workspace_0, + const torch::Tensor& backup_workspace_1, bool is_h20_device, - cudaStream_t stream) { + cudaStream_t stream, + cudaStream_t backup_stream_0, + cudaStream_t backup_stream_1) { using LowMGemmH20Traits = ExpertSpecializationSm90FP8BlockwiseGroupedGemmTraits; using LowMGemmHx00Traits = @@ -184,6 +189,40 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype( using HighMGemmHx00Traits = ExpertSpecializationSm90FP8BlockwiseGroupedGemmTraits; + if (!is_h20_device) { + launch_sm90_fp8_blockwise_scaled_group_mm( + out_ptrs, + a_ptrs, + b_ptrs, + a_scales_ptrs, + b_scales_ptrs, + stride_a, + stride_b, + stride_d, + layout_sfa, + layout_sfb, + hm_problem_sizes, + workspace, + stream, + 132); + } else { + launch_sm90_fp8_blockwise_scaled_group_mm( + out_ptrs, + a_ptrs, + b_ptrs, + a_scales_ptrs, + b_scales_ptrs, + stride_a, + stride_b, + stride_d, + layout_sfa, + layout_sfb, + hm_problem_sizes, + workspace, + stream, + 78); + } + if (!is_h20_device) { launch_sm90_fp8_blockwise_scaled_group_mm( out_ptrs, @@ -197,8 +236,9 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype( layout_sfb, layout_sfa, lm_problem_sizes, - workspace, - stream); + backup_workspace_1, + backup_stream_1, + 132); } else { launch_sm90_fp8_blockwise_scaled_group_mm( out_ptrs, @@ -212,8 +252,9 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype( layout_sfb, layout_sfa, lm_problem_sizes, - workspace, - stream); + backup_workspace_1, + backup_stream_1, + 78); } if (!is_h20_device) { @@ -229,8 +270,9 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype( layout_sfb, layout_sfa, mm_problem_sizes, - workspace, - stream); + backup_workspace_0, + backup_stream_0, + 132); } else { launch_sm90_fp8_blockwise_scaled_group_mm( out_ptrs, @@ -244,40 +286,9 @@ void es_sm90_fp8_blockwise_scaled_group_mm_distpatch_out_dtype( layout_sfa, layout_sfb, mm_problem_sizes, - workspace, - stream); - } - - if (!is_h20_device) { - launch_sm90_fp8_blockwise_scaled_group_mm( - out_ptrs, - a_ptrs, - b_ptrs, - a_scales_ptrs, - b_scales_ptrs, - stride_a, - stride_b, - stride_d, - layout_sfa, - layout_sfb, - hm_problem_sizes, - workspace, - stream); - } else { - launch_sm90_fp8_blockwise_scaled_group_mm( - out_ptrs, - a_ptrs, - b_ptrs, - a_scales_ptrs, - b_scales_ptrs, - stride_a, - stride_b, - stride_d, - layout_sfa, - layout_sfb, - hm_problem_sizes, - workspace, - stream); + backup_workspace_0, + backup_stream_0, + 78); } } diff --git a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_traits.cuh b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_traits.cuh index 3bc7d929a..31106f67d 100644 --- a/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_traits.cuh +++ b/sgl-kernel/csrc/expert_specialization/es_fp8_blockwise_traits.cuh @@ -30,10 +30,10 @@ using namespace cute; struct PerfConfigLowMH20 { // Swap A/B using ElementA = cutlass::float_e4m3_t; - using MmaTileShape = Shape<_128, _32, _128>; + using MmaTileShape = Shape<_256, _32, _128>; using ClusterShape = Shape<_2, _1, _1>; - using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8Blockwise; - using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong; + using KernelSchedule = cutlass::gemm::KernelPtrArrayTmaWarpSpecializedCooperativeFP8Blockwise; + using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecializedCooperative; using ScaleConfig = cutlass::detail::Sm90BlockwiseScaleConfig<128, 1, 128, cute::GMMA::Major::K, cute::GMMA::Major::K>; using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());