diff --git a/sgl-kernel/csrc/common_extension_musa.cc b/sgl-kernel/csrc/common_extension_musa.cc index 4cb3a5f04..bcb3a8893 100644 --- a/sgl-kernel/csrc/common_extension_musa.cc +++ b/sgl-kernel/csrc/common_extension_musa.cc @@ -72,6 +72,12 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { m.def("concat_mla_k(Tensor! k, Tensor k_nope, Tensor k_rope) -> ()"); m.impl("concat_mla_k", torch::kMUSA, &concat_mla_k); + m.def( + "rotary_embedding(Tensor positions, Tensor! query," + " Tensor!? key, int head_size," + " Tensor cos_sin_cache, bool is_neox) -> ()"); + m.impl("rotary_embedding", torch::kMUSA, &rotary_embedding); + /* * From csrc/gemm */ @@ -272,6 +278,42 @@ TORCH_LIBRARY_EXPAND(sgl_kernel, m) { m.def("top_p_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_p_arr, float top_p_val) -> ()"); m.impl("top_p_renorm_probs", torch::kMUSA, &top_p_renorm_probs); + /* + * From csrc/musa + */ + m.def( + "musa_batched_rotary_embedding_contiguous(Tensor! positions, Tensor! query, Tensor! key, " + "int head_size, Tensor! cos_sin_cache, bool is_neox, int rot_dim, Tensor! cos_sin_cache_offsets) -> ()"); + m.impl("musa_batched_rotary_embedding_contiguous", torch::kMUSA, &batched_rotary_embedding_contiguous); + + m.def( + "musa_rotary_embedding_contiguous(Tensor! positions, Tensor! query, Tensor! key, " + "int head_size, Tensor! cos_sin_cache, bool is_neox) -> ()"); + m.impl("musa_rotary_embedding_contiguous", torch::kMUSA, &rotary_embedding_contiguous); + + m.def( + "musa_fused_moe_gemv(Tensor! A, Tensor! B, Tensor! C, Tensor? A_scale, Tensor? B_scale," + "Tensor! topk_weights, Tensor! topk_ids, bool mul_routed_weight, int topk, bool use_int4_w4a16," + "bool use_swigelu) -> ()"); + m.impl("fused_moe_gemv", torch::kMUSA, &fused_moe_gemv); + + m.def( + "musa_fused_gemv(Tensor! A, Tensor! B, Tensor! C, Tensor? A_scale, Tensor? B_scale," + "bool use_int4_w4a16, bool use_swigelu, bool use_rms_norm, Tensor? gamma," + "float eps) -> ()"); + m.impl("musa_fused_gemv", torch::kMUSA, &musa_fused_gemv); + + m.def( + "musa_fused_mul_add(Tensor! output, Tensor! self, Tensor! bias," + "float scale) -> ()"); + m.impl("musa_fused_mul_add", torch::kMUSA, &fused_mul_add); + + m.def( + "musa_top_k_top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? " + "maybe_top_k_arr, " + "float top_k_val, Tensor? maybe_top_p_arr, float top_p_val, bool deterministic, Generator? gen) -> ()"); + m.impl("musa_top_k_top_p_sampling_from_probs", torch::kMUSA, &musa_top_k_top_p_sampling_from_probs); + /* * From csrc/memory */ diff --git a/sgl-kernel/csrc/elementwise/fused_add_rms_norm_kernel.mu b/sgl-kernel/csrc/elementwise/fused_add_rms_norm_kernel.mu index c010a2868..4142a8aae 100644 --- a/sgl-kernel/csrc/elementwise/fused_add_rms_norm_kernel.mu +++ b/sgl-kernel/csrc/elementwise/fused_add_rms_norm_kernel.mu @@ -127,6 +127,7 @@ GEN_VECTYPE(float16_t, Half8, 16, 8); GEN_VECTYPE(bfloat16_t, Bhalf8, 16, 8); + GEN_VECTYPE(float, float8, 32, 8); template class Dtype; @@ -309,6 +310,7 @@ DEF_VECT(bfloat16_t, Bhalf8); DEF_VECT(float16_t, Half8); DEF_VECT(float, float4); + DEF_VECT(float, float8); enum class VarUpdateMode { WELFORD, WELFORD_ONLY_MEAN, CHAN, CHAN_ONLY_MEAN }; @@ -396,7 +398,10 @@ size_t n_idx = tx * vlen; size_t n_step = (size_t)blockDim.x * vlen; + extern __shared__ ComputeType smem[]; + using SrcVec = VecType; + using ComputeVec = VecType; ComputeType var = 0; const SrcDtype* __restrict p_src = input + m_idx * N; @@ -406,6 +411,7 @@ bool m_valid = m_idx < M; if (m_valid) { for (size_t j = n_idx; j < N; j += n_step) { + ComputeVec x_vec; SrcVec curr, res_vec, fused_vec; #if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 220)) curr = *(SrcVec *)(p_src+j); @@ -416,10 +422,13 @@ #endif #pragma unroll for (int k = 0; k < vlen; k++) { - fused_vec.val_.elem[k] = curr.val_.elem[k] + res_vec.val_.elem[k]; - var += (ComputeType)fused_vec.val_.elem[k] * (ComputeType)fused_vec.val_.elem[k]; + ComputeType x = (ComputeType)curr.val_.elem[k] + (ComputeType)res_vec.val_.elem[k]; + var += x * x; + fused_vec.val_.elem[k] = (SrcDtype)x; + x_vec.val_.elem[k] = x; } *(SrcVec*)(p_res + j) = fused_vec; + *(ComputeVec*)(smem + j) = x_vec; } } AllReduceOp all_reduce_op; @@ -430,17 +439,17 @@ bool with_weight = (weight != NULL); if (with_weight) { for (size_t j = n_idx; j < N; j += n_step) { - SrcVec fused_vec, weight_val, dst; + SrcVec weight_val, dst; + ComputeVec x_vec; + x_vec = *(ComputeVec *)(smem + j); #if ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 220)) - fused_vec = *(SrcVec *)(p_res+j); - weight_val = *(SrcVec *)(weight+j); + weight_val = *(SrcVec *)(weight + j); #elif ((defined __MUSA_ARCH__) && (__MUSA_ARCH__ == 310)) - fused_vec = SrcVec::load_byp_slc(p_res, j); weight_val = SrcVec::load_byp_slc(weight, j); #endif #pragma unroll for (int k = 0; k < vlen; k++) { - dst.val_.elem[k] = (SrcDtype)((ComputeType)fused_vec.val_.elem[k] * inv_var * + dst.val_.elem[k] = (SrcDtype)(x_vec.val_.elem[k] * inv_var * (ComputeType)weight_val.val_.elem[k]); } *(SrcVec*)(p_dst + j) = dst; @@ -458,7 +467,7 @@ dim3 grid_size{nr_blocks, 1, 1}; \ LayerNorm##_KERN##KernelVlen<_SRC_DTYPE, float, \ block_x, block_y, _VLEN> \ - <<>>( \ + <<>>( \ static_cast<_SRC_DTYPE*>(input), \ static_cast<_SRC_DTYPE*>(residual), \ static_cast<_SRC_DTYPE*>(weight), \ diff --git a/sgl-kernel/csrc/musa/common.muh b/sgl-kernel/csrc/musa/common.muh new file mode 100644 index 000000000..3f8d1d137 --- /dev/null +++ b/sgl-kernel/csrc/musa/common.muh @@ -0,0 +1,42 @@ + /* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once +#include "dtype.muh" +#include +#include +#include +#include +#include + + +/// Returns the ceiling of (a / b) +__device__ __host__ __forceinline__ constexpr int ceil_div(int a, int b) { + return (a + b - 1) / b; +} + +template __device__ __host__ inline T sigmoid(T x) { + return 1.f / (1.f + __expf(-x)); +} + +#define MUDNN_ATTR_INLINE inline __attribute__((always_inline)) + +#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ == 310 // MUSIFY_EXCL_LINE +#define __SYNCTHREADS_LM __syncthreads_lm() +#else +#define __SYNCTHREADS_LM __syncthreads() +#endif diff --git a/sgl-kernel/csrc/musa/dtype.muh b/sgl-kernel/csrc/musa/dtype.muh new file mode 100644 index 000000000..587774ed9 --- /dev/null +++ b/sgl-kernel/csrc/musa/dtype.muh @@ -0,0 +1,446 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once +#include +#include + +// #include "mudnn_config.h" +// #include "mudnn/api/impl_base.h" +#include "musa/integer_subbyte.h" + +#ifdef __MTGPU__ +#include +#endif + +#include "musa_fp16.h" +#include "musa_bf16.h" +#include "musa_fp8.h" +#include + +namespace musa { +namespace dnn { + +template struct sizeof_bits { + static constexpr size_t value = sizeof(T) * 8; +}; + + +template +struct sizeof_bits> { + static constexpr size_t value = Bits; +}; + +template +static constexpr int sizeof_bits_v = sizeof_bits::value; + +typedef __half float16_t; +#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__)) +typedef __mt_bfloat16 bfloat16_t; +#else +typedef int16_t bfloat16_t; +#endif + +#define MACRO_UNROLL _Pragma("unroll") + +#define ATTR_ALIGNED(v) __attribute__((aligned(v))) +#define SELF_VEC_DEF(BASE_TYPE, VEC_TYPE_V2, VEC_TYPE_V4) \ + struct ATTR_ALIGNED(sizeof(BASE_TYPE) * 2) VEC_TYPE_V2 { \ + __device__ VEC_TYPE_V2() {} \ + __device__ VEC_TYPE_V2(const VEC_TYPE_V2 &t) { \ + this->x = t.x; \ + this->y = t.y; \ + } \ + BASE_TYPE x, y; \ + }; \ + \ + __device__ __forceinline__ VEC_TYPE_V2 make_##VEC_TYPE_V2(BASE_TYPE x, \ + BASE_TYPE y) { \ + VEC_TYPE_V2 t; \ + t.x = x, t.y = y; \ + return t; \ + } \ + \ + struct ATTR_ALIGNED(sizeof(BASE_TYPE) * 4) VEC_TYPE_V4 { \ + __device__ VEC_TYPE_V4() {} \ + __device__ VEC_TYPE_V4(const VEC_TYPE_V4 &t) { \ + this->x = t.x; \ + this->y = t.y; \ + this->z = t.z; \ + this->w = t.w; \ + } \ + BASE_TYPE x, y, z, w; \ + }; \ + \ + __device__ __forceinline__ VEC_TYPE_V4 make_##VEC_TYPE_V4( \ + BASE_TYPE x, BASE_TYPE y, BASE_TYPE z, BASE_TYPE w) { \ + VEC_TYPE_V4 t; \ + t.x = x, t.y = y, t.z = z, t.w = w; \ + return t; \ + } + +SELF_VEC_DEF(float16_t, Half2, Half4) +#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__)) +SELF_VEC_DEF(bfloat16_t, Bhalf2, Bhalf4) +#endif + +#define GEN_VECTYPE(_CTYPE, _VECTYPE, _BYTES, _VLEN) \ + struct ATTR_ALIGNED(_BYTES) _VECTYPE { \ + __device__ _VECTYPE() {} \ + __device__ _VECTYPE(const _VECTYPE &t) { \ + MACRO_UNROLL \ + for (int i = 0; i < _VLEN; i++) { \ + this->arr[i] = t.arr[i]; \ + } \ + } \ + _CTYPE arr[_VLEN]; \ + } +GEN_VECTYPE(float16_t, Half8, 16, 8); +GEN_VECTYPE(signed char, Char8, 8, 8); +GEN_VECTYPE(uint8_t, Uint8, 8, 8); +GEN_VECTYPE(int16_t, Short8, 16, 8); +GEN_VECTYPE(int16_t, Short16, 32, 16); +GEN_VECTYPE(uint16_t, Ushort8, 16, 8); +GEN_VECTYPE(uint32_t, UInt8, 32, 8); +GEN_VECTYPE(uint32_t, UInt16, 64, 16); +GEN_VECTYPE(uint64_t, ULong8, 64, 8); +GEN_VECTYPE(uint64_t, ULong16, 128, 16); +GEN_VECTYPE(float, Float8, 32, 8); +#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__)) +GEN_VECTYPE(bfloat16_t, Bhalf8, 16, 8); +GEN_VECTYPE(bfloat16_t, Bhalf16, 32, 16); +GEN_VECTYPE(bfloat16_t, Bhalf32, 64, 32); +#endif +GEN_VECTYPE(int32_t, Int8, 32, 8); +GEN_VECTYPE(int64_t, Long8, 64, 8); +GEN_VECTYPE(int64_t, Long16, 128, 16); +GEN_VECTYPE(signed char, Char16, 16, 16); +GEN_VECTYPE(float16_t, Half16, 32, 16); +GEN_VECTYPE(float, Float16, 64, 16); +GEN_VECTYPE(int32_t, Int16, 64, 16); + +template class Dtype; + +#define INST(_type, _vec2, _vec4) \ + template <> class Dtype<_type> { \ + public: \ + using Scalar = _type; \ + using Vec2 = _vec2; \ + using Vec4 = _vec4; \ + static __device__ __forceinline__ Vec2 make_vec2(_type x, _type y) { \ + return make_##_vec2(x, y); \ + } \ + static __device__ __forceinline__ Vec4 make_vec4(_type x, _type y, \ + _type z, _type w) { \ + return make_##_vec4(x, y, z, w); \ + } \ + } + +INST(float, float2, float4); +INST(float16_t, Half2, Half4); +#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__)) +INST(bfloat16_t, Bhalf2, Bhalf4); +#endif +INST(int32_t, int2, int4); +INST(uint32_t, uint2, uint4); +INST(int8_t, char2, char4); +INST(uint8_t, uchar2, uchar4); +INST(int16_t, short2, short4); +INST(uint16_t, ushort2, ushort4); +INST(int64_t, long2, long4); +INST(uint64_t, ulong2, ulong4); +INST(double, double2, double4); + +#undef INST + +template struct DeduceVectorizedType { + using Type = T; +}; + +template <> struct DeduceVectorizedType { + using Type = int8_t; +}; + +template <> struct DeduceVectorizedType { + using Type = _Float16; +}; + +#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__)) +template <> struct DeduceVectorizedType { + using Type = _Float16; +}; +#endif + +#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ == 310 // MUSIFY_EXCL_LINE +#define LD_BYP_SLC(_BITS, _BYTES) \ + VecType dst; \ + const BaseType *addr = ptr + idx; \ + __lsu_ld_cache_hint(addr,0,2,1,1); \ + return dst; +#else +#define LD_BYP_SLC(_BITS, _BYTES) return *(VecType *)(ptr + idx); +#endif + +template struct VecType; + +#define DEF_VECT(_CTYPE, _VECTYPE) \ + template <> struct VecType<_CTYPE, sizeof(_VECTYPE) * 8> { \ + static constexpr int vec_bytes = sizeof(_VECTYPE); \ + static constexpr int bit_per_byte = 8; \ + using BaseType = _CTYPE; \ + using RobustTypePtr = __musa::robust_ptr<_CTYPE>; \ + using Ttype = _VECTYPE; \ + static constexpr int bits = vec_bytes * bit_per_byte; \ + static constexpr int vlen = bits / (sizeof(BaseType) * bit_per_byte); \ + using VectorizedType = typename DeduceVectorizedType::Type; \ + typedef VectorizedType VxTy __attribute__((vector_size(vec_bytes))); \ + template \ + static __device__ __forceinline__ VecType load(const BaseType *ptr, \ + OffsetType idx) { \ + return *(VecType *)(ptr + idx); \ + } \ + template \ + static __device__ __forceinline__ VecType \ + load_byp_slc(const BaseType *ptr, OffsetType idx) { \ + if constexpr (vec_bytes == 16) { \ + LD_BYP_SLC(128, 16); \ + } else if constexpr (vec_bytes == 8) { \ + LD_BYP_SLC(64, 8); \ + } else if constexpr (vec_bytes == 4) { \ + LD_BYP_SLC(32, 4); \ + } else if constexpr (vec_bytes == 2) { \ + LD_BYP_SLC(16, 2); \ + } else { \ + LD_BYP_SLC(8, 1); \ + } \ + } \ + template \ + static __device__ __forceinline__ VecType \ + robust_load(const RobustTypePtr ptr, OffsetType idx) { \ + return __musa::robust_load(ptr, idx); \ + } \ + \ + template \ + static __device__ __forceinline__ void \ + store(BaseType *ptr, OffsetType idx, const VecType &dst) { \ + *(VecType *)(ptr + idx) = dst; \ + } \ + template \ + static __device__ __forceinline__ void \ + robust_store(RobustTypePtr ptr, OffsetType idx, const VecType &dst) { \ + __musa::robust_store(dst, ptr, idx); \ + } \ + \ + __device__ VecType() { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + this->val_.elem[i] = 0; \ + } \ + } \ + __device__ VecType(const VecType &t) { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + this->val_.elem[i] = t.val_.elem[i]; \ + } \ + } \ + __device__ VecType &operator=(const VecType &t) { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + this->val_.elem[i] = t.val_.elem[i]; \ + } \ + return *this; \ + } \ + __device__ VecType(_CTYPE val) { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + this->val_.elem[i] = val; \ + } \ + } \ + template \ + friend __device__ VecType operator+(VecType lhs, const SrcVecType &rhs) { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + lhs.val_.elem[i] += static_cast(rhs.val_.elem[i]); \ + } \ + return lhs; \ + } \ + friend __device__ VecType operator+(VecType lhs, const _CTYPE &rhs) { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + lhs.val_.elem[i] += rhs; \ + } \ + return lhs; \ + } \ + friend __device__ VecType operator-(VecType lhs, const VecType &rhs) { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + lhs.val_.elem[i] -= rhs.val_.elem[i]; \ + } \ + return lhs; \ + } \ + friend __device__ VecType operator*(VecType lhs, const VecType &rhs) { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + lhs.val_.elem[i] *= rhs.val_.elem[i]; \ + } \ + return lhs; \ + } \ + template __device__ VecType &apply() { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + this->val_.elem[i] = Func::apply(this->val_.elem[i]); \ + } \ + return *this; \ + } \ + template \ + static __device__ VecType cvt(const SrcVecType &src) { \ + VecType dst; \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + dst.val_.elem[i] = (BaseType)(src.val_.elem[i]); \ + } \ + return dst; \ + } \ + union U { \ + __device__ U() { \ + MACRO_UNROLL \ + for (int i = 0; i < sizeof(Ttype) / sizeof(BaseType); i++) { \ + this->elem[i] = 0; \ + } \ + } \ + Ttype storage; \ + BaseType elem[sizeof(Ttype) / sizeof(BaseType)]; \ + VxTy vt_elem; \ + }; \ + U val_{}; \ + } +DEF_VECT(float16_t, float16_t); +DEF_VECT(float16_t, Half2); +DEF_VECT(float16_t, Half4); +DEF_VECT(float16_t, Half8); +DEF_VECT(float16_t, Half16); +#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__)) +DEF_VECT(bfloat16_t, bfloat16_t); +DEF_VECT(bfloat16_t, Bhalf2); +DEF_VECT(bfloat16_t, Bhalf4); +DEF_VECT(bfloat16_t, Bhalf8); +DEF_VECT(bfloat16_t, Bhalf16); +DEF_VECT(bfloat16_t, Bhalf32); +#endif +DEF_VECT(bool, char); +DEF_VECT(bool, char2); +DEF_VECT(bool, char3); +DEF_VECT(bool, char4); +DEF_VECT(bool, Char8); +DEF_VECT(bool, Char16); +DEF_VECT(int8_t, int8_t); +DEF_VECT(int8_t, char2); +DEF_VECT(int8_t, char3); +DEF_VECT(int8_t, char4); +DEF_VECT(int8_t, Char8); +DEF_VECT(int8_t, Char16); +DEF_VECT(uint8_t, uint8_t); +DEF_VECT(uint8_t, uchar2); +DEF_VECT(uint8_t, uchar3); +DEF_VECT(uint8_t, uchar4); +DEF_VECT(uint8_t, uint4); +DEF_VECT(uint8_t, Uint8); +DEF_VECT(int16_t, int16_t); +DEF_VECT(int16_t, short2); +DEF_VECT(int16_t, short3); +DEF_VECT(int16_t, short4); +DEF_VECT(int16_t, Short8); +DEF_VECT(int16_t, Short16); +DEF_VECT(uint16_t, ushort); +DEF_VECT(uint16_t, ushort2); +DEF_VECT(uint16_t, ushort3); +DEF_VECT(uint16_t, ushort4); +DEF_VECT(uint16_t, Ushort8); +DEF_VECT(int32_t, int); +DEF_VECT(int32_t, int2); +DEF_VECT(int32_t, int3); +DEF_VECT(int32_t, int4); +DEF_VECT(int32_t, Int8); +DEF_VECT(int32_t, Int16); +DEF_VECT(uint32_t, uint); +DEF_VECT(uint32_t, uint2); +DEF_VECT(uint32_t, uint3); +DEF_VECT(uint32_t, uint4); +DEF_VECT(uint32_t, UInt8); +DEF_VECT(uint32_t, UInt16); +DEF_VECT(uint64_t, uint64_t); +DEF_VECT(uint64_t, ulong2); +DEF_VECT(uint64_t, ulong3); +DEF_VECT(uint64_t, ulong4); +DEF_VECT(uint64_t, ULong8); +DEF_VECT(uint64_t, ULong16); +DEF_VECT(int64_t, int64_t); +DEF_VECT(int64_t, long2); +DEF_VECT(int64_t, long3); +DEF_VECT(int64_t, long4); +DEF_VECT(int64_t, Long8); +DEF_VECT(int64_t, Long16); +DEF_VECT(float, float); +DEF_VECT(float, float2); +DEF_VECT(float, float3); +DEF_VECT(float, float4); +DEF_VECT(float, Float8); +DEF_VECT(float, Float16); + +DEF_VECT(double, double); +DEF_VECT(double, double2); +DEF_VECT(double, double3); +DEF_VECT(double, double4); + +#undef DEF_VECT +#undef MACRO_UNROLL +template struct ComputeDType { + using Type = typename std::conditional< + (sizeof(T) >= 4), T, + typename std::conditional< + std::is_integral::value, + typename std::conditional::value, uint32_t, + int32_t>::type, + float>::type>::type; +}; +template <> struct ComputeDType { + using Type = bool; +}; + +// static inline bool check_qint8_only_scale(CTR x) { +// if (x.type == TensorImpl::Type::QINT8) { +// size_t zp_size = x.quant_desc.zero_point.size(); +// return zp_size == 0 || (zp_size == 1 && x.quant_desc.zero_point[0] == 0); +// } +// return true; +// } + +template +bool check_qint8_only_scale(T x0, Types... xn) { + bool ok = check_qint8_only_scale(x0); + return ok && check_qint8_only_scale(xn...); +} +#define CHECK_QINT8(ARGS...) \ + { \ + if (!check_qint8_only_scale(ARGS)) { \ + return fail::NOT_SUPPORTED() << "qint8 only support zero_point==0"; \ + } \ + } + + +} // namespace dnn +} // namespace musa diff --git a/sgl-kernel/csrc/musa/moe_gemv_swiglu.mu b/sgl-kernel/csrc/musa/moe_gemv_swiglu.mu new file mode 100644 index 000000000..4d408df38 --- /dev/null +++ b/sgl-kernel/csrc/musa/moe_gemv_swiglu.mu @@ -0,0 +1,846 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include +#include +#include + +#include "common.muh" +#include "dtype.muh" +#include "torch_musa/csrc/core/MUSAGuard.h" +#include "torch_musa/csrc/core/MUSAStream.h" +#include "torch_musa/csrc/aten/musa/MUSAContext.h" + +using namespace musa::dnn; + +#if defined(__MUSA_ARCH__) && __MUSA_ARCH__ == 310 +#define ThreadNumPerWarp 32 +#else +#define ThreadNumPerWarp 128 +#endif + +#define SYNC_IF_NEEDED() \ + if constexpr (BLOCK_N * BLOCK_K > ThreadNumPerWarp) { \ + __SYNCTHREADS_LM; \ + } + +#define CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, _IS_FP8, _IS_RMSNORM) \ + if (scale_k_group_tile == 128) { \ + musa_gemv_kernel<_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, false, false, _IS_FP8, 128, _IS_RMSNORM> \ + <<>>( \ + static_cast<_CDTYPE*>(C.data_ptr()), \ + static_cast<_ADTYPE*>(A.data_ptr()), \ + static_cast<_BDTYPE*>(B.data_ptr()), \ + static_cast(topk_ids_ptr), \ + static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \ + static_cast<_SCALE_DTYPE*>(a_scale_ptr), \ + static_cast<_SCALE_DTYPE*>(b_scale_ptr), \ + topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \ + static_cast(rms_gamma_ptr), static_cast(rms_sum_out_ptr), static_cast(rms_count_ptr), eps); \ + } else { \ + musa_gemv_kernel<_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, false, false, _IS_FP8, 64, _IS_RMSNORM> \ + <<>>( \ + static_cast<_CDTYPE*>(C.data_ptr()), \ + static_cast<_ADTYPE*>(A.data_ptr()), \ + static_cast<_BDTYPE*>(B.data_ptr()), \ + static_cast(topk_ids_ptr), \ + static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \ + static_cast<_SCALE_DTYPE*>(a_scale_ptr), \ + static_cast<_SCALE_DTYPE*>(b_scale_ptr), \ + topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \ + static_cast(rms_gamma_ptr), static_cast(rms_sum_out_ptr), static_cast(rms_count_ptr), eps); \ + } \ + return; + +#define RUN_SCALE_ROUTE_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _IS_FP8) \ + if (mul_routed_weight) { \ + if (use_swigelu) { \ + CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, true, _IS_FP8, false) \ + } else if(use_rms_norm) { \ + CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, false, _IS_FP8, true) \ + } else { \ + CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, false, _IS_FP8, false) \ + } \ + } else { \ + if (use_swigelu) { \ + CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, true, _IS_FP8, false) \ + } else if (use_rms_norm) { \ + CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, false, _IS_FP8, true) \ + } else { \ + CAL_MOE_GEMV_FP8(_ADTYPE, _BDTYPE, _CDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, false, _IS_FP8, false) \ + } \ + } + +#define CAL_MOE_GEMV_W4A16(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, _IS_RMS_NROM) \ + if (is_pergroup_scale) { \ + if (scale_k_group_tile == 128) { \ + musa_gemv_kernel<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, true, true, false, 128, _IS_RMS_NROM> \ + <<>>( \ + static_cast<_ADTYPE*>(C.data_ptr()), \ + static_cast<_ADTYPE*>(A.data_ptr()), \ + static_cast<_BDTYPE*>(B.data_ptr()), \ + static_cast(topk_ids_ptr), \ + static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \ + static_cast<_SCALE_DTYPE*>(a_scale_ptr), \ + static_cast<_SCALE_DTYPE*>(b_scale_ptr), \ + topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \ + static_cast(rms_gamma_ptr), static_cast(rms_sum_out_ptr), static_cast(rms_count_ptr), eps); \ + return; \ + } else { \ + musa_gemv_kernel<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, true, true, false, 64, _IS_RMS_NROM> \ + <<>>( \ + static_cast<_ADTYPE*>(C.data_ptr()), \ + static_cast<_ADTYPE*>(A.data_ptr()), \ + static_cast<_BDTYPE*>(B.data_ptr()), \ + static_cast(topk_ids_ptr), \ + static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \ + static_cast<_SCALE_DTYPE*>(a_scale_ptr), \ + static_cast<_SCALE_DTYPE*>(b_scale_ptr), \ + topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \ + static_cast(rms_gamma_ptr), static_cast(rms_sum_out_ptr), static_cast(rms_count_ptr), eps); \ + return; \ + } \ + } else { \ + musa_gemv_kernel<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, true, false, false, 1, _IS_RMS_NROM> \ + <<>>( \ + static_cast<_ADTYPE*>(C.data_ptr()), \ + static_cast<_ADTYPE*>(A.data_ptr()), \ + static_cast<_BDTYPE*>(B.data_ptr()), \ + static_cast(topk_ids_ptr), \ + static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \ + static_cast<_SCALE_DTYPE*>(a_scale_ptr), \ + static_cast<_SCALE_DTYPE*>(b_scale_ptr), \ + topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \ + static_cast(rms_gamma_ptr), static_cast(rms_sum_out_ptr), static_cast(rms_count_ptr), eps); \ + return; \ + } + +#define CAL_MOE_GEMV(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, _IS_RMS_NROM) \ + musa_gemv_kernel<_ADTYPE, _BDTYPE, _ADTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, block_n, block_k, iobit, _IS_MUL_ROUTED_WEIGHT, _IS_SWGELU, false, false, false, 1, _IS_RMS_NROM> \ + <<>>( \ + static_cast<_ADTYPE*>(C.data_ptr()), \ + static_cast<_ADTYPE*>(A.data_ptr()), \ + static_cast<_BDTYPE*>(B.data_ptr()), \ + static_cast(topk_ids_ptr), \ + static_cast<_TOPK_WEIGHT_DTYPE*>(topk_weights_ptr), \ + nullptr, \ + nullptr, \ + topk, expert_offset_stride, nr_n, hidden_size, num_experts, half_n_idx, scale_k_len, \ + static_cast(rms_gamma_ptr), static_cast(rms_sum_out_ptr), static_cast(rms_count_ptr), eps); \ + return; + +#define RUN_SCALE_ROUTE(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, _CAL_FUNC) \ + if (mul_routed_weight) { \ + if (use_swigelu) { \ + _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, true, false) \ + } else if (use_rms_norm) { \ + _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, false, true) \ + } else { \ + _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, true, false, false) \ + } \ + } else { \ + if (use_swigelu) { \ + _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, true, false) \ + } else if (use_rms_norm) { \ + _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, false, true) \ + } else { \ + _CAL_FUNC(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _SCALE_DTYPE, false, false, false) \ + } \ + } + +#define RUN_ROUNTE_WEIGHT(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, _CAL_FUNC) \ + if (!B_scale.has_value() || B_scale->scalar_type() == at::ScalarType::Float) { \ + RUN_SCALE_ROUTE(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, float, _CAL_FUNC) \ + } else if (B_scale.has_value() && B_scale->scalar_type() == at::ScalarType::BFloat16) { \ + RUN_SCALE_ROUTE(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, bfloat16_t, _CAL_FUNC) \ + } else if (B_scale.has_value() && B_scale->scalar_type() == at::ScalarType::Half) { \ + RUN_SCALE_ROUTE(_ADTYPE, _BDTYPE, _TOPK_WEIGHT_DTYPE, float16_t, _CAL_FUNC) \ + } + +#define GEN_LAUNCH_KERN_GEMV(_BLK_N, _BLK_K) \ + { \ + launch_kernel = [&]() { \ + constexpr int block_n = _BLK_N; \ + constexpr int block_k = _BLK_K; \ + TORCH_CHECK(nr_n % block_n == 0, "gemv n need align"); \ + TORCH_CHECK(hidden_size % block_k == 0, "gemv k need align"); \ + dim3 block_size{block_n * block_k, 1, 1}; \ + dim3 grid_size{(uint32_t)ceil_div(reduce_size, block_n), (uint32_t)topk, (uint32_t)bseqlen}; \ + int shmem_size = block_n * sizeof(float) * block_k; \ + if (use_int4_w4a16) { \ + if (A.scalar_type() == at::ScalarType::BFloat16) { \ + RUN_ROUNTE_WEIGHT(bfloat16_t, int8_t, float, CAL_MOE_GEMV_W4A16) \ + } else if (A.scalar_type() == at::ScalarType::Half) { \ + RUN_ROUNTE_WEIGHT(float16_t, int8_t, float, CAL_MOE_GEMV_W4A16) \ + } \ + } else if (is_fp8) { \ + if (A.dtype() == at::ScalarType::BFloat16) { \ + RUN_SCALE_ROUTE_FP8(bfloat16_t, __mt_fp8_e4m3, bfloat16_t, float, float, true) \ + } else { \ + RUN_SCALE_ROUTE_FP8(__mt_fp8_e4m3, __mt_fp8_e4m3, bfloat16_t, float, float, true) \ + } \ + } else { \ + if (A.scalar_type() == at::ScalarType::BFloat16) { \ + RUN_ROUNTE_WEIGHT(bfloat16_t, bfloat16_t, float, CAL_MOE_GEMV) \ + } else if (A.scalar_type() == at::ScalarType::Half) { \ + RUN_ROUNTE_WEIGHT(float16_t, float16_t, float, CAL_MOE_GEMV) \ + } \ + } \ + TORCH_CHECK(false, "no support on moe gemv"); \ + }; \ + } + +#define GEN_LAUNCH_KERN(_BLK_N, _BLK_K) \ + { \ + launch_kernel = [&]() { \ + constexpr int block_n = _BLK_N; \ + constexpr int block_k = _BLK_K; \ + TORCH_CHECK(nr_n % block_n == 0, "gemv n need align"); \ + TORCH_CHECK(hidden_size % block_k == 0, "gemv k need align"); \ + dim3 block_size{block_n * block_k, 1, 1}; \ + dim3 grid_size{(uint32_t)ceil_div(reduce_size, block_n), (uint32_t)topk, (uint32_t)bseqlen}; \ + int shmem_size = block_n * sizeof(float) * block_k; \ + if (use_int4_w4a16) { \ + if (A.scalar_type() == at::ScalarType::BFloat16) { \ + if (topk_weights.scalar_type() == at::ScalarType::Float) { \ + RUN_ROUNTE_WEIGHT(bfloat16_t, int8_t, float, CAL_MOE_GEMV_W4A16) \ + } else if (topk_weights.scalar_type() == at::ScalarType::BFloat16) { \ + RUN_ROUNTE_WEIGHT(bfloat16_t, int8_t, bfloat16_t, CAL_MOE_GEMV_W4A16) \ + } \ + } else if (A.scalar_type() == at::ScalarType::Half) { \ + if (topk_weights.scalar_type() == at::ScalarType::Float) { \ + RUN_ROUNTE_WEIGHT(float16_t, int8_t, float, CAL_MOE_GEMV_W4A16) \ + } else if (topk_weights.scalar_type() == at::ScalarType::Half) { \ + RUN_ROUNTE_WEIGHT(float16_t, int8_t, float16_t, CAL_MOE_GEMV_W4A16) \ + } \ + } \ + } else if (is_fp8) { \ + if (A.dtype() == at::ScalarType::BFloat16) { \ + RUN_SCALE_ROUTE_FP8(bfloat16_t, __mt_fp8_e4m3, bfloat16_t, float, float, true) \ + } else { \ + RUN_SCALE_ROUTE_FP8(__mt_fp8_e4m3, __mt_fp8_e4m3, bfloat16_t, float, float, true) \ + } \ + } else { \ + if (A.scalar_type() == at::ScalarType::BFloat16) { \ + if (topk_weights.scalar_type() == at::ScalarType::Float) { \ + RUN_ROUNTE_WEIGHT(bfloat16_t, bfloat16_t, float, CAL_MOE_GEMV) \ + } else if (topk_weights.scalar_type() == at::ScalarType::BFloat16) { \ + RUN_ROUNTE_WEIGHT(bfloat16_t, bfloat16_t, bfloat16_t, CAL_MOE_GEMV) \ + } \ + } else if (A.scalar_type() == at::ScalarType::Half) { \ + if (topk_weights.scalar_type() == at::ScalarType::Float) { \ + RUN_ROUNTE_WEIGHT(float16_t, float16_t, float, CAL_MOE_GEMV) \ + } else if (topk_weights.scalar_type() == at::ScalarType::Half) { \ + RUN_ROUNTE_WEIGHT(float16_t, float16_t, float16_t, CAL_MOE_GEMV) \ + } \ + } \ + } \ + TORCH_CHECK(false, "no support on moe gemv"); \ + }; \ + } + +template +__global__ void musa_gemv_kernel( + CType *c_ptr, + const AType *a_ptr, + const BType *b_ptr, + int *expert_idx_table, + ScoreType *score_ptr, + ScaleType *scale_a, + ScaleType *scale_b, + int topk, + int expert_offset_stride, + int n, int k, + int nr_expert, + int half_n_idx, + int scale_k_len, + bfloat16_t *gamma, + float* sum_out, + volatile int *count, + float eps) { + + constexpr int bits_of_byte = 8; + constexpr int half_blockn = BLOCK_N / 2; + constexpr int b_vec_bits = 128; + constexpr int Vlen = is_w4a16 ? b_vec_bits / 4 : b_vec_bits / (sizeof(BType) * bits_of_byte); + constexpr int w4a16_shift = is_w4a16 ? 2 : 1; + constexpr int scale_k_load_cntdown_init = ceil_div(scale_block, (BLOCK_K * Vlen)); + constexpr bool fuse_castfp8 = (is_fp8 && !std::is_same_v); + + using AVecSType = std::conditional_t, uint8_t, AType>; + using BVecSType = std::conditional_t; + using v16f32_t = float __attribute__((vector_size(64))); + using v8f32_t = float __attribute__((vector_size(32))); + using AVec = typename std::conditional_t< + is_w4a16, + v16f32_t, + typename std::conditional_t< + fuse_castfp8, + v8f32_t, + typename VecType::Ttype + > + >; + using BVec = typename VecType::Ttype; + using fp8x4_vec = unsigned char __attribute__((vector_size(4))); + + int token_idx = blockIdx.z; + int expert_idx = blockIdx.y; + int real_expert_idx = 0; + int t_n_idx = threadIdx.x / BLOCK_K; + int t_k_idx = threadIdx.x % BLOCK_K; + int n_idx = blockIdx.x * BLOCK_N + t_n_idx; + + if (expert_idx_table != nullptr) { + real_expert_idx = expert_idx_table[token_idx * topk + expert_idx]; + if (real_expert_idx < 0 || real_expert_idx >= nr_expert) { + if constexpr (is_swigelu) { + if (n_idx < half_n_idx) { + int offsets = (token_idx * topk + expert_idx) * half_n_idx + n_idx; + c_ptr[offsets] = 0; + } + } else { + int offsets = (token_idx * topk + expert_idx) * half_n_idx * 2 + n_idx; + c_ptr[offsets] = 0; + } + return; + } + } + + constexpr int thread_sum_len = is_fp8 ? Vlen / 4 : Vlen; + float cur_thread_sum[thread_sum_len]; + int scale_k_load_cntdown = scale_k_load_cntdown_init; + + extern __shared__ float shared_array[]; + + if constexpr (is_swigelu) { + if (t_n_idx < half_blockn) { + n_idx = blockIdx.x * half_blockn + t_n_idx; + } else { + n_idx = blockIdx.x * half_blockn + t_n_idx - half_blockn + half_n_idx; + } + } + + + #pragma unroll + for (int i = 0; i < thread_sum_len; i++) { + cur_thread_sum[i] = 0.0f; + } + + float scale_a_val = 1.0f; + float scale_b_val = 1.0f; + int scale_a_offset = 0; + int scale_b_offset = 0; + + if constexpr (is_w4a16) { + scale_b_offset = (real_expert_idx * n + n_idx) * scale_k_len + t_k_idx * Vlen / scale_block; + if constexpr (is_swigelu) { + scale_b_offset = (real_expert_idx * 2 * n + n_idx) * scale_k_len + t_k_idx * Vlen / scale_block; + } + scale_b_val = scale_b[scale_b_offset]; + scale_k_load_cntdown -= 1; + } else if (is_fp8) { + scale_a_offset = token_idx * scale_k_len + t_k_idx * Vlen / scale_block; + scale_b_offset = (real_expert_idx * n + n_idx) / scale_block * scale_k_len + t_k_idx * Vlen / scale_block; + if constexpr (is_swigelu) { + scale_b_offset = (real_expert_idx * 2 * n + n_idx) / scale_block * scale_k_len + t_k_idx * Vlen / scale_block; + } + if constexpr (!fuse_castfp8) { + scale_a_val = scale_a[scale_a_offset]; + } + scale_b_val = scale_b[scale_b_offset]; + scale_k_load_cntdown -= 1; + } + + const BType *b_base_ptr = b_ptr + ((size_t)real_expert_idx * expert_offset_stride + n_idx * k + t_k_idx * Vlen) / w4a16_shift; + + for (int k_idx = 0; k_idx < k; k_idx += Vlen * BLOCK_K) { + AType a_reg[Vlen]; + BType b_reg[Vlen / w4a16_shift]; + *(AVec *)(a_reg) = *(AVec *)(a_ptr + token_idx * k + t_k_idx * Vlen + k_idx); + *(BVec *)(b_reg) = *(BVec *)(b_base_ptr + k_idx / w4a16_shift); + + if constexpr (is_w4a16 && !is_fp8) { + float b_reg_float[Vlen]; + #pragma unroll + for (int i = 0; i < Vlen / 2; i++) { + if constexpr (is_per_group_scale) { + uint8_t read_u8 = b_reg[i]; + b_reg_float[i * 2 + 0] = scale_b_val * ((float)(read_u8 & 0xF) - 8.f); + b_reg_float[i * 2 + 1] = scale_b_val * ((float)(read_u8 >> 4) - 8.f); + } else { + int8_t read_s8 = b_reg[i]; + b_reg_float[i * 2 + 0] = scale_b_val * (float)((int8_t)(read_s8 << 4)); + b_reg_float[i * 2 + 1] = scale_b_val * (float)((int8_t)(read_s8 & 0xF0)); + } + } + if constexpr (is_per_group_scale) { + if (scale_k_load_cntdown == 0 && (k_idx + Vlen * BLOCK_K) < k) { + scale_b_offset += ceil_div(BLOCK_K * Vlen, scale_block); + scale_b_val = scale_b[scale_b_offset]; + scale_k_load_cntdown = scale_k_load_cntdown_init; + } + scale_k_load_cntdown -= 1; + } + #pragma unroll + for (int i = 0; i < thread_sum_len; i++) { + cur_thread_sum[i] += b_reg_float[i] * (float)a_reg[i]; + } + } else if constexpr (is_fp8) { + float scale_val = scale_a_val * scale_b_val; + if (scale_k_load_cntdown == 0 && (k_idx + Vlen * BLOCK_K) < k) { + scale_a_offset += ceil_div(BLOCK_K * Vlen, scale_block); + scale_b_offset += ceil_div(BLOCK_K * Vlen, scale_block); + if constexpr (!fuse_castfp8) { + scale_a_val = scale_a[scale_a_offset]; + } + scale_b_val = scale_b[scale_b_offset]; + scale_k_load_cntdown = scale_k_load_cntdown_init; + } + scale_k_load_cntdown -= 1; + for (int i = 0; i < thread_sum_len; i++) { + typedef _Float16 _half_v4 __attribute__((ext_vector_type(4))); + typedef _Float32 _float_v4 __attribute__((ext_vector_type(4))); + _half_v4 a_halfv4; + _half_v4 b_halfv4; + _float_v4 a_float4; + _float_v4 b_float4; + if constexpr (fuse_castfp8) { + b_halfv4 = __musa_e4m32f16_rn_bst4(reinterpret_cast(b_reg)[i]); + #pragma unroll + for (int j = 0; j < 4; j++) { + cur_thread_sum[i] += scale_val * float(a_reg[i * 4 + j]) * (b_halfv4[j]); + } + } else { + a_halfv4 = __musa_e4m32f16_rn_bst4(reinterpret_cast(a_reg)[i]); + b_halfv4 = __musa_e4m32f16_rn_bst4(reinterpret_cast(b_reg)[i]); + #pragma unroll + for (int j = 0; j < 4; j++) { + cur_thread_sum[i] += scale_val * (a_halfv4[j]) * (b_halfv4[j]); + } + } + } + } else { + #pragma unroll + for (int i = 0; i < thread_sum_len; i++) { + cur_thread_sum[i] += (float)b_reg[i] * (float)a_reg[i]; + } + } + } + + float rst = 0; + #pragma unroll + for (int i = 0; i < thread_sum_len; i++) { + rst += cur_thread_sum[i]; + } + + if constexpr (is_w4a16 && !is_per_group_scale) { + rst = rst / 16.f; + } + + if constexpr (BLOCK_K > 1) { + shared_array[threadIdx.x] = rst; + SYNC_IF_NEEDED() + if (threadIdx.x < BLOCK_N) { + rst = 0; + #pragma unroll + for (int i = 0; i < BLOCK_K; i++) { + rst += shared_array[threadIdx.x * BLOCK_K + i]; + } + } + if constexpr (is_swigelu) { + SYNC_IF_NEEDED() + } + } + + if constexpr (BLOCK_N > ThreadNumPerWarp) { + return; + } + + if (threadIdx.x < BLOCK_N) { + int dst_n_idx = blockIdx.x * BLOCK_N + threadIdx.x; + if constexpr (is_swigelu) { + dst_n_idx = blockIdx.x * half_blockn + threadIdx.x; + } + + if constexpr (mul_routed_weight) { + float score = (float)score_ptr[token_idx * topk + expert_idx]; + rst = rst * score; + } + + if constexpr (is_swigelu) { + shared_array[threadIdx.x] = rst; + if (threadIdx.x < half_blockn) { + float b = shared_array[threadIdx.x + half_blockn]; + rst = rst * sigmoid(rst) * b; + c_ptr[token_idx * topk * n + expert_idx * n + dst_n_idx] = rst; + } + } else if constexpr (use_rms_norm) { + float rms = rst * rst; + int count_val = 0; + + for (int offset = 1; offset < BLOCK_N; offset *= 2) { + float peer = __shfl_xor_sync(BLOCK_N, rms, offset); + rms += peer; + } + + if (threadIdx.x == 0) { + atomicAdd(sum_out, rms); + __threadfence_block(); + atomicAdd((int*)(count), 1); + } + + while (count_val < gridDim.x) { + count_val = count[0]; + } + + rms = sum_out[0]; + rst = rst * rsqrtf(rms / n + eps) * float(gamma[dst_n_idx]); + c_ptr[token_idx * topk * n + expert_idx * n + dst_n_idx] = rst; + } else { + c_ptr[token_idx * topk * n + expert_idx * n + dst_n_idx] = rst; + } + } +} + +struct BlockConfig { + int block_n; + int block_k; + float score; + bool valid; +}; + +void musa_fused_gemv( + torch::Tensor &A, + torch::Tensor &B, + torch::Tensor &C, + const c10::optional &A_scale, + const c10::optional &B_scale, + bool use_int4_w4a16, + bool use_swigelu, + bool use_rms_norm, + const c10::optional &gamma, + double eps) { + + TORCH_CHECK(A.dim() == 2, "A must be dim 2.") + TORCH_CHECK(B.dim() == 2, "B must be dim 2.") + + bool mul_routed_weight = false; + int topk = 1; + int32_t bseqlen = A.size(0); + int32_t hidden_size = A.size(1); + int32_t num_experts = 1; + int32_t reduce_size = B.size(0); + bool is_fp8 = false; + + if (B.dtype() == torch::kFloat8_e4m3fn) { + is_fp8 = true; + } + + int current_arch = at::musa::getMUSAArch(); + if (current_arch < 300) { + if (is_fp8) { + TORCH_CHECK(false, "gemv moe not support Float8_e4m3fn on MUSA arch ", current_arch); + } + } + + const at::musa::OptionalMUSAGuard device_guard(device_of(A)); + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + void *topk_ids_ptr = nullptr; + void *topk_weights_ptr = nullptr; + void *a_scale_ptr = nullptr; + void *b_scale_ptr = nullptr; + + if (A_scale.has_value()) { + a_scale_ptr = A_scale.value().data_ptr(); + } + if (B_scale.has_value()) { + b_scale_ptr = B_scale.value().data_ptr(); + } + + void *rms_gamma_ptr = nullptr; + void *rms_sum_out_ptr = nullptr; + void *rms_count_ptr = nullptr; + + if (use_rms_norm && gamma.has_value()) { + torch::Tensor sum_out = torch::zeros({1}, A.options().dtype(torch::kFloat)); + torch::Tensor count = torch::zeros({1}, A.options().dtype(torch::kInt)); + rms_gamma_ptr = gamma.value().data_ptr(); + rms_sum_out_ptr = sum_out.data_ptr(); + rms_count_ptr = count.data_ptr(); + } + + int device; + musaGetDevice(&device); + musaDeviceProp device_prop; + musaGetDeviceProperties(&device_prop, device); + int num_mp = device_prop.multiProcessorCount; + int expert_offset_stride = reduce_size * hidden_size; + int half_n_idx = reduce_size / 2; + int scale_k_len = 1; + int scale_k_group_tile = 128; + + if (use_int4_w4a16 || is_fp8) { + scale_k_len = B_scale->size(1); + if (scale_k_len != 1) { + scale_k_group_tile = ceil_div(hidden_size, scale_k_len); + TORCH_CHECK(scale_k_group_tile == 128 || scale_k_group_tile == 64, "scale_k_group_tile only support 128 or 64"); + } + } + + bool is_pergroup_scale = scale_k_len != 1; + int nr_n = use_swigelu ? reduce_size / 2 : reduce_size; + + std::function launch_kernel; + + BlockConfig configs[] = { + {8, 16, 0.f, false}, + {16, 8, 0.f, false}, + {32, 4, 0.f, false}, + {4, 32, 0.f, false}, + }; + + constexpr int iobit = 128; + const int bits_of_byte = 8; + const int vlen = use_int4_w4a16 ? + (iobit / 4): + (iobit / (torch::elementSize(B.scalar_type()) * bits_of_byte)); + + float target_ratio = static_cast(reduce_size) / hidden_size; + + for (auto& config : configs) { + int load_size = config.block_k * vlen; + config.valid = (reduce_size % config.block_n == 0) && (hidden_size % load_size == 0) && (load_size % scale_k_group_tile == 0); + + if (config.valid) { + float block_ratio = static_cast(config.block_n) / config.block_k; + config.score = 1.0f / (1.0f + fabsf(block_ratio - target_ratio)); + } + } + + BlockConfig best_config_storage; + if (current_arch < 300) { + best_config_storage = {128, 1, -1.0f, false}; + } else { + best_config_storage = {32, 1, -1.0f, false}; + } + BlockConfig* best_config = &best_config_storage; + for (auto& config : configs) { + if (config.valid && config.score > best_config->score) { + best_config = &config; + } + } + + switch (best_config->block_n) { + case 4: + switch (best_config->block_k) { + case 32: GEN_LAUNCH_KERN_GEMV(4, 32); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=4"); + } + break; + case 8: + switch (best_config->block_k) { + case 16: GEN_LAUNCH_KERN_GEMV(8, 16); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=8"); + } + break; + case 16: + switch (best_config->block_k) { + case 8: GEN_LAUNCH_KERN_GEMV(16, 8); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=16"); + } + break; + case 32: + switch (best_config->block_k) { + case 4: GEN_LAUNCH_KERN_GEMV(32, 4); break; + case 1: GEN_LAUNCH_KERN_GEMV(32, 1); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=32"); + } + break; + case 128: + switch (best_config->block_k) { + case 1: GEN_LAUNCH_KERN_GEMV(128, 1); + break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=128"); + } + break; + default: + TORCH_CHECK(false, "Unsupported block configuration"); + } + + launch_kernel(); +} + +void fused_moe_gemv( + torch::Tensor &A, + torch::Tensor &B, + torch::Tensor &C, + const c10::optional &A_scale, + const c10::optional &B_scale, + torch::Tensor &topk_weights, + torch::Tensor &topk_ids, + bool mul_routed_weight, + int64_t topk, + bool use_int4_w4a16, + bool use_swigelu) { + + TORCH_CHECK(A.dim() == 2, "A must be dim 2.") + TORCH_CHECK(B.dim() == 3, "B must be dim 3.") + + int32_t bseqlen = A.size(0); + bool is_fp8 = false; + if (B.dtype() == torch::kFloat8_e4m3fn) { + is_fp8 = true; + } + + bool use_rms_norm = false; + void *rms_gamma_ptr = nullptr; + void *rms_sum_out_ptr = nullptr; + void *rms_count_ptr = nullptr; + float eps = 1e-6; + + int current_arch = at::musa::getMUSAArch(); + if (current_arch < 300) { + if (is_fp8) { + TORCH_CHECK(false, "gemv moe not support Float8_e4m3fn on MUSA arch ", current_arch); + } + } + + int32_t hidden_size = A.size(1); + int32_t num_experts = B.size(0); + int32_t reduce_size = B.size(1); + + const at::musa::OptionalMUSAGuard device_guard(device_of(A)); + musaStream_t stream = at::musa::getCurrentMUSAStream(); + + void *topk_ids_ptr = topk_ids.data_ptr(); + void *topk_weights_ptr = topk_weights.data_ptr(); + void *a_scale_ptr = nullptr; + void *b_scale_ptr = nullptr; + + if (A_scale.has_value()) { + a_scale_ptr = A_scale.value().data_ptr(); + } + if (B_scale.has_value()) { + b_scale_ptr = B_scale.value().data_ptr(); + } + + int device; + musaGetDevice(&device); + musaDeviceProp device_prop; + musaGetDeviceProperties(&device_prop, device); + int num_mp = device_prop.multiProcessorCount; + int expert_offset_stride = reduce_size * hidden_size; + int half_n_idx = reduce_size / 2; + int scale_k_len = 1; + int scale_k_group_tile = 128; + + bool is_pergroup_scale = false; + if (use_int4_w4a16 || is_fp8) { + scale_k_len = B_scale->size(2); + if (scale_k_len != 1) { + is_pergroup_scale = true; + scale_k_group_tile = ceil_div(hidden_size, scale_k_len); + TORCH_CHECK(scale_k_group_tile == 128 || scale_k_group_tile == 64, "scale_k_group_tile only support 128 or 64"); + } + } + + int nr_n = use_swigelu ? reduce_size / 2 : reduce_size; + + std::function launch_kernel; + + BlockConfig configs[] = { + {8, 16, 0.f, false}, + {16, 8, 0.f, false}, + {32, 4, 0.f, false}, + {4, 32, 0.f, false}, + }; + + constexpr int iobit = 128; + const int bits_of_byte = 8; + const int vlen = use_int4_w4a16 ? + (iobit / 4): + (iobit / (torch::elementSize(B.scalar_type()) * bits_of_byte)); + + float target_ratio = static_cast(reduce_size) / hidden_size; + + for (auto& config : configs) { + int load_size = config.block_k * vlen; + config.valid = (reduce_size % config.block_n == 0) && (hidden_size % load_size == 0) && (load_size % scale_k_group_tile == 0); + + if (config.valid) { + float block_ratio = static_cast(config.block_n) / config.block_k; + config.score = 1.0f / (1.0f + fabsf(block_ratio - target_ratio)); + } + } + + BlockConfig best_config_storage; + if (current_arch < 300) { + best_config_storage = {128, 1, -1.0f, false}; + } else { + best_config_storage = {32, 1, -1.0f, false}; + } + BlockConfig* best_config = &best_config_storage; + + for (auto& config : configs) { + if (config.valid && config.score > best_config->score) { + best_config = &config; + } + } + + switch (best_config->block_n) { + case 4: + switch (best_config->block_k) { + case 32: GEN_LAUNCH_KERN(4, 32); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=4"); + } + break; + case 8: + switch (best_config->block_k) { + case 16: GEN_LAUNCH_KERN(8, 16); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=8"); + } + break; + case 16: + switch (best_config->block_k) { + case 8: GEN_LAUNCH_KERN(16, 8); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=16"); + } + break; + case 32: + switch (best_config->block_k) { + case 4: GEN_LAUNCH_KERN(32, 4); break; + case 1: GEN_LAUNCH_KERN(32, 1); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=32"); + } + break; + case 128: + switch (best_config->block_k) { + case 1: GEN_LAUNCH_KERN(128, 1); break; + default: TORCH_CHECK(false, "Unsupported block_k for block_n=128"); + } + break; + default: + TORCH_CHECK(false, "Unsupported block configuration"); + } + + launch_kernel(); +} diff --git a/sgl-kernel/csrc/musa/pos_encoding_contiguous.mu b/sgl-kernel/csrc/musa/pos_encoding_contiguous.mu new file mode 100644 index 000000000..bde7c0515 --- /dev/null +++ b/sgl-kernel/csrc/musa/pos_encoding_contiguous.mu @@ -0,0 +1,264 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include "musa.h" +#include "musa/dispatch_utils.h" +#include "torch_musa/csrc/core/MUSAGuard.h" +#include "torch_musa/csrc/core/MUSAStream.h" + +template +inline __device__ void apply_token_rotary_embedding_contiguous( + scalar_t* __restrict__ arr, const scalar_t* __restrict__ cos_ptr, + const scalar_t* __restrict__ sin_ptr, int rot_offset, int embed_dim, + int64_t head_stride) { + int x_index, y_index; + scalar_t cos, sin; + if (IS_NEOX) { + // GPT-NeoX style rotary embedding. + x_index = rot_offset; + y_index = embed_dim + rot_offset; + cos = MUSA_LDG(cos_ptr + x_index); + sin = MUSA_LDG(sin_ptr + x_index); + } else { + // GPT-J style rotary embedding. + x_index = 2 * rot_offset; + y_index = 2 * rot_offset + 1; + cos = MUSA_LDG(cos_ptr + x_index / 2); + sin = MUSA_LDG(sin_ptr + x_index / 2); + } + + scalar_t* x_ptr = arr + x_index * head_stride; + scalar_t* y_ptr = arr + y_index * head_stride; + + const scalar_t x = *x_ptr; + const scalar_t y = *y_ptr; + *x_ptr = x * cos - y * sin; + *y_ptr = y * cos + x * sin; +} + +template +inline __device__ void apply_rotary_embedding_contiguous( + scalar_t* __restrict__ query, // [num_tokens, num_heads, head_size] + scalar_t* __restrict__ key, // [num_tokens, num_kv_heads, head_size] + const scalar_t* cache_ptr, const int head_size, const int num_heads, + const int num_kv_heads, const int rot_dim, const int token_idx, + const int64_t query_token_stride, const int64_t query_head_stride, + const int64_t query_dim_stride, + const int64_t key_token_stride, const int64_t key_head_stride, + const int64_t key_dim_stride) { + const int embed_dim = rot_dim / 2; + const scalar_t* cos_ptr = cache_ptr; + const scalar_t* sin_ptr = cache_ptr + embed_dim; + + const int nq = num_heads * embed_dim; + for (int i = threadIdx.x; i < nq; i += blockDim.x) { + const int head_idx = i / embed_dim; + const int rot_offset = i % embed_dim; + + scalar_t* head_query = query + + token_idx * query_token_stride + + head_idx * query_head_stride; + + apply_token_rotary_embedding_contiguous( + head_query, cos_ptr, sin_ptr, rot_offset, embed_dim, query_dim_stride); + } + + const int nk = num_kv_heads * embed_dim; + for (int i = threadIdx.x; i < nk; i += blockDim.x) { + const int head_idx = i / embed_dim; + const int rot_offset = i % embed_dim; + + scalar_t* head_key = key + + token_idx * key_token_stride + + head_idx * key_head_stride; + + apply_token_rotary_embedding_contiguous( + head_key, cos_ptr, sin_ptr, rot_offset, embed_dim, key_dim_stride); + } +} + +template +__global__ void rotary_embedding_kernel_contiguous( + const int64_t* __restrict__ positions, // [num_tokens] + scalar_t* __restrict__ query, // [num_tokens, num_heads, head_size] + scalar_t* __restrict__ key, // [num_tokens, num_kv_heads, head_size] + const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim // 2] + const int rot_dim, + const int64_t query_token_stride, const int64_t query_head_stride, + const int64_t query_dim_stride, + const int64_t key_token_stride, const int64_t key_head_stride, + const int64_t key_dim_stride, + const int num_heads, const int num_kv_heads, const int head_size) { + // Each thread block is responsible for one token. + const int token_idx = blockIdx.x; + int64_t pos = positions[token_idx]; + const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim; + + apply_rotary_embedding_contiguous( + query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim, + token_idx, + query_token_stride, query_head_stride, query_dim_stride, + key_token_stride, key_head_stride, key_dim_stride); +} + +template +__global__ void batched_rotary_embedding_kernel_contiguous( + const int64_t* __restrict__ positions, // [num_tokens] + scalar_t* __restrict__ query, // [num_tokens, num_heads, head_size] + scalar_t* __restrict__ key, // [num_tokens, num_kv_heads, head_size] + const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2, rot_dim // 2] + const int64_t* __restrict__ cos_sin_cache_offsets, // [num_tokens] + const int rot_dim, + const int64_t query_token_stride, const int64_t query_head_stride, + const int64_t query_dim_stride, // stride for each dimension + const int64_t key_token_stride, const int64_t key_head_stride, + const int64_t key_dim_stride, + const int num_heads, const int num_kv_heads, const int head_size) { + // Each thread block is responsible for one token. + const int token_idx = blockIdx.x; + int64_t pos = positions[token_idx]; + int64_t cos_sin_cache_offset = cos_sin_cache_offsets[token_idx]; + const scalar_t* cache_ptr = + cos_sin_cache + (cos_sin_cache_offset + pos) * rot_dim; + + apply_rotary_embedding_contiguous( + query, key, cache_ptr, head_size, num_heads, num_kv_heads, rot_dim, + token_idx, + query_token_stride, query_head_stride, query_dim_stride, + key_token_stride, key_head_stride, key_dim_stride); +} + + +void rotary_embedding_contiguous( + torch::Tensor& positions, // [num_tokens] + torch::Tensor& query, // [num_tokens, num_heads, head_size] + torch::Tensor& key, // [num_tokens, num_kv_heads, head_size] + int64_t head_size, + torch::Tensor& cos_sin_cache, // [max_position, rot_dim] + bool is_neox) { + int64_t num_tokens = positions.size(0); + + TORCH_CHECK(query.dim() == 3, "query must be 3D [num_tokens, num_heads, head_size]"); + TORCH_CHECK(key.dim() == 3, "key must be 3D [num_tokens, num_kv_heads, head_size]"); + TORCH_CHECK(query.size(0) == num_tokens && key.size(0) == num_tokens, + "query, key and positions must have the same number of tokens"); + + int64_t query_token_stride = query.stride(0); + int64_t query_head_stride = query.stride(1); + int64_t query_dim_stride = query.stride(2); + + int64_t key_token_stride = key.stride(0); + int64_t key_head_stride = key.stride(1); + int64_t key_dim_stride = key.stride(2); + + int num_heads = query.size(1); + int num_kv_heads = key.size(1); + int rot_dim = cos_sin_cache.size(1); + + dim3 grid(num_tokens); + dim3 block(std::min(num_heads * rot_dim / 2, 512)); + + const at::musa::OptionalMUSAGuard device_guard(device_of(query)); + const musaStream_t stream = at::musa::getCurrentMUSAStream(); + + MUSA_DISPATCH_FLOATING_TYPES(query.scalar_type(), "rotary_embedding_contiguous", [&] { + if (is_neox) { + rotary_embedding_kernel_contiguous<<>>( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + rot_dim, + query_token_stride, query_head_stride, query_dim_stride, + key_token_stride, key_head_stride, key_dim_stride, + num_heads, num_kv_heads, head_size); + } else { + rotary_embedding_kernel_contiguous<<>>( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + rot_dim, + query_token_stride, query_head_stride, query_dim_stride, + key_token_stride, key_head_stride, key_dim_stride, + num_heads, num_kv_heads, head_size); + } + }); +} + +void batched_rotary_embedding_contiguous( + torch::Tensor& positions, // [num_tokens] + torch::Tensor& query, // [num_tokens, num_heads, head_size] + torch::Tensor& key, // [num_tokens, num_kv_heads, head_size] + int64_t head_size, + torch::Tensor& cos_sin_cache, // [max_position, rot_dim] + bool is_neox, int64_t rot_dim, + torch::Tensor& cos_sin_cache_offsets // [num_tokens] +) { + int64_t num_tokens = cos_sin_cache_offsets.size(0); + + + TORCH_CHECK(positions.size(0) == num_tokens, + "positions must have the same num_tokens as cos_sin_cache_offsets"); + TORCH_CHECK(query.dim() == 3, "query must be 3D [num_tokens, num_heads, head_size]"); + TORCH_CHECK(key.dim() == 3, "key must be 3D [num_tokens, num_kv_heads, head_size]"); + + int64_t query_token_stride = query.stride(0); + int64_t query_head_stride = query.stride(1); + int64_t query_dim_stride = query.stride(2); + + int64_t key_token_stride = key.stride(0); + int64_t key_head_stride = key.stride(1); + int64_t key_dim_stride = key.stride(2); + + int num_heads = query.size(1); + int num_kv_heads = key.size(1); + + dim3 grid(num_tokens); + dim3 block(std::min(num_heads * rot_dim / 2, 512)); + + const at::musa::OptionalMUSAGuard device_guard(device_of(query)); + const musaStream_t stream = at::musa::getCurrentMUSAStream(); + + MUSA_DISPATCH_FLOATING_TYPES(query.scalar_type(), "batched_rotary_embedding_contiguous", [&] { + if (is_neox) { + batched_rotary_embedding_kernel_contiguous<<>>( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + cos_sin_cache_offsets.data_ptr(), + rot_dim, + query_token_stride, query_head_stride, query_dim_stride, + key_token_stride, key_head_stride, key_dim_stride, + num_heads, num_kv_heads, head_size); + } else { + batched_rotary_embedding_kernel_contiguous<<>>( + positions.data_ptr(), + query.data_ptr(), + key.data_ptr(), + cos_sin_cache.data_ptr(), + cos_sin_cache_offsets.data_ptr(), + rot_dim, + query_token_stride, query_head_stride, query_dim_stride, + key_token_stride, key_head_stride, key_dim_stride, + num_heads, num_kv_heads, head_size); + } + }); +} diff --git a/sgl-kernel/csrc/musa/ternary.mu b/sgl-kernel/csrc/musa/ternary.mu new file mode 100644 index 000000000..e67e0aeae --- /dev/null +++ b/sgl-kernel/csrc/musa/ternary.mu @@ -0,0 +1,123 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include "musa/dispatch_utils.h" +#include "musa.h" +#include "torch_musa/csrc/aten/musa/MUSADtype.muh" +#include "torch_musa/csrc/core/MUSAGuard.h" +#include "torch_musa/csrc/core/MUSAStream.h" + + +typedef __half float16_t; +typedef __mt_bfloat16 bfloat16_t; + +__device__ __host__ __forceinline__ constexpr int64_t ceil_div(int64_t a, + int64_t b) { + return (a + b - 1) / b; +} + +template +__global__ void FusedMulAdd(scalar_t *out, scalar_t *self, scalar_t *bias, + const scalar_t scale, const int64_t elem) { + constexpr int bits_of_byte = 8; + using Vec = + at::musa::VecType; + + int64_t tid = (int64_t)blockIdx.x * blockDim.x + threadIdx.x; + int64_t grid_stride = (int64_t)gridDim.x * blockDim.x; + int64_t grid_stride_vec = grid_stride * vlen; + + for (int64_t offset = tid * vlen; offset < elem; offset += grid_stride_vec) { + Vec t = Vec::load(self, offset); + Vec b = Vec::load(bias, offset); + Vec o; +#pragma unroll + for (int k = 0; k < vlen; ++k) { + o.val_.elem[k] = b.val_.elem[k] + t.val_.elem[k] * scale; + } + Vec::store(out, offset, o); + } +} + +void fused_mul_add(torch::Tensor &output, torch::Tensor &self, + torch::Tensor &bias, const double scale) { + TORCH_CHECK(self.sizes() == output.sizes(), + "self and output shape don't match"); + TORCH_CHECK(self.sizes() == bias.sizes(), "self and bias shape don't match"); + TORCH_CHECK(self.scalar_type() == bias.scalar_type(), + "self and bias should be same type"); + TORCH_CHECK(self.scalar_type() == output.scalar_type(), + "self and output should be same type"); + TORCH_CHECK(self.scalar_type() == at::ScalarType::Float || + self.scalar_type() == at::ScalarType::BFloat16 || + self.scalar_type() == at::ScalarType::Half, + "self's dtype should be in float/half/bfloat16"); + + // device guard + const at::musa::OptionalMUSAGuard device_guard(device_of(self)); + + // Suppose the uncontiguous elementwise is much slower + if C10_UNLIKELY (!self.is_contiguous()) { + self = self.contiguous(); + } + if C10_UNLIKELY (!bias.is_contiguous()) { + bias = bias.contiguous(); + } + + // follow mudnn config for arch==22 + const int64_t max_load_vec = + self.scalar_type() == at::ScalarType::Float ? 4 : 8; + const int64_t numel = self.numel(); + size_t thread_per_block = 512; + if (ceil_div(numel, max_load_vec) <= 128) { + thread_per_block = 128; + } else if (ceil_div(numel, max_load_vec) <= 256) { + thread_per_block = 256; + } + size_t nr_block = ceil_div(numel, max_load_vec * thread_per_block); + + const musaStream_t stream = at::musa::getCurrentMUSAStream(); + + switch (self.scalar_type()) { + case at::ScalarType::Float: + FusedMulAdd<<>>( + static_cast(output.data_ptr()), + static_cast(self.data_ptr()), + static_cast(bias.data_ptr()), scale, numel); + break; + case at::ScalarType::Half: + FusedMulAdd + <<>>( + static_cast(output.data_ptr()), + static_cast(self.data_ptr()), + static_cast(bias.data_ptr()), + static_cast(scale), numel); + break; + case at::ScalarType::BFloat16: + FusedMulAdd + <<>>( + static_cast(output.data_ptr()), + static_cast(self.data_ptr()), + static_cast(bias.data_ptr()), + static_cast(scale), numel); + break; + default: + break; + } +} diff --git a/sgl-kernel/csrc/musa/top_k_top_p_sampling.mu b/sgl-kernel/csrc/musa/top_k_top_p_sampling.mu new file mode 100644 index 000000000..4e149086f --- /dev/null +++ b/sgl-kernel/csrc/musa/top_k_top_p_sampling.mu @@ -0,0 +1,382 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include +#include "torch_musa/csrc/aten/musa/MUSAGeneratorImpl.h" +#include + +#include "torch_musa/csrc/aten/musa/UnpackRaw.muh" +#include +#include + +#include "musa.h" +#include "musa/dispatch_utils.h" +#include "pytorch_extension_utils.h" +#include "torch_musa/csrc/core/MUSAGuard.h" +#include "torch_musa/csrc/core/MUSAStream.h" + +namespace musa { +namespace sampling { + +#define kRmemEles 16 // should not too large, otherwise it will cause register spilling +// reuse code of flashinfer +using namespace flashinfer; +using namespace flashinfer::sampling; + +template +__device__ __forceinline__ void DeviceSamplingFromProbWithOffset( + uint32_t i, uint32_t d, Predicate pred, float u, vec_t prob_vec, + float& aggregate, + SamplingTempStorage* temp_storage, + int offset = 0) { + const uint32_t tx = threadIdx.x; + float prob_greater_than_threshold[VEC_SIZE]; + float inclusive_cdf[VEC_SIZE]; + bool greater_than_u[VEC_SIZE], valid[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + prob_greater_than_threshold[j] = pred(prob_vec[j]) ? prob_vec[j] : 0; + valid[j] = pred(prob_vec[j]) && (offset + (i * BLOCK_THREADS + tx) * VEC_SIZE + j < d); + } + float aggregate_local = + BlockReduce(temp_storage->block_prim.reduce) + .template Sum(prob_greater_than_threshold); + if (tx == 0) { + temp_storage->block_aggregate.value = aggregate_local; + } + __syncthreads(); + aggregate_local = temp_storage->block_aggregate.value; + + if (aggregate + aggregate_local > u) { + if constexpr (DETERMINISTIC) { + DeterministicInclusiveSum( + prob_greater_than_threshold, inclusive_cdf, temp_storage); + } else { + BlockScan(temp_storage->block_prim.scan) + .template InclusiveSum(prob_greater_than_threshold, inclusive_cdf); + + __syncthreads(); + } + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + greater_than_u[j] = (inclusive_cdf[j] + aggregate > u) && valid[j]; + } + + bool greater_than_u_diff[VEC_SIZE]; +#ifdef FLASHINFER_CUB_SUBTRACTLEFT_DEFINED + BlockAdjacentDifference(temp_storage->block_prim.adj_diff) + .SubtractLeft(greater_than_u, greater_than_u_diff, BoolDiffOp()); +#else + BlockAdjacentDifference(temp_storage->block_prim.adj_diff) + .template FlagHeads(greater_than_u_diff, greater_than_u, BoolDiffOp(), 0); +#endif + __syncthreads(); + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + if (greater_than_u_diff[j]) { + atomicMin(&(temp_storage->sampled_id), offset + (i * BLOCK_THREADS + tx) * VEC_SIZE + j); + } + } + __syncthreads(); + } + + // update the last valid index + int valid_index[VEC_SIZE]; +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + if (valid[j]) { + valid_index[j] = offset + (i * BLOCK_THREADS + tx) * VEC_SIZE + j; + } else { + valid_index[j] = -1; + } + } + int max_valid_index = + BlockReduce(temp_storage->block_prim.reduce_int) + .Reduce(valid_index, MaxReduceOp{}); + if (tx == 0 && max_valid_index != -1) { + temp_storage->last_valid_id = max_valid_index; + } + __syncthreads(); + aggregate += aggregate_local; +} + +template +__global__ void TopKTopPSamplingFromProbKernel(DType* probs, IdType* top_k_arr, float* top_p_arr, + IdType* output, IdType* indices, IdType top_k_val, + float top_p_val, uint32_t d, uint64_t philox_seed, + uint64_t philox_offset) { + const uint32_t batch_size = gridDim.x; + const uint32_t bx = blockIdx.x, tx = threadIdx.x; + murandStatePhilox4_32_10_t state; + murand_init(philox_seed, bx, philox_offset, &state); + const uint32_t row_idx = indices == nullptr ? bx : indices[bx]; + const uint32_t k = top_k_arr == nullptr ? top_k_val : top_k_arr[row_idx]; + const float p = top_p_arr == nullptr ? top_p_val : top_p_arr[row_idx]; + extern __shared__ __align__(alignof(SamplingTempStorage)) + uint8_t smem_sampling[]; + auto& temp_storage = + reinterpret_cast&>(smem_sampling); + + vec_t rprobs; // persistent + vec_t rprobs_d; + // elements of one row will be split into three stages by the address it will be stored + // [rprobs, probs_vec] + // rprobs will be persistent to reuse those elements instead of reload + int real_r_size = d > (BLOCK_THREADS * kRmemEles) ? BLOCK_THREADS * kRmemEles : d; + int real_rd_size = d > (BLOCK_THREADS * kRmemEles) ? d - BLOCK_THREADS * kRmemEles : 0; + + probs = probs + row_idx * d; + DType* global_rprobs = probs; + DType* global_rprobs_d = probs + real_r_size; + + // fisrtly, we load some elements to register + rprobs.fill(0); + if (tx * kRmemEles < real_r_size) { + rprobs.cast_load(global_rprobs + tx * kRmemEles); + int valid_size = real_r_size - tx * kRmemEles; + // mask invalid data + for (int i = valid_size; i < kRmemEles; ++i) { + rprobs[i] = 0.0; + } + } + + float aggregate; + float q = 1; + double low = 0, high = 1.f; + int sampled_id; + + // sample and check + do { + temp_storage.sampled_id = d; + __syncthreads(); + float u = murand_uniform(&state) * q; + aggregate = 0; + + // fisrtly, sample from rprobs + DeviceSamplingFromProbWithOffset( + 0, d, [&](float x) { return x > low; }, u, rprobs, aggregate, &temp_storage); + + if (aggregate <= u) { + // secondly, sample from rprobs_d + for (uint32_t i = 0; i < ceil_div(real_rd_size, BLOCK_THREADS * VEC_SIZE); ++i) { + rprobs_d.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE + real_r_size < d) { + rprobs_d.cast_load(global_rprobs_d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + + DeviceSamplingFromProbWithOffset( + i, d, [&](float x) { return x > low; }, u, rprobs_d, aggregate, &temp_storage, real_r_size); + if (aggregate > u) { + break; + } + } + } + + // id is sampled + __syncthreads(); + sampled_id = temp_storage.sampled_id; + if (sampled_id == d) { + // NOTE(Zihao): this would happen when u is very close to 1 + // and the sum of probabilities is smaller than u + // In this case, we use the last valid index as the sampled id + sampled_id = temp_storage.last_valid_id; + } + double pivot_0 = probs[sampled_id]; + double pivot_1 = (pivot_0 + high) / 2; + + ValueCount aggregate_gt_pivot_0{0, 0}, aggregate_gt_pivot_1{0, 0}; + // check in rprobs + ValueCount probs_gt_pivot_0[VEC_SIZE], probs_gt_pivot_1[VEC_SIZE]; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + probs_gt_pivot_0[i] = {0, 0}; + probs_gt_pivot_1[i] = {0, 0}; + } + + // init to 0 + for (uint32_t j = 0; j < kRmemEles; j += VEC_SIZE) { +#pragma unroll + for (uint32_t k = 0; k < VEC_SIZE; ++k) { + probs_gt_pivot_0[k] += + {(rprobs[j + k] > pivot_0) ? rprobs[j + k] : 0, (rprobs[j + k] > pivot_0 && (tx)*kRmemEles + j + k < d)}; + probs_gt_pivot_1[k] += + {(rprobs[j + k] > pivot_1) ? rprobs[j + k] : 0, (rprobs[j + k] > pivot_1 && (tx)*kRmemEles + j + k < d)}; + } + } + + // check in rprobs_d + for (uint32_t i = 0; i < ceil_div(real_rd_size, BLOCK_THREADS * VEC_SIZE); ++i) { + rprobs_d.fill(0); + if ((i * BLOCK_THREADS + tx) * VEC_SIZE + real_r_size < d) { + rprobs_d.cast_load(global_rprobs_d + (i * BLOCK_THREADS + tx) * VEC_SIZE); + } + +#pragma unroll + for (uint32_t j = 0; j < VEC_SIZE; ++j) { + probs_gt_pivot_0[j] += + {(rprobs_d[j] > pivot_0) ? rprobs_d[j] : 0, + (rprobs_d[j] > pivot_0 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j + real_r_size < d)}; + probs_gt_pivot_1[j] += + {(rprobs_d[j] > pivot_1) ? rprobs_d[j] : 0, + (rprobs_d[j] > pivot_1 && (i * BLOCK_THREADS + tx) * VEC_SIZE + j + real_r_size < d)}; + } + } + + aggregate_gt_pivot_0 = BlockReduce, BLOCK_THREADS>(temp_storage.block_prim.reduce_value_count) + .template Sum(probs_gt_pivot_0); + if (tx == 0) { + temp_storage.block_aggregate.pair = aggregate_gt_pivot_0; + } + __syncthreads(); + aggregate_gt_pivot_0 = temp_storage.block_aggregate.pair; + + aggregate_gt_pivot_1 = BlockReduce, BLOCK_THREADS>(temp_storage.block_prim.reduce_value_count) + .template Sum(probs_gt_pivot_1); + if (tx == 0) { + temp_storage.block_aggregate.pair = aggregate_gt_pivot_1; + } + __syncthreads(); + aggregate_gt_pivot_1 = temp_storage.block_aggregate.pair; + + if (aggregate_gt_pivot_0.count < k && aggregate_gt_pivot_0.value < p) { + // case 1: pivot_0 accepted + break; + } + if (aggregate_gt_pivot_1.count < k && aggregate_gt_pivot_1.value < p) { + // case 2: pivot_0 rejected, pivot_1 accepted + low = pivot_0; + high = pivot_1; + q = aggregate_gt_pivot_0.value; + } else { + // case 3: pivot_0 rejected, pivot_1 rejected + low = pivot_1; + q = aggregate_gt_pivot_1.value; + } + } while (low < high); + __syncthreads(); + if (tx == 0) { + output[bx] = sampled_id; + } +} +} // namespace sampling +} // namespace musa + +template +musaError_t MusaTopKTopPSamplingFromProb(T* probs, IdType* top_k_arr, T* top_p_arr, IdType* output, + IdType* indices, uint32_t batch_size, IdType top_k_val, + T top_p_val, uint32_t d, bool deterministic, + uint64_t philox_seed, uint64_t philox_offset, + musaStream_t stream = 0) { + const uint32_t vec_size = std::gcd(16 / sizeof(T), d); + using namespace flashinfer; + using namespace flashinfer::sampling; + auto compute_capacity = GetCudaComputeCapability(); + + DISPATCH_COMPUTE_CAP_NUM_THREADS(compute_capacity, BLOCK_THREADS, { + const uint32_t smem_size = sizeof(SamplingTempStorage); + dim3 nblks(batch_size); + dim3 nthrs(BLOCK_THREADS); + void* args[] = { + &probs, &top_k_arr, &top_p_arr, &output, &indices, &top_k_val, &top_p_val, &d, &philox_seed, &philox_offset}; + // fall back to flashinfer implementation + if (d < BLOCK_THREADS * kRmemEles) { + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = TopKTopPSamplingFromProbKernel< + BLOCK_THREADS, + SCAN_ALGO, + REDUCE_ALGO, + VEC_SIZE, + DETERMINISTIC, + T, + IdType>; + FLASHINFER_CUDA_CALL(musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(musaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + } else { + DISPATCH_ALIGNED_VEC_SIZE( + vec_size, VEC_SIZE, {DISPATCH_DETERMINISTIC(deterministic, DETERMINISTIC, { + auto kernel = musa::sampling::TopKTopPSamplingFromProbKernel< + BLOCK_THREADS, + SCAN_ALGO, + REDUCE_ALGO, + VEC_SIZE, + DETERMINISTIC, + T, + IdType>; + FLASHINFER_CUDA_CALL(musaFuncSetAttribute(kernel, musaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + FLASHINFER_CUDA_CALL(musaLaunchKernel((void*)kernel, nblks, nthrs, args, smem_size, stream)); + })}); + } + return musaSuccess; + }); +} + +void musa_top_k_top_p_sampling_from_probs( + at::Tensor probs, + at::Tensor output, + std::optional maybe_indices, + std::optional maybe_top_k_arr, + double top_k_val, + std::optional maybe_top_p_arr, + double top_p_val, + bool deterministic, + std::optional gen_) { + CHECK_INPUT(probs); + CHECK_INPUT(output); + auto device = probs.device(); + CHECK_EQ(output.device(), device); + CHECK_EQ(probs.dtype(), torch::kFloat); + CHECK_DIM(2, probs); // probs: (batch_size, vocab_size) + CHECK_DIM(1, output); // output: (batch_size) + unsigned int batch_size = output.size(0); + unsigned int vocab_size = probs.size(1); + bool has_top_k_arr = maybe_top_k_arr.has_value(); + bool has_top_p_arr = maybe_top_p_arr.has_value(); + uint64_t philox_seed, philox_offset; + auto gen = at::get_generator_or_default(gen_, at::musa::detail::getDefaultMUSAGenerator()); + std::lock_guard lock(gen->mutex_); + at::PhiloxMusaState rng_engine_inputs = gen->philox_musa_state(32 * batch_size); + philox_seed = rng_engine_inputs.seed_.val; + philox_offset = rng_engine_inputs.offset_.val; + + const c10::musa::OptionalMUSAGuard device_guard(device); + auto stream = at::musa::getCurrentMUSAStream(); + musaError_t status = MusaTopKTopPSamplingFromProb( + static_cast(probs.data_ptr()), + has_top_k_arr ? static_cast(maybe_top_k_arr->data_ptr()) : nullptr, + has_top_p_arr ? static_cast(maybe_top_p_arr->data_ptr()) : nullptr, + static_cast(output.data_ptr()), + maybe_indices.has_value() ? static_cast(maybe_indices->data_ptr()) : nullptr, + batch_size, + top_k_val, + top_p_val, + vocab_size, + deterministic, + philox_seed, + philox_offset, + stream); + TORCH_CHECK( + status == musaSuccess, + "MusaTopKTopPSamplingFromProb failed with error code " + std::string(musaGetErrorString(status))); +} diff --git a/sgl-kernel/include/musa/dispatch_utils.h b/sgl-kernel/include/musa/dispatch_utils.h new file mode 100644 index 000000000..42eaa7024 --- /dev/null +++ b/sgl-kernel/include/musa/dispatch_utils.h @@ -0,0 +1,34 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +/* + * Adapted from + * https://github.com/pytorch/pytorch/blob/v2.0.1/aten/src/ATen/Dispatch.h + */ +#pragma once + +#include + +#define MUSA_LDG(arg) __ldg(arg) + +#define MUSA_DISPATCH_CASE_FLOATING_TYPES(...) \ + AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \ + AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__) + +#define MUSA_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \ + AT_DISPATCH_SWITCH(TYPE, NAME, MUSA_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__)) diff --git a/sgl-kernel/include/musa/integer_subbyte.h b/sgl-kernel/include/musa/integer_subbyte.h new file mode 100644 index 000000000..3c7b4c62b --- /dev/null +++ b/sgl-kernel/include/musa/integer_subbyte.h @@ -0,0 +1,49 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +namespace musa::dnn { + +// cutlass integer_subbyte class +template +struct integer_subbyte { + using Storage = uint8_t; + + static_assert(Bits <= 8 * sizeof(Storage), "Require a subbyte of bits in integer_subbyte"); + + using xint_t = typename std::conditional::type; + + static constexpr Storage bits_mask_ = Storage((1 << Bits) - 1); + + static constexpr Storage sign_mask_ = Storage((Signed ? 1 : 0) << (Bits - 1)); + + Storage storage; + + __host__ __device__ constexpr integer_subbyte() {} + + __host__ __device__ constexpr integer_subbyte(int value) + : storage(reinterpret_cast(value) & bits_mask_) {} + + __host__ __device__ constexpr integer_subbyte(unsigned value) + : storage(reinterpret_cast(value) & bits_mask_) {} +}; + +} // namespace musa::dnn diff --git a/sgl-kernel/include/sgl_kernel_musa_ops.h b/sgl-kernel/include/sgl_kernel_musa_ops.h new file mode 100644 index 000000000..a50cd7fc4 --- /dev/null +++ b/sgl-kernel/include/sgl_kernel_musa_ops.h @@ -0,0 +1,80 @@ +/* + * Copyright (c) 2020-2026, Moore Threads Technology Co., Ltd("Moore Threads"). + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include +#include + +#include + +void batched_rotary_embedding_contiguous( + torch::Tensor& positions, + torch::Tensor& query, + torch::Tensor& key, + int64_t head_size, + torch::Tensor& cos_sin_cache, + bool is_neox, + int64_t rot_dim, + torch::Tensor& cos_sin_cache_offsets); + +void rotary_embedding_contiguous( + torch::Tensor& positions, + torch::Tensor& query, + torch::Tensor& key, + int64_t head_size, + torch::Tensor& cos_sin_cache, + bool is_neox); + +void fused_moe_gemv( + torch::Tensor& A, + torch::Tensor& B, + torch::Tensor& C, + const c10::optional& A_scale, + const c10::optional& B_scale, + torch::Tensor& topk_weights, + torch::Tensor& topk_ids, + bool mul_routed_weight, + int64_t topk, + bool use_int4_w4a16, + bool use_swigelu); + +void musa_fused_gemv( + torch::Tensor& A, + torch::Tensor& B, + torch::Tensor& C, + const c10::optional& A_scale, + const c10::optional& B_scale, + bool use_int4_w4a16, + bool use_swigelu, + bool use_rms_norm, + const c10::optional& gamma, + double eps); + +void fused_mul_add(torch::Tensor& output, torch::Tensor& self, torch::Tensor& bias, double scale); + +void musa_top_k_top_p_sampling_from_probs( + at::Tensor probs, + at::Tensor output, + std::optional maybe_indices, + std::optional maybe_top_k_arr, + double top_k_val, + std::optional maybe_top_p_arr, + double top_p_val, + bool deterministic, + std::optional gen); diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 3891dc900..5c0261b64 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/sgl-kernel/include/sgl_kernel_ops.h @@ -27,6 +27,10 @@ limitations under the License. #include "scalar_type.hpp" +#ifdef USE_MUSA +#include "sgl_kernel_musa_ops.h" +#endif + #define _CONCAT(A, B) A##B #define CONCAT(A, B) _CONCAT(A, B) diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index ed08e3f0f..c0a2af611 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/sgl-kernel/python/sgl_kernel/__init__.py @@ -111,6 +111,15 @@ from sgl_kernel.version import __version__ if torch.version.hip is not None: from sgl_kernel.elementwise import gelu_quick +if hasattr(torch.version, "musa") and torch.version.musa is not None: + from sgl_kernel.musa import ( + musa_batched_rotary_embedding_contiguous, + musa_fused_gemv, + musa_fused_moe_gemv, + musa_fused_mul_add, + musa_rotary_embedding_contiguous, + ) + _DEBUG_EXPORT_NAMES = [ "apply_shuffle_mul_sum", diff --git a/sgl-kernel/python/sgl_kernel/musa.py b/sgl-kernel/python/sgl_kernel/musa.py new file mode 100644 index 000000000..b1bd5eb23 --- /dev/null +++ b/sgl-kernel/python/sgl_kernel/musa.py @@ -0,0 +1,169 @@ +from typing import Optional + +import torch + + +def musa_batched_rotary_embedding_contiguous( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox: bool, + rot_dim: int, + cos_sin_cache_offsets: torch.Tensor, +) -> None: + return torch.ops.sgl_kernel.musa_batched_rotary_embedding_contiguous( + positions, + query, + key, + head_size, + cos_sin_cache, + is_neox, + rot_dim, + cos_sin_cache_offsets, + ) + + +def musa_rotary_embedding_contiguous( + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + head_size: int, + cos_sin_cache: torch.Tensor, + is_neox: bool, +) -> None: + return torch.ops.sgl_kernel.musa_rotary_embedding_contiguous( + positions, + query, + key, + head_size, + cos_sin_cache, + is_neox, + ) + + +def musa_fused_moe_gemv( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + A_scale, + B_scale, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + mul_routed_weight: bool, + topk: int, + use_int4_w4a16: bool, + use_swigelu: bool, +) -> None: + return torch.ops.sgl_kernel.musa_fused_moe_gemv( + A, + B, + C, + A_scale, + B_scale, + topk_weights, + topk_ids, + mul_routed_weight, + topk, + use_int4_w4a16, + use_swigelu, + ) + + +def musa_fused_gemv( + x: torch.Tensor, + qweight: torch.Tensor, + x_scales: Optional[torch.Tensor] = None, + qweight_scales: Optional[torch.Tensor] = None, + use_swigelu: bool = False, + use_rms_norm: bool = False, + gamma: Optional[torch.Tensor] = None, + eps: float = 1e-6, +): + use_int4_w4a16 = False + out_shape = x.shape[:-1] + ( + qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, + ) + assert not ( + use_swigelu and use_rms_norm + ), "gemv only fused one activation (swigelu or rms_norm)!" + + if use_rms_norm: + if gamma is None: + assert False, "rms_norm gamma is None!" + + # fp8 grouped matmul + if qweight.dtype == torch.float8_e4m3fn: + assert qweight_scales is not None, "FP8 grouped matmul weight scales is None!" + output = torch.empty(out_shape, device=x.device, dtype=torch.bfloat16) + torch.ops.sgl_kernel.musa_fused_gemv( + x, + qweight, + output, + x_scales, + qweight_scales, + use_int4_w4a16, + use_swigelu, + use_rms_norm, + gamma, + eps, + ) + return output + # w4a16 gemv + elif qweight_scales is not None: + assert ( + x.dtype == torch.bfloat16 or x.dtype == torch.float16 + ), "W4A16 gemv only support bfloat16 or float16!" + use_int4_w4a16 = True + out_shape = x.shape[:-1] + ( + qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2, + ) + output = torch.empty(out_shape, device=x.device, dtype=x.dtype) + torch.ops.sgl_kernel.musa_fused_gemv( + x, + qweight, + output, + None, + qweight_scales, + use_int4_w4a16, + use_swigelu, + use_rms_norm, + gamma, + eps, + ) + return output + # general gemv + else: + output = torch.empty(out_shape, device=x.device, dtype=x.dtype) + torch.ops.sgl_kernel.musa_fused_gemv( + x, + qweight, + output, + None, + None, + use_int4_w4a16, + use_swigelu, + use_rms_norm, + gamma, + eps, + ) + return output + + +def musa_fused_mul_add( + self: torch.Tensor, + bias: Optional[torch.Tensor], + scale: Optional[float], + accurate: bool = True, +): + # if accurate == False, then we call inplace op: bias += (self * scale) + if not accurate: + bias.add_(self, alpha=scale) + return bias + + # otherwise, we call custom outplace op, act: output = self * scale + bias + output = torch.empty_like(self) + torch.ops.sgl_kernel.musa_fused_mul_add(output, self, bias, scale) + + return output diff --git a/sgl-kernel/setup_musa.py b/sgl-kernel/setup_musa.py index cdfcf87b4..4f204dccb 100644 --- a/sgl-kernel/setup_musa.py +++ b/sgl-kernel/setup_musa.py @@ -81,6 +81,7 @@ sources = [ "csrc/common_extension_musa.cc", "csrc/elementwise/activation.cu", "csrc/elementwise/concat_mla.cu", + "csrc/elementwise/pos_enc.cu", "csrc/elementwise/fused_add_rms_norm_kernel.mu", "csrc/grammar/apply_token_bitmask_inplace_cuda.cu", "csrc/moe/moe_align_kernel.cu", @@ -107,6 +108,11 @@ sources = [ "csrc/memory/weak_ref_tensor.cpp", str(_FLASHINFER_REPO.source_dir / "csrc/norm.cu"), str(_FLASHINFER_REPO.source_dir / "csrc/renorm.cu"), + # XXX (MUSA): The following files contain MUSA-specific implementations. + "csrc/musa/pos_encoding_contiguous.mu", + "csrc/musa/moe_gemv_swiglu.mu", + "csrc/musa/ternary.mu", + "csrc/musa/top_k_top_p_sampling.mu", ] cxx_flags = ["force_mcc"]