diff --git a/python/sglang/jit_kernel/benchmark/bench_qknorm.py b/python/sglang/jit_kernel/benchmark/bench_qknorm.py index 70f61a4ca..76e217f09 100644 --- a/python/sglang/jit_kernel/benchmark/bench_qknorm.py +++ b/python/sglang/jit_kernel/benchmark/bench_qknorm.py @@ -74,8 +74,6 @@ def torch_impl_qknorm( k.copy_(k.float() * k_norm * k_weight.float()) -HEAD_DIM = 128 - BS_RANGE = get_benchmark_range( full_range=[2**n for n in range(0, 14)], ci_range=[16], @@ -88,17 +86,21 @@ KV_HEAD_RANGE = get_benchmark_range( full_range=[1, 2, 4, 8], ci_range=[1], ) +HEAD_DIM_RANGE = get_benchmark_range( + full_range=[128, 256, 512, 1024], + ci_range=[128], +) LINE_VALS = ["aot", "jit", "fi", "torch"] LINE_NAMES = ["SGL AOT Kernel", "SGL JIT Kernel", "FlashInfer", "PyTorch"] STYLES = [("orange", "-"), ("blue", "--"), ("green", "-."), ("red", ":")] -configs = list(itertools.product(GQA_RANGE, KV_HEAD_RANGE, BS_RANGE)) +configs = list(itertools.product(HEAD_DIM_RANGE, GQA_RANGE, KV_HEAD_RANGE, BS_RANGE)) @triton.testing.perf_report( triton.testing.Benchmark( - x_names=["GQA", "num_kv_heads", "batch_size"], + x_names=["head_dim", "GQA", "num_kv_heads", "batch_size"], x_vals=configs, line_arg="provider", line_vals=LINE_VALS, @@ -109,16 +111,18 @@ configs = list(itertools.product(GQA_RANGE, KV_HEAD_RANGE, BS_RANGE)) args={}, ) ) -def benchmark(batch_size: int, GQA: int, num_kv_heads: int, provider: str): +def benchmark( + head_dim: int, GQA: int, num_kv_heads: int, batch_size: int, provider: str +): num_qo_heads = GQA * num_kv_heads q = torch.randn( - (batch_size, num_qo_heads, HEAD_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE + (batch_size, num_qo_heads, head_dim), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) k = torch.randn( - (batch_size, num_kv_heads, HEAD_DIM), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE + (batch_size, num_kv_heads, head_dim), dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE ) - q_weight = torch.randn(HEAD_DIM, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) - k_weight = torch.randn(HEAD_DIM, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) + q_weight = torch.randn(head_dim, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) + k_weight = torch.randn(head_dim, dtype=DEFAULT_DTYPE, device=DEFAULT_DEVICE) FN_MAP = { "aot": sglang_aot_qknorm, "jit": sglang_jit_qknorm, diff --git a/python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh b/python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh index 9e6cec8d2..789e796b0 100644 --- a/python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh +++ b/python/sglang/jit_kernel/csrc/elementwise/qknorm.cuh @@ -33,8 +33,9 @@ struct QKNormParams { constexpr uint32_t kWarpsPerBlock = 4; constexpr uint32_t kThreadsPerBlock = kWarpsPerBlock * device::kWarpThreads; +// Warp-level kernel for head_dim <= 256 template -__global__ void fused_qknorm(const QKNormParams __grid_constant__ params) { +__global__ void fused_qknorm_warp(const QKNormParams __grid_constant__ params) { using namespace device; using Storage = norm::StorageType; @@ -66,11 +67,47 @@ __global__ void fused_qknorm(const QKNormParams __grid_constant__ params) { PDLTriggerSecondary(); // launch secondary kernel } +// For CTA level, used for head_dim > 256 (512,1024) +template +__global__ void fused_qknorm_cta(const QKNormParams __grid_constant__ params) { + using namespace device; + using Storage = norm::StorageType; + + constexpr auto kNumThreads = host::norm::get_cta_threads(); + constexpr auto kNumWarps = kNumThreads / kWarpThreads; + + static_assert(sizeof(Float) == 2, "Only support FP16/BF16"); + const auto& [q, k, q_stride, k_stride, num_qo_heads, num_kv_heads, eps, q_weight, k_weight, num_tokens] = params; + + const auto num_q_and_k_heads = num_qo_heads + num_kv_heads; + const auto num_works = num_q_and_k_heads * num_tokens; + const auto gmem = tile::Memory::cta(kNumThreads); + __shared__ float smem[norm::kSmemBufferSize]; + + PDLWaitPrimary(); // wait for primary kernel + + for (auto idx = blockIdx.x; idx < num_works; idx += gridDim.x) { + const int64_t token_id = idx / num_q_and_k_heads; + const int64_t head_id = idx % num_q_and_k_heads; + const auto load_q = head_id < num_qo_heads; + const auto input = load_q ? pointer::offset(q, 2 * (token_id * q_stride + head_id * kHeadDim)) + : pointer::offset(k, 2 * (token_id * k_stride + head_id * kHeadDim)); + const auto weight = load_q ? q_weight : k_weight; + const auto input_vec = gmem.load(input); + const auto weight_vec = gmem.load(weight); + const auto output_vec = norm::apply_norm_cta(input_vec, weight_vec, eps, smem, kNumWarps); + gmem.store(input, output_vec); + } + + PDLTriggerSecondary(); // launch secondary kernel +} + +// Warp-level kernel struct for head_dim <= 256 template -struct QKNormKernel { +struct QKNormKernelWarp { static_assert(std::is_same_v || std::is_same_v); - static_assert(!host::norm::should_use_cta(), "Head dim too large for QKNorm"); - static constexpr auto kernel = fused_qknorm; + static_assert(!host::norm::should_use_cta(), "Use QKNormKernelCTA for head_dim > 256"); + static constexpr auto kernel = fused_qknorm_warp; static void run(const tvm::ffi::TensorView q, @@ -138,4 +175,83 @@ struct QKNormKernel { } }; +// This goes with fused_qknorm_cta +template +struct QKNormKernelCTA { + static_assert(std::is_same_v || std::is_same_v); + static_assert(host::norm::should_use_cta(), "Use QKNormKernelWarp for head_dim <= 256"); + static constexpr auto kernel = fused_qknorm_cta; + static constexpr auto kNumThreads = host::norm::get_cta_threads(); + + 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 Q = SymbolicSize{"num_qo_heads"}; + auto K = SymbolicSize{"num_kv_heads"}; + auto D = SymbolicSize{"head_dim"}; + auto Sq = SymbolicSize{"q_stride"}; + auto Sk = SymbolicSize{"k_stride"}; + auto device = SymbolicDevice{}; + D.set_value(kHeadDim); + device.set_options(); + + TensorMatcher({N, Q, D}) // q input + .with_strides({Sq, D, 1}) + .with_dtype() + .with_device(device) + .verify(q); + TensorMatcher({N, K, D}) // k input + .with_strides({Sk, D, 1}) + .with_dtype() + .with_device(device) + .verify(k); + TensorMatcher({D}) // weight + .with_dtype() + .with_device(device) + .verify(q_weight) + .verify(k_weight); + + const auto num_tokens = static_cast(N.unwrap()); + const auto num_qo_heads = static_cast(Q.unwrap()); + const auto num_kv_heads = static_cast(K.unwrap()); + + // NOTE: we offset the k here to reduce computation cost in the kernel + const auto params = QKNormParams{ + .q = q.data_ptr(), + .k = pointer::offset(k.data_ptr(), -2 * static_cast(num_qo_heads) * kHeadDim), + .q_stride = static_cast(Sq.unwrap()), + .k_stride = static_cast(Sk.unwrap()), + .num_qo_heads = num_qo_heads, + .num_kv_heads = num_kv_heads, + .eps = eps, + .q_weight = q_weight.data_ptr(), + .k_weight = k_weight.data_ptr(), + .num_tokens = num_tokens, + }; + + static const uint32_t max_occupancy = runtime::get_blocks_per_sm(kernel, kNumThreads); + static const uint32_t kNumSM = runtime::get_sm_count(device.unwrap().device_id); + + const auto num_works = (num_qo_heads + num_kv_heads) * num_tokens; + + // we use persistent kernel, which limit the number of blocks to reduce overhead + const auto num_blocks = std::min(num_works, max_occupancy * kNumSM); + LaunchKernel(num_blocks, kNumThreads, device.unwrap()) // + .enable_pdl(kUsePDL)(kernel, params); + } +}; + +// Unified dispatch: select warp or CTA kernel based on head_dim +template +using QKNormKernel = std::conditional_t< + host::norm::should_use_cta(), + QKNormKernelCTA, + QKNormKernelWarp>; + } // namespace diff --git a/python/sglang/jit_kernel/tests/test_qknorm.py b/python/sglang/jit_kernel/tests/test_qknorm.py index 2938341d7..4dd1963d8 100644 --- a/python/sglang/jit_kernel/tests/test_qknorm.py +++ b/python/sglang/jit_kernel/tests/test_qknorm.py @@ -63,7 +63,7 @@ BS_LIST = [2**n for n in range(0, 14)] BS_LIST += [x + 1 + i for i, x in enumerate(BS_LIST)] N_K_LIST = [2, 4] N_Q_LIST = [8, 16] -HEAD_DIM_LIST = [64, 128, 256] +HEAD_DIM_LIST = [64, 128, 256, 512, 1024] DEVICE = "cuda" DTYPE = torch.bfloat16