Skip to content

GitLab

"...git@developer.sourcefind.cn:2222/OpenDAS/vllm_cscc.git" did not exist on "8f0d5167155247934d247eb10ae086108db8d473"
Unverified Commit 1485396a authored by Дзержи́нский's avatar Дзержи́нский Committed by GitHub
Browse files

[Kernel] Apply 256bit LDG/STG To Activation Kernels (#33022) 


Signed-off-by: default avatarDzerzhinsky <256908701+AstroVoyager7@users.noreply.github.com>
Signed-off-by: default avatarДзержи́нский <256908701+AstroVoyager7@users.noreply.github.com>
Co-authored-by: default avatarWentao Ye <44945378+yewentao256@users.noreply.github.com>
parent 5ee5c86e
Hide whitespace changes
Inline Side-by-side
Showing with 401 additions and 123 deletions
csrc/activation_kernels.cu
View file @ 1485396a
...@@ -9,6 +9,111 @@ ...@@ -9,6 +9,111 @@
namespace vllm { namespace vllm {
struct alignas(32) u32x8_t {
uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
};
__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
: "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
"=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
: "l"(ptr));
#else
const uint4* uint_ptr = reinterpret_cast<const uint4*>(ptr);
uint4 top_half = __ldg(&uint_ptr[0]);
uint4 bottom_half = __ldg(&uint_ptr[1]);
val.u0 = top_half.x;
val.u1 = top_half.y;
val.u2 = top_half.z;
val.u3 = top_half.w;
val.u4 = bottom_half.x;
val.u5 = bottom_half.y;
val.u6 = bottom_half.z;
val.u7 = bottom_half.w;
#endif
}
__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
:
: "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
"r"(val.u4), "r"(val.u5), "r"(val.u6), "r"(val.u7)
: "memory");
#else
uint4* uint_ptr = reinterpret_cast<uint4*>(ptr);
uint_ptr[0] = make_uint4(val.u0, val.u1, val.u2, val.u3);
uint_ptr[1] = make_uint4(val.u4, val.u5, val.u6, val.u7);
#endif
}
template <bool support_256>
struct VecTraits;
template <>
struct VecTraits<true> {
static constexpr int ARCH_MAX_VEC_SIZE = 32;
using vec_t = u32x8_t;
};
template <>
struct VecTraits<false> {
static constexpr int ARCH_MAX_VEC_SIZE = 16;
using vec_t = int4;
};
template <typename T>
struct PackedTraits;
template <>
struct PackedTraits<c10::BFloat16> {
using packed_t = __nv_bfloat162;
};
template <>
struct PackedTraits<c10::Half> {
using packed_t = __half2;
};
template <>
struct PackedTraits<float> {
using packed_t = float2;
};
template <typename packed_t>
__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
return __bfloat1622float2(val);
} else if constexpr (std::is_same_v<packed_t, __half2>) {
return __half22float2(val);
} else if constexpr (std::is_same_v<packed_t, float2>) {
return float2(val);
}
}
template <typename packed_t>
__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
return __float22bfloat162_rn(val);
} else if constexpr (std::is_same_v<packed_t, __half2>) {
return __float22half2_rn(val);
} else if constexpr (std::is_same_v<packed_t, float2>) {
return float2(val);
}
}
template <typename packed_t>
__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
const packed_t& y) {
if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
std::is_same_v<packed_t, __half2>) {
return __hmul2(x, y);
} else if constexpr (std::is_same_v<packed_t, float2>) {
return make_float2(x.x * y.x, x.y * y.y);
}
}
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
bool act_first> bool act_first>
__device__ __forceinline__ scalar_t compute(const scalar_t& x, __device__ __forceinline__ scalar_t compute(const scalar_t& x,
...@@ -16,52 +121,69 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x, ...@@ -16,52 +121,69 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
return act_first ? ACT_FN(x) * y : x * ACT_FN(y); return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
} }
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
bool act_first>
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
const packed_t& y) {
return act_first ? packed_mul(PACKED_ACT_FN(x), y)
: packed_mul(x, PACKED_ACT_FN(y));
}
// Check if all pointers are 16-byte aligned for int4 vectorized access // Check if all pointers are 16-byte aligned for int4 vectorized access
__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) { __host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0; return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
} }
// Check if all pointers are 16-byte aligned for longlong4_32a vectorized access
__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
}
// Activation and gating kernel template. // Activation and gating kernel template.
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), template <typename scalar_t, typename packed_t,
bool act_first> scalar_t (*ACT_FN)(const scalar_t&),
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
bool use_vec, bool use_256b = false>
__global__ void act_and_mul_kernel( __global__ void act_and_mul_kernel(
scalar_t* __restrict__ out, // [..., d] scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., 2, d] const scalar_t* __restrict__ input, // [..., 2, d]
const int d) { const int d) {
constexpr int VEC_SIZE = 16 / sizeof(scalar_t); const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
const int64_t token_idx = blockIdx.x;
const scalar_t* x_ptr = input + token_idx * 2 * d;
const scalar_t* y_ptr = x_ptr + d; const scalar_t* y_ptr = x_ptr + d;
scalar_t* out_ptr = out + token_idx * d; scalar_t* out_ptr = out + blockIdx.x * d;
// Check alignment for 128-bit vectorized access. if constexpr (use_vec) {
// All three pointers must be 16-byte aligned for safe int4 operations. // Fast path: 128-bit/256-bit vectorized loop
const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) && using vec_t = typename VecTraits<use_256b>::vec_t;
is_16byte_aligned(out_ptr); constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
if (aligned && d >= VEC_SIZE) { const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
// Fast path: 128-bit vectorized loop const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
const int4* x_vec = reinterpret_cast<const int4*>(x_ptr); vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
const int4* y_vec = reinterpret_cast<const int4*>(y_ptr); const int num_vecs = d / 2 / VEC_SIZE;
int4* out_vec = reinterpret_cast<int4*>(out_ptr);
const int num_vecs = d / VEC_SIZE;
const int vec_end = num_vecs * VEC_SIZE;
for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r; vec_t x, y;
auto* xp = reinterpret_cast<scalar_t*>(&x); if constexpr (use_256b) {
auto* yp = reinterpret_cast<scalar_t*>(&y); ld256(x, &x_vec[i]);
auto* rp = reinterpret_cast<scalar_t*>(&r); ld256(y, &y_vec[i]);
} else {
x = VLLM_LDG(&x_vec[i]);
y = VLLM_LDG(&y_vec[i]);
}
auto* xp = reinterpret_cast<packed_t*>(&x);
auto* yp = reinterpret_cast<packed_t*>(&y);
#pragma unroll #pragma unroll
for (int j = 0; j < VEC_SIZE; j++) { for (int j = 0; j < VEC_SIZE; j++) {
rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]); xp[j] =
packed_compute<packed_t, PACKED_ACT_FN, act_first>(xp[j], yp[j]);
}
if constexpr (use_256b) {
st256(x, &out_vec[i]);
} else {
out_vec[i] = x;
} }
out_vec[i] = r;
}
// Scalar cleanup for remaining elements
for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
out_ptr[i] = compute<scalar_t, ACT_FN, act_first>(VLLM_LDG(&x_ptr[i]),
VLLM_LDG(&y_ptr[i]));
} }
} else { } else {
// Scalar fallback for unaligned data or small d // Scalar fallback for unaligned data or small d
...@@ -79,6 +201,15 @@ __device__ __forceinline__ T silu_kernel(const T& x) { ...@@ -79,6 +201,15 @@ __device__ __forceinline__ T silu_kernel(const T& x) {
return (T)(((float)x) / (1.0f + expf((float)-x))); return (T)(((float)x) / (1.0f + expf((float)-x)));
} }
template <typename packed_t>
__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
// x * sigmoid(x)
float2 fval = cast_to_float2(val);
fval.x = fval.x / (1.0f + expf(-fval.x));
fval.y = fval.y / (1.0f + expf(-fval.y));
return cast_to_packed<packed_t>(fval);
}
template <typename T> template <typename T>
__device__ __forceinline__ T gelu_kernel(const T& x) { __device__ __forceinline__ T gelu_kernel(const T& x) {
// Equivalent to PyTorch GELU with 'none' approximation. // Equivalent to PyTorch GELU with 'none' approximation.
...@@ -89,6 +220,18 @@ __device__ __forceinline__ T gelu_kernel(const T& x) { ...@@ -89,6 +220,18 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA))); return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA)));
} }
template <typename packed_t>
__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
// Equivalent to PyTorch GELU with 'none' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
constexpr float ALPHA = M_SQRT1_2;
float2 fval = cast_to_float2(val);
fval.x = fval.x * 0.5f * (1.0f + ::erf(fval.x * ALPHA));
fval.y = fval.y * 0.5f * (1.0f + ::erf(fval.y * ALPHA));
return cast_to_packed<packed_t>(fval);
}
template <typename T> template <typename T>
__device__ __forceinline__ T gelu_tanh_kernel(const T& x) { __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
// Equivalent to PyTorch GELU with 'tanh' approximation. // Equivalent to PyTorch GELU with 'tanh' approximation.
...@@ -102,32 +245,83 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) { ...@@ -102,32 +245,83 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
return (T)(0.5f * f * (1.0f + ::tanhf(inner))); return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
} }
template <typename packed_t>
__device__ __forceinline__ packed_t
packed_gelu_tanh_kernel(const packed_t& val) {
// Equivalent to PyTorch GELU with 'tanh' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
float2 fval = cast_to_float2(val);
constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
constexpr float KAPPA = 0.044715;
float x_cube = fval.x * fval.x * fval.x;
float inner = BETA * (fval.x + KAPPA * x_cube);
fval.x = 0.5f * fval.x * (1.0f + ::tanhf(inner));
x_cube = fval.y * fval.y * fval.y;
inner = BETA * (fval.y + KAPPA * x_cube);
fval.y = 0.5f * fval.y * (1.0f + ::tanhf(inner));
return cast_to_packed<packed_t>(fval);
}
} // namespace vllm } // namespace vllm
// Launch activation and gating kernel. // Launch activation and gating kernel.
// Use ACT_FIRST (bool) indicating whether to apply the activation function // Use ACT_FIRST (bool) indicating whether to apply the activation function
// first. // first.
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST) \ #define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
int d = input.size(-1) / 2; \ auto dtype = input.scalar_type(); \
int64_t num_tokens = input.numel() / input.size(-1); \ int d = input.size(-1) / 2; \
dim3 grid(num_tokens); \ int64_t num_tokens = input.numel() / input.size(-1); \
dim3 block(std::min(d, 1024)); \ if (num_tokens == 0) { \
if (num_tokens == 0) { \ return; \
return; \ } \
} \ dim3 grid(num_tokens); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \ int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \ int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
VLLM_DISPATCH_FLOATING_TYPES( \ int vec_size = support_vec / at::elementSize(dtype); \
input.scalar_type(), "act_and_mul_kernel", [&] { \ const bool use_vec = (d % vec_size == 0); \
vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST> \ const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \ const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
input.data_ptr<scalar_t>(), d); \ if (use_vec) { \
}); dim3 block(std::min(d / vec_size, 1024)); \
if (cc_major >= 10 && num_tokens > 128) { \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
vllm::act_and_mul_kernel< \
scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
} else { \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
vllm::act_and_mul_kernel< \
scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
} \
} else { \
dim3 block(std::min(d, 1024)); \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
vllm::act_and_mul_kernel< \
scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
ACT_FIRST, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
}); \
}
void silu_and_mul(torch::Tensor& out, // [..., d] void silu_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d] torch::Tensor& input) // [..., 2 * d]
{ {
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, true); LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
true);
} }
void mul_and_silu(torch::Tensor& out, // [..., d] void mul_and_silu(torch::Tensor& out, // [..., d]
...@@ -135,19 +329,22 @@ void mul_and_silu(torch::Tensor& out, // [..., d] ...@@ -135,19 +329,22 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
{ {
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
// applies the silu to the latter half of the input. // applies the silu to the latter half of the input.
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, false); LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
false);
} }
void gelu_and_mul(torch::Tensor& out, // [..., d] void gelu_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d] torch::Tensor& input) // [..., 2 * d]
{ {
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, true); LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
true);
} }
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d] void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d] torch::Tensor& input) // [..., 2 * d]
{ {
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel, true); LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
vllm::packed_gelu_tanh_kernel, true);
} }
namespace vllm { namespace vllm {
...@@ -158,42 +355,57 @@ __device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) { ...@@ -158,42 +355,57 @@ __device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) {
return (T)(f > threshold ? f : 0.0f); return (T)(f > threshold ? f : 0.0f);
} }
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float)> template <typename packed_t>
__device__ __forceinline__ packed_t
packed_fatrelu_kernel(const packed_t& val, const float threshold) {
float2 fval = cast_to_float2(val);
fval.x = fval.x > threshold ? fval.x : 0.0f;
fval.y = fval.y > threshold ? fval.y : 0.0f;
return cast_to_packed<packed_t>(fval);
}
template <typename scalar_t, typename packed_t,
scalar_t (*ACT_FN)(const scalar_t&, const float),
packed_t (*PACKED_ACT_FN)(const packed_t&, const float), bool use_vec,
bool use_256b = false>
__global__ void act_and_mul_kernel_with_param( __global__ void act_and_mul_kernel_with_param(
scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d, scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
const float param) { const float param) {
constexpr int VEC_SIZE = 16 / sizeof(scalar_t); const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
const int64_t token_idx = blockIdx.x;
const scalar_t* x_ptr = input + token_idx * 2 * d;
const scalar_t* y_ptr = x_ptr + d; const scalar_t* y_ptr = x_ptr + d;
scalar_t* out_ptr = out + token_idx * d; scalar_t* out_ptr = out + blockIdx.x * d;
// Check alignment for 128-bit vectorized access if constexpr (use_vec) {
const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) && // Fast path: 128-bit/256-bit vectorized loop
is_16byte_aligned(out_ptr); using vec_t = typename VecTraits<use_256b>::vec_t;
constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
if (aligned && d >= VEC_SIZE) { const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
// Fast path: 128-bit vectorized loop const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
const int4* x_vec = reinterpret_cast<const int4*>(x_ptr); vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
const int4* y_vec = reinterpret_cast<const int4*>(y_ptr); const int num_vecs = d / 2 / VEC_SIZE;
int4* out_vec = reinterpret_cast<int4*>(out_ptr);
const int num_vecs = d / VEC_SIZE;
const int vec_end = num_vecs * VEC_SIZE;
for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r; vec_t x, y;
auto* xp = reinterpret_cast<scalar_t*>(&x); if constexpr (use_256b) {
auto* yp = reinterpret_cast<scalar_t*>(&y); ld256(x, &x_vec[i]);
auto* rp = reinterpret_cast<scalar_t*>(&r); ld256(y, &y_vec[i]);
} else {
x = VLLM_LDG(&x_vec[i]);
y = VLLM_LDG(&y_vec[i]);
}
auto* xp = reinterpret_cast<packed_t*>(&x);
auto* yp = reinterpret_cast<packed_t*>(&y);
#pragma unroll #pragma unroll
for (int j = 0; j < VEC_SIZE; j++) { for (int j = 0; j < VEC_SIZE; j++) {
rp[j] = ACT_FN(xp[j], param) * yp[j]; xp[j] = packed_mul(PACKED_ACT_FN(xp[j], param), yp[j]);
}
if constexpr (use_256b) {
st256(x, &out_vec[i]);
} else {
out_vec[i] = x;
} }
out_vec[i] = r;
}
// Scalar cleanup for remaining elements
for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
out_ptr[i] = ACT_FN(VLLM_LDG(&x_ptr[i]), param) * VLLM_LDG(&y_ptr[i]);
} }
} else { } else {
// Scalar fallback for unaligned data or small d // Scalar fallback for unaligned data or small d
...@@ -276,20 +488,58 @@ __global__ void swigluoai_and_mul_kernel( ...@@ -276,20 +488,58 @@ __global__ void swigluoai_and_mul_kernel(
} // namespace vllm } // namespace vllm
#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PARAM) \ #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
int d = input.size(-1) / 2; \ auto dtype = input.scalar_type(); \
int64_t num_tokens = input.numel() / input.size(-1); \ int d = input.size(-1) / 2; \
dim3 grid(num_tokens); \ int64_t num_tokens = input.numel() / input.size(-1); \
dim3 block(std::min(d, 1024)); \ if (num_tokens == 0) { \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \ return; \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \ } \
VLLM_DISPATCH_FLOATING_TYPES( \ dim3 grid(num_tokens); \
input.scalar_type(), "act_and_mul_kernel_with_param", [&] { \ int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
vllm::act_and_mul_kernel_with_param<scalar_t, KERNEL<scalar_t>> \ int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \ int vec_size = support_vec / at::elementSize(dtype); \
input.data_ptr<scalar_t>(), d, \ const bool use_vec = (d % vec_size == 0); \
PARAM); \ const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
}); const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
if (use_vec) { \
dim3 block(std::min(d / vec_size, 1024)); \
if (cc_major >= 10 && num_tokens > 128) { \
VLLM_DISPATCH_FLOATING_TYPES( \
dtype, "act_and_mul_kernel_with_param", [&] { \
vllm::act_and_mul_kernel_with_param< \
scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
KERNEL<scalar_t>, \
PACKED_KERNEL< \
typename vllm::PackedTraits<scalar_t>::packed_t>, \
true, true><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, \
PARAM); \
}); \
} else { \
VLLM_DISPATCH_FLOATING_TYPES( \
dtype, "act_and_mul_kernel_with_param", [&] { \
vllm::act_and_mul_kernel_with_param< \
scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
KERNEL<scalar_t>, \
PACKED_KERNEL< \
typename vllm::PackedTraits<scalar_t>::packed_t>, \
true, false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, \
PARAM); \
}); \
} \
} else { \
dim3 block(std::min(d, 1024)); \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
vllm::act_and_mul_kernel_with_param< \
scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
KERNEL<scalar_t>, \
PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
false><<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM); \
}); \
}
#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \ #define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \
int d = input.size(-1) / 2; \ int d = input.size(-1) / 2; \
...@@ -309,7 +559,8 @@ __global__ void swigluoai_and_mul_kernel( ...@@ -309,7 +559,8 @@ __global__ void swigluoai_and_mul_kernel(
void fatrelu_and_mul(torch::Tensor& out, // [..., d], void fatrelu_and_mul(torch::Tensor& out, // [..., d],
torch::Tensor& input, // [..., 2 * d] torch::Tensor& input, // [..., 2 * d]
double threshold) { double threshold) {
LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(vllm::fatrelu_kernel, threshold); LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
} }
void swigluoai_and_mul(torch::Tensor& out, // [..., d] void swigluoai_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input, // [..., 2 * d] torch::Tensor& input, // [..., 2 * d]
...@@ -319,39 +570,41 @@ void swigluoai_and_mul(torch::Tensor& out, // [..., d] ...@@ -319,39 +570,41 @@ void swigluoai_and_mul(torch::Tensor& out, // [..., d]
namespace vllm { namespace vllm {
// Element-wise activation kernel template. // Element-wise activation kernel template.
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&)> template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), bool use_vec,
bool use_256b = false>
__global__ void activation_kernel( __global__ void activation_kernel(
scalar_t* __restrict__ out, // [..., d] scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., d] const scalar_t* __restrict__ input, // [..., d]
const int d) { const int d) {
constexpr int VEC_SIZE = 16 / sizeof(scalar_t); const scalar_t* in_ptr = input + blockIdx.x * d;
const int64_t token_idx = blockIdx.x; scalar_t* out_ptr = out + blockIdx.x * d;
const scalar_t* in_ptr = input + token_idx * d;
scalar_t* out_ptr = out + token_idx * d; if constexpr (use_vec) {
// Fast path: 128-bit/256-bit vectorized loop
// Check alignment for 128-bit vectorized access using vec_t = typename VecTraits<use_256b>::vec_t;
const bool aligned = is_16byte_aligned(in_ptr) && is_16byte_aligned(out_ptr); constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(scalar_t);
if (aligned && d >= VEC_SIZE) { const vec_t* in_vec = reinterpret_cast<const vec_t*>(in_ptr);
// Fast path: 128-bit vectorized loop vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
const int4* in_vec = reinterpret_cast<const int4*>(in_ptr);
int4* out_vec = reinterpret_cast<int4*>(out_ptr);
const int num_vecs = d / VEC_SIZE; const int num_vecs = d / VEC_SIZE;
const int vec_end = num_vecs * VEC_SIZE;
for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) { for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
int4 v = VLLM_LDG(&in_vec[i]), r; vec_t v;
if constexpr (use_256b) {
ld256(v, &in_vec[i]);
} else {
v = VLLM_LDG(&in_vec[i]);
}
auto* vp = reinterpret_cast<scalar_t*>(&v); auto* vp = reinterpret_cast<scalar_t*>(&v);
auto* rp = reinterpret_cast<scalar_t*>(&r);
#pragma unroll #pragma unroll
for (int j = 0; j < VEC_SIZE; j++) { for (int j = 0; j < VEC_SIZE; j++) {
rp[j] = ACT_FN(vp[j]); vp[j] = ACT_FN(vp[j]);
}
if constexpr (use_256b) {
st256(v, &out_vec[i]);
} else {
out_vec[i] = v;
} }
out_vec[i] = r;
}
// Scalar cleanup for remaining elements
for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
out_ptr[i] = ACT_FN(VLLM_LDG(&in_ptr[i]));
} }
} else { } else {
// Scalar fallback for unaligned data or small d // Scalar fallback for unaligned data or small d
...@@ -365,18 +618,43 @@ __global__ void activation_kernel( ...@@ -365,18 +618,43 @@ __global__ void activation_kernel(
} // namespace vllm } // namespace vllm
// Launch element-wise activation kernel. // Launch element-wise activation kernel.
#define LAUNCH_ACTIVATION_KERNEL(KERNEL) \ #define LAUNCH_ACTIVATION_KERNEL(KERNEL) \
int d = input.size(-1); \ auto dtype = input.scalar_type(); \
int64_t num_tokens = input.numel() / d; \ int d = input.size(-1); \
dim3 grid(num_tokens); \ int64_t num_tokens = input.numel() / input.size(-1); \
dim3 block(std::min(d, 1024)); \ if (num_tokens == 0) { \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \ return; \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \ } \
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \ dim3 grid(num_tokens); \
vllm::activation_kernel<scalar_t, KERNEL<scalar_t>> \ int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \ int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
input.data_ptr<scalar_t>(), d); \ int vec_size = support_vec / at::elementSize(dtype); \
}); const bool use_vec = (d % vec_size == 0); \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
if (use_vec) { \
dim3 block(std::min(d / vec_size, 1024)); \
if (cc_major >= 10 && num_tokens > 128) { \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true> \
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
input.data_ptr<scalar_t>(), d); \
}); \
} else { \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
input.data_ptr<scalar_t>(), d); \
}); \
} \
} else { \
dim3 block(std::min(d, 1024)); \
VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false> \
<<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
input.data_ptr<scalar_t>(), d); \
}); \
}
namespace vllm { namespace vllm {
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment