Optimize and Extend FusedQKNormRope to support FusedQKNorm and more general inputs

include/sgl_kernel_ops.h

@@ -165,6 +165,8 @@ void fused_qk_norm_rope(
     double high,
     double attention_factor,
     int64_t rotary_dim);
+void fused_inplace_qknorm(
+    torch::Tensor& q, torch::Tensor& k, torch::Tensor& q_weight, torch::Tensor& k_weight, double eps);
 void fused_qk_rope(
     torch::Tensor& qkv,
     int64_t num_heads_q,

python/sgl_kernel/__init__.py

@@ -22,6 +22,7 @@
 from sgl_kernel.elementwise import (
     apply_rope_with_cos_sin_cache_inplace,
     fused_add_rmsnorm,
+    fused_inplace_qknorm,
     fused_qk_norm_rope,
     fused_qk_rope,
     fused_qk_rope_with_cos_sin_cache_inplace,

python/sgl_kernel/elementwise.py

@@ -140,6 +140,22 @@ def gemma_fused_add_rmsnorm(
     torch.ops.sgl_kernel.gemma_fused_add_rmsnorm(input, residual, weight, eps)
 
 
+def fused_inplace_qknorm(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    eps: float = 1e-6,
+) -> None:
+    r"""Fused in-place Q/K RMS normalization.
+
+    ``q`` and ``k`` are expected to be 3D tensors with shape
+    ``(num_tokens, num_heads, head_dim)`` and contiguous last dimension. Only
+    the Q and K tensors are modified in-place.
+    """
+    torch.ops.sgl_kernel.fused_inplace_qknorm(q, k, q_weight, k_weight, eps)
+
+
 def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None:
     assert input.ndim == output.ndim, f"{input.ndim} != {output.ndim}"
     assert (

src/sycl/MemoryAccess.h

@@ -84,7 +84,7 @@ struct alignas(sizeof(scalar_t) * vec_size) aligned_vector_loop {
 };
 
 template <typename scalar_t>
-inline int can_vectorize_up_to(at::DeviceIndex dev_id, char* pointer) {
+inline int can_vectorize_up_to(at::DeviceIndex dev_id, const char* pointer) {
   int elem_size = sizeof(scalar_t);
   int preferred_width = preferred_vector_width(dev_id, elem_size);
   uint64_t address = reinterpret_cast<uint64_t>(pointer);
@@ -108,8 +108,8 @@ template <typename... Args>
 int get_min_vec_size(int vec_size, Args*... args) {
   auto limit_func = [](int vec_size, auto* data) {
     if (!data) return vec_size;
-    return can_vectorize_up_to<std::remove_pointer_t<decltype(data)>>(
-        dpcppGetDeviceIdOfCurrentQueue(), reinterpret_cast<char*>(data));
+    using data_t = std::remove_cv_t<std::remove_pointer_t<decltype(data)>>;
+    return can_vectorize_up_to<data_t>(dpcppGetDeviceIdOfCurrentQueue(), reinterpret_cast<const char*>(data));
   };
   return get_min(limit_func, vec_size, args...);
 }