[Diffsuion & JIT_kernel] QKNorm cross heads kernel (#18073)

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)

python/sglang/jit_kernel/csrc/elementwise/qknorm_across_heads.cuh

@@ -0,0 +1,232 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/runtime.cuh>
+#include <sgl_kernel/tile.cuh>
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <cooperative_groups/reduce.h>
+#include <tvm/ffi/container/tensor.h>
+
+#include <cooperative_groups.h>
+#include <type_traits>
+
+namespace {
+
+template <typename T, int VEC_SIZE_IN_BYTE>
+struct VecTypeTrait;
+
+template <>
+struct VecTypeTrait<bf16_t, 16> {
+  using packed_t = packed_t<bf16_t>;
+  using vec_t = device::AlignedVector<packed_t, 4>;
+};
+
+template <>
+struct VecTypeTrait<fp16_t, 16> {
+  using packed_t = packed_t<fp16_t>;
+  using vec_t = device::AlignedVector<packed_t, 4>;
+};
+
+template <>
+struct VecTypeTrait<bf16_t, 32> {
+  using packed_t = packed_t<bf16_t>;
+  using vec_t = device::AlignedVector<packed_t, 8>;
+};
+
+template <>
+struct VecTypeTrait<fp16_t, 32> {
+  using packed_t = packed_t<fp16_t>;
+  using vec_t = device::AlignedVector<packed_t, 8>;
+};
+
+template <typename packed_t>
+SGL_DEVICE packed_t rms(packed_t& val, packed_t& weight, float rsqrt_square_sum) {
+  float2 valf = device::cast<fp32x2_t, packed_t>(val);
+  float2 weightf = device::cast<fp32x2_t, packed_t>(weight);
+  return device::cast<packed_t, fp32x2_t>(
+      make_float2(valf.x * weightf.x * rsqrt_square_sum, valf.y * weightf.y * rsqrt_square_sum));
+}
+
+template <typename T, int VEC_SIZE_IN_BYTE>
+__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<T, VEC_SIZE_IN_BYTE>::vec_t;
+  using packed_t = typename VecTypeTrait<T, VEC_SIZE_IN_BYTE>::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<vec_t*>(q) + token_id * vec_hidden_size;
+    vec_t* p_k = reinterpret_cast<vec_t*>(k) + token_id * vec_hidden_size;
+    const vec_t* p_q_weight = reinterpret_cast<const vec_t*>(q_weight);
+    const vec_t* p_k_weight = reinterpret_cast<const vec_t*>(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<fp32x2_t, packed_t>(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<fp32x2_t, packed_t>(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<float>());
+  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<float>());
+  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<float>());
+    buffer_q[threadIdx.x] =
+        rsqrtf(eps + cta_sum_q * (1.0f / static_cast<float>(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<float>());
+    buffer_k[threadIdx.x] =
+        rsqrtf(eps + cta_sum_k * (1.0f / static_cast<float>(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<vec_t*>(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<vec_t*>(k) + token_id * vec_hidden_size;
+    p_k_out[threadIdx.x] = v_k_out;
+  }
+}
+
+template <typename DType>
+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<kDLCUDA>();
+
+    TensorMatcher({N, D})  // q
+        .with_strides({D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q);
+    TensorMatcher({N, D})  // k
+        .with_strides({D, 1})
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k);
+    TensorMatcher({D})  // q_weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(q_weight);
+    TensorMatcher({D})  // k_weight
+        .with_dtype<DType>()
+        .with_device(device)
+        .verify(k_weight);
+
+    auto cc_major = host::runtime::get_cc_major(device.unwrap().device_id);
+    int hidden_size = static_cast<int>(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<DType, 32>
+                                            : qknorm_across_heads_reg_kernel<DType, 16>;
+
+      LaunchKernel(static_cast<uint>(N.unwrap()), threads, device.unwrap())
+          .enable_pdl(false)(
+              kernel,
+              reinterpret_cast<DType*>(q.data_ptr()),
+              reinterpret_cast<DType*>(k.data_ptr()),
+              reinterpret_cast<DType*>(q_weight.data_ptr()),
+              reinterpret_cast<DType*>(k_weight.data_ptr()),
+              vec_hidden_size,
+              eps);
+    } else {
+      host::RuntimeCheck(false, "Large hidden_sizes are not supported for now.");
+    }
+  }
+};
+
+}  // namespace

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)

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__])