Qwen3.5 moe optimizations

vllm_musa/jit_kernel/csrc/norm.py

@@ -150,3 +150,138 @@ def _fused_add_rmsnorm_custom_fake(
     mutates_args=["input", "residual"],
     fake_impl=_fused_add_rmsnorm_custom_fake,
 )
+
+
+@cache_once
+def _qk_mrope_module():
+    import tilelang
+
+    tilelang_dir = Path(tilelang.__file__).resolve().parent
+    return load_musa_jit(
+        "vllm_musa_norm_qk_mrope",
+        ("norm/qk_mrope.mu",),
+        extra_musa_cflags=(
+            f"-I{(tilelang_dir / 'src').resolve()}",
+            f"-I{(tilelang_dir / '3rdparty' / 'mutlass' / 'include').resolve()}",
+            "-Wno-error=address-of-temporary",
+            "-fmusa-flush-denormals-to-zero",
+            "-fno-signed-zeros",
+            "-D__MUSA_ARCH_LIST__=310",
+            "-mllvm",
+            "-mtgpu-opt-level=1",
+            "-mllvm",
+            "-mtgpu-load-store-opt=1",
+            "-mllvm",
+            "-mtgpu-fold-global-ldst=1",
+            "-mllvm",
+            "-mtgpu-load-cluster-mutation=1",
+            "-mllvm",
+            "-mtgpu-store-cluster-mutation=1",
+            "-mllvm",
+            "-mtgpu-memory-sched-mutation=1",
+            "-mllvm",
+            "-mtgpu-alloc-shared-memory-from-zero=1",
+        ),
+    )
+
+
+def fused_qk_rmsnorm_mrope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    eps: float = 1e-6,
+    gemma: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """QK-RMSNorm + MRoPE in one kernel. q/k are (tokens, heads, head_dim)."""
+    q_out = torch.empty_like(q)
+    k_out = torch.empty_like(k)
+    torch.ops.vllm.musa_csrc_fused_qk_rmsnorm_mrope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        positions,
+        cos_sin_cache,
+        q_out,
+        k_out,
+        bool(is_neox),
+        int(mrope_section_t),
+        int(mrope_section_h),
+        int(mrope_section_w),
+        bool(is_interleaved),
+        float(eps),
+        bool(gemma),
+    )
+    return q_out, k_out
+
+
+def _fused_qk_rmsnorm_mrope_custom(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    eps: float,
+    gemma: bool,
+) -> None:
+    _qk_mrope_module().sgl_musa_fused_qk_rmsnorm_mrope(
+        q,
+        k,
+        q_weight,
+        k_weight,
+        positions,
+        cos_sin_cache,
+        q_out,
+        k_out,
+        bool(is_neox),
+        int(mrope_section_t),
+        int(mrope_section_h),
+        int(mrope_section_w),
+        bool(is_interleaved),
+        float(eps),
+        bool(gemma),
+    )
+
+
+def _fused_qk_rmsnorm_mrope_custom_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    q_weight: torch.Tensor,
+    k_weight: torch.Tensor,
+    positions: torch.Tensor,
+    cos_sin_cache: torch.Tensor,
+    q_out: torch.Tensor,
+    k_out: torch.Tensor,
+    is_neox: bool,
+    mrope_section_t: int,
+    mrope_section_h: int,
+    mrope_section_w: int,
+    is_interleaved: bool,
+    eps: float,
+    gemma: bool,
+) -> None:
+    return
+
+
+direct_register_custom_op(
+    op_name="musa_csrc_fused_qk_rmsnorm_mrope",
+    op_func=_fused_qk_rmsnorm_mrope_custom,
+    mutates_args=["q_out", "k_out"],
+    fake_impl=_fused_qk_rmsnorm_mrope_custom_fake,
+)

vllm_musa/jit_kernel/csrc/norm/common.mu

