diff --git a/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py b/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py new file mode 100644 index 000000000..e4c0a7e8e --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py @@ -0,0 +1,73 @@ +import itertools + +import torch +import triton +import triton.testing +from flashinfer import fused_add_rmsnorm as fi_fused_add_rmsnorm + +from sglang.jit_kernel.benchmark.utils import is_in_ci +from sglang.jit_kernel.norm import fused_add_rmsnorm as jit_fused_add_rmsnorm + +IS_CI = is_in_ci() + + +def sglang_jit_fused_add_rmsnorm( + input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float +) -> None: + jit_fused_add_rmsnorm(input, residual, weight, eps) + + +def flashinfer_fused_add_rmsnorm( + input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float +) -> None: + fi_fused_add_rmsnorm(input, residual, weight, eps=eps) + + +DTYPE = torch.bfloat16 +DEVICE = "cuda" + +if IS_CI: + BS_LIST = [16] + HIDDEN_SIZE_LIST = [512, 2048] +else: + BS_LIST = [2**n for n in range(0, 14)] + HIDDEN_SIZE_LIST = [1536, 3072, 4096, 5120, 8192] + +LINE_VALS = ["jit", "fi"] +LINE_NAMES = ["SGL JIT Kernel", "FlashInfer"] +STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")] + +configs = list(itertools.product(HIDDEN_SIZE_LIST, BS_LIST)) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["hidden_size", "batch_size"], + x_vals=configs, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="fused-add-rmsnorm-performance", + args={}, + ) +) +def benchmark(hidden_size: int, batch_size: int, provider: str): + input = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE) + residual = torch.randn((batch_size, hidden_size), dtype=DTYPE, device=DEVICE) + weight = torch.randn(hidden_size, dtype=DTYPE, device=DEVICE) + FN_MAP = { + "jit": sglang_jit_fused_add_rmsnorm, + "fi": flashinfer_fused_add_rmsnorm, + } + fn = lambda: FN_MAP[provider]( + input.clone(), residual.clone(), weight, torch.finfo(torch.bfloat16).eps + ) + quantiles = [0.5, 0.2, 0.8] + ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(fn, quantiles=quantiles) # type: ignore + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + +if __name__ == "__main__": + benchmark.run(print_data=True) diff --git a/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh b/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh new file mode 100644 index 000000000..5455796af --- /dev/null +++ b/python/sglang/jit_kernel/csrc/elementwise/fused_add_rmsnorm.cuh @@ -0,0 +1,186 @@ +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +namespace { + +template +struct VecTypeTrait; + +template <> +struct VecTypeTrait { + using packed_t = packed_t; + using vec_t = device::AlignedVector; +}; + +template <> +struct VecTypeTrait { + using packed_t = packed_t; + using vec_t = device::AlignedVector; +}; + +template <> +struct VecTypeTrait { + using packed_t = packed_t; + using vec_t = device::AlignedVector; +}; + +template <> +struct VecTypeTrait { + using packed_t = packed_t; + using vec_t = device::AlignedVector; +}; + +template +SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) { + float2 valf = device::cast(val); + float2 weightf = device::cast(weight); + return device::cast( + make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum)); +} + +template +__global__ void fused_add_rmsnorm_reg_kernel( + T* __restrict__ input, T* __restrict__ residual, const T* __restrict__ weight, int vec_hidden_size, float eps) { + constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8; + + __shared__ float shared_memory[32]; // Used for CTA reduce + + using vec_t = typename VecTypeTrait::vec_t; + using packed_t = typename VecTypeTrait::packed_t; + vec_t v; // Save input + vec_t v_res; // Save residual + vec_t v_weight; // Save weight + vec_t v_out; // Save output + + auto token_id = blockIdx.x; + float2 acc_square = make_float2(0.0f, 0.0f); // Sum of squares for each thread + + if (threadIdx.x < vec_hidden_size) { + // Compute address + vec_t* p = reinterpret_cast(input) + token_id * vec_hidden_size; + vec_t* p_res = reinterpret_cast(residual) + token_id * vec_hidden_size; + const vec_t* p_weight = reinterpret_cast(weight); + + // Load data + v = p[threadIdx.x]; + v_res = p_res[threadIdx.x]; + v_weight = p_weight[threadIdx.x]; + + for (int i = 0; i < inner_loop; i++) { + float2 val = device::cast(v[i]); + float2 res = device::cast(v_res[i]); + float2 inp_res = make_float2(val.x + res.x, val.y + res.y); + acc_square.x += inp_res.x * inp_res.x; + acc_square.y += inp_res.y * inp_res.y; + v[i] = device::cast(inp_res); + } + + // Store inp+res to residual + p_res[threadIdx.x] = v; + } + + // CTA Reduce + // Step 0: Warp Reduce + auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block()); + float warp_sum = cooperative_groups::reduce(cg_warp, acc_square.x + acc_square.y, cooperative_groups::plus()); + + float* buffer = shared_memory; + if (threadIdx.x % 32 == 0) { + buffer[threadIdx.x / 32] = warp_sum; // Write warp_sum to buffer + } + + // Step 1: CTA Reduce + __syncthreads(); + if (threadIdx.x < 32) { + float cta_sum = cooperative_groups::reduce( + cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer[threadIdx.x] : 0.0f, cooperative_groups::plus()); + buffer[threadIdx.x] = + rsqrtf(eps + cta_sum * (1.0f / static_cast(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T))))); + } + __syncthreads(); + + // Compute RMSNorm + if (threadIdx.x < vec_hidden_size) { + float rsqrt_square_sum = buffer[threadIdx.x / 32]; // Read rsqrt from Shared Memory(Broadcast) + for (int i = 0; i < inner_loop; i++) { + v_out[i] = rms(v[i], v_weight[i], rsqrt_square_sum); + } + vec_t* p_out = reinterpret_cast(input) + token_id * vec_hidden_size; + p_out[threadIdx.x] = v_out; + } +} + +template +struct FusedAddRMSNormKernel { + static void + run(const tvm::ffi::TensorView input, + const tvm::ffi::TensorView residual, + const tvm::ffi::TensorView weight, + float eps) { + using namespace host; + auto N = SymbolicSize{"num_tokens"}; + auto D = SymbolicSize{"hidden_size"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, D}) // input + .with_strides({D, 1}) + .with_dtype() + .with_device(device) + .verify(input); + TensorMatcher({D}) // weight + .with_dtype() + .with_device(device) + .verify(weight); + TensorMatcher({N, D}) // residual + .with_strides({D, 1}) + .with_dtype() + .with_device(device) + .verify(residual); + + auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id); + int hidden_size = static_cast(D.unwrap()); + if ((cc_major <= 9 && hidden_size <= 8192) || (cc_major >= 10 && hidden_size <= 12288)) { + int max_vec_size_byte = cc_major >= 10 ? 32 : 16; + int elements_in_vec = max_vec_size_byte / sizeof(DType); + int vec_hidden_size = hidden_size / elements_in_vec; + uint threads = (vec_hidden_size + 31) / 32 * 32; + + // Runtime check + host::RuntimeCheck( + hidden_size % elements_in_vec == 0, + "hidden_size", + hidden_size, + " can not align to elements_in_vec ", + elements_in_vec); + + // Launch kernel + auto kernel = + max_vec_size_byte == 32 ? fused_add_rmsnorm_reg_kernel : fused_add_rmsnorm_reg_kernel; + LaunchKernel(static_cast(N.unwrap()), threads, device.unwrap()) + .enable_pdl(false)( + kernel, + reinterpret_cast(input.data_ptr()), + reinterpret_cast(residual.data_ptr()), + reinterpret_cast(weight.data_ptr()), + vec_hidden_size, + eps); + } else { + host::RuntimeCheck(false, "Large hidden_sizes are not supported for now."); + } + } +}; + +} // namespace diff --git a/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh b/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh index 5510b4474..c1150d428 100644 --- a/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh +++ b/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh @@ -45,7 +45,7 @@ template struct AlignedVector { private: /// NOTE: 1. must be pow of two 2. 16 * 8 = 128 byte, which is the max vector size supported by most devices - static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 16, "CUDA only support at most 128B vector op"); + static_assert((N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= 32, "CUDA only support at most 256B vector op"); using element_t = typename details::sized_int; using storage_t = AlignedStorage; diff --git a/python/sglang/jit_kernel/norm.py b/python/sglang/jit_kernel/norm.py index 4e8082c34..ef3f93681 100644 --- a/python/sglang/jit_kernel/norm.py +++ b/python/sglang/jit_kernel/norm.py @@ -38,6 +38,17 @@ def _jit_rmsnorm_module(hidden_size: int, dtype: torch.dtype) -> Module: ) +@cache_once +def _jit_fused_add_rmsnorm_module(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + return load_jit( + "fused_add_rmsnorm", + *args, + cuda_files=["elementwise/fused_add_rmsnorm.cuh"], + cuda_wrappers=[("fused_add_rmsnorm", f"FusedAddRMSNormKernel<{args}>::run")], + ) + + @cache_once def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool: logger = logging.getLogger(__name__) @@ -76,3 +87,13 @@ def rmsnorm( hidden_size = input.size(-1) module = _jit_rmsnorm_module(hidden_size, input.dtype) module.rmsnorm(input, weight, output, eps) + + +def fused_add_rmsnorm( + input: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + eps: float = 1e-6, +) -> None: + module = _jit_fused_add_rmsnorm_module(input.dtype) + module.fused_add_rmsnorm(input, residual, weight, eps) diff --git a/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py b/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py new file mode 100644 index 000000000..a763408d0 --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py @@ -0,0 +1,55 @@ +import itertools + +import pytest +import torch + + +def sglang_jit_fused_add_rmsnorm( + input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float +) -> None: + from sglang.jit_kernel.norm import fused_add_rmsnorm + + fused_add_rmsnorm(input, residual, weight, eps) + + +def flashinfer_fused_add_rmsnorm( + input: torch.Tensor, residual: torch.Tensor, weight: torch.Tensor, eps: float +) -> None: + from flashinfer.norm import fused_add_rmsnorm + + fused_add_rmsnorm(input, residual, weight, eps=eps) + + +BS_LIST = [2**n for n in range(0, 14)] +BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] +HIDDEN_SIZE_LIST = [512, 1024, 1536, 2048, 3072, 4096, 5120, 6144, 7168, 8192] +DEVICE = "cuda" +DTYPE = torch.bfloat16 + + +@pytest.mark.parametrize( + "batch_size,hidden_size", list(itertools.product(BS_LIST, HIDDEN_SIZE_LIST)) +) +def test_fused_add_rmsnorm(batch_size: int, hidden_size: int) -> None: + input = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE) + residual = torch.randn(batch_size, hidden_size, device=DEVICE, dtype=DTYPE) + weight = torch.randn(hidden_size, device=DEVICE, dtype=DTYPE) + + input_sglang = input.clone() + residual_sglang = residual.clone() + input_flashinfer = input.clone() + residual_flashinfer = residual.clone() + sglang_jit_fused_add_rmsnorm( + input_sglang, residual_sglang, weight, torch.finfo(torch.bfloat16).eps + ) + flashinfer_fused_add_rmsnorm( + input_flashinfer, residual_flashinfer, weight, torch.finfo(torch.bfloat16).eps + ) + torch.testing.assert_close(input_sglang, input_flashinfer, atol=1e-2, rtol=1e-2) + torch.testing.assert_close( + residual_sglang, residual_flashinfer, atol=1e-2, rtol=1e-2 + ) + + +if __name__ == "__main__": + pytest.main([__file__])