diff --git a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py b/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py deleted file mode 100644 index a44c8ffc10c5..000000000000 --- a/benchmark/kernels/deepseek/benchmark_deepgemm_dsv3_router_gemm_blackwell.py +++ /dev/null @@ -1,250 +0,0 @@ -import argparse -import os -from typing import List - -import torch -import triton -from flashinfer.gemm import mm_M1_16_K7168_N256 -from sgl_kernel import dsv3_router_gemm - -N = 256 -K = 7168 - - -def create_benchmark_configs(tp_sizes: List[int]): - configs = [] - for tp_size in tp_sizes: - for m in range(1, 17): - configs.append((m, N, K, tp_size)) - return configs - - -def dsv3_router_gemm_flashinfer( - hidden_states: torch.Tensor, - router_weights: torch.Tensor, -): - """Flashinfer implementation of dsv3 router gemm""" - output = torch.empty( - hidden_states.shape[0], - router_weights.shape[0], - device="cuda", - dtype=torch.float32, - ) - mm_M1_16_K7168_N256( - hidden_states, router_weights.t(), output, launch_with_pdl=args.use_pdl - ) - return output - - -def dsv3_router_gemm_sgl( - hidden_states: torch.Tensor, - router_weights: torch.Tensor, -): - """SGLang implementation of dsv3 router gemm""" - output = dsv3_router_gemm( - hidden_states, - router_weights, - out_dtype=torch.float32, - ) - return output - - -def check_accuracy(a, b, atol, rtol, percent): - """Unified accuracy checking function with detailed error reporting.""" - if not torch.isfinite(a).all(): - print("Non-finite values in reference output") - return False - if not torch.isfinite(b).all(): - print("Non-finite values in actual output") - return False - assert a.shape == b.shape, f"Shape mismatch: {a.shape} vs {b.shape}" - - close = torch.isclose(a, b, atol=atol, rtol=rtol) - match_ratio = close.float().mean() - if match_ratio >= percent: - return True - - mismatch_percent = 1.0 - match_ratio.item() - if mismatch_percent > 1 - percent: - print( - f"Mismatch percentage is {mismatch_percent:.4f} for rtol {rtol} " - f"(threshold: {1 - percent:.4f})" - ) - return False - - -def calculate_diff(m: int, n: int, k: int): - hidden_states = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) - router_weights = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) - - out_flashinfer = dsv3_router_gemm_flashinfer( - hidden_states.clone(memory_format=torch.contiguous_format), - router_weights.clone(memory_format=torch.contiguous_format), - ) - - out_sgl = dsv3_router_gemm_sgl( - hidden_states.clone(memory_format=torch.contiguous_format), - router_weights.clone(memory_format=torch.contiguous_format), - ) - - print(f"Shape m={m}, n={n}, k={k}:") - print(f"Using PDL={args.use_pdl}") - print(f"Flashinfer output: {out_flashinfer[0, 0:5]}") - print(f"SGLang output: {out_sgl[0, 0:5]}") - - flashinfer_sgl_match = check_accuracy(out_flashinfer, out_sgl, 0.1, 0.6, 0.95) - print("Correctness check:") - print(f" - Flashinfer vs SGLang: {'✅' if flashinfer_sgl_match else '❌'}") - - -def _benchmark(m, n, k, tp_size, provider): - print(f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), Provider: {provider}") - hidden_states = torch.randn( - (m, k), device="cuda", dtype=torch.bfloat16 - ).contiguous() - router_weights = torch.randn( - (n, k), device="cuda", dtype=torch.bfloat16 - ).contiguous() - - quantiles = [0.5, 0.2, 0.8] - - if provider == "sglang": - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: dsv3_router_gemm_sgl( - hidden_states.clone(memory_format=torch.contiguous_format), - router_weights.clone(memory_format=torch.contiguous_format), - ), - quantiles=quantiles, - ) - elif provider == "flashinfer": - ms, min_ms, max_ms = triton.testing.do_bench( - lambda: dsv3_router_gemm_flashinfer( - hidden_states.clone(memory_format=torch.contiguous_format), - router_weights.clone(memory_format=torch.contiguous_format), - ), - quantiles=quantiles, - ) - - # Calculate TFLOPS - flops = 2 * m * n * k # multiply-adds - tflops = flops / (ms * 1e-3) / 1e12 - - # Print shape-specific results with TFLOPS - print(f"Time: {ms*1000:.2f} us, TFLOPS: {tflops:.2f}") - return ms, max_ms, min_ms - - -def get_benchmark_plot_friendly(tp_sizes): - all_configs = create_benchmark_configs(tp_sizes) - x_vals = list(range(len(all_configs))) - - @triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["cfg_id"], - x_vals=x_vals, - line_arg="provider", - line_vals=["sglang", "flashinfer"], - line_names=["SGLang", "Flashinfer"], - styles=[("blue", "-"), ("red", "-")], - ylabel="us", - plot_name=f"fp8-gemm-performance-comparison-tp-{'-'.join(str(tp) for tp in tp_sizes)}", - args={}, - ) - ) - def benchmark(cfg_id, provider): - m, n, k, tp_size = all_configs[cfg_id] - ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, provider) - return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms - - return benchmark - - -def get_benchmark(tp_sizes): - all_configs = create_benchmark_configs(tp_sizes) - - @triton.testing.perf_report( - triton.testing.Benchmark( - x_names=[ - "m", - "n", - "k", - "tp_size", - ], - x_vals=[list(config) for config in all_configs], - line_arg="provider", - line_vals=["sglang", "flashinfer"], - line_names=["SGLang", "Flashinfer"], - styles=[("blue", "-"), ("red", "-")], - ylabel="us", - plot_name=f"fp8-gemm-performance-comparison-tp-{'-'.join(str(tp) for tp in tp_sizes)}", - args={}, - ) - ) - def benchmark(m, n, k, tp_size, provider): - ms, min_ms, max_ms = _benchmark(m, n, k, tp_size, provider) - return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms - - return benchmark - - -if __name__ == "__main__": - if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10: - print("Skipping benchmark because the device is not supported") - exit(0) - - parser = argparse.ArgumentParser() - parser.add_argument( - "--save-path", - type=str, - default="./configs/benchmark_ops/dsv3_router_gemm/", - help="Path to save dsv3 router gemm benchmark results", - ) - parser.add_argument( - "--run-correctness", - action="store_true", - default=True, - help="Whether to run correctness test", - ) - parser.add_argument( - "--tp-sizes", - type=int, - nargs="+", - default=[1], - help="List of tensor parallelism sizes to benchmark", - ) - parser.add_argument( - "--plot-friendly", - action="store_true", - default=False, - help="Plot x axis as the config index instead of the m", - ) - parser.add_argument( - "--use-pdl", - action="store_true", - default=False, - help="Use PDL if true.", - ) - args = parser.parse_args() - - # Set random seed for reproducibility - torch.manual_seed(0) - torch.cuda.manual_seed(0) - - if args.use_pdl: - os.environ["TRTLLM_ENABLE_PDL"] = "1" - - # Run correctness tests on a few examples - if args.run_correctness: - print("Running correctness tests...") - for m, n, k, _ in create_benchmark_configs(args.tp_sizes): - calculate_diff(m, n, k) - - # Get the benchmark function with the specified tp_size - benchmark = ( - get_benchmark_plot_friendly(args.tp_sizes) - if args.plot_friendly - else get_benchmark(args.tp_sizes) - ) - - print(f"Running performance benchmark for TP sizes = {args.tp_sizes}...") - benchmark.run(print_data=True, save_path=args.save_path) diff --git a/python/sglang/jit_kernel/dsv3_fused_a_gemm.py b/python/sglang/jit_kernel/dsv3_fused_a_gemm.py index c718a5131925..3649beb20a3c 100644 --- a/python/sglang/jit_kernel/dsv3_fused_a_gemm.py +++ b/python/sglang/jit_kernel/dsv3_fused_a_gemm.py @@ -1,7 +1,7 @@ """ JIT kernel for DeepSeek V3 fused QKV-A GEMM (min-latency). -Replaces the AOT sgl_kernel.dsv3_fused_a_gemm for SM90+ (Hopper) GPUs. +Runtime-compiled CUDA C++ kernel for SM90+ (Hopper) GPUs. Shapes: hd_in a multiple of 256, hd_out a multiple of 16, num_tokens 1-16, bfloat16. """ diff --git a/python/sglang/jit_kernel/dsv3_router_gemm.py b/python/sglang/jit_kernel/dsv3_router_gemm.py index f33d883ddbd3..a7fadae5293e 100644 --- a/python/sglang/jit_kernel/dsv3_router_gemm.py +++ b/python/sglang/jit_kernel/dsv3_router_gemm.py @@ -1,7 +1,7 @@ """ JIT kernel for DeepSeek V3 router GEMM. -Replaces the AOT sgl_kernel.dsv3_router_gemm for SM90+ (Hopper) GPUs. +Runtime-compiled CUDA C++ kernel for SM90+ (Hopper) GPUs. Supports num_experts in {256, 384}, hidden_dim a multiple of 1024, num_tokens 1-16. """ diff --git a/python/sglang/jit_kernel/fused_a_gemm.py b/python/sglang/jit_kernel/fused_a_gemm.py index f5e3446e3518..70b69bb49634 100644 --- a/python/sglang/jit_kernel/fused_a_gemm.py +++ b/python/sglang/jit_kernel/fused_a_gemm.py @@ -1,8 +1,7 @@ """Unified entry point for the DeepSeek-V3 fused QKV-A GEMM. -Dispatches to one of three interchangeable implementations via ``backend``: +Dispatches to one of two interchangeable implementations via ``backend``: -- ``"aot"``: prebuilt ``sgl_kernel.dsv3_fused_a_gemm`` (CUDA C++). - ``"jit"``: runtime-compiled CUDA C++ (``sglang.jit_kernel.dsv3_fused_a_gemm``). - ``"cutedsl"``: CuTe DSL (``sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm``). - ``"auto"``: CuTe DSL on SM120+, otherwise the JIT kernel. @@ -21,7 +20,6 @@ class FusedAGemmBackend(str, Enum): AUTO = "auto" - AOT = "aot" JIT = "jit" CUTEDSL = "cutedsl" @@ -70,9 +68,7 @@ def dsv3_fused_a_gemm( if backend == FusedAGemmBackend.AUTO: backend = _AUTO_BACKEND - if backend == FusedAGemmBackend.AOT: - from sgl_kernel import dsv3_fused_a_gemm as impl - elif backend == FusedAGemmBackend.JIT: + if backend == FusedAGemmBackend.JIT: from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm as impl else: from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import ( diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index d4382e073114..d76c00880368 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/sgl-kernel/CMakeLists.txt @@ -263,7 +263,6 @@ set(SOURCES "csrc/expert_specialization/es_sm100_mxfp8_blockscaled_group_quant.cu" "csrc/gemm/awq_kernel.cu" - "csrc/gemm/dsv3_fused_a_gemm.cu" "csrc/gemm/fp8_gemm_kernel.cu" "csrc/gemm/int8_gemm_kernel.cu" "csrc/gemm/per_token_group_quant_8bit.cu" diff --git a/sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py b/sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py deleted file mode 100644 index 43f0961f7a96..000000000000 --- a/sgl-kernel/benchmark/bench_dsv3_fused_a_gemm.py +++ /dev/null @@ -1,73 +0,0 @@ -import argparse - -import torch -import torch.nn.functional as F -import triton -import triton.testing -from sgl_kernel import dsv3_fused_a_gemm - -from sglang.utils import is_in_ci - -IS_CI = is_in_ci() - -# CI environment uses simplified parameters -if IS_CI: - num_tokens_vals = [1] # Only test 1 value in CI - line_vals = ["sgl-kernel"] # Only test sgl-kernel implementation in CI -else: - num_tokens_vals = [i + 1 for i in range(16)] # Test 1-16 in full mode - line_vals = ["torch", "sgl-kernel"] - - -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["num_tokens"], - x_vals=num_tokens_vals, - x_log=False, - line_arg="impl", - line_vals=line_vals, - line_names=( - ["torch (bf16)", "dsv3_fused_a_gemm"] - if not IS_CI - else ["dsv3_fused_a_gemm"] - ), - styles=[("blue", "-"), ("orange", "-")] if not IS_CI else [("orange", "-")], - ylabel="TFLOPs", - plot_name="bf16 dsv3 fused a GEMM throughput", - args={}, - ) -) -def benchmark(num_tokens, impl): - kHdIn = 7168 - kHdOut = 2112 - M, K, N = num_tokens, kHdIn, kHdOut - - mat_a = torch.randn((M, K), dtype=torch.bfloat16, device="cuda").contiguous() - mat_b = torch.randn((N, K), dtype=torch.bfloat16, device="cuda").transpose(0, 1) - - quantiles = [0.5, 0.2, 0.8] - - if impl == "torch": - - def runner(): - F.linear(mat_a, mat_b.T) - - elif impl == "sgl-kernel": - - def runner(): - dsv3_fused_a_gemm(mat_a, mat_b) - - ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(runner, quantiles=quantiles) - - def tflops(t_ms): - flops = 2 * M * K * N - return flops / (t_ms * 1e-3) / 1e12 - - return tflops(ms), tflops(max_ms), tflops(min_ms) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - args = parser.parse_args() - - benchmark.run(print_data=True) diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index bd84206599da..4748c3cd1843 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/sgl-kernel/csrc/common_extension.cc @@ -136,9 +136,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { m.def("sgl_per_token_quant_fp8(Tensor input, Tensor! output_q, Tensor! output_s) -> ()"); m.impl("sgl_per_token_quant_fp8", torch::kCUDA, &sgl_per_token_quant_fp8); - m.def("dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); - m.impl("dsv3_fused_a_gemm", torch::kCUDA, &dsv3_fused_a_gemm); - /* * From csrc/gemm/gptq */ diff --git a/sgl-kernel/csrc/common_extension_musa.cc b/sgl-kernel/csrc/common_extension_musa.cc index bf83cba3c703..1c8d714979e0 100644 --- a/sgl-kernel/csrc/common_extension_musa.cc +++ b/sgl-kernel/csrc/common_extension_musa.cc @@ -97,9 +97,6 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { m.def("sgl_per_token_quant_fp8(Tensor input, Tensor output_q, Tensor output_s) -> ()"); m.impl("sgl_per_token_quant_fp8", torch::kMUSA, &sgl_per_token_quant_fp8); - m.def("dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); - m.impl("dsv3_fused_a_gemm", torch::kMUSA, &dsv3_fused_a_gemm); - /* * From csrc/moe */ diff --git a/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu b/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu deleted file mode 100644 index c393b5a587f1..000000000000 --- a/sgl-kernel/csrc/gemm/dsv3_fused_a_gemm.cu +++ /dev/null @@ -1,677 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/619709fc33bd5dc268f19d6a741fe7ed51c0f8f5/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3FusedAGemm.cu - * - * Copyright (c) 2019-2024, NVIDIA CORPORATION. All rights reserved. - * Copyright (c) 2021, NAVER Corp. Authored by CLOVA. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include -#include -#include - -#include "utils.h" - -using bf16_t = __nv_bfloat16; - -__device__ void hmma_16_8_16_f32acc_bf16ab( - float (&d_reg)[4], const bf16_t (&a_reg)[8], const bf16_t (&b_reg)[4], float const (&c_reg)[4]) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - uint32_t a0 = *reinterpret_cast(a_reg + 0); - uint32_t a1 = *reinterpret_cast(a_reg + 2); - uint32_t a2 = *reinterpret_cast(a_reg + 4); - uint32_t a3 = *reinterpret_cast(a_reg + 6); - uint32_t b0 = *reinterpret_cast(b_reg + 0); - uint32_t b1 = *reinterpret_cast(b_reg + 2); - asm volatile( - "mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 " - "{%0, %1, %2, %3}," - "{%4, %5, %6, %7}," - "{%8, %9}," - "{%10, %11, %12, %13};\n" - : "=f"(d_reg[0]), "=f"(d_reg[1]), "=f"(d_reg[2]), "=f"(d_reg[3]) - : "r"(a0), - "r"(a1), - "r"(a2), - "r"(a3), - "r"(b0), - "r"(b1), - "f"(d_reg[0]), - "f"(d_reg[1]), - "f"(d_reg[2]), - "f"(d_reg[3])); -#endif -} - -extern "C" { -__device__ uint32_t __nvvm_get_smem_pointer(void*); -} - -__device__ void ldgsts_128(void const* gPtr, void* sPtr, uint32_t pred) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - if (pred) { - uint32_t smemPtrAsUint32 = __nvvm_get_smem_pointer(sPtr); - asm volatile("cp.async.cg.shared.global.L2::128B [%0], [%1], %2;\n" ::"r"(smemPtrAsUint32), "l"(gPtr), "n"(16)); - } -#endif -} - -__device__ void ldsm_x4(void* smem_ptr, uint32_t* reg_ptr) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n" - : "=r"(reg_ptr[0]), "=r"(reg_ptr[1]), "=r"(reg_ptr[2]), "=r"(reg_ptr[3]) - : "r"(__nvvm_get_smem_pointer(smem_ptr))); -#endif -} - -template -__device__ int apply_swizzle_343_on_elem_row_col(int row_idx_, int col_idx_) { - uint32_t row_idx = *reinterpret_cast(&row_idx_); - uint32_t col_idx = *reinterpret_cast(&col_idx_); - row_idx = row_idx % 8; - row_idx = row_idx * (16 / sizeof(Type)); - col_idx = col_idx ^ row_idx; - return *reinterpret_cast(&col_idx); -} - -__device__ void initialize_barrier( - uint64_t* smem_barrier, // 64 bits user-manged barrier in smem - int thread_count = 1) // Thread count expected to arrive/wait on this barrier -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); - asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;\n" ::"r"(smem_int_ptr), "r"(thread_count)); -#endif -} - -// Barrier wait -__device__ void wait_barrier( - uint64_t* smem_barrier, // 64 bits user-manged barrier in smem - int phase_bit) // Current phase bit the barrier waiting to flip -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); - asm volatile( - "{\n" - ".reg .pred P1;\n" - "LAB_WAIT:\n" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%0], %1;\n" - "@P1 bra DONE;\n" - "bra LAB_WAIT;\n" - "DONE:\n" - "}\n" ::"r"(smem_int_ptr), - "r"(phase_bit)); -#endif -} - -__device__ bool try_wait_barrier(uint64_t* smem_ptr, int phase_bit) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - uint32_t wait_complete; - uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_ptr); - asm volatile( - "{\n\t" - ".reg .pred P1; \n\t" - "mbarrier.try_wait.parity.shared::cta.b64 P1, [%1], %2; \n\t" - "selp.b32 %0, 1, 0, P1; \n\t" - "}" - : "=r"(wait_complete) - : "r"(smem_int_ptr), "r"(phase_bit)); - return static_cast(wait_complete); -#endif - return false; -} - -// Barrier arrive -__device__ void arrive_barrier(uint64_t* smem_barrier) // 64 bits user-manged barrier in smem -{ -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); - asm volatile( - "{\n" - ".reg .b64 state; \n" - "mbarrier.arrive.shared::cta.b64 state, [%0];\n" - "}\n" ::"r"(smem_int_ptr)); -#endif -} - -__device__ void ldgsts_arrive(uint64_t* smem_barrier) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - uint32_t smem_int_ptr = __nvvm_get_smem_pointer(smem_barrier); - asm volatile("cp.async.mbarrier.arrive.noinc.shared.b64 [%0];" : : "r"(smem_int_ptr)); -#endif -} - -template -struct GmemLoaderA { - static constexpr int elem_bytes = 2; - static constexpr int vec_bytes = 16; - static constexpr int vec_elems = vec_bytes / elem_bytes; - static constexpr int thread_cnt = 64; - static_assert((tile_m * tile_k) % (vec_elems * thread_cnt) == 0); - static constexpr int a_inst_cnt_per_iter = (tile_m * tile_k) / (vec_elems * thread_cnt); - static_assert(gemm_k % tile_k == 0); - static constexpr int k_iter_cnt = gemm_k / tile_k; - - // Extra params to keep the order of k reduction... - static constexpr int mma_warp_cnt = 4; - static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt; - static constexpr int k_each_chunk = gemm_k / mma_warp_cnt; - - private: - __device__ int k_project(int tile_k_idx) { - return (tile_k_idx / per_mma_warp_k * k_each_chunk) + (tile_k_idx % per_mma_warp_k); - } - - public: - __device__ GmemLoaderA(bf16_t const* gmem_a_local_, bf16_t* smem_a_, uint64_t* smem_barrier_) - : gmem_a(gmem_a_local_), smem_a(smem_a_), smem_barrier(smem_barrier_), local_tid(threadIdx.x % thread_cnt) {} - - __device__ void prepare() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 -// swizzle, that's what we want. -#pragma unroll - for (int i = 0; i < a_inst_cnt_per_iter; i++) { - int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems; - int m_idx = linear_idx / tile_k; - int k_idx = linear_idx % tile_k; - k_idx = apply_swizzle_343_on_elem_row_col(m_idx, k_idx); - a_smem_offsets[i] = m_idx * tile_k + k_idx; - } -#endif - } - - __device__ void issue_mainloop() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 -#pragma unroll 1 - for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) { - if (need_wait) { - wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit); - } - int next_stage_idx = stage_idx + 1; - int next_phase_bit = next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit; - next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx; - if (loop_idx != k_iter_cnt - 1) { - need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2, next_phase_bit); - } - -#pragma unroll - for (int i = 0; i < a_inst_cnt_per_iter; i++) { - int smem_offset = a_smem_offsets[i]; - bf16_t* smem_ptr_this_iter = smem_a + stage_idx * tile_m * tile_k + smem_offset; - int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems; - int m_idx = linear_idx / tile_k; - int k_idx = linear_idx % tile_k; - int gmem_offset = m_idx * gemm_k + k_project(k_idx); - bf16_t const* gmem_ptr_this_iter = gmem_a + gmem_offset; - ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, true); - } - ldgsts_arrive(smem_barrier + stage_idx * 2); - - stage_idx = next_stage_idx; - phase_bit = next_phase_bit; - gmem_a += per_mma_warp_k; - } -#endif - } - - bf16_t const* gmem_a; - bf16_t* smem_a; - uint64_t* smem_barrier; - int local_tid; - int stage_idx = 0; - int phase_bit = 1; - bool need_wait = true; - - // per smem_stage, store with swizzle information - int a_smem_offsets[a_inst_cnt_per_iter]; -}; - -template -struct GmemLoaderB { - static constexpr int elem_bytes = 2; - static constexpr int vec_bytes = 16; - static constexpr int vec_elems = vec_bytes / elem_bytes; - static constexpr int thread_cnt = 64; - static_assert((tile_n * tile_k) % (vec_elems * thread_cnt) == 0); - static constexpr int b_inst_cnt_per_iter = (tile_n * tile_k) / (vec_elems * thread_cnt); - static_assert(gemm_k % tile_k == 0); - static constexpr int k_iter_cnt = gemm_k / tile_k; - - // Extra params to keep the order of k reduction... - static constexpr int mma_warp_cnt = 4; - static constexpr int per_mma_warp_k = tile_k / mma_warp_cnt; - static constexpr int k_each_chunk = gemm_k / mma_warp_cnt; - - private: - __device__ int k_project(int tile_k_idx) { - return (tile_k_idx / per_mma_warp_k * k_each_chunk) + (tile_k_idx % per_mma_warp_k); - } - - public: - __device__ GmemLoaderB(bf16_t const* gmem_b_local_, bf16_t* smem_b_, uint64_t* smem_barrier_, int gemm_n_) - : gmem_b(gmem_b_local_), - smem_b(smem_b_), - smem_barrier(smem_barrier_), - gemm_n(gemm_n_), - local_tid(threadIdx.x % thread_cnt) {} - - __device__ void prepare() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 -// swizzle, that's what we want. -#pragma unroll - for (int i = 0; i < b_inst_cnt_per_iter; i++) { - int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems; - int n_idx = linear_idx / tile_k; - int k_idx = linear_idx % tile_k; - k_idx = apply_swizzle_343_on_elem_row_col(n_idx, k_idx); - b_smem_offsets[i] = n_idx * tile_k + k_idx; - preds[i] = n_idx < gemm_n; - } -#endif - } - - __device__ void issue_mainloop() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - cudaGridDependencySynchronize(); -#pragma unroll 1 - for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) { - if (need_wait) { - wait_barrier(smem_barrier + 1 + stage_idx * 2, phase_bit); - } - int next_stage_idx = stage_idx + 1; - int next_phase_bit = next_stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit; - next_stage_idx = next_stage_idx == stage_cnt ? 0 : next_stage_idx; - if (loop_idx != k_iter_cnt - 1) { - need_wait = !try_wait_barrier(smem_barrier + 1 + next_stage_idx * 2, next_phase_bit); - } -#pragma unroll - for (int i = 0; i < b_inst_cnt_per_iter; i++) { - int smem_offset = b_smem_offsets[i]; - bf16_t* smem_ptr_this_iter = smem_b + stage_idx * tile_n * tile_k + smem_offset; - int linear_idx = local_tid * vec_elems + i * thread_cnt * vec_elems; - int n_idx = linear_idx / tile_k; - int k_idx = linear_idx % tile_k; - int gmem_offset = n_idx * gemm_k + k_project(k_idx); - bf16_t const* gmem_ptr_this_iter = gmem_b + gmem_offset; - ldgsts_128(gmem_ptr_this_iter, smem_ptr_this_iter, preds[i]); - } - ldgsts_arrive(smem_barrier + stage_idx * 2); - - stage_idx = next_stage_idx; - phase_bit = next_phase_bit; - gmem_b += per_mma_warp_k; - } -#endif - } - - bf16_t const* gmem_b; - bf16_t* smem_b; - uint64_t* smem_barrier; - int gemm_n; - int local_tid; - int stage_idx = 0; - int phase_bit = 1; - bool need_wait = true; - - // per smem_stage, store with swizzle information - int b_smem_offsets[b_inst_cnt_per_iter]; - uint32_t preds[b_inst_cnt_per_iter]; -}; - -template -struct MmaComputer { - static constexpr int elem_bytes = 2; - static constexpr int thread_cnt = 128; - static_assert(gemm_k % tile_k == 0); - static_assert(tile_k % (thread_cnt / 32) == 0); - static constexpr int per_warp_tile_k = tile_k / (thread_cnt / 32); - static constexpr int k_iter_cnt = gemm_k / tile_k; - static constexpr int k_phase_cnt = per_warp_tile_k / 16; - static constexpr int m_iter_cnt = (tile_m + 15) / 16; - static constexpr int n_iter_cnt = (tile_n + 7) / 8; // Possible to have non-1 n_iter_cnt for ab_swap m16 case. - static_assert(m_iter_cnt == 1); - static_assert(n_iter_cnt == 1 || n_iter_cnt == 2); - - __device__ MmaComputer( - bf16_t* gmem_c_local_, bf16_t* smem_a_, bf16_t* smem_b_, uint64_t* smem_barrier_, int warp_idx_, int gemm_n_) - : gmem_c(gmem_c_local_), - smem_a(smem_a_), - smem_b(smem_b_), - smem_barrier(smem_barrier_), - warp_idx(warp_idx_ - (thread_cnt / 32)), - gemm_n(gemm_n_) {} - - private: - __device__ constexpr int internal_b_atom_func(int tid) { - if constexpr (tile_n < 8) { - return (tid % tile_n) + ((tid % 8) / tile_n * 0) + tid / 8 * 8 * tile_n; - } else { - return (tid % 8) + ((tid % 32) / 8 * (tile_n * 8)); - } - } - - public: - __device__ void prepare() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 -#pragma unroll - for (int i = 0; i < k_phase_cnt; i++) { - int linear_idx = (lane_idx % 16) + (lane_idx / 16) * 128 + i * 256; - int m_idx = linear_idx % tile_m; - int k_idx = linear_idx / tile_m + warp_k_offset_in_tile_k; - k_idx = apply_swizzle_343_on_elem_row_col(m_idx, k_idx); - a_smem_offsets[0][i] = m_idx * tile_k + k_idx; - } -#pragma unroll - for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { -#pragma unroll - for (int i = 0; i < k_phase_cnt; i += 2) { // Special i+=2 for B. - int linear_idx = internal_b_atom_func(lane_idx) + i * tile_n * 16 + n_iter_idx * 8; - int n_idx = linear_idx % tile_n; - int k_idx = linear_idx / tile_n + warp_k_offset_in_tile_k; - k_idx = apply_swizzle_343_on_elem_row_col(n_idx, k_idx); - b_smem_offsets[n_iter_idx][i] = n_idx * tile_k + k_idx; - } - } -#endif - } - - __device__ void issue_mainloop() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 -#pragma unroll 1 - for (int loop_idx = 0; loop_idx < k_iter_cnt; loop_idx++) { - wait_barrier(smem_barrier + 0 + stage_idx * 2, phase_bit); - -#pragma unroll - for (int i = 0; i < k_phase_cnt; i++) { - int smem_offset = a_smem_offsets[0][i]; - bf16_t* smem_ptr_this_iter = smem_a + stage_idx * tile_m * tile_k + smem_offset; - ldsm_x4(smem_ptr_this_iter, reinterpret_cast(a_reg[0][i])); - } - -#pragma unroll - for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { -#pragma unroll - for (int i = 0; i < k_phase_cnt; i += 2) { - int smem_offset = b_smem_offsets[n_iter_idx][i]; - bf16_t* smem_ptr_this_iter = smem_b + stage_idx * tile_n * tile_k + smem_offset; - ldsm_x4(smem_ptr_this_iter, reinterpret_cast(b_reg[n_iter_idx][i])); - } - } - -#pragma unroll - for (int k_iter_idx = 0; k_iter_idx < k_phase_cnt; k_iter_idx++) { -#pragma unroll - for (int n_iter_idx = 0; n_iter_idx < n_iter_cnt; n_iter_idx++) { - hmma_16_8_16_f32acc_bf16ab( - acc_reg[0][n_iter_idx], a_reg[0][k_iter_idx], b_reg[n_iter_idx][k_iter_idx], acc_reg[0][n_iter_idx]); - } - } - ::arrive_barrier(smem_barrier + 1 + stage_idx * 2); - stage_idx += 1; - phase_bit = stage_idx == stage_cnt ? phase_bit ^ 1 : phase_bit; - stage_idx = stage_idx == stage_cnt ? 0 : stage_idx; - } -#endif - } - - __device__ void epi() { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt)); - // reorganize the acc_reg - constexpr int thread_m = 2; - constexpr int thread_n = 2 * n_iter_cnt; - constexpr int cta_mma_n = n_iter_cnt * 8; - float acc_reg_reorg[thread_m][thread_n]; - - for (int i = 0; i < thread_m; i++) { - for (int j = 0; j < thread_n; j++) { - acc_reg_reorg[i][j] = acc_reg[0][j / 2][(j % 2) + (i * 2)]; - } - } - - // 4 x cosize(smem_c_layout) - float* smem_c = reinterpret_cast(smem_a); - // coord -> index - auto smem_c_index_func = [&](int m_idx, int n_idx) { - int group_rows = 32 / cta_mma_n; - int group_cnt = 2; - return (m_idx % group_rows * cta_mma_n) + (m_idx / group_rows * (32 + group_cnt)) + n_idx; - }; - constexpr int cosize_smem_c = ((tile_m * cta_mma_n) / 32) * (32 + 2); - -// This should be optimized to STS.64 but can not be STS.128 due to the bank index. -#pragma unroll - for (int m_idx_thread = 0; m_idx_thread < thread_m; m_idx_thread++) { -#pragma unroll - for (int n_idx_thread = 0; n_idx_thread < thread_n; n_idx_thread++) { - int m_idx = (lane_idx / 4) + m_idx_thread * 8; - int n_idx = ((lane_idx % 4) * 2) + (n_idx_thread % 2) + (n_idx_thread / 2) * 8; - smem_c[cosize_smem_c * warp_idx + smem_c_index_func(m_idx, n_idx)] = acc_reg_reorg[m_idx_thread][n_idx_thread]; - } - } - asm volatile("bar.sync %0, %1;" : : "r"(1), "r"(thread_cnt)); - - if (warp_idx == 0) { - constexpr int final_acc_reg_cnt = (tile_m * tile_n + 31) / 32; - float acc_final[final_acc_reg_cnt]{}; - -#pragma unroll - for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) { - int linear_idx = reg_idx * 32 + lane_idx; - int m_idx = linear_idx % tile_m; - int n_idx = linear_idx / tile_m; - acc_final[reg_idx] += smem_c[smem_c_index_func(m_idx, n_idx) + 0 * cosize_smem_c] + - smem_c[smem_c_index_func(m_idx, n_idx) + 1 * cosize_smem_c] + - smem_c[smem_c_index_func(m_idx, n_idx) + 2 * cosize_smem_c] + - smem_c[smem_c_index_func(m_idx, n_idx) + 3 * cosize_smem_c]; - } - -#pragma unroll - for (int reg_idx = 0; reg_idx < final_acc_reg_cnt; reg_idx++) { - int linear_idx = reg_idx * 32 + lane_idx; - int m_idx = linear_idx % tile_m; - int n_idx = linear_idx / tile_m; - if (m_idx < tile_m && n_idx < gemm_n) { - gmem_c[n_idx * gemm_m + m_idx] = acc_final[reg_idx]; - } - } - } -#endif - } - - bf16_t* gmem_c; - bf16_t* smem_a; - bf16_t* smem_b; - uint64_t* smem_barrier; - int warp_idx; - int gemm_n; - int stage_idx = 0; - int phase_bit = 0; - int lane_idx = threadIdx.x % 32; - int warp_k_offset_in_tile_k = warp_idx * per_warp_tile_k; - - int a_smem_offsets[m_iter_cnt][k_phase_cnt]; - int b_smem_offsets[n_iter_cnt][k_phase_cnt]; - - bf16_t a_reg[m_iter_cnt][k_phase_cnt][8]; - bf16_t b_reg[n_iter_cnt][k_phase_cnt][4]; - float acc_reg[m_iter_cnt][n_iter_cnt][4]{}; -}; - -// AB swapped, kernel is k-major, k-major, m-major -template -__global__ __launch_bounds__(256, 1) void fused_a_gemm_kernel( - bf16_t* output, bf16_t const* mat_a, bf16_t const* mat_b, int gemm_n) { -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 - constexpr int load_thread_cnt = 128; - constexpr int compute_thread_cnt = 128; - constexpr int thread_cnt = load_thread_cnt + compute_thread_cnt; - (void)thread_cnt; - static_assert(gemm_m % 16 == 0); - static_assert(gemm_k % tile_k == 0); - static_assert(gemm_m % tile_m == 0); - static_assert( - tile_k == 128 || tile_k == 256 || tile_k == 512 || - tile_k == 1024); // tile_k must be larger than 64 since 4 warp splitK. - static_assert(tile_m == 16); - constexpr int g2s_vec_bytes = 16; - constexpr int a_elem_bytes = 2; - constexpr int b_elem_bytes = 2; - // constexpr int c_elem_bytes = 2; - static_assert((tile_m * a_elem_bytes + tile_n * b_elem_bytes) * tile_k * stage_cnt <= 225 * 1024); - static_assert((tile_m * tile_k * a_elem_bytes) % (load_thread_cnt * g2s_vec_bytes) == 0); - static_assert((tile_n * tile_k * b_elem_bytes) % (load_thread_cnt * g2s_vec_bytes) == 0); - - extern __shared__ char smem[]; - uint64_t* smem_barrier = reinterpret_cast(smem); // producer,consumer; producer,consumer; ... - bf16_t* smem_a = reinterpret_cast(smem + (stage_cnt * 8 * 2 + 1024) / 1024 * 1024); - bf16_t* smem_b = smem_a + tile_m * tile_k * stage_cnt; - - int cta_m_idx = tile_m * blockIdx.x; - int cta_n_idx = tile_n * blockIdx.y; - bf16_t const* gmem_a_local = mat_a + cta_m_idx * gemm_k; - bf16_t const* gmem_b_local = mat_b + cta_n_idx * gemm_k; - bf16_t* gmem_c_local = output + cta_n_idx * gemm_m + cta_m_idx; - - int warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); - - if (warp_idx == 4) { - for (int i = 0; i < stage_cnt; i++) { - initialize_barrier(smem_barrier + i * 2 + 0, load_thread_cnt); // producer - initialize_barrier(smem_barrier + i * 2 + 1, compute_thread_cnt); // consumer - } - } - __syncthreads(); - - if (warp_idx < 2) { - GmemLoaderA a_loader(gmem_a_local, smem_a, smem_barrier); - a_loader.prepare(); - a_loader.issue_mainloop(); - } else if (warp_idx < 4) { - GmemLoaderB b_loader(gmem_b_local, smem_b, smem_barrier, gemm_n); - b_loader.prepare(); - b_loader.issue_mainloop(); - } else { - MmaComputer mma_computer( - gmem_c_local, smem_a, smem_b, smem_barrier, warp_idx, gemm_n); - mma_computer.prepare(); - mma_computer.issue_mainloop(); - mma_computer.epi(); - } - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -template -void invokeFusedAGemm(T* output, T const* mat_a, T const* mat_b, int num_tokens, cudaStream_t const stream) { - constexpr int gemm_m = kHdOut; // 2112 - int const gemm_n = num_tokens; // 16 - constexpr int gemm_k = kHdIn; // 7168 - constexpr int batch_size = 1; - std::swap(mat_a, mat_b); - constexpr int tile_m = 16; - constexpr int tile_n = kTileN; // 8 or 16 - constexpr int tile_k = std::max(256, 1024 / tile_n); // 256 - constexpr int max_stage_cnt = 1024 * 192 / ((tile_m + tile_n) * tile_k * sizeof(bf16_t)); - constexpr int k_iter_cnt = gemm_k / tile_k; - constexpr int stage_cnt = - k_iter_cnt > max_stage_cnt ? max_stage_cnt : k_iter_cnt; // possible tunable for smallK > 1 wave n. // 22 - int cta_m_cnt = gemm_m / tile_m; - int cta_n_cnt = (gemm_n + tile_n - 1) / tile_n; - constexpr int barrier_bytes = (stage_cnt * 16 + 1023) / 1024 * 1024; // 4096 - constexpr int smem_bytes = ((tile_m * 2 + tile_n * 2) * tile_k * stage_cnt + barrier_bytes + 1023) / 1024 * 1024; - - dim3 grid(cta_m_cnt, cta_n_cnt, 1); - dim3 block_size(256); - cudaLaunchConfig_t config; - config.gridDim = grid; - config.blockDim = block_size; - config.dynamicSmemBytes = smem_bytes; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL(); - config.numAttrs = 1; - config.attrs = attrs; - if (smem_bytes >= (48 * 1024)) { - cudaFuncSetAttribute( - fused_a_gemm_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, - smem_bytes); - } - cudaLaunchKernelEx( - &config, - fused_a_gemm_kernel, - output, - mat_a, - mat_b, - gemm_n); -} - -template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 8>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, cudaStream_t); - -template void invokeFusedAGemm<__nv_bfloat16, 7168, 2112, 16>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, int num_tokens, cudaStream_t); - -void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b) { - TORCH_CHECK(mat_a.dim() == 2 && mat_b.dim() == 2 && output.dim() == 2); - int const num_tokens = mat_a.size(0); - int const hd_in = mat_a.size(1); - int const hd_out = mat_b.size(1); - - constexpr int kHdIn = 7168; - constexpr int kHdOut = 2112; - TORCH_CHECK(num_tokens >= 1 && num_tokens <= 16, "required 1 <= mat_a.shape[0] <= 16") - TORCH_CHECK(hd_in == kHdIn, "required mat_a.shape[1] == 7168") - TORCH_CHECK(hd_out == kHdOut, "required mat_b.shape[1] == 2112") - TORCH_CHECK(output.size(0) == num_tokens, "required output.shape[0] == mat_a.shape[0]") - TORCH_CHECK(output.size(1) == hd_out, "required output.shape[1] == mat_b.shape[1]") - - TORCH_CHECK(mat_a.stride(1) == 1, "mat_a must be a row major tensor"); // Row-major - TORCH_CHECK(output.stride(1) == 1, "output must be a row major tensor"); // Row-major - TORCH_CHECK(mat_b.stride(0) == 1, "mat_b must be a column major tensor"); // Column-major - - auto const data_type = mat_a.scalar_type(); - TORCH_CHECK( - mat_a.scalar_type() == torch::kBFloat16 && mat_b.scalar_type() == torch::kBFloat16, - "Only BFloat16 input dtype is supported") - TORCH_CHECK(output.scalar_type() == torch::kBFloat16, "Only BFloat16 output dtype is supported") - - auto const sm = getSMVersion(); -#ifndef USE_MUSA - TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90"); -#else - TORCH_CHECK(sm >= 22, "required MUSA ARCH >= MP_22"); -#endif - - auto stream = at::cuda::getCurrentCUDAStream(mat_a.get_device()); - if (num_tokens <= 8) { - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 8>( - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), - num_tokens, - stream); - } else { - invokeFusedAGemm<__nv_bfloat16, kHdIn, kHdOut, 16>( - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), - num_tokens, - stream); - } -} diff --git a/sgl-kernel/csrc/gemm/dsv3_router_gemm_bf16_out.cu b/sgl-kernel/csrc/gemm/dsv3_router_gemm_bf16_out.cu deleted file mode 100644 index e60a83db6f61..000000000000 --- a/sgl-kernel/csrc/gemm/dsv3_router_gemm_bf16_out.cu +++ /dev/null @@ -1,284 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp - * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include - -#include "cuda_bf16.h" -#include "cuda_runtime.h" -#include "utils.h" - -// Custom FMA implementation using PTX assembly instructions -__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b, float2 const& c) { - asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" - : "=l"(reinterpret_cast(d)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b)), - "l"(reinterpret_cast(c))); -} - -// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion -template -__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec, float* dst) { - __nv_bfloat16* bf16_ptr = reinterpret_cast<__nv_bfloat16*>(const_cast(&vec)); - -#pragma unroll - for (int i = 0; i < VPT; i++) { - dst[i] = __bfloat162float(bf16_ptr[i]); - } -} - -template -__global__ -__launch_bounds__(128, 1) void router_gemm_kernel_bf16_output(__nv_bfloat16* out, T const* mat_a, T const* mat_b) { - // Each block handles one expert column - int const n_idx = blockIdx.x; - int const tid = threadIdx.x; - constexpr int kWarpSize = 32; - constexpr int kNumWarps = kBlockSize / kWarpSize; - // Constants for this kernel - constexpr int k_elems_per_k_iteration = VPT * kBlockSize; - constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; // Total K iterations - - // Initialize accumulators for all M rows - float acc[kNumTokens] = {}; - - // Shared memory for warp-level reduction - __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps - - // B matrix is in column-major order, so we can directly load a column for the n_idx expert - T const* b_col = mat_b + n_idx * kHiddenDim; - - // Pre-compute k_base values for each iteration to help compiler optimize - // int k_bases[k_iterations]; - int k_bases[k_iterations]; -#pragma unroll - for (int ki = 0; ki < k_iterations; ki++) { - k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); -#endif - - // Process the GEMM in chunks - for (int ki = 0; ki < k_iterations; ki++) { - int const k_base = k_bases[ki]; - - // Load B matrix values using vector load (8 bf16 values) - uint4 b_vec = *reinterpret_cast(b_col + k_base); - - // Convert B values to float - float b_float[VPT]; - bf16_uint4_to_float8(b_vec, b_float); - -// Process each token -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { - // Load both rows of A matrix using vector loads - uint4 a_vec = *reinterpret_cast(mat_a + (m_idx * kHiddenDim) + k_base); - - // Convert A values to float - float a_float[VPT]; - bf16_uint4_to_float8(a_vec, a_float); - -// Process elements in this chunk -#pragma unroll - for (int k = 0; k < VPT; k++) { - float a = a_float[k]; - float b = b_float[k]; - acc[m_idx] += a * b; - } - } - } - - // Perform warp-level reduction - int const warpSize = 32; - int const warpId = tid / warpSize; - int const laneId = tid % warpSize; - - // Register for warp-level reduction results - float warp_result[kNumTokens]; - -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { - warp_result[m_idx] = acc[m_idx]; - } - -// Perform warp-level reduction using optimized butterfly pattern -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float sum = warp_result[m]; - - // Butterfly reduction pattern - sum += __shfl_xor_sync(0xffffffff, sum, 16); - sum += __shfl_xor_sync(0xffffffff, sum, 8); - sum += __shfl_xor_sync(0xffffffff, sum, 4); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 1); - - // Only the first thread in each warp stores to shared memory - if (laneId == 0) { - sm_reduction[m][warpId] = sum; - } - } - - __syncthreads(); - - // Final reduction across warps (only first thread) - if (tid == 0) { -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float final_sum = 0.0f; - -// Sum across the kNumWarps -#pragma unroll - for (int w = 0; w < kNumWarps; w++) { - final_sum += sm_reduction[m][w]; - } - - // Write final result - out[m * kNumExperts + n_idx] = __float2bfloat16(final_sum); - } - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -template -void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a, T const* mat_b, cudaStream_t stream) { - constexpr int VPT = 16 / sizeof(T); - constexpr int kBlockSize = 128; - cudaLaunchConfig_t config; - config.gridDim = kNumExperts; - config.blockDim = kBlockSize; - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL(); - config.numAttrs = 1; - config.attrs = attrs; - cudaLaunchKernelEx( - &config, - router_gemm_kernel_bf16_output, - output, - mat_a, - mat_b); -} - -// Template instantiations for DEFAULT_NUM_EXPERTS experts -template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 256, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -// Template instantiations for KIMI_K2_NUM_EXPERTS experts -template void invokeRouterGemmBf16Output<__nv_bfloat16, 1, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 2, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 3, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 4, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 5, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 6, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 7, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 8, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 9, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 10, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 11, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 12, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 13, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 14, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 15, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmBf16Output<__nv_bfloat16, 16, 384, 7168>( - __nv_bfloat16*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); diff --git a/sgl-kernel/csrc/gemm/dsv3_router_gemm_entry.cu b/sgl-kernel/csrc/gemm/dsv3_router_gemm_entry.cu deleted file mode 100644 index a3b6b272c82e..000000000000 --- a/sgl-kernel/csrc/gemm/dsv3_router_gemm_entry.cu +++ /dev/null @@ -1,165 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp - * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include - -#include "cuda_bf16.h" -#include "cuda_runtime.h" -#include "utils.h" - -static constexpr int DEFAULT_NUM_EXPERTS = 256; -static constexpr int KIMI_K2_NUM_EXPERTS = 384; -static constexpr int DEFAULT_HIDDEN_DIM = 7168; - -template -void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b, cudaStream_t stream); - -template -void invokeRouterGemmBf16Output(__nv_bfloat16* output, T const* mat_a, T const* mat_b, cudaStream_t stream); - -template -struct LoopUnroller { - static void unroll_float_output( - int num_tokens, float* output, __nv_bfloat16 const* input, __nv_bfloat16 const* weights, cudaStream_t stream) { - if (num_tokens == kBegin) { - invokeRouterGemmFloatOutput<__nv_bfloat16, kBegin, kNumExperts, kHiddenDim>(output, input, weights, stream); - } else { - LoopUnroller::unroll_float_output( - num_tokens, output, input, weights, stream); - } - } - - static void unroll_bf16_output( - int num_tokens, - __nv_bfloat16* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kBegin) { - invokeRouterGemmBf16Output<__nv_bfloat16, kBegin, kNumExperts, kHiddenDim>(output, input, weights, stream); - } else { - LoopUnroller::unroll_bf16_output( - num_tokens, output, input, weights, stream); - } - } -}; - -template -struct LoopUnroller { - static void unroll_float_output( - int num_tokens, float* output, __nv_bfloat16 const* input, __nv_bfloat16 const* weights, cudaStream_t stream) { - if (num_tokens == kEnd) { - invokeRouterGemmFloatOutput<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(output, input, weights, stream); - } else { - throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); - } - } - - static void unroll_bf16_output( - int num_tokens, - __nv_bfloat16* output, - __nv_bfloat16 const* input, - __nv_bfloat16 const* weights, - cudaStream_t stream) { - if (num_tokens == kEnd) { - invokeRouterGemmBf16Output<__nv_bfloat16, kEnd, kNumExperts, kHiddenDim>(output, input, weights, stream); - } else { - throw std::invalid_argument("Invalid num_tokens, only supports 1 to 16"); - } - } -}; - -void dsv3_router_gemm( - torch::Tensor& output, // [num_tokens, num_experts] - const torch::Tensor& mat_a, // [num_tokens, hidden_dim] - const torch::Tensor& mat_b // [num_experts, hidden_dim] -) { - TORCH_CHECK(output.dim() == 2 && mat_a.dim() == 2 && mat_b.dim() == 2); - - const int num_tokens = mat_a.size(0); - const int num_experts = mat_b.size(0); - const int hidden_dim = mat_a.size(1); - - TORCH_CHECK(mat_a.size(1) == mat_b.size(1), "mat_a and mat_b must have the same hidden_dim"); - TORCH_CHECK( - hidden_dim == DEFAULT_HIDDEN_DIM, - "Expected hidden_dim=", - DEFAULT_HIDDEN_DIM, - ", but got hidden_dim=", - hidden_dim); - TORCH_CHECK( - num_experts == DEFAULT_NUM_EXPERTS || num_experts == KIMI_K2_NUM_EXPERTS, - "Expected num_experts=", - DEFAULT_NUM_EXPERTS, - " or num_experts=", - KIMI_K2_NUM_EXPERTS, - ", but got num_experts=", - num_experts); - TORCH_CHECK( - num_tokens >= 1 && num_tokens <= 16, "currently num_tokens must be less than or equal to 16 for router_gemm"); - TORCH_CHECK(mat_a.dtype() == torch::kBFloat16, "mat_a must be bf16"); - TORCH_CHECK(mat_b.dtype() == torch::kBFloat16, "mat_b must be bf16"); - TORCH_CHECK( - output.dtype() == torch::kFloat32 || output.dtype() == torch::kBFloat16, "output must be float32 or bf16"); - - auto const sm = getSMVersion(); -#ifndef USE_MUSA - TORCH_CHECK(sm >= 90, "required CUDA ARCH >= SM_90"); -#else - TORCH_CHECK(sm >= 22, "required MUSA ARCH >= MP_22"); -#endif - - const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); - - if (output.dtype() == torch::kFloat32) { - if (num_experts == DEFAULT_NUM_EXPERTS) { - LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::unroll_float_output( - num_tokens, - reinterpret_cast(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), - stream); - } else if (num_experts == KIMI_K2_NUM_EXPERTS) { - LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::unroll_float_output( - num_tokens, - reinterpret_cast(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), - stream); - } - } else if (output.dtype() == torch::kBFloat16) { - if (num_experts == DEFAULT_NUM_EXPERTS) { - LoopUnroller<1, 16, DEFAULT_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::unroll_bf16_output( - num_tokens, - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), - stream); - } else if (num_experts == KIMI_K2_NUM_EXPERTS) { - LoopUnroller<1, 16, KIMI_K2_NUM_EXPERTS, DEFAULT_HIDDEN_DIM>::unroll_bf16_output( - num_tokens, - reinterpret_cast<__nv_bfloat16*>(output.mutable_data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_a.data_ptr()), - reinterpret_cast<__nv_bfloat16 const*>(mat_b.data_ptr()), - stream); - } - } -} diff --git a/sgl-kernel/csrc/gemm/dsv3_router_gemm_float_out.cu b/sgl-kernel/csrc/gemm/dsv3_router_gemm_float_out.cu deleted file mode 100644 index 0abcaf1c1e35..000000000000 --- a/sgl-kernel/csrc/gemm/dsv3_router_gemm_float_out.cu +++ /dev/null @@ -1,283 +0,0 @@ -/* - * Adapted from - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/kernels/dsv3MinLatencyKernels/dsv3RouterGemm.cu - * https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/thop/dsv3RouterGemmOp.cpp - * - * Copyright (c) 2019-2023, NVIDIA CORPORATION. All rights reserved. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -#include -#include - -#include "cuda_bf16.h" -#include "cuda_runtime.h" -#include "utils.h" - -// Custom FMA implementation using PTX assembly instructions -__device__ __forceinline__ void fma(float2& d, float2 const& a, float2 const& b, float2 const& c) { - asm volatile("fma.rn.f32x2 %0, %1, %2, %3;\n" - : "=l"(reinterpret_cast(d)) - : "l"(reinterpret_cast(a)), - "l"(reinterpret_cast(b)), - "l"(reinterpret_cast(c))); -} - -// Convert 8 bfloat16 values from a uint4 to float array - optimized conversion -template -__device__ __forceinline__ void bf16_uint4_to_float8(uint4 const& vec, float* dst) { - __nv_bfloat16* bf16_ptr = reinterpret_cast<__nv_bfloat16*>(const_cast(&vec)); - -#pragma unroll - for (int i = 0; i < VPT; i++) { - dst[i] = __bfloat162float(bf16_ptr[i]); - } -} - -template -__global__ __launch_bounds__(128, 1) void router_gemm_kernel_float_output(float* out, T const* mat_a, T const* mat_b) { - // Each block handles one expert column - int const n_idx = blockIdx.x; - int const tid = threadIdx.x; - constexpr int kWarpSize = 32; - constexpr int kNumWarps = kBlockSize / kWarpSize; - // Constants for this kernel - constexpr int k_elems_per_k_iteration = VPT * kBlockSize; - constexpr int k_iterations = kHiddenDim / k_elems_per_k_iteration; // Total K iterations - - // Initialize accumulators for all M rows - float acc[kNumTokens] = {}; - - // Shared memory for warp-level reduction - __shared__ float sm_reduction[kNumTokens][kNumWarps]; // kNumWarps - - // B matrix is in column-major order, so we can directly load a column for the n_idx expert - T const* b_col = mat_b + n_idx * kHiddenDim; - - // Pre-compute k_base values for each iteration to help compiler optimize - // int k_bases[k_iterations]; - int k_bases[k_iterations]; -#pragma unroll - for (int ki = 0; ki < k_iterations; ki++) { - k_bases[ki] = ki * k_elems_per_k_iteration + tid * VPT; - } - -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaGridDependencySynchronize(); -#endif - - // Process the GEMM in chunks - for (int ki = 0; ki < k_iterations; ki++) { - int const k_base = k_bases[ki]; - - // Load B matrix values using vector load (8 bf16 values) - uint4 b_vec = *reinterpret_cast(b_col + k_base); - - // Convert B values to float - float b_float[VPT]; - bf16_uint4_to_float8(b_vec, b_float); - -// Process each token -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { - // Load both rows of A matrix using vector loads - uint4 a_vec = *reinterpret_cast(mat_a + (m_idx * kHiddenDim) + k_base); - - // Convert A values to float - float a_float[VPT]; - bf16_uint4_to_float8(a_vec, a_float); - -// Process elements in this chunk -#pragma unroll - for (int k = 0; k < VPT; k++) { - float a = a_float[k]; - float b = b_float[k]; - acc[m_idx] += a * b; - } - } - } - - // Perform warp-level reduction - int const warpSize = 32; - int const warpId = tid / warpSize; - int const laneId = tid % warpSize; - - // Register for warp-level reduction results - float warp_result[kNumTokens]; - -#pragma unroll - for (int m_idx = 0; m_idx < kNumTokens; m_idx++) { - warp_result[m_idx] = acc[m_idx]; - } - -// Perform warp-level reduction using optimized butterfly pattern -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float sum = warp_result[m]; - - // Butterfly reduction pattern - sum += __shfl_xor_sync(0xffffffff, sum, 16); - sum += __shfl_xor_sync(0xffffffff, sum, 8); - sum += __shfl_xor_sync(0xffffffff, sum, 4); - sum += __shfl_xor_sync(0xffffffff, sum, 2); - sum += __shfl_xor_sync(0xffffffff, sum, 1); - - // Only the first thread in each warp stores to shared memory - if (laneId == 0) { - sm_reduction[m][warpId] = sum; - } - } - - __syncthreads(); - - // Final reduction across warps (only first thread) - if (tid == 0) { -#pragma unroll - for (int m = 0; m < kNumTokens; m++) { - float final_sum = 0.0f; - -// Sum across the kNumWarps -#pragma unroll - for (int w = 0; w < kNumWarps; w++) { - final_sum += sm_reduction[m][w]; - } - - // Write final result - out[m * kNumExperts + n_idx] = final_sum; - } - } -#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) - cudaTriggerProgrammaticLaunchCompletion(); -#endif -} - -template -void invokeRouterGemmFloatOutput(float* output, T const* mat_a, T const* mat_b, cudaStream_t stream) { - constexpr int VPT = 16 / sizeof(T); - constexpr int kBlockSize = 128; - cudaLaunchConfig_t config; - config.gridDim = kNumExperts; - config.blockDim = kBlockSize; - config.dynamicSmemBytes = 0; - config.stream = stream; - cudaLaunchAttribute attrs[1]; - attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; - attrs[0].val.programmaticStreamSerializationAllowed = getEnvEnablePDL(); - config.numAttrs = 1; - config.attrs = attrs; - cudaLaunchKernelEx( - &config, - router_gemm_kernel_float_output, - output, - mat_a, - mat_b); -} - -// Template instantiations for DEFAULT_NUM_EXPERTS experts -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 256, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -// Template instantiations for KIMI_K2_NUM_EXPERTS experts -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 1, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 2, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 3, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 4, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 5, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 6, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 7, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 8, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 9, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 10, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 11, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 12, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 13, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 14, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 15, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); - -template void invokeRouterGemmFloatOutput<__nv_bfloat16, 16, 384, 7168>( - float*, __nv_bfloat16 const*, __nv_bfloat16 const*, cudaStream_t); diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 2e63f5fecc05..69dd943ea62f 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -256,7 +256,6 @@ void sgl_per_token_group_quant_8bit_v2( bool fuse_silu_and_mul, const std::optional& masked_m); void sgl_per_token_quant_fp8(at::Tensor input, at::Tensor output_q, at::Tensor output_s); -void dsv3_fused_a_gemm(torch::Tensor& output, torch::Tensor const& mat_a, torch::Tensor const& mat_b); torch::Tensor gptq_gemm( torch::Tensor a, diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index f4874cd73660..f725f62c2ed3 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -55,7 +55,6 @@ ) from sgl_kernel.gemm import ( awq_dequantize, - dsv3_fused_a_gemm, fp8_scaled_mm, gptq_gemm, gptq_shuffle, @@ -160,8 +159,6 @@ "copy_to_gpu_no_ce", "cutlass_mla_decode", "cutlass_mla_get_workspace_size", - "dsv3_fused_a_gemm", - "dsv3_router_gemm", "dsv4_fused_k_norm_rope_flashmla", "dsv4_fused_q_indexer_rope_hadamard_quant", "dsv4_fused_q_norm_rope", diff --git a/sgl-kernel/python/sgl_kernel/gemm.py b/sgl-kernel/python/sgl_kernel/gemm.py index e6a9b50d1a67..fa896625dd8b 100644 --- a/sgl-kernel/python/sgl_kernel/gemm.py +++ b/sgl-kernel/python/sgl_kernel/gemm.py @@ -31,21 +31,6 @@ def fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None): ) -def dsv3_fused_a_gemm( - mat_a: torch.Tensor, - mat_b: torch.Tensor, - output: Optional[torch.Tensor] = None, -) -> torch.Tensor: - if output is None: - output = torch.empty( - (mat_a.shape[0], mat_b.shape[1]), - device=mat_a.device, - dtype=mat_a.dtype, - ) - torch.ops.sgl_kernel.dsv3_fused_a_gemm.default(output, mat_a, mat_b) - return output - - def sgl_per_token_group_quant_8bit( input: torch.Tensor, output_q: torch.Tensor, diff --git a/sgl-kernel/setup_musa.py b/sgl-kernel/setup_musa.py index db8cb3fd76fc..504d6be22449 100644 --- a/sgl-kernel/setup_musa.py +++ b/sgl-kernel/setup_musa.py @@ -95,10 +95,6 @@ def _get_version(): "csrc/speculative/speculative_sampling.cu", "csrc/kvcacheio/transfer.cu", "csrc/gemm/awq_kernel.cu", - "csrc/gemm/dsv3_fused_a_gemm.cu", - "csrc/gemm/dsv3_router_gemm_bf16_out.cu", - "csrc/gemm/dsv3_router_gemm_entry.cu", - "csrc/gemm/dsv3_router_gemm_float_out.cu", "csrc/gemm/per_token_quant_fp8.cu", "csrc/gemm/per_token_group_quant_8bit.cu", "csrc/gemm/per_token_group_quant_8bit_v2.cu", diff --git a/sgl-kernel/tests/test_dsv3_fused_a_gemm.py b/sgl-kernel/tests/test_dsv3_fused_a_gemm.py deleted file mode 100644 index 66110383736e..000000000000 --- a/sgl-kernel/tests/test_dsv3_fused_a_gemm.py +++ /dev/null @@ -1,34 +0,0 @@ -import sys - -import pytest -import torch -import torch.nn.functional as F -from sgl_kernel import dsv3_fused_a_gemm - - -@pytest.mark.parametrize("num_tokens", [1, 8, 15, 16]) -def test_dsv3_fused_a_gemm(num_tokens): - kHdIn = 7168 - kHdOut = 2112 - - mat_a = torch.randn( - (num_tokens, kHdIn), dtype=torch.bfloat16, device="cuda" - ).contiguous() - mat_b = torch.randn((kHdOut, kHdIn), dtype=torch.bfloat16, device="cuda").transpose( - 0, 1 - ) - output = torch.empty( - (num_tokens, kHdOut), dtype=torch.bfloat16, device="cuda" - ).contiguous() - - ref = F.linear(mat_a, mat_b.T) - - output = dsv3_fused_a_gemm(mat_a, mat_b) - - assert torch.allclose( - output, ref, rtol=1e-2, atol=1e-3 - ), "Fused GEMM output mismatch with torch.nn.functional.linear reference" - - -if __name__ == "__main__": - sys.exit(pytest.main([__file__])) diff --git a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py b/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py index ef084d79d3a5..9fa7462ec449 100644 --- a/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py +++ b/test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py @@ -1,5 +1,4 @@ -"""Benchmark for DeepSeek V3 fused QKV-A GEMM: CuTe DSL vs CUDA JIT vs -sgl_kernel AOT vs torch. +"""Benchmark for DeepSeek V3 fused QKV-A GEMM: CuTe DSL vs CUDA JIT vs torch. Run on SM90+ (Hopper or later): python test/registered/jit/benchmark/bench_dsv3_fused_a_gemm.py @@ -8,7 +7,6 @@ import torch import torch.nn.functional as F import triton.testing -from sgl_kernel import dsv3_fused_a_gemm as sgl_kernel_dsv3_fused_a_gemm from sglang.jit_kernel.benchmark import marker from sglang.jit_kernel.cutedsl_dsv3_fused_a_gemm import ( @@ -16,7 +14,6 @@ ) from sglang.jit_kernel.dsv3_fused_a_gemm import dsv3_fused_a_gemm from sglang.jit_kernel.utils import get_jit_cuda_arch, is_hip_runtime -from sglang.srt.utils.common import is_sm120_supported from sglang.test.ci.ci_register import register_cuda_ci from sglang.utils import is_in_ci @@ -31,14 +28,11 @@ HD_OUT = 2112 HD_IN_LIST = [6144, 7168] -AOT_HD_IN = 7168 -HAS_AOT = not is_sm120_supported() - NUM_TOKENS_LIST = [1, 8, 16] if IS_CI else list(range(1, 17)) -LINE_VALS = ["cutedsl", "jit", "sgl_kernel", "torch"] -LINE_NAMES = ["CuTe DSL", "CUDA JIT", "sgl_kernel AOT", "torch F.linear"] -STYLES = [("blue", "-"), ("orange", "--"), ("red", ":"), ("green", "-.")] +LINE_VALS = ["cutedsl", "jit", "torch"] +LINE_NAMES = ["CuTe DSL", "CUDA JIT", "torch F.linear"] +STYLES = [("blue", "-"), ("orange", "--"), ("green", "-.")] def _median_us(fn, *args) -> float: @@ -53,15 +47,11 @@ def _median_us(fn, *args) -> float: def _bench(num_tokens, provider, hd_in): - if provider == "sgl_kernel" and not (HAS_AOT and hd_in == AOT_HD_IN): - return float("nan") - mat_a = torch.randn((num_tokens, hd_in), dtype=DTYPE, device=DEVICE) mat_b = torch.randn((HD_OUT, hd_in), dtype=DTYPE, device=DEVICE).transpose(0, 1) fn_map = { "cutedsl": cutedsl_dsv3_fused_a_gemm, "jit": dsv3_fused_a_gemm, - "sgl_kernel": sgl_kernel_dsv3_fused_a_gemm, "torch": lambda a, b: F.linear(a, b.T), } return _median_us(fn_map[provider], mat_a, mat_b) diff --git a/test/registered/jit/benchmark/bench_dsv3_router_gemm.py b/test/registered/jit/benchmark/bench_dsv3_router_gemm.py index 91018638c5f4..dd3d0adf06e7 100644 --- a/test/registered/jit/benchmark/bench_dsv3_router_gemm.py +++ b/test/registered/jit/benchmark/bench_dsv3_router_gemm.py @@ -1,4 +1,4 @@ -"""Benchmark for DeepSeek V3 router GEMM (JIT kernel vs sgl_kernel AOT vs torch). +"""Benchmark for DeepSeek V3 router GEMM (JIT kernel vs torch). Run on a Hopper (SM90+) GPU: python -m sglang.jit_kernel.benchmark.bench_dsv3_router_gemm @@ -7,11 +7,6 @@ import torch import torch.nn.functional as F -try: - from sgl_kernel import dsv3_router_gemm as sgl_kernel_dsv3_router_gemm -except ImportError: - sgl_kernel_dsv3_router_gemm = None - from sglang.jit_kernel.benchmark import marker from sglang.jit_kernel.benchmark.utils import create_random from sglang.jit_kernel.dsv3_router_gemm import dsv3_router_gemm @@ -23,9 +18,6 @@ ) register_amd_ci(est_time=5, stage="jit-kernel-benchmark", runner_config="amd") -# sgl_kernel AOT kernel is specialized for hidden_dim=7168 only. -SGL_KERNEL_HIDDEN_DIM = 7168 - def _torch(mat_a, mat_b, out_dtype): return F.linear(mat_a, mat_b).to(out_dtype) @@ -33,7 +25,6 @@ def _torch(mat_a, mat_b, out_dtype): FN_MAP = { "jit": dsv3_router_gemm, - "sgl_kernel": sgl_kernel_dsv3_router_gemm, "torch": _torch, } @@ -42,13 +33,8 @@ def _torch(mat_a, mat_b, out_dtype): @marker.parametrize("hidden_dim", [6144, 7168], [7168]) @marker.parametrize("num_tokens", list(range(1, 17)), [1, 8, 16]) @marker.parametrize("out_dtype", [torch.bfloat16, torch.float32]) -@marker.benchmark("provider", ["jit", "sgl_kernel", "torch"]) +@marker.benchmark("provider", ["jit", "torch"]) def benchmark(num_experts, hidden_dim, num_tokens, out_dtype, provider): - if provider == "sgl_kernel": - if sgl_kernel_dsv3_router_gemm is None: - marker.skip("sgl_kernel dsv3_router_gemm not available in this build") - if hidden_dim != SGL_KERNEL_HIDDEN_DIM: - marker.skip("sgl_kernel AOT only supports hidden_dim=7168") mat_a = create_random(num_tokens, hidden_dim) mat_b = create_random(num_experts, hidden_dim) return marker.do_bench(