diff --git a/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py b/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py new file mode 100644 index 000000000..64d6bd921 --- /dev/null +++ b/python/sglang/jit_kernel/benchmark/bench_qknorm_across_heads.py @@ -0,0 +1,121 @@ +import itertools +from typing import Tuple + +import torch +import triton +import triton.testing +from sgl_kernel import rmsnorm + +from sglang.jit_kernel.benchmark.utils import is_in_ci +from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads +from sglang.srt.utils import get_current_device_stream_fast + +IS_CI = is_in_ci() + +alt_stream = torch.cuda.Stream() + + +def sglang_jit_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, +) -> None: + + fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight) + + +def sglang_aot_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, +) -> None: + + current_stream = get_current_device_stream_fast() + alt_stream.wait_stream(current_stream) + rmsnorm(q, q_weight, out=q) + with torch.cuda.stream(alt_stream): + rmsnorm(k, k_weight, out=k) + current_stream.wait_stream(alt_stream) + + +def flashinfer_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, +) -> None: + from flashinfer import rmsnorm + + rmsnorm(q, q_weight, out=q) + rmsnorm(k, k_weight, out=k) + + +@torch.compile() +def torch_impl_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float = 1e-6, +) -> None: + q_mean = q.float().pow(2).mean(dim=-1, keepdim=True) + k_mean = k.float().pow(2).mean(dim=-1, keepdim=True) + q_norm = (q_mean + eps).rsqrt() + k_norm = (k_mean + eps).rsqrt() + q.copy_(q.float() * q_norm * q_weight.float()) + k.copy_(k.float() * k_norm * k_weight.float()) + + +DTYPE = torch.bfloat16 +DEVICE = "cuda" + +if IS_CI: + BS_RANGE = [16] + HIDDEN_DIM_RANGE = [1024] +else: + BS_RANGE = [2**n for n in range(0, 14)] + HIDDEN_DIM_RANGE = [512, 1024, 2048, 4096, 8192] + +LINE_VALS = ["jit", "aot", "fi", "torch"] +LINE_NAMES = ["SGL JIT Kernel", "SGL AOT Kernel", "FlashInfer", "PyTorch"] +STYLES = [("blue", "-"), ("orange", "--"), ("green", "-."), ("red", ":")] + +configs = list(itertools.product(BS_RANGE, HIDDEN_DIM_RANGE)) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "hidden_dim"], + x_vals=configs, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="qknorm-across-heads-performance", + args={}, + ) +) +def benchmark( + batch_size: int, hidden_dim: int, provider: str +) -> Tuple[float, float, float]: + q = torch.randn((batch_size, hidden_dim), dtype=DTYPE, device=DEVICE) + k = torch.randn((batch_size, hidden_dim), dtype=DTYPE, device=DEVICE) + q_weight = torch.randn(hidden_dim, dtype=DTYPE, device=DEVICE) + k_weight = torch.randn(hidden_dim, dtype=DTYPE, device=DEVICE) + FN_MAP = { + "jit": sglang_jit_qknorm_across_heads, + "aot": sglang_aot_qknorm_across_heads, + "fi": flashinfer_qknorm_across_heads, + "torch": torch_impl_qknorm_across_heads, + } + fn = lambda: FN_MAP[provider](q, k, q_weight, k_weight) + 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/qknorm_across_heads.cuh b/python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh new file mode 100644 index 000000000..1c231390b --- /dev/null +++ b/python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh @@ -0,0 +1,232 @@ +#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 qknorm_across_heads_reg_kernel( + T* __restrict__ q, + T* __restrict__ k, + const T* __restrict__ q_weight, + const T* __restrict__ k_weight, + int vec_hidden_size, + float eps) { + constexpr int inner_loop = VEC_SIZE_IN_BYTE == 16 ? 4 : 8; + + __shared__ float shared_memory[64]; // Used for CTA reduce, store both Q and K rsqrt + + using vec_t = typename VecTypeTrait::vec_t; + using packed_t = typename VecTypeTrait::packed_t; + vec_t v_q; // Save q + vec_t v_k; // Save k + vec_t v_q_weight; // Save q_weight + vec_t v_k_weight; // Save k_weight + vec_t v_q_out; // Save q output + vec_t v_k_out; // Save k output + + auto token_id = blockIdx.x; + float2 acc_square_q = make_float2(0.0f, 0.0f); // Sum of squares for q + float2 acc_square_k = make_float2(0.0f, 0.0f); // Sum of squares for k + + if (threadIdx.x < vec_hidden_size) { + // Compute address for q and k + vec_t* p_q = reinterpret_cast(q) + token_id * vec_hidden_size; + vec_t* p_k = reinterpret_cast(k) + token_id * vec_hidden_size; + const vec_t* p_q_weight = reinterpret_cast(q_weight); + const vec_t* p_k_weight = reinterpret_cast(k_weight); + + // Load data + v_q = p_q[threadIdx.x]; + v_k = p_k[threadIdx.x]; + v_q_weight = p_q_weight[threadIdx.x]; + v_k_weight = p_k_weight[threadIdx.x]; + + // Compute sum of squares for q + for (int i = 0; i < inner_loop; i++) { + float2 val = device::cast(v_q[i]); + acc_square_q.x += val.x * val.x; + acc_square_q.y += val.y * val.y; + } + + // Compute sum of squares for k + for (int i = 0; i < inner_loop; i++) { + float2 val = device::cast(v_k[i]); + acc_square_k.x += val.x * val.x; + acc_square_k.y += val.y * val.y; + } + } + + auto cg_warp = cooperative_groups::tiled_partition<32>(cooperative_groups::this_thread_block()); + float* buffer_q = shared_memory; // [0, 31] for Q + float* buffer_k = shared_memory + 32; // [32, 63] for K + + // ========== Reduction phase: Compute rsqrt for both Q and K ========== + + // Step 0: Warp Reduce for Q + float warp_sum_q = + cooperative_groups::reduce(cg_warp, acc_square_q.x + acc_square_q.y, cooperative_groups::plus()); + if (threadIdx.x % 32 == 0) { + buffer_q[threadIdx.x / 32] = warp_sum_q; + } + + // Step 0: Warp Reduce for K + float warp_sum_k = + cooperative_groups::reduce(cg_warp, acc_square_k.x + acc_square_k.y, cooperative_groups::plus()); + if (threadIdx.x % 32 == 0) { + buffer_k[threadIdx.x / 32] = warp_sum_k; + } + + // Step 1: CTA Reduce for both Q and K + __syncthreads(); + if (threadIdx.x < 32) { + // CTA Reduce for Q + float cta_sum_q = cooperative_groups::reduce( + cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer_q[threadIdx.x] : 0.0f, cooperative_groups::plus()); + buffer_q[threadIdx.x] = + rsqrtf(eps + cta_sum_q * (1.0f / static_cast(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T))))); + + // CTA Reduce for K + float cta_sum_k = cooperative_groups::reduce( + cg_warp, (threadIdx.x < blockDim.x / 32) ? buffer_k[threadIdx.x] : 0.0f, cooperative_groups::plus()); + buffer_k[threadIdx.x] = + rsqrtf(eps + cta_sum_k * (1.0f / static_cast(vec_hidden_size * (VEC_SIZE_IN_BYTE / sizeof(T))))); + } + __syncthreads(); + + // ========== Apply normalization phase: Compute and write back Q and K ========== + + if (threadIdx.x < vec_hidden_size) { + // Apply RMSNorm for Q + float rsqrt_q = buffer_q[threadIdx.x / 32]; + for (int i = 0; i < inner_loop; i++) { + v_q_out[i] = rms(v_q[i], v_q_weight[i], rsqrt_q); + } + vec_t* p_q_out = reinterpret_cast(q) + token_id * vec_hidden_size; + p_q_out[threadIdx.x] = v_q_out; + + // Apply RMSNorm for K + float rsqrt_k = buffer_k[threadIdx.x / 32]; + for (int i = 0; i < inner_loop; i++) { + v_k_out[i] = rms(v_k[i], v_k_weight[i], rsqrt_k); + } + vec_t* p_k_out = reinterpret_cast(k) + token_id * vec_hidden_size; + p_k_out[threadIdx.x] = v_k_out; + } +} + +template +struct QKNormAcrossHeadsKernel { + static void + run(const tvm::ffi::TensorView q, + const tvm::ffi::TensorView k, + const tvm::ffi::TensorView q_weight, + const tvm::ffi::TensorView k_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}) // q + .with_strides({D, 1}) + .with_dtype() + .with_device(device) + .verify(q); + TensorMatcher({N, D}) // k + .with_strides({D, 1}) + .with_dtype() + .with_device(device) + .verify(k); + TensorMatcher({D}) // q_weight + .with_dtype() + .with_device(device) + .verify(q_weight); + TensorMatcher({D}) // k_weight + .with_dtype() + .with_device(device) + .verify(k_weight); + + 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 single kernel for both q and k + auto kernel = max_vec_size_byte == 32 ? qknorm_across_heads_reg_kernel + : qknorm_across_heads_reg_kernel; + + LaunchKernel(static_cast(N.unwrap()), threads, device.unwrap()) + .enable_pdl(false)( + kernel, + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + reinterpret_cast(q_weight.data_ptr()), + reinterpret_cast(k_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/norm.py b/python/sglang/jit_kernel/norm.py index ef3f93681..e3b2aee1b 100644 --- a/python/sglang/jit_kernel/norm.py +++ b/python/sglang/jit_kernel/norm.py @@ -49,6 +49,19 @@ def _jit_fused_add_rmsnorm_module(dtype: torch.dtype) -> Module: ) +@cache_once +def _jit_qknorm_across_heads_module(dtype: torch.dtype) -> Module: + args = make_cpp_args(dtype) + return load_jit( + "qknorm_across_heads", + *args, + cuda_files=["elementwise/qknorm_across_heads.cuh"], + cuda_wrappers=[ + ("qknorm_across_heads", f"QKNormAcrossHeadsKernel<{args}>::run") + ], + ) + + @cache_once def can_use_fused_inplace_qknorm(head_dim: int, dtype: torch.dtype) -> bool: logger = logging.getLogger(__name__) @@ -97,3 +110,24 @@ def fused_add_rmsnorm( ) -> None: module = _jit_fused_add_rmsnorm_module(input.dtype) module.fused_add_rmsnorm(input, residual, weight, eps) + + +def fused_inplace_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float = 1e-6, +) -> None: + """ + Fused inplace QK normalization across all heads. + + Args: + q: Query tensor of shape [batch_size, num_heads * head_dim] + k: Key tensor of shape [batch_size, num_heads * head_dim] + q_weight: Query weight tensor of shape [num_heads * head_dim] + k_weight: Key weight tensor of shape [num_heads * head_dim] + eps: Epsilon for numerical stability + """ + module = _jit_qknorm_across_heads_module(q.dtype) + module.qknorm_across_heads(q, k, q_weight, k_weight, eps) diff --git a/python/sglang/jit_kernel/tests/test_qknorm_across_heads.py b/python/sglang/jit_kernel/tests/test_qknorm_across_heads.py new file mode 100644 index 000000000..d00c713de --- /dev/null +++ b/python/sglang/jit_kernel/tests/test_qknorm_across_heads.py @@ -0,0 +1,75 @@ +import itertools + +import pytest +import torch +import triton + + +def sglang_jit_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, +) -> None: + from sglang.jit_kernel.norm import fused_inplace_qknorm_across_heads + + fused_inplace_qknorm_across_heads(q, k, q_weight, k_weight) + + +def sglang_aot_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, +) -> None: + from sgl_kernel import rmsnorm + + rmsnorm(q, q_weight, out=q) + rmsnorm(k, k_weight, out=k) + + +@torch.compile() +def torch_impl_qknorm_across_heads( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + eps: float = 1e-6, +) -> None: + q_mean = q.float().pow(2).mean(dim=-1, keepdim=True) + k_mean = k.float().pow(2).mean(dim=-1, keepdim=True) + q_norm = (q_mean + eps).rsqrt() + k_norm = (k_mean + eps).rsqrt() + q.copy_(q.float() * q_norm * q_weight.float()) + k.copy_(k.float() * k_norm * k_weight.float()) + + +BS_LIST = [2**n for n in range(0, 14)] +BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] +HIDDEN_DIM_LIST = [512, 1024, 2048, 4096] +DEVICE = "cuda" +DTYPE = torch.bfloat16 + + +@pytest.mark.parametrize( + "batch_size,hidden_dim", + list(itertools.product(BS_LIST, HIDDEN_DIM_LIST)), +) +def test_qknorm_across_heads(batch_size: int, hidden_dim: int) -> None: + q = torch.randn(batch_size, hidden_dim, device=DEVICE, dtype=DTYPE) + k = torch.randn(batch_size, hidden_dim, device=DEVICE, dtype=DTYPE) + q_weight = torch.randn(hidden_dim, device=DEVICE, dtype=DTYPE) + k_weight = torch.randn(hidden_dim, device=DEVICE, dtype=DTYPE) + + q_k_jit = (q.clone(), k.clone()) + q_k_aot = (q.clone(), k.clone()) + + sglang_jit_qknorm_across_heads(q_k_jit[0], q_k_jit[1], q_weight, k_weight) + sglang_aot_qknorm_across_heads(q_k_aot[0], q_k_aot[1], q_weight, k_weight) + + triton.testing.assert_close(q_k_jit[0], q_k_aot[0], atol=1e-2, rtol=1e-2) + triton.testing.assert_close(q_k_jit[1], q_k_aot[1], atol=1e-2, rtol=1e-2) + + +if __name__ == "__main__": + pytest.main([__file__])