minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[MUSA][18/N] Add MUSA-optimized kernel implementations for hot ops (#23255)

Browse Source 
Signed-off-by: Joey-gvwal <joey_gvwal@yeah.net>
Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com>
This commit is contained in:
Joey
2026-05-07 20:38:33 -07:00
committed by GitHub
co-authored by R0CKSTAR
parent 461bc8af49
commit 15e6572f21
15 changed files with 2513 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-kernel/csrc/common_extension_musa.cc
+42
View File
@@ -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
*/
sgl-kernel/csrc/elementwise/fused_add_rms_norm_kernel.mu
+17 -8
View File
@@ -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 <typename type>
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<SrcDtype, vlen * sizeof(SrcDtype) * 8>;
using ComputeVec = VecType<ComputeType, vlen * sizeof(ComputeType) * 8>;
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<ComputeType, BLOCK_X, BLOCK_Y, 1> 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> \
<<<grid_size, block_size, 0, stream>>>( \
<<<grid_size, block_size, n * sizeof(float), stream>>>( \
static_cast<_SRC_DTYPE*>(input), \
static_cast<_SRC_DTYPE*>(residual), \
static_cast<_SRC_DTYPE*>(weight), \
sgl-kernel/csrc/musa/common.muh
+42
View File
@@ -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 <functional>
#include <iostream>
#include <musa_runtime.h>
#include <sstream>
#include <string>
/// Returns the ceiling of (a / b)
__device__ __host__ __forceinline__ constexpr int ceil_div(int a, int b) {
return (a + b - 1) / b;
}
template <typename T> __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
sgl-kernel/csrc/musa/dtype.muh
+446
View File
@@ -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 <musa_runtime.h>
#include <stdint.h>
// #include "mudnn_config.h"
// #include "mudnn/api/impl_base.h"
#include "musa/integer_subbyte.h"
#ifdef __MTGPU__
#include <musa_robust.h>
#endif
#include "musa_fp16.h"
#include "musa_bf16.h"
#include "musa_fp8.h"
#include <torch/all.h>
namespace musa {
namespace dnn {
template <typename T> struct sizeof_bits {
static constexpr size_t value = sizeof(T) * 8;
};
template <int Bits, bool Signed>
struct sizeof_bits<integer_subbyte<Bits, Signed>> {
static constexpr size_t value = Bits;
};
template <typename T>
static constexpr int sizeof_bits_v = sizeof_bits<T>::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 <typename type> 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 <typename T> struct DeduceVectorizedType {
using Type = T;
};
template <> struct DeduceVectorizedType<bool> {
using Type = int8_t;
};
template <> struct DeduceVectorizedType<half> {
using Type = _Float16;
};
#if defined(__MUSACC__) && (__MUSA_ARCH__ >= 220 || !defined(__MUSA_ARCH__))
template <> struct DeduceVectorizedType<bfloat16_t> {
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 <typename T, int bits = 16 * 8> 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<BaseType>::Type; \
typedef VectorizedType VxTy __attribute__((vector_size(vec_bytes))); \
template <typename OffsetType> \
static __device__ __forceinline__ VecType load(const BaseType *ptr, \
OffsetType idx) { \
return *(VecType *)(ptr + idx); \
} \
template <typename OffsetType> \
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 <typename OffsetType> \
static __device__ __forceinline__ VecType \
robust_load(const RobustTypePtr ptr, OffsetType idx) { \
return __musa::robust_load<VecType, BaseType>(ptr, idx); \
} \
\
template <typename OffsetType> \
static __device__ __forceinline__ void \
store(BaseType *ptr, OffsetType idx, const VecType &dst) { \
*(VecType *)(ptr + idx) = dst; \
} \
template <typename OffsetType> \
static __device__ __forceinline__ void \
robust_store(RobustTypePtr ptr, OffsetType idx, const VecType &dst) { \
__musa::robust_store<VecType, BaseType>(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 <typename SrcVecType> \
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<BaseType>(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 <typename Func> __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 <typename SrcVecType> \
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 <typename T> struct ComputeDType {
using Type = typename std::conditional<
(sizeof(T) >= 4), T,
typename std::conditional<
std::is_integral<T>::value,
typename std::conditional<std::is_unsigned<T>::value, uint32_t,
int32_t>::type,
float>::type>::type;
};
template <> struct ComputeDType<bool> {
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 <typename T, typename... Types>
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
sgl-kernel/csrc/musa/moe_gemv_swiglu.mu
+846
View File
@@ -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 <musa_runtime.h>
#include <cassert>
#include <mutex>
#include <musa_bf16.h>
#include <musa_fp16.h>
#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> \
<<<grid_size, block_size, shmem_size, stream>>>( \
static_cast<_CDTYPE*>(C.data_ptr()), \
static_cast<_ADTYPE*>(A.data_ptr()), \
static_cast<_BDTYPE*>(B.data_ptr()), \
static_cast<int*>(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<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(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> \
<<<grid_size, block_size, shmem_size, stream>>>( \
static_cast<_CDTYPE*>(C.data_ptr()), \
static_cast<_ADTYPE*>(A.data_ptr()), \
static_cast<_BDTYPE*>(B.data_ptr()), \
static_cast<int*>(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<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(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> \
<<<grid_size, block_size, shmem_size, stream>>>( \
static_cast<_ADTYPE*>(C.data_ptr()), \
static_cast<_ADTYPE*>(A.data_ptr()), \
static_cast<_BDTYPE*>(B.data_ptr()), \
static_cast<int*>(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<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(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> \
<<<grid_size, block_size, shmem_size, stream>>>( \
static_cast<_ADTYPE*>(C.data_ptr()), \
static_cast<_ADTYPE*>(A.data_ptr()), \
static_cast<_BDTYPE*>(B.data_ptr()), \
static_cast<int*>(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<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(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> \
<<<grid_size, block_size, shmem_size, stream>>>( \
static_cast<_ADTYPE*>(C.data_ptr()), \
static_cast<_ADTYPE*>(A.data_ptr()), \
static_cast<_BDTYPE*>(B.data_ptr()), \
static_cast<int*>(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<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(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> \
<<<grid_size, block_size, shmem_size, stream>>>( \
static_cast<_ADTYPE*>(C.data_ptr()), \
static_cast<_ADTYPE*>(A.data_ptr()), \
static_cast<_BDTYPE*>(B.data_ptr()), \
static_cast<int*>(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<bfloat16_t*>(rms_gamma_ptr), static_cast<float*>(rms_sum_out_ptr), static_cast<int*>(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 <typename AType, typename BType, typename CType, typename ScoreType, typename ScaleType,
int BLOCK_N, int BLOCK_K, int iobit, bool mul_routed_weight, bool is_swigelu,
bool is_w4a16, bool is_per_group_scale, bool is_fp8, int scale_block, bool use_rms_norm>
__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<AType, __mt_fp8_e4m3>);
using AVecSType = std::conditional_t<std::is_same_v<AType, __mt_fp8_e4m3>, uint8_t, AType>;
using BVecSType = std::conditional_t<is_fp8, uint8_t, BType>;
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<AVecSType, 128>::Ttype
>
>;
using BVec = typename VecType<BVecSType, b_vec_bits>::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<const fp8x4_vec*>(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<const fp8x4_vec*>(a_reg)[i]);
b_halfv4 = __musa_e4m32f16_rn_bst4(reinterpret_cast<const fp8x4_vec*>(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<torch::Tensor> &A_scale,
const c10::optional<torch::Tensor> &B_scale,
bool use_int4_w4a16,
bool use_swigelu,
bool use_rms_norm,
const c10::optional<torch::Tensor> &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<void()> 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<float>(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<float>(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<torch::Tensor> &A_scale,
const c10::optional<torch::Tensor> &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<void()> 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<float>(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<float>(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();
}
sgl-kernel/csrc/musa/pos_encoding_contiguous.mu
+264
View File
@@ -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 <torch/all.h>
#include "musa.h"
#include "musa/dispatch_utils.h"
#include "torch_musa/csrc/core/MUSAGuard.h"
#include "torch_musa/csrc/core/MUSAStream.h"
template <typename scalar_t, bool IS_NEOX>
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 <typename scalar_t, bool IS_NEOX>
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<scalar_t, IS_NEOX>(
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<scalar_t, IS_NEOX>(
head_key, cos_ptr, sin_ptr, rot_offset, embed_dim, key_dim_stride);
}
}
template <typename scalar_t, bool IS_NEOX>
__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<scalar_t, IS_NEOX>(
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 <typename scalar_t, bool IS_NEOX>
__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<scalar_t, IS_NEOX>(
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<int64_t>(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<scalar_t, true><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.data_ptr<scalar_t>(),
cos_sin_cache.data_ptr<scalar_t>(),
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<scalar_t, false><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.data_ptr<scalar_t>(),
cos_sin_cache.data_ptr<scalar_t>(),
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<int64_t>(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<scalar_t, true><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.data_ptr<scalar_t>(),
cos_sin_cache.data_ptr<scalar_t>(),
cos_sin_cache_offsets.data_ptr<int64_t>(),
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<scalar_t, false><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.data_ptr<scalar_t>(),
cos_sin_cache.data_ptr<scalar_t>(),
cos_sin_cache_offsets.data_ptr<int64_t>(),
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);
}
});
}
sgl-kernel/csrc/musa/ternary.mu
+123
View File
@@ -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 <torch/all.h>
#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 <typename scalar_t, int64_t vlen, int iobit = 128>
__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<scalar_t, vlen * sizeof(scalar_t) * bits_of_byte>;
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<float, 4><<<nr_block, thread_per_block, 0, stream>>>(
static_cast<float *>(output.data_ptr()),
static_cast<float *>(self.data_ptr()),
static_cast<float *>(bias.data_ptr()), scale, numel);
break;
case at::ScalarType::Half:
FusedMulAdd<float16_t, 8>
<<<nr_block, thread_per_block, 0, stream>>>(
static_cast<float16_t *>(output.data_ptr()),
static_cast<float16_t *>(self.data_ptr()),
static_cast<float16_t *>(bias.data_ptr()),
static_cast<float16_t>(scale), numel);
break;
case at::ScalarType::BFloat16:
FusedMulAdd<bfloat16_t, 8>
<<<nr_block, thread_per_block, 0, stream>>>(
static_cast<bfloat16_t *>(output.data_ptr()),
static_cast<bfloat16_t *>(self.data_ptr()),
static_cast<bfloat16_t *>(bias.data_ptr()),
static_cast<bfloat16_t>(scale), numel);
break;
default:
break;
}
}
sgl-kernel/csrc/musa/top_k_top_p_sampling.mu
+382
View File
@@ -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 <ATen/Utils.h>
#include <ATen/core/Generator.h>
#include "torch_musa/csrc/aten/musa/MUSAGeneratorImpl.h"
#include <torch/all.h>
#include "torch_musa/csrc/aten/musa/UnpackRaw.muh"
#include <flashinfer/sampling.muh>
#include <mutex>
#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 <uint32_t VEC_SIZE, uint32_t BLOCK_THREADS, BlockScanAlgorithm SCAN_ALGORITHM,
BlockReduceAlgorithm REDUCE_ALGORITHM, bool DETERMINISTIC, typename Predicate>
__device__ __forceinline__ void DeviceSamplingFromProbWithOffset(
uint32_t i, uint32_t d, Predicate pred, float u, vec_t<float, VEC_SIZE> prob_vec,
float& aggregate,
SamplingTempStorage<BLOCK_THREADS, SCAN_ALGORITHM, REDUCE_ALGORITHM>* 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<float, BLOCK_THREADS, REDUCE_ALGORITHM>(temp_storage->block_prim.reduce)
.template Sum<VEC_SIZE>(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<VEC_SIZE, BLOCK_THREADS, SCAN_ALGORITHM, REDUCE_ALGORITHM>(
prob_greater_than_threshold, inclusive_cdf, temp_storage);
} else {
BlockScan<float, BLOCK_THREADS, SCAN_ALGORITHM>(temp_storage->block_prim.scan)
.template InclusiveSum<VEC_SIZE>(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<bool, BLOCK_THREADS>(temp_storage->block_prim.adj_diff)
.SubtractLeft<VEC_SIZE>(greater_than_u, greater_than_u_diff, BoolDiffOp());
#else
BlockAdjacentDifference<bool, BLOCK_THREADS>(temp_storage->block_prim.adj_diff)
.template FlagHeads<VEC_SIZE>(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<int, BLOCK_THREADS, REDUCE_ALGORITHM>(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 <uint32_t BLOCK_THREADS, BlockScanAlgorithm SCAN_ALGORITHM,
BlockReduceAlgorithm REDUCE_ALGORITHM, uint32_t VEC_SIZE, bool DETERMINISTIC,
typename DType, typename IdType>
__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<BLOCK_THREADS, SCAN_ALGORITHM, REDUCE_ALGORITHM>))
uint8_t smem_sampling[];
auto& temp_storage =
reinterpret_cast<SamplingTempStorage<BLOCK_THREADS, SCAN_ALGORITHM, REDUCE_ALGORITHM>&>(smem_sampling);
vec_t<float, kRmemEles> rprobs; // persistent
vec_t<float, VEC_SIZE> 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<kRmemEles, BLOCK_THREADS, SCAN_ALGORITHM, REDUCE_ALGORITHM, DETERMINISTIC>(
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<VEC_SIZE, BLOCK_THREADS, SCAN_ALGORITHM, REDUCE_ALGORITHM, DETERMINISTIC>(
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<float> aggregate_gt_pivot_0{0, 0}, aggregate_gt_pivot_1{0, 0};
// check in rprobs
ValueCount<float> 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<ValueCount<float>, BLOCK_THREADS>(temp_storage.block_prim.reduce_value_count)
.template Sum<VEC_SIZE>(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<ValueCount<float>, BLOCK_THREADS>(temp_storage.block_prim.reduce_value_count)
.template Sum<VEC_SIZE>(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 <typename T, typename IdType>
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<BLOCK_THREADS, SCAN_ALGO, REDUCE_ALGO>);
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<at::Tensor> maybe_indices,
std::optional<at::Tensor> maybe_top_k_arr,
double top_k_val,
std::optional<at::Tensor> maybe_top_p_arr,
double top_p_val,
bool deterministic,
std::optional<at::Generator> 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<at::MUSAGeneratorImpl>(gen_, at::musa::detail::getDefaultMUSAGenerator());
std::lock_guard<std::mutex> 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<float, int>(
static_cast<float*>(probs.data_ptr()),
has_top_k_arr ? static_cast<int*>(maybe_top_k_arr->data_ptr()) : nullptr,
has_top_p_arr ? static_cast<float*>(maybe_top_p_arr->data_ptr()) : nullptr,
static_cast<int*>(output.data_ptr()),
maybe_indices.has_value() ? static_cast<int*>(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)));
}
sgl-kernel/include/musa/dispatch_utils.h
+34
View File
@@ -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 <torch/all.h>
#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__))
sgl-kernel/include/musa/integer_subbyte.h
+49
View File
@@ -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 <limits>
#include <type_traits>
namespace musa::dnn {
// cutlass integer_subbyte class
template <int Bits, bool Signed = true>
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<Signed, int, unsigned>::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<Storage const&>(value) & bits_mask_) {}
__host__ __device__ constexpr integer_subbyte(unsigned value)
: storage(reinterpret_cast<Storage const&>(value) & bits_mask_) {}
};
} // namespace musa::dnn
sgl-kernel/include/sgl_kernel_musa_ops.h
+80
View File
@@ -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 <ATen/ATen.h>
#include <ATen/Tensor.h>
#include <torch/torch.h>
#include <optional>
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<torch::Tensor>& A_scale,
const c10::optional<torch::Tensor>& 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<torch::Tensor>& A_scale,
const c10::optional<torch::Tensor>& B_scale,
bool use_int4_w4a16,
bool use_swigelu,
bool use_rms_norm,
const c10::optional<torch::Tensor>& 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<at::Tensor> maybe_indices,
std::optional<at::Tensor> maybe_top_k_arr,
double top_k_val,
std::optional<at::Tensor> maybe_top_p_arr,
double top_p_val,
bool deterministic,
std::optional<at::Generator> gen);
sgl-kernel/include/sgl_kernel_ops.h
+4
View File
@@ -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)
sgl-kernel/python/sgl_kernel/__init__.py
+9
View File
@@ -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",
sgl-kernel/python/sgl_kernel/musa.py
+169
View File
@@ -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
sgl-kernel/setup_musa.py
+6
View File
@@ -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"]