@@ -0,0 +1,158 @@
+template <typename T>
+struct __align__(16) Vec8Storage {
+  T elem[8];
+};
+
+struct __align__(32) Float8Storage {
+  float elem[8];
+};
+
+template <typename T>
+struct __align__(16) Vec8 {
+  union {
+    Vec8Storage<T> storage;
+    T elem[8];
+  } val;
+
+  __device__ __forceinline__ Vec8() {}
+
+  template <typename Offset>
+  static __device__ __forceinline__ Vec8 load(const T* ptr, Offset idx) {
+    return *(const Vec8*)(ptr + idx);
+  }
+
+  template <typename Offset>
+  static __device__ __forceinline__ Vec8 load_byp_slc(const T* ptr, Offset idx) {
+#if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310))
+    Vec8 dst;
+    const T* addr = ptr + idx;
+    asm volatile(
+        "LSU.LD.B128 %0, %1, _, 16, 1, 1, inner_persist=0, outer_persist=2, "
+        "chrnt=l2_l3, slc=byp, persist=0, stride_add_first=0"
+        : "=R"(dst)
+        : "R"(addr));
+    return dst;
+#else
+    return *(const Vec8*)(ptr + idx);
+#endif
+  }
+};
+
+struct __align__(32) Float8 {
+  union {
+    Float8Storage storage;
+    float elem[8];
+  } val;
+
+  __device__ __forceinline__ Float8() {}
+};
+
+__device__ __forceinline__ int mrope_24_20_20_interleaved_axis(int rot_offset) {
+  constexpr unsigned long long axis1_mask = 0x492492492492492ULL;
+  constexpr unsigned long long axis2_mask = 0x924924924924924ULL;
+  const unsigned long long bit = 1ULL << rot_offset;
+  return ((axis1_mask & bit) != 0ULL) + (((axis2_mask & bit) != 0ULL) << 1);
+}
+
+__device__ __forceinline__ int mrope_11_11_10_interleaved_axis(int rot_offset) {
+  constexpr unsigned int axis1_mask = 0x92492492U;
+  constexpr unsigned int axis2_mask = 0x24924924U;
+  const unsigned int bit = 1U << rot_offset;
+  return ((axis1_mask & bit) != 0U) + (((axis2_mask & bit) != 0U) << 1);
+}
+
+__device__ __forceinline__ float fast_rsqrt(float value) {
+#if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310))
+  const float half_value = 0.5f * value;
+  float y = __frsqrt_rn(value);
+  y = y * (1.5f - half_value * y * y);
+  return y;
+#else
+  return rsqrtf(value);
+#endif
+}
+
+__device__ __forceinline__ float block_sum(float value, float* warp_sums) {
+  const int tid = (int)threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+  const int num_warps = ((int)blockDim.x + 31) >> 5;
+
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    value += __shfl_down_sync(0xffffffff, value, offset, 32);
+  }
+  if (lane == 0) {
+    warp_sums[warp] = value;
+  }
+  __syncthreads_lm();
+
+  value = tid < num_warps ? warp_sums[lane] : 0.0f;
+  if (warp == 0) {
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_down_sync(0xffffffff, value, offset, 32);
+    }
+    if (lane == 0) {
+      warp_sums[0] = value;
+    }
+  }
+  __syncthreads_lm();
+  return warp_sums[0];
+}
+
+__device__ __forceinline__ float block_sum_8warps(float value, float* warp_sums) {
+  const int tid = (int)threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    value += __shfl_down_sync(0xffffffff, value, offset, 32);
+  }
+  if (lane == 0) {
+    warp_sums[warp] = value;
+  }
+  __syncthreads_lm();
+
+  value = lane < 8 ? warp_sums[lane] : 0.0f;
+  if (warp == 0) {
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_down_sync(0xffffffff, value, offset, 32);
+    }
+    if (lane == 0) {
+      warp_sums[0] = value;
+    }
+  }
+  __syncthreads_lm();
+  return warp_sums[0];
+}
+
+__device__ __forceinline__ float block_sum_4warps(float value, float* warp_sums) {
+  const int tid = (int)threadIdx.x;
+  const int lane = tid & 31;
+  const int warp = tid >> 5;
+
+#pragma unroll
+  for (int offset = 16; offset > 0; offset >>= 1) {
+    value += __shfl_down_sync(0xffffffff, value, offset, 32);
+  }
+  if (lane == 0) {
+    warp_sums[warp] = value;
+  }
+  __syncthreads_lm();
+
+  value = lane < 4 ? warp_sums[lane] : 0.0f;
+  if (warp == 0) {
+#pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+      value += __shfl_down_sync(0xffffffff, value, offset, 32);
+    }
+    if (lane == 0) {
+      warp_sums[0] = value;
+    }
+  }
+  __syncthreads_lm();
+  return warp_sums[0];
+}