Skip to content

GitLab

"scripts/generation.sh" did not exist on "870fd841a8879d8394910f76b89eaeaa11d1c6ed"
Unverified Commit 37a27712 authored by Muyang Li's avatar Muyang Li Committed by GitHub
Browse files

Merge pull request #340 from mit-han-lab/dev 

feat: support PuLID, Double FBCache and TeaCache; better linter
parents c1d6fc84 760ab022
Hide whitespace changes
Inline Side-by-side
Showing with 2260 additions and 2216 deletions
src/kernels/gemm_f16.cu
View file @ 37a27712
......@@ -14,10 +14,9 @@ using spdlog::fmt_lib::format;
Tensor gemm_f16(Tensor input, // FP16
Tensor weight, // FP16
Tensor out, // FP16
Tensor out, // FP16
Tensor bias,
float alpha
) {
float alpha) {
auto N = weight.size(0);
auto K = input.size(-1);
auto M = input.numel() / K;
......@@ -26,30 +25,34 @@ Tensor gemm_f16(Tensor input, // FP16
spdlog::debug("gemm_f16: M={} K={} N={}", M, K, N);
dispatchF16(weight.dtype(), [&]<typename half_t>() {
using ElementOutput = half_t;
using ElementAccumulator = float;
using ElementOutput = half_t;
using ElementAccumulator = float;
using ElementComputeEpilogue = half_t;
using ElementInputA = half_t; // <- data type of elements in input matrix A
using ElementInputB = half_t; // <- data type of elements in input matrix B
using ElementInputA = half_t; // <- data type of elements in input matrix A
using ElementInputB = half_t; // <- data type of elements in input matrix B
using LayoutInputA = cutlass::layout::RowMajor;
using LayoutInputB = cutlass::layout::ColumnMajor;
using LayoutOutput = cutlass::layout::RowMajor;
// #if CUDA_ARCH >= 800
using Gemm = cutlass::gemm::device::Gemm<
ElementInputA, cutlass::layout::RowMajor, ElementInputB, cutlass::layout::ColumnMajor,
ElementOutput, cutlass::layout::RowMajor, ElementAccumulator,
cutlass::arch::OpClassTensorOp, cutlass::arch::Sm75>;
// cutlass::gemm::GemmShape<128, 128, 64>,
// cutlass::gemm::GemmShape<32, 64, 64>, cutlass::gemm::GemmShape<16, 8, 16>,
// cutlass::epilogue::thread::LinearCombination<
// ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
// ElementAccumulator, ElementComputeEpilogue>,
// cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>;
auto input_size = cutlass::MatrixCoord(M, K);
// #if CUDA_ARCH >= 800
using Gemm = cutlass::gemm::device::Gemm<ElementInputA,
cutlass::layout::RowMajor,
ElementInputB,
cutlass::layout::ColumnMajor,
ElementOutput,
cutlass::layout::RowMajor,
ElementAccumulator,
cutlass::arch::OpClassTensorOp,
cutlass::arch::Sm75>;
// cutlass::gemm::GemmShape<128, 128, 64>,
// cutlass::gemm::GemmShape<32, 64, 64>, cutlass::gemm::GemmShape<16, 8, 16>,
// cutlass::epilogue::thread::LinearCombination<
// ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
// ElementAccumulator, ElementComputeEpilogue>,
// cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>;
auto input_size = cutlass::MatrixCoord(M, K);
auto weight_size = cutlass::MatrixCoord(K, N);
auto output_size = cutlass::MatrixCoord(M, N);
......@@ -59,8 +62,8 @@ Tensor gemm_f16(Tensor input, // FP16
if (!out.valid()) {
auto out_shape = TensorShape(input.shape.dataExtent);
out_shape[-1] = N;
out = Tensor::empty(out_shape, input.scalar_type(), input.device());
out_shape[-1] = N;
out = Tensor::empty(out_shape, input.scalar_type(), input.device());
}
// FIXME: check contiguous of input if dims >= 3
......@@ -83,23 +86,22 @@ Tensor gemm_f16(Tensor input, // FP16
// constexpr int kMetaSizeInBits = Gemm::kMetaSizeInBits;
cutlass::gemm::GemmCoord problem_size(M, N, K);
cutlass::TensorRef<ElementInputA, LayoutInputA> input_ref(
input.data_ptr<ElementInputA>(), LayoutInputA(input.stride(-2)));
cutlass::TensorRef<ElementInputB, LayoutInputB> weight_ref(
weight.data_ptr<ElementInputB>(), LayoutInputB::packed(weight_size));
cutlass::TensorRef<ElementInputA, LayoutInputA> input_ref(input.data_ptr<ElementInputA>(),
LayoutInputA(input.stride(-2)));
cutlass::TensorRef<ElementInputB, LayoutInputB> weight_ref(weight.data_ptr<ElementInputB>(),
LayoutInputB::packed(weight_size));
cutlass::TensorRef<ElementOutput, LayoutOutput> bias_ref(
bias.valid() ? bias.data_ptr<ElementOutput>() : out.data_ptr<ElementOutput>(), LayoutOutput(0));
cutlass::TensorRef<ElementOutput, LayoutOutput> out_ref(
out.data_ptr<ElementOutput>(), LayoutOutput(out.stride(-2)));
typename Gemm::Arguments arguments{
problem_size, // <- problem size of matrix multiplication
input_ref, // <- reference to matrix A on device
weight_ref, // <- reference to matrix B on device
bias_ref, // <- reference to matrix C on device
out_ref, // <- reference to matrix D on device
{ElementOutput(alpha), ElementOutput(bias.valid() ? 1.0f : 0.0f)},
1};
cutlass::TensorRef<ElementOutput, LayoutOutput> out_ref(out.data_ptr<ElementOutput>(),
LayoutOutput(out.stride(-2)));
typename Gemm::Arguments arguments{problem_size, // <- problem size of matrix multiplication
input_ref, // <- reference to matrix A on device
weight_ref, // <- reference to matrix B on device
bias_ref, // <- reference to matrix C on device
out_ref, // <- reference to matrix D on device
{ElementOutput(alpha), ElementOutput(bias.valid() ? 1.0f : 0.0f)},
1};
Gemm gemm_op;
// Using the arguments, query for extra workspace required for matrix
......@@ -127,9 +129,7 @@ Tensor gemm_f16(Tensor input, // FP16
if (status != cutlass::Status::kSuccess) {
throw std::runtime_error("cutlass cannot run");
}
});
return out;
}
src/kernels/gemm_f16.h
View file @ 37a27712
......@@ -3,10 +3,8 @@
#include "common.h"
#include "Tensor.h"
Tensor gemm_f16(
Tensor input, // FP16
Tensor weight, // FP16
Tensor out, // FP16
Tensor bias,
float alpha
);
\ No newline at end of file
Tensor gemm_f16(Tensor input, // FP16
Tensor weight, // FP16
Tensor out, // FP16
Tensor bias,
float alpha);
src/kernels/gemm_w8a8.cu
View file @ 37a27712
......@@ -11,10 +11,10 @@
using spdlog::fmt_lib::format;
Tensor gemm_w8a8_fp16(Tensor input, // INT8
Tensor weight, // INT8
Tensor out, // FP16
half alpha,
half beta // FP16
Tensor weight, // INT8
Tensor out, // FP16
half alpha,
half beta // FP16
) {
auto N = weight.size(0);
auto K = input.size(-1);
......@@ -23,57 +23,66 @@ Tensor gemm_w8a8_fp16(Tensor input, // INT8
spdlog::debug("gemm_w8a8: M={} K={} N={}", M, K, N);
using ElementOutput = cutlass::half_t;
using ElementAccumulator = int32_t;
using ElementOutput = cutlass::half_t;
using ElementAccumulator = int32_t;
using ElementComputeEpilogue = cutlass::half_t;
using ElementInputA = int8_t; // <- data type of elements in input matrix A
using ElementInputB = int8_t; // <- data type of elements in input matrix B
using ElementInputA = int8_t; // <- data type of elements in input matrix A
using ElementInputB = int8_t; // <- data type of elements in input matrix B
using LayoutInputA = cutlass::layout::RowMajor;
using LayoutInputB = cutlass::layout::ColumnMajor;
using LayoutOutput = cutlass::layout::RowMajor;
// #if CUDA_ARCH >= 800
// #if CUDA_ARCH >= 800
using Gemm = cutlass::gemm::device::Gemm<
int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor,
ElementOutput, cutlass::layout::RowMajor, ElementAccumulator,
cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80,
int8_t,
cutlass::layout::RowMajor,
int8_t,
cutlass::layout::ColumnMajor,
ElementOutput,
cutlass::layout::RowMajor,
ElementAccumulator,
cutlass::arch::OpClassTensorOp,
cutlass::arch::Sm80,
cutlass::gemm::GemmShape<128, 128, 64>,
cutlass::gemm::GemmShape<32, 64, 64>, cutlass::gemm::GemmShape<16, 8, 32>,
cutlass::epilogue::thread::LinearCombination<
ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
ElementAccumulator, ElementComputeEpilogue>,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, 3>;
// #elif CUDA_ARCH >= 750
// using DefaultGemmCfg = cutlass::gemm::device::DefaultGemmConfiguration<
// cutlass::arch::OpClassTensorOp, cutlass::arch::Sm75,
// ElementInputA, ElementInputB, ElementOutput, ElementAccumulator>;
// using Gemm = cutlass::gemm::device::Gemm<
// int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor,
// ElementOutput, cutlass::layout::RowMajor, ElementAccumulator,
// cutlass::arch::OpClassTensorOp, cutlass::arch::Sm75,
// DefaultGemmCfg::ThreadblockShape, DefaultGemmCfg::WarpShape,
// DefaultGemmCfg::InstructionShape,
// cutlass::epilogue::thread::LinearCombination<
// ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
// ElementAccumulator, ElementComputeEpilogue>>;
// #elif CUDA_ARCH >= 700
// using DefaultGemmCfg = cutlass::gemm::device::DefaultGemmConfiguration<
// cutlass::arch::OpClassSimt, cutlass::arch::Sm70,
// ElementInputA, ElementInputB, ElementOutput, ElementAccumulator>;
// using Gemm = cutlass::gemm::device::Gemm<
// int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor,
// ElementOutput, cutlass::layout::RowMajor, ElementAccumulator,
// cutlass::arch::OpClassSimt, cutlass::arch::Sm70,
// DefaultGemmCfg::ThreadblockShape, DefaultGemmCfg::WarpShape,
// DefaultGemmCfg::InstructionShape,
// cutlass::epilogue::thread::LinearCombination<
// ElementOutput, 1, ElementAccumulator, ElementComputeEpilogue>>;
// #else
// #error "Unsupported cuda arch"
// #endif
auto input_size = cutlass::MatrixCoord(M, K);
cutlass::gemm::GemmShape<32, 64, 64>,
cutlass::gemm::GemmShape<16, 8, 32>,
cutlass::epilogue::thread::LinearCombination<ElementOutput,
128 / cutlass::sizeof_bits<ElementOutput>::value,
ElementAccumulator,
ElementComputeEpilogue>,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>,
3>;
// #elif CUDA_ARCH >= 750
// using DefaultGemmCfg = cutlass::gemm::device::DefaultGemmConfiguration<
// cutlass::arch::OpClassTensorOp, cutlass::arch::Sm75,
// ElementInputA, ElementInputB, ElementOutput, ElementAccumulator>;
// using Gemm = cutlass::gemm::device::Gemm<
// int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor,
// ElementOutput, cutlass::layout::RowMajor, ElementAccumulator,
// cutlass::arch::OpClassTensorOp, cutlass::arch::Sm75,
// DefaultGemmCfg::ThreadblockShape, DefaultGemmCfg::WarpShape,
// DefaultGemmCfg::InstructionShape,
// cutlass::epilogue::thread::LinearCombination<
// ElementOutput, 128 / cutlass::sizeof_bits<ElementOutput>::value,
// ElementAccumulator, ElementComputeEpilogue>>;
// #elif CUDA_ARCH >= 700
// using DefaultGemmCfg = cutlass::gemm::device::DefaultGemmConfiguration<
// cutlass::arch::OpClassSimt, cutlass::arch::Sm70,
// ElementInputA, ElementInputB, ElementOutput, ElementAccumulator>;
// using Gemm = cutlass::gemm::device::Gemm<
// int8_t, cutlass::layout::RowMajor, int8_t, cutlass::layout::ColumnMajor,
// ElementOutput, cutlass::layout::RowMajor, ElementAccumulator,
// cutlass::arch::OpClassSimt, cutlass::arch::Sm70,
// DefaultGemmCfg::ThreadblockShape, DefaultGemmCfg::WarpShape,
// DefaultGemmCfg::InstructionShape,
// cutlass::epilogue::thread::LinearCombination<
// ElementOutput, 1, ElementAccumulator, ElementComputeEpilogue>>;
// #else
// #error "Unsupported cuda arch"
// #endif
auto input_size = cutlass::MatrixCoord(M, K);
auto weight_size = cutlass::MatrixCoord(K, N);
auto output_size = cutlass::MatrixCoord(M, N);
......@@ -83,8 +92,8 @@ Tensor gemm_w8a8_fp16(Tensor input, // INT8
if (!out.valid()) {
auto out_shape = TensorShape(input.shape.dataExtent);
out_shape[-1] = N;
out = Tensor::empty(out_shape, Tensor::FP16, input.device());
out_shape[-1] = N;
out = Tensor::empty(out_shape, Tensor::FP16, input.device());
}
// FIXME: check contiguous of input if dims >= 3
......@@ -105,21 +114,20 @@ Tensor gemm_w8a8_fp16(Tensor input, // INT8
// constexpr int kMetaSizeInBits = Gemm::kMetaSizeInBits;
cutlass::gemm::GemmCoord problem_size(M, N, K);
cutlass::TensorRef<ElementInputA, LayoutInputA> input_ref(
input.data_ptr<ElementInputA>(), LayoutInputA(input.stride(-2)));
cutlass::TensorRef<ElementInputB, LayoutInputB> weight_ref(
weight.data_ptr<ElementInputB>(), LayoutInputB::packed(weight_size));
cutlass::TensorRef<ElementOutput, LayoutOutput> out_ref(
out.data_ptr<ElementOutput>(), LayoutOutput(out.stride(-2)));
typename Gemm::Arguments arguments{
problem_size, // <- problem size of matrix multiplication
input_ref, // <- reference to matrix A on device
weight_ref, // <- reference to matrix B on device
out_ref, // <- reference to matrix C on device
out_ref, // <- reference to matrix D on device
{ElementOutput(alpha), ElementOutput(beta)},
1};
cutlass::TensorRef<ElementInputA, LayoutInputA> input_ref(input.data_ptr<ElementInputA>(),
LayoutInputA(input.stride(-2)));
cutlass::TensorRef<ElementInputB, LayoutInputB> weight_ref(weight.data_ptr<ElementInputB>(),
LayoutInputB::packed(weight_size));
cutlass::TensorRef<ElementOutput, LayoutOutput> out_ref(out.data_ptr<ElementOutput>(),
LayoutOutput(out.stride(-2)));
typename Gemm::Arguments arguments{problem_size, // <- problem size of matrix multiplication
input_ref, // <- reference to matrix A on device
weight_ref, // <- reference to matrix B on device
out_ref, // <- reference to matrix C on device
out_ref, // <- reference to matrix D on device
{ElementOutput(alpha), ElementOutput(beta)},
1};
Gemm gemm_op;
// Using the arguments, query for extra workspace required for matrix
......
src/kernels/gemm_w8a8.h
View file @ 37a27712
......@@ -7,5 +7,4 @@ Tensor gemm_w8a8_fp16(Tensor input, // INT8
Tensor weight, // INT8
Tensor out,
half scale,
half bias
);
\ No newline at end of file
half bias);
src/kernels/layernorm_kernels.cu
View file @ 37a27712
......@@ -6,175 +6,212 @@ void rms_norm(Tensor &out, // [..., hidden_size]
Tensor &weight, // [hidden_size]
float epsilon,
bool use_quant) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
if (use_quant) {
vllm::rms_norm_kernel<scalar_t, int8_t, true><<<grid, block, 0, stream>>>(
out.data_ptr<int8_t>(), input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
} else {
vllm::rms_norm_kernel<scalar_t, scalar_t, false><<<grid, block, 0, stream>>>(
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
}
});
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
if (use_quant) {
vllm::rms_norm_kernel<scalar_t, int8_t, true><<<grid, block, 0, stream>>>(out.data_ptr<int8_t>(),
input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
epsilon,
num_tokens,
hidden_size);
} else {
vllm::rms_norm_kernel<scalar_t, scalar_t, false><<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
weight.data_ptr<scalar_t>(),
epsilon,
num_tokens,
hidden_size);
}
});
}
void layernorm_general(Tensor out, Tensor input, Tensor weight, Tensor bias, float epsilon) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 256));
block.x = 32 * ((block.x + 31) / 32);
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 256));
block.x = 32 * ((block.x + 31) / 32);
size_t size_shmem = input.scalar_size() * hidden_size;
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "generalLayerNorm", [&] {
using T = typename packed_as<scalar_t, 2>::type;
vllm::generalLayerNorm<T, half, true><<<grid, block, size_shmem, stream>>>(
reinterpret_cast<T*>(input.data_ptr<scalar_t>()),
weight.valid() ? reinterpret_cast<T*>(weight.data_ptr<scalar_t>()) : nullptr,
bias.valid() ? reinterpret_cast<T*>(bias.data_ptr<scalar_t>()) : nullptr,
reinterpret_cast<T*>(out.data_ptr<scalar_t>()),
epsilon, num_tokens, hidden_size, nullptr, nullptr, nullptr, true
);
});
}
size_t size_shmem = input.scalar_size() * hidden_size;
void rms_norm_general(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &weight, // [hidden_size]
Tensor &scaling, // [tokens] or [1]
float epsilon,
bool use_per_token_quant) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
block.x = 32 * ((block.x + 31) / 32);
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "generalLayerNorm", [&] {
using T = scalar_t;
if (use_per_token_quant) {
// per-token
vllm::generalLayerNorm<T, half><<<grid, block, 0, stream>>>(
reinterpret_cast<T*>(input.data_ptr<scalar_t>()),
reinterpret_cast<T*>(weight.data_ptr<scalar_t>()), nullptr,
nullptr, epsilon, num_tokens, hidden_size, nullptr, scaling.data_ptr<half>(),
out.data_ptr<int8_t>(), false
);
// input, gamma, beta, normed_output, eps, tokens, hidden_dim, per_tensor_scale, per_token_scale
// normed_output_quant, use_shmem
// out.data_ptr<int8_t>(), input.data_ptr<scalar_t>(),
// weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
} else {
// per-tensor
vllm::generalLayerNorm<T, half><<<grid, block, 0, stream>>>(
reinterpret_cast<T*>(input.data_ptr<scalar_t>()),
reinterpret_cast<T*>(weight.data_ptr<scalar_t>()), nullptr,
nullptr, epsilon, num_tokens, hidden_size, scaling.data_ptr<half>(), nullptr,
out.data_ptr<int8_t>(), false
);
}
});
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "generalLayerNorm", [&] {
using T = typename packed_as<scalar_t, 2>::type;
vllm::generalLayerNorm<T, half, true><<<grid, block, size_shmem, stream>>>(
reinterpret_cast<T *>(input.data_ptr<scalar_t>()),
weight.valid() ? reinterpret_cast<T *>(weight.data_ptr<scalar_t>()) : nullptr,
bias.valid() ? reinterpret_cast<T *>(bias.data_ptr<scalar_t>()) : nullptr,
reinterpret_cast<T *>(out.data_ptr<scalar_t>()),
epsilon,
num_tokens,
hidden_size,
nullptr,
nullptr,
nullptr,
true);
});
}
void rms_norm_general_fuse_sum(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &weight, // [hidden_size]
Tensor &input_sum, // [tokens] or [1]
Tensor &scaling, // [tokens] or [1]
float epsilon,
bool use_per_token_quant) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
block.x = 32 * ((block.x + 31) / 32);
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "generalLayerNorm_fuse_sum", [&] {
using T = scalar_t;
if (use_per_token_quant) {
// per-token
vllm::generalLayerNorm_fuse_sum<T, half><<<grid, block, 0, stream>>>(
reinterpret_cast<T*>(input.data_ptr<scalar_t>()),
reinterpret_cast<T*>(weight.data_ptr<scalar_t>()), nullptr,
nullptr, epsilon, num_tokens, hidden_size, input_sum.data_ptr<half>(), nullptr, scaling.data_ptr<half>(),
out.data_ptr<int8_t>(), false
);
// input, gamma, beta, normed_output, eps, tokens, hidden_dim, per_tensor_scale, per_token_scale
// normed_output_quant, use_shmem
// out.data_ptr<int8_t>(), input.data_ptr<scalar_t>(),
// weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
} else {
// per-tensor
// Rasing error here
// Not implemented per-tensor input_sum
assert(false);
vllm::generalLayerNorm_fuse_sum<T, half><<<grid, block, 0, stream>>>(
reinterpret_cast<T*>(input.data_ptr<scalar_t>()),
reinterpret_cast<T*>(weight.data_ptr<scalar_t>()), nullptr,
nullptr, epsilon, num_tokens, hidden_size, nullptr, scaling.data_ptr<half>(), nullptr,
out.data_ptr<int8_t>(), false
);
}
});
void rms_norm_general(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &weight, // [hidden_size]
Tensor &scaling, // [tokens] or [1]
float epsilon,
bool use_per_token_quant) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
block.x = 32 * ((block.x + 31) / 32);
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "generalLayerNorm", [&] {
using T = scalar_t;
if (use_per_token_quant) {
// per-token
vllm::generalLayerNorm<T, half>
<<<grid, block, 0, stream>>>(reinterpret_cast<T *>(input.data_ptr<scalar_t>()),
reinterpret_cast<T *>(weight.data_ptr<scalar_t>()),
nullptr,
nullptr,
epsilon,
num_tokens,
hidden_size,
nullptr,
scaling.data_ptr<half>(),
out.data_ptr<int8_t>(),
false);
// input, gamma, beta, normed_output, eps, tokens, hidden_dim, per_tensor_scale, per_token_scale
// normed_output_quant, use_shmem
// out.data_ptr<int8_t>(), input.data_ptr<scalar_t>(),
// weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
} else {
// per-tensor
vllm::generalLayerNorm<T, half>
<<<grid, block, 0, stream>>>(reinterpret_cast<T *>(input.data_ptr<scalar_t>()),
reinterpret_cast<T *>(weight.data_ptr<scalar_t>()),
nullptr,
nullptr,
epsilon,
num_tokens,
hidden_size,
scaling.data_ptr<half>(),
nullptr,
out.data_ptr<int8_t>(),
false);
}
});
}
void rms_norm_general_fuse_sum(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &weight, // [hidden_size]
Tensor &input_sum, // [tokens] or [1]
Tensor &scaling, // [tokens] or [1]
float epsilon,
bool use_per_token_quant) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
block.x = 32 * ((block.x + 31) / 32);
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "generalLayerNorm_fuse_sum", [&] {
using T = scalar_t;
if (use_per_token_quant) {
// per-token
vllm::generalLayerNorm_fuse_sum<T, half>
<<<grid, block, 0, stream>>>(reinterpret_cast<T *>(input.data_ptr<scalar_t>()),
reinterpret_cast<T *>(weight.data_ptr<scalar_t>()),
nullptr,
nullptr,
epsilon,
num_tokens,
hidden_size,
input_sum.data_ptr<half>(),
nullptr,
scaling.data_ptr<half>(),
out.data_ptr<int8_t>(),
false);
// input, gamma, beta, normed_output, eps, tokens, hidden_dim, per_tensor_scale, per_token_scale
// normed_output_quant, use_shmem
// out.data_ptr<int8_t>(), input.data_ptr<scalar_t>(),
// weight.data_ptr<scalar_t>(), epsilon, num_tokens, hidden_size);
} else {
// per-tensor
// Rasing error here
// Not implemented per-tensor input_sum
assert(false);
vllm::generalLayerNorm_fuse_sum<T, half>
<<<grid, block, 0, stream>>>(reinterpret_cast<T *>(input.data_ptr<scalar_t>()),
reinterpret_cast<T *>(weight.data_ptr<scalar_t>()),
nullptr,
nullptr,
epsilon,
num_tokens,
hidden_size,
nullptr,
scaling.data_ptr<half>(),
nullptr,
out.data_ptr<int8_t>(),
false);
}
});
}
void invoke_dequant_add_residual_rms_norm_quant(
Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &residual, // [..., hidden_size]
Tensor &gamma, // [hidden_size]
half scale,
float epsilon) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(
residual.scalar_type(), "dequant_add_residual_rms_norm_quant_kernel",
[&] {
vllm::dequant_add_residual_rms_norm_quant_kernel<scalar_t, half, false>
<<<grid, block, 0, stream>>>(
input.data_ptr<int32_t>(), residual.data_ptr<scalar_t>(),
out.data_ptr<int8_t>(), gamma.data_ptr<scalar_t>(), epsilon,
scale, num_tokens, hidden_size);
});
void invoke_dequant_add_residual_rms_norm_quant(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &residual, // [..., hidden_size]
Tensor &gamma, // [hidden_size]
half scale,
float epsilon) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(residual.scalar_type(), "dequant_add_residual_rms_norm_quant_kernel", [&] {
vllm::dequant_add_residual_rms_norm_quant_kernel<scalar_t, half, false>
<<<grid, block, 0, stream>>>(input.data_ptr<int32_t>(),
residual.data_ptr<scalar_t>(),
out.data_ptr<int8_t>(),
gamma.data_ptr<scalar_t>(),
epsilon,
scale,
num_tokens,
hidden_size);
});
}
void invoke_dequant_add_residual_rms_norm_quant(
Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &residual, // [..., hidden_size]
Tensor &gamma, // [hidden_size]
Tensor &scale, // [num_tokens]
float epsilon) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
void invoke_dequant_add_residual_rms_norm_quant(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &residual, // [..., hidden_size]
Tensor &gamma, // [hidden_size]
Tensor &scale, // [num_tokens]
float epsilon) {
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(
residual.scalar_type(), "dequant_add_residual_rms_norm_quant_kernel",
[&] {
vllm::dequant_add_residual_rms_norm_quant_kernel<scalar_t, half*, true>
<<<grid, block, 0, stream>>>(
input.data_ptr<int32_t>(), residual.data_ptr<scalar_t>(),
out.data_ptr<int8_t>(), gamma.data_ptr<scalar_t>(), epsilon,
scale.data_ptr<half>(), num_tokens, hidden_size);
});
const cudaStream_t stream = getCurrentCUDAStream();
VLLM_DISPATCH_FLOATING_TYPES(residual.scalar_type(), "dequant_add_residual_rms_norm_quant_kernel", [&] {
vllm::dequant_add_residual_rms_norm_quant_kernel<scalar_t, half *, true>
<<<grid, block, 0, stream>>>(input.data_ptr<int32_t>(),
residual.data_ptr<scalar_t>(),
out.data_ptr<int8_t>(),
gamma.data_ptr<scalar_t>(),
epsilon,
scale.data_ptr<half>(),
num_tokens,
hidden_size);
});
}
src/kernels/layernorm_kernels.h
View file @ 37a27712
......@@ -7,36 +7,36 @@
void rms_norm(Tensor &out, // [num_tokens, hidden_size]
Tensor &input, // [num_tokens, hidden_size]
Tensor &weight, // [hidden_size]
float epsilon, bool use_quant);
float epsilon,
bool use_quant);
void layernorm_general(Tensor out, Tensor input, Tensor weight, Tensor bias, float epsilon);
void rms_norm_general(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &weight, // [hidden_size]
Tensor &scaling, // [tokens] or [1]
float epsilon,
bool use_per_token_quant);
void rms_norm_general(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &weight, // [hidden_size]
Tensor &scaling, // [tokens] or [1]
float epsilon,
bool use_per_token_quant);
void rms_norm_general_fuse_sum(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &weight, // [hidden_size]
Tensor &input_sum, // [tokens] or [1]
Tensor &scaling, // [tokens] or [1]
float epsilon,
bool use_per_token_quant);
void rms_norm_general_fuse_sum(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &weight, // [hidden_size]
Tensor &input_sum, // [tokens] or [1]
Tensor &scaling, // [tokens] or [1]
float epsilon,
bool use_per_token_quant);
void invoke_dequant_add_residual_rms_norm_quant(
Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &residual, // [..., hidden_size]
Tensor &gamma, // [hidden_size]
half scale, float epsilon);
void invoke_dequant_add_residual_rms_norm_quant(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &residual, // [..., hidden_size]
Tensor &gamma, // [hidden_size]
half scale,
float epsilon);
void invoke_dequant_add_residual_rms_norm_quant(
Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &residual, // [..., hidden_size]
Tensor &gamma, // [hidden_size]
Tensor &scale, // [num_tokens]
float epsilon);
\ No newline at end of file
void invoke_dequant_add_residual_rms_norm_quant(Tensor &out, // [..., hidden_size]
Tensor &input, // [..., hidden_size]
Tensor &residual, // [..., hidden_size]
Tensor &gamma, // [hidden_size]
Tensor &scale, // [num_tokens]
float epsilon);
src/kernels/layernorm_kernels_impl.cuh
View file @ 37a27712
......@@ -5,13 +5,12 @@
#include "utils.cuh"
#include "reduction_utils.cuh"
namespace vllm {
// from TRTLLM
template <typename Tf, typename T>
__inline__ __device__ Tf compute_layernorm(Tf val, float s_mean, float s_variance, const T* gamma, const T* beta, int i)
{
template<typename Tf, typename T>
__inline__ __device__ Tf
compute_layernorm(Tf val, float s_mean, float s_variance, const T *gamma, const T *beta, int i) {
Tf ret = (val - s_mean) * s_variance;
if (gamma != nullptr) {
ret = ret * cuda_cast<Tf>(gamma[i]);
......@@ -44,353 +43,320 @@ __inline__ __device__ Tf compute_layernorm(Tf val, float s_mean, float s_varianc
* amax per row. A final pass scales to int8 accordingly, and writes output to
* normed_output_quant.
*/
template <typename T, typename scale_type, bool USE_DIFF_OF_SQUARES = false>
__global__ void generalLayerNorm(const T* input, const T* gamma, const T* beta, T* normed_output, const float eps,
int tokens, int hidden_dim, const scale_type* scale_orig_quant_per_tensor, scale_type* scale_orig_quant_per_token,
int8_t* normed_output_quant, bool use_shmem)
{
template<typename T, typename scale_type, bool USE_DIFF_OF_SQUARES = false>
__global__ void generalLayerNorm(const T *input,
const T *gamma,
const T *beta,
T *normed_output,
const float eps,
int tokens,
int hidden_dim,
const scale_type *scale_orig_quant_per_tensor,
scale_type *scale_orig_quant_per_token,
int8_t *normed_output_quant,
bool use_shmem) {
constexpr auto num_elems_T = num_elems<T>::value;
using int8_packed_t = typename packed_as<int8_t, num_elems_T>::type;
using float_packed_t = typename packed_as<float, num_elems_T>::type;
using T_scalar = typename packed_as<T, 1>::type;
using int8_packed_t = typename packed_as<int8_t, num_elems_T>::type;
using float_packed_t = typename packed_as<float, num_elems_T>::type;
using T_scalar = typename packed_as<T, 1>::type;
extern __shared__ __align__(sizeof(float)) char _shmem[];
T* shmem = reinterpret_cast<T*>(_shmem);
T *shmem = reinterpret_cast<T *>(_shmem);
__shared__ float s_mean;
__shared__ float s_variance;
const int tidx = threadIdx.x;
const int bidx = blockIdx.x;
float mean = 0.0f;
float variance = 0.0f;
float local_sum = 0.0f;
float mean = 0.0f;
float variance = 0.0f;
float local_sum = 0.0f;
float local_var_sum = 0.0f;
const int n_elems = hidden_dim / num_elems_T;
for (int i = tidx; i < n_elems; i += blockDim.x)
{
for (int i = tidx; i < n_elems; i += blockDim.x) {
const T val = input[bidx * n_elems + i];
if (use_shmem)
{
if (use_shmem) {
shmem[i] = val;
}
const float_packed_t val_f = cuda_cast<float_packed_t>(val);
local_sum += cuda_sum<float>(val_f);
if (USE_DIFF_OF_SQUARES)
{
if (USE_DIFF_OF_SQUARES) {
local_var_sum += cuda_sum<float>(val_f * val_f);
}
}
if (USE_DIFF_OF_SQUARES)
{
if (USE_DIFF_OF_SQUARES) {
float packed[2] = {local_sum, local_var_sum};
blockReduceSumV2<float, 2>(packed);
mean = packed[0];
mean = packed[0];
variance = packed[1];
}
else
{
} else {
mean = blockReduceSum(local_sum);
}
if (threadIdx.x == 0)
{
mean = mean / hidden_dim;
if (threadIdx.x == 0) {
mean = mean / hidden_dim;
s_mean = mean;
if (USE_DIFF_OF_SQUARES)
{
variance = (variance / hidden_dim) - (mean * mean); // Var[x] = E[x²] - E[x]²
if (USE_DIFF_OF_SQUARES) {
variance = (variance / hidden_dim) - (mean * mean); // Var[x] = E[x²] - E[x]²
s_variance = rsqrtf(variance + eps);
}
}
__syncthreads();
if (!USE_DIFF_OF_SQUARES)
{
for (int i = tidx; i < n_elems; i += blockDim.x)
{
const T val = use_shmem ? shmem[i] : input[bidx * n_elems + i];
if (!USE_DIFF_OF_SQUARES) {
for (int i = tidx; i < n_elems; i += blockDim.x) {
const T val = use_shmem ? shmem[i] : input[bidx * n_elems + i];
float_packed_t diff = cuda_cast<float_packed_t>(val) - s_mean;
local_var_sum += cuda_sum<float>(diff * diff);
}
variance = blockReduceSum(local_var_sum);
if (threadIdx.x == 0)
{
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_dim + eps);
}
__syncthreads();
}
const bool with_per_token_scaling = scale_orig_quant_per_token != nullptr;
const bool with_per_token_scaling = scale_orig_quant_per_token != nullptr;
const bool with_per_tensor_scaling = scale_orig_quant_per_tensor != nullptr;
const float_packed_t scale_orig_quant
= cuda_cast<float_packed_t>(with_per_tensor_scaling ? __half2float(*scale_orig_quant_per_tensor) : 0.0f);
const float_packed_t scale_orig_quant =
cuda_cast<float_packed_t>(with_per_tensor_scaling ? __half2float(*scale_orig_quant_per_tensor) : 0.0f);
T_scalar amax = 1e-6f;
for (int i = tidx; i < n_elems; i += blockDim.x)
{
const int index = bidx * n_elems + i;
for (int i = tidx; i < n_elems; i += blockDim.x) {
const int index = bidx * n_elems + i;
const float_packed_t val_f = cuda_cast<float_packed_t>(use_shmem ? shmem[i] : input[index]);
const T val = cuda_cast<T>(compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i));
const T val = cuda_cast<T>(compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i));
if (with_per_token_scaling)
{
if (with_per_token_scaling) {
amax = cuda_max(cuda_max<T_scalar, T>(cuda_abs(val)), amax);
if (use_shmem)
{
if (use_shmem) {
shmem[i] = val;
}
}
else if (with_per_tensor_scaling)
{
reinterpret_cast<int8_packed_t*>(normed_output_quant)[index]
= cuda_cast<int8_packed_t>(cuda_cast<float_packed_t>(val) * scale_orig_quant);
}
else
{
} else if (with_per_tensor_scaling) {
reinterpret_cast<int8_packed_t *>(normed_output_quant)[index] =
cuda_cast<int8_packed_t>(cuda_cast<float_packed_t>(val) * scale_orig_quant);
} else {
normed_output[index] = val;
}
}
if (with_per_token_scaling)
{
float abs_max_f = blockAllReduceMax(cuda_cast<float>(amax));
if (with_per_token_scaling) {
float abs_max_f = blockAllReduceMax(cuda_cast<float>(amax));
const float dynamic_per_token_scale = 127.f / abs_max_f;
for (int i = tidx; i < n_elems; i += blockDim.x)
{
const int index = bidx * n_elems + i;
for (int i = tidx; i < n_elems; i += blockDim.x) {
const int index = bidx * n_elems + i;
float_packed_t val_f = cuda_cast<float_packed_t>(use_shmem ? shmem[i] : input[index]);
if (!use_shmem)
{
if (!use_shmem) {
val_f = compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i);
}
reinterpret_cast<int8_packed_t*>(normed_output_quant)[index]
= cuda_cast<int8_packed_t>(val_f * cuda_cast<float_packed_t>(dynamic_per_token_scale));
reinterpret_cast<int8_packed_t *>(normed_output_quant)[index] =
cuda_cast<int8_packed_t>(val_f * cuda_cast<float_packed_t>(dynamic_per_token_scale));
}
if (tidx == 0)
{
if (tidx == 0) {
scale_orig_quant_per_token[bidx] = abs_max_f / 127.f;
}
}
}
template <typename T, typename scale_type, bool USE_DIFF_OF_SQUARES = false>
__global__ void generalLayerNorm_fuse_sum(const T* input, const T* gamma, const T* beta, T* normed_output, const float eps,
int tokens, int hidden_dim, scale_type* input_sum, const scale_type* scale_orig_quant_per_tensor, scale_type* scale_orig_quant_per_token,
int8_t* normed_output_quant, bool use_shmem)
{
template<typename T, typename scale_type, bool USE_DIFF_OF_SQUARES = false>
__global__ void generalLayerNorm_fuse_sum(const T *input,
const T *gamma,
const T *beta,
T *normed_output,
const float eps,
int tokens,
int hidden_dim,
scale_type *input_sum,
const scale_type *scale_orig_quant_per_tensor,
scale_type *scale_orig_quant_per_token,
int8_t *normed_output_quant,
bool use_shmem) {
constexpr auto num_elems_T = num_elems<T>::value;
using int8_packed_t = typename packed_as<int8_t, num_elems_T>::type;
using float_packed_t = typename packed_as<float, num_elems_T>::type;
using T_scalar = typename packed_as<T, 1>::type;
using int8_packed_t = typename packed_as<int8_t, num_elems_T>::type;
using float_packed_t = typename packed_as<float, num_elems_T>::type;
using T_scalar = typename packed_as<T, 1>::type;
extern __shared__ __align__(sizeof(float)) char _shmem[];
T* shmem = reinterpret_cast<T*>(_shmem);
T *shmem = reinterpret_cast<T *>(_shmem);
__shared__ float s_mean;
__shared__ float s_variance;
const int tidx = threadIdx.x;
const int bidx = blockIdx.x;
float mean = 0.0f;
float variance = 0.0f;
float local_sum = 0.0f;
float mean = 0.0f;
float variance = 0.0f;
float local_sum = 0.0f;
float local_var_sum = 0.0f;
const int n_elems = hidden_dim / num_elems_T;
for (int i = tidx; i < n_elems; i += blockDim.x)
{
for (int i = tidx; i < n_elems; i += blockDim.x) {
const T val = input[bidx * n_elems + i];
if (use_shmem)
{
if (use_shmem) {
shmem[i] = val;
}
const float_packed_t val_f = cuda_cast<float_packed_t>(val);
local_sum += cuda_sum<float>(val_f);
if (USE_DIFF_OF_SQUARES)
{
if (USE_DIFF_OF_SQUARES) {
local_var_sum += cuda_sum<float>(val_f * val_f);
}
}
if (USE_DIFF_OF_SQUARES)
{
if (USE_DIFF_OF_SQUARES) {
float packed[2] = {local_sum, local_var_sum};
blockReduceSumV2<float, 2>(packed);
mean = packed[0];
mean = packed[0];
variance = packed[1];
}
else
{
} else {
mean = blockReduceSum(local_sum);
}
if (threadIdx.x == 0)
{
mean = mean / hidden_dim;
if (threadIdx.x == 0) {
mean = mean / hidden_dim;
s_mean = mean;
if (USE_DIFF_OF_SQUARES)
{
variance = (variance / hidden_dim) - (mean * mean); // Var[x] = E[x²] - E[x]²
if (USE_DIFF_OF_SQUARES) {
variance = (variance / hidden_dim) - (mean * mean); // Var[x] = E[x²] - E[x]²
s_variance = rsqrtf(variance + eps);
}
}
__syncthreads();
if (!USE_DIFF_OF_SQUARES)
{
for (int i = tidx; i < n_elems; i += blockDim.x)
{
const T val = use_shmem ? shmem[i] : input[bidx * n_elems + i];
if (!USE_DIFF_OF_SQUARES) {
for (int i = tidx; i < n_elems; i += blockDim.x) {
const T val = use_shmem ? shmem[i] : input[bidx * n_elems + i];
float_packed_t diff = cuda_cast<float_packed_t>(val) - s_mean;
local_var_sum += cuda_sum<float>(diff * diff);
}
variance = blockReduceSum(local_var_sum);
if (threadIdx.x == 0)
{
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_dim + eps);
}
__syncthreads();
}
const bool with_per_token_scaling = scale_orig_quant_per_token != nullptr;
const bool with_per_token_scaling = scale_orig_quant_per_token != nullptr;
const bool with_per_tensor_scaling = scale_orig_quant_per_tensor != nullptr;
const float_packed_t scale_orig_quant
= cuda_cast<float_packed_t>(with_per_tensor_scaling ? __half2float(*scale_orig_quant_per_tensor) : 0.0f);
const float_packed_t scale_orig_quant =
cuda_cast<float_packed_t>(with_per_tensor_scaling ? __half2float(*scale_orig_quant_per_tensor) : 0.0f);
T_scalar amax = 1e-6f;
T_scalar sum = 0.0f;
T_scalar sum = 0.0f;
for (int i = tidx; i < n_elems; i += blockDim.x)
{
const int index = bidx * n_elems + i;
for (int i = tidx; i < n_elems; i += blockDim.x) {
const int index = bidx * n_elems + i;
const float_packed_t val_f = cuda_cast<float_packed_t>(use_shmem ? shmem[i] : input[index]);
const T val = cuda_cast<T>(compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i));
const T val = cuda_cast<T>(compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i));
if (with_per_token_scaling)
{
if (with_per_token_scaling) {
amax = cuda_max(cuda_max<T_scalar, T>(cuda_abs(val)), amax);
sum += cuda_sum<float>(val);
if (use_shmem)
{
if (use_shmem) {
shmem[i] = val;
}
}
else if (with_per_tensor_scaling)
{
reinterpret_cast<int8_packed_t*>(normed_output_quant)[index]
= cuda_cast<int8_packed_t>(cuda_cast<float_packed_t>(val) * scale_orig_quant);
}
else
{
} else if (with_per_tensor_scaling) {
reinterpret_cast<int8_packed_t *>(normed_output_quant)[index] =
cuda_cast<int8_packed_t>(cuda_cast<float_packed_t>(val) * scale_orig_quant);
} else {
normed_output[index] = val;
}
}
if (with_per_token_scaling)
{
float abs_max_f = blockAllReduceMax(cuda_cast<float>(amax));
float sum_f = blockAllReduceSum(cuda_cast<float>(sum));
if (with_per_token_scaling) {
float abs_max_f = blockAllReduceMax(cuda_cast<float>(amax));
float sum_f = blockAllReduceSum(cuda_cast<float>(sum));
const float dynamic_per_token_scale = 127.f / abs_max_f;
for (int i = tidx; i < n_elems; i += blockDim.x)
{
const int index = bidx * n_elems + i;
for (int i = tidx; i < n_elems; i += blockDim.x) {
const int index = bidx * n_elems + i;
float_packed_t val_f = cuda_cast<float_packed_t>(use_shmem ? shmem[i] : input[index]);
if (!use_shmem)
{
if (!use_shmem) {
val_f = compute_layernorm(val_f, s_mean, s_variance, gamma, beta, i);
}
reinterpret_cast<int8_packed_t*>(normed_output_quant)[index]
= cuda_cast<int8_packed_t>(val_f * cuda_cast<float_packed_t>(dynamic_per_token_scale));
reinterpret_cast<int8_packed_t *>(normed_output_quant)[index] =
cuda_cast<int8_packed_t>(val_f * cuda_cast<float_packed_t>(dynamic_per_token_scale));
}
if (tidx == 0)
{
if (tidx == 0) {
scale_orig_quant_per_token[bidx] = abs_max_f / 127.f;
input_sum[bidx] = sum_f;
input_sum[bidx] = sum_f;
}
}
}
// TODO(woosuk): Further optimize this kernel.
template <typename scalar_t, typename out_type, bool use_quant>
__global__ void
rms_norm_kernel(out_type *__restrict__ out, // [..., hidden_size]
const scalar_t *__restrict__ input, // [..., hidden_size]
const scalar_t *__restrict__ weight, // [hidden_size]
const float epsilon, const int num_tokens,
const int hidden_size) {
__shared__ float s_variance;
float variance = 0.0f;
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
const float x = (float)input[blockIdx.x * hidden_size + idx];
variance += x * x;
}
variance = blockReduceSum<float>(variance);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_size + epsilon);
}
__syncthreads();
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)input[blockIdx.x * hidden_size + idx];
if constexpr (use_quant) {
out[blockIdx.x * hidden_size + idx] = float_to_int8_rn(
((float)(x * s_variance)) * (float)(weight[idx]));
} else {
out[blockIdx.x * hidden_size + idx] =
((scalar_t)(x * s_variance)) * weight[idx];
template<typename scalar_t, typename out_type, bool use_quant>
__global__ void rms_norm_kernel(out_type *__restrict__ out, // [..., hidden_size]
const scalar_t *__restrict__ input, // [..., hidden_size]
const scalar_t *__restrict__ weight, // [hidden_size]
const float epsilon,
const int num_tokens,
const int hidden_size) {
__shared__ float s_variance;
float variance = 0.0f;
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
const float x = (float)input[blockIdx.x * hidden_size + idx];
variance += x * x;
}
}
}
variance = blockReduceSum<float>(variance);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_size + epsilon);
}
__syncthreads();
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = (float)input[blockIdx.x * hidden_size + idx];
if constexpr (use_quant) {
out[blockIdx.x * hidden_size + idx] = float_to_int8_rn(((float)(x * s_variance)) * (float)(weight[idx]));
} else {
out[blockIdx.x * hidden_size + idx] = ((scalar_t)(x * s_variance)) * weight[idx];
}
}
}
template<typename T, typename scale_type, bool use_per_token_dequant>
__global__ void dequant_add_residual_rms_norm_quant_kernel(const int32_t *__restrict__ input,
T *__restrict__ residual,
int8_t *__restrict__ output,
const T *__restrict__ gamma,
const float layernorm_eps,
const scale_type scale,
int num_tokens,
int hidden_size) {
// layernorm module in the T5 style No bias and no subtraction of mean.
const int tid = threadIdx.x;
__shared__ float s_variance;
float variance = 0.0f;
template <typename T, typename scale_type, bool use_per_token_dequant>
__global__ void dequant_add_residual_rms_norm_quant_kernel(
const int32_t *__restrict__ input, T *__restrict__ residual,
int8_t *__restrict__ output, const T *__restrict__ gamma,
const float layernorm_eps, const scale_type scale, int num_tokens, int hidden_size) {
// layernorm module in the T5 style No bias and no subtraction of mean.
const int tid = threadIdx.x;
float local_var_sum = 0.0f;
for (int i = tid; i < hidden_size; i += blockDim.x) {
float diff = 0.0f;
if constexpr (use_per_token_dequant) {
diff = ((((float)input[blockIdx.x * hidden_size + i]) * __half2float(scale[blockIdx.x])) +
(float)residual[blockIdx.x * hidden_size + i]);
} else {
diff = ((((float)input[blockIdx.x * hidden_size + i]) * __half2float(scale)) +
(float)residual[blockIdx.x * hidden_size + i]);
}
residual[blockIdx.x * hidden_size + i] = (T)diff;
local_var_sum += diff * diff;
}
variance = blockReduceSum(local_var_sum);
__shared__ float s_variance;
float variance = 0.0f;
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / (float)hidden_size + layernorm_eps);
}
__syncthreads();
float local_var_sum = 0.0f;
for (int i = tid; i < hidden_size; i += blockDim.x) {
float diff = 0.0f;
if constexpr (use_per_token_dequant) {
diff = ((((float)input[blockIdx.x * hidden_size + i]) * __half2float(scale[blockIdx.x])) +
(float)residual[blockIdx.x * hidden_size + i]);
} else {
diff = ((((float)input[blockIdx.x * hidden_size + i]) * __half2float(scale)) +
(float)residual[blockIdx.x * hidden_size + i]);
for (int i = tid; i < hidden_size; i += blockDim.x) {
output[blockIdx.x * hidden_size + i] =
float_to_int8_rn((((float)(residual[blockIdx.x * hidden_size + i])) * s_variance) * (float)(gamma[i]));
}
residual[blockIdx.x * hidden_size + i] = (T)diff;
local_var_sum += diff * diff;
}
variance = blockReduceSum(local_var_sum);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / (float)hidden_size + layernorm_eps);
}
__syncthreads();
for (int i = tid; i < hidden_size; i += blockDim.x) {
output[blockIdx.x * hidden_size + i] = float_to_int8_rn(
(((float)(residual[blockIdx.x * hidden_size + i])) * s_variance) *
(float)(gamma[i]));
}
}
} // namespace vllm
src/kernels/misc_kernels.cu
View file @ 37a27712
......@@ -11,7 +11,7 @@ Tensor add(Tensor a, Tensor b) {
assert(b.is_contiguous());
int threadsPerBlock = 1024;
int blocksPerGrid = (a.numel() + threadsPerBlock - 1) / threadsPerBlock;
int blocksPerGrid = (a.numel() + threadsPerBlock - 1) / threadsPerBlock;
auto stream = getCurrentCUDAStream();
......@@ -44,14 +44,23 @@ void mul_add(Tensor x, Tensor scale, Tensor bias) {
assert(bias.numel() % unroll == 0);
int threadsPerBlock = 1024;
int blocksPerGrid = (x.numel() + threadsPerBlock * unroll - 1) / (threadsPerBlock * unroll);
int blocksPerGrid = (x.numel() + threadsPerBlock * unroll - 1) / (threadsPerBlock * unroll);
auto stream = getCurrentCUDAStream();
dispatch(x.scalar_type(), [&]<typename scalar_t>() {
if (scale.valid()) {
mul_add_kernel<scalar_t, unroll, false><<<blocksPerGrid, threadsPerBlock, 0, stream>>>(
x.data_ptr<scalar_t>(), scale.data_ptr<scalar_t>(), bias.data_ptr<scalar_t>(), 0, x.numel(), scale.numel(), bias.numel(), 0, 0, 0);
mul_add_kernel<scalar_t, unroll, false>
<<<blocksPerGrid, threadsPerBlock, 0, stream>>>(x.data_ptr<scalar_t>(),
scale.data_ptr<scalar_t>(),
bias.data_ptr<scalar_t>(),
0,
x.numel(),
scale.numel(),
bias.numel(),
0,
0,
0);
} else {
mul_add_kernel<scalar_t, unroll, true><<<blocksPerGrid, threadsPerBlock, 0, stream>>>(
x.data_ptr<scalar_t>(), nullptr, bias.data_ptr<scalar_t>(), 0, x.numel(), 1, bias.numel(), 0, 0, 0);
......@@ -65,7 +74,7 @@ void mul_add_batch(Tensor x, Tensor scale, bool batch_scale, double scale_shift,
assert(!batch_scale || scale.shape[0] == batch_size);
assert(!batch_bias || bias.shape[0] == batch_size);
const int numel = x.numel() / batch_size;
const int numel = x.numel() / batch_size;
const int numel_scale = scale.valid() ? (scale.numel() / (batch_scale ? batch_size : 1)) : 1;
const int numel_bias = bias.numel() / (batch_bias ? batch_size : 1);
......@@ -91,17 +100,29 @@ void mul_add_batch(Tensor x, Tensor scale, bool batch_scale, double scale_shift,
dispatch(x.scalar_type(), [&]<typename scalar_t>() {
if (scale.valid()) {
mul_add_kernel<scalar_t, unroll, false><<<grid, threadsPerBlock, 0, stream>>>(
x.data_ptr<scalar_t>(), scale.data_ptr<scalar_t>(), bias.data_ptr<scalar_t>(),
(scalar_t)scale_shift,
numel, numel_scale, numel_bias,
x.stride(0), batch_scale ? scale.stride(0) : 0, batch_bias ? bias.stride(0) : 0);
mul_add_kernel<scalar_t, unroll, false>
<<<grid, threadsPerBlock, 0, stream>>>(x.data_ptr<scalar_t>(),
scale.data_ptr<scalar_t>(),
bias.data_ptr<scalar_t>(),
(scalar_t)scale_shift,
numel,
numel_scale,
numel_bias,
x.stride(0),
batch_scale ? scale.stride(0) : 0,
batch_bias ? bias.stride(0) : 0);
} else {
mul_add_kernel<scalar_t, unroll, true><<<grid, threadsPerBlock, 0, stream>>>(
x.data_ptr<scalar_t>(), nullptr, bias.data_ptr<scalar_t>(),
(scalar_t)scale_shift,
numel, 1, numel_bias,
x.stride(0), 0, batch_bias ? bias.stride(0) : 0);
mul_add_kernel<scalar_t, unroll, true>
<<<grid, threadsPerBlock, 0, stream>>>(x.data_ptr<scalar_t>(),
nullptr,
bias.data_ptr<scalar_t>(),
(scalar_t)scale_shift,
numel,
1,
numel_bias,
x.stride(0),
0,
batch_bias ? bias.stride(0) : 0);
}
});
}
......@@ -134,8 +155,7 @@ Tensor argmax_sample(Tensor logits) {
dispatch(logits.scalar_type(), [&]<typename scalar_t>() {
argmax_sample_kernel<<<logits.shape[0], std::min(logits.shape[1], 1024), 0, stream>>>(
logits.data_ptr<scalar_t>(), out.data_ptr<int32_t>(), logits.shape[1]
);
logits.data_ptr<scalar_t>(), out.data_ptr<int32_t>(), logits.shape[1]);
});
return out;
......@@ -155,20 +175,17 @@ void splitqkv(Tensor qkv, Tensor q, Tensor k, Tensor v) {
assert(dim_k == dim_v);
assert(dim_q + dim_k + dim_v == qkv.shape[-1]);
int num_tokens = qkv.numel() / qkv.shape[-1];
dispatch(qkv.scalar_type(), [&]<typename scalar_t>() {
splitqkv_kernel<<<num_tokens, std::min(qkv.shape[-1], 1024), 0, stream>>>(
qkv.data_ptr<scalar_t>(),
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
dim_q,
dim_k
);
splitqkv_kernel<<<num_tokens, std::min(qkv.shape[-1], 1024), 0, stream>>>(qkv.data_ptr<scalar_t>(),
q.data_ptr<scalar_t>(),
k.data_ptr<scalar_t>(),
v.data_ptr<scalar_t>(),
dim_q,
dim_k);
});
}
template<size_t N>
......@@ -176,7 +193,7 @@ std::array<Tensor, N> split_mod(Tensor input) {
assert(input.shape[-1] % N == 0);
int threadsPerBlock = 1024;
int blocksPerGrid = (input.numel() + threadsPerBlock - 1) / threadsPerBlock;
int blocksPerGrid = (input.numel() + threadsPerBlock - 1) / threadsPerBlock;
auto stream = getCurrentCUDAStream();
......@@ -194,8 +211,7 @@ std::array<Tensor, N> split_mod(Tensor input) {
outPtr[k] = out[k].template data_ptr<scalar_t>();
}
split_mod_kernel<<<blocksPerGrid, threadsPerBlock, 0, stream>>>(
input.data_ptr<scalar_t>(),
outPtr, input.numel());
input.data_ptr<scalar_t>(), outPtr, input.numel());
});
return out;
......@@ -209,7 +225,7 @@ Tensor quant_static(Tensor x, float scale) {
assert((uintptr_t)x.data_ptr() % (x.scalar_size() * unroll) == 0);
int threadsPerBlock = 1024;
int blocksPerGrid = (x.numel() + threadsPerBlock * unroll - 1) / (threadsPerBlock * unroll);
int blocksPerGrid = (x.numel() + threadsPerBlock * unroll - 1) / (threadsPerBlock * unroll);
auto stream = getCurrentCUDAStream();
......@@ -228,9 +244,8 @@ Tensor quant_static_fuse_gelu(Tensor x, float scale) {
assert((uintptr_t)x.data_ptr() % (x.scalar_size() * unroll) == 0);
int threadsPerBlock = 1024;
int blocksPerGrid = (x.numel() + threadsPerBlock * unroll - 1) / (threadsPerBlock * unroll);
int blocksPerGrid = (x.numel() + threadsPerBlock * unroll - 1) / (threadsPerBlock * unroll);
auto stream = getCurrentCUDAStream();
......@@ -258,7 +273,7 @@ void cast(Tensor input, Tensor output) {
constexpr int unroll = 16 / std::max(sizeof(input_t), sizeof(output_t));
int threadsPerBlock = 1024;
int blocksPerGrid = (int)ceilDiv<int64_t>(input.numel(), threadsPerBlock * unroll);
int blocksPerGrid = (int)ceilDiv<int64_t>(input.numel(), threadsPerBlock * unroll);
cast_kernel<input_t, output_t, unroll><<<blocksPerGrid, threadsPerBlock, 0, stream>>>(
input.data_ptr<input_t>(), output.data_ptr<output_t>(), input.numel());
......@@ -271,17 +286,16 @@ void cast(Tensor input, Tensor output) {
Tensor topk(Tensor x, int k) {
constexpr int MAXK = 64 + 4;
const int N = x.shape[-1];
const int N = x.shape[-1];
const int batch = x.numel() / N;
assert(k <= N);
assert(k <= MAXK);
auto outShape = TensorShape(x.shape.dataExtent);
outShape[-1] = k;
outShape[-1] = k;
outShape.dataStride.clear();
Tensor out = Tensor::empty(outShape, Tensor::INT32, x.device());
auto stream = getCurrentCUDAStream();
......@@ -294,10 +308,7 @@ Tensor topk(Tensor x, int k) {
if constexpr (K > 0) {
dispatch(x.scalar_type(), [&]<typename scalar_t>() {
topk_kernel<scalar_t, K><<<ceilDiv(batch, 32), 32, 0, stream>>>(
x.data_ptr<scalar_t>(),
out.data_ptr<int>(),
N, x.stride(-2), batch
);
x.data_ptr<scalar_t>(), out.data_ptr<int>(), N, x.stride(-2), batch);
checkCUDA(cudaGetLastError());
});
}
......@@ -312,4 +323,4 @@ template std::array<Tensor, 4> split_mod<4>(Tensor input);
template std::array<Tensor, 5> split_mod<5>(Tensor input);
template std::array<Tensor, 6> split_mod<6>(Tensor input);
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/misc_kernels.h
View file @ 37a27712
......@@ -22,4 +22,4 @@ Tensor topk(Tensor x, int k);
template<size_t N>
std::array<Tensor, N> split_mod(Tensor input);
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/misc_kernels_impl.cuh
View file @ 37a27712
......@@ -9,7 +9,6 @@
namespace nunchaku::kernels {
template<typename T>
__global__ void add_kernel(T *a, T *b, T *c, size_t length) {
int i = threadIdx.x + blockIdx.x * blockDim.x;
......@@ -24,12 +23,21 @@ struct alignas(sizeof(T) * unroll) Tvec {
};
template<typename T, int unroll, bool no_scale>
__global__ void mul_add_kernel(T *x, T *scale, T *bias, T scale_shift, size_t length, int mod_scale, int mod_bias, int64_t batch_stride_x, int64_t batch_stride_scale, int64_t batch_stride_bias) {
__global__ void mul_add_kernel(T *x,
T *scale,
T *bias,
T scale_shift,
size_t length,
int mod_scale,
int mod_bias,
int64_t batch_stride_x,
int64_t batch_stride_scale,
int64_t batch_stride_bias) {
const int batch_id = blockIdx.y;
int thread = threadIdx.x + blockIdx.x * blockDim.x;
int i = thread * unroll;
int i_scale = i % mod_scale;
int i_bias = i % mod_bias;
int thread = threadIdx.x + blockIdx.x * blockDim.x;
int i = thread * unroll;
int i_scale = i % mod_scale;
int i_bias = i % mod_bias;
if (i >= length) {
return;
......@@ -37,9 +45,9 @@ __global__ void mul_add_kernel(T *x, T *scale, T *bias, T scale_shift, size_t le
using Tvec = nunchaku::kernels::Tvec<T, unroll>;
Tvec rx = *reinterpret_cast<Tvec *>(&x[i + batch_stride_x * batch_id]);
Tvec rx = *reinterpret_cast<Tvec *>(&x[i + batch_stride_x * batch_id]);
Tvec rscale = *reinterpret_cast<Tvec *>(&scale[i_scale + batch_stride_scale * batch_id]);
Tvec rbias = *reinterpret_cast<Tvec *>(&bias[i_bias + batch_stride_bias * batch_id]);
Tvec rbias = *reinterpret_cast<Tvec *>(&bias[i_bias + batch_stride_bias * batch_id]);
#pragma unroll
for (int k = 0; k < unroll; k++) {
......@@ -58,16 +66,16 @@ __global__ void mul_add_kernel(T *x, T *scale, T *bias, T scale_shift, size_t le
*reinterpret_cast<Tvec *>(&x[i + batch_stride_x * batch_id]) = rx;
// #pragma unroll
// for (int k = 0; k < unroll; k++) {
// // assert(i < length);
// x[i] = x[i] * scale[i_scale] + bias[i_bias];
// i++;
// i_scale++;
// i_bias++;
// // assert(i_scale < mod_scale);
// // assert(i_bias < mod_bias);
// }
// #pragma unroll
// for (int k = 0; k < unroll; k++) {
// // assert(i < length);
// x[i] = x[i] * scale[i_scale] + bias[i_bias];
// i++;
// i_scale++;
// i_bias++;
// // assert(i_scale < mod_scale);
// // assert(i_bias < mod_bias);
// }
}
template<typename T, size_t N>
......@@ -82,12 +90,13 @@ __global__ void split_mod_kernel(T *input, std::array<T *, N> output, size_t len
}
template<typename T>
__global__ void EmbeddingKernel(int32_t *__restrict__ input_id, T *__restrict__ output, T *__restrict__ lookup, int embed_dim) {
__global__ void
EmbeddingKernel(int32_t *__restrict__ input_id, T *__restrict__ output, T *__restrict__ lookup, int embed_dim) {
int i = blockIdx.x;
int32_t token_id = input_id[i];
int32_t token_id = input_id[i];
T *output_sample_ptr = output + i * embed_dim;
T *target_embed = lookup + token_id * embed_dim;
T *target_embed = lookup + token_id * embed_dim;
for (int j = threadIdx.x; j < embed_dim; j += blockDim.x) {
output_sample_ptr[j] = target_embed[j];
......@@ -97,7 +106,7 @@ __global__ void EmbeddingKernel(int32_t *__restrict__ input_id, T *__restrict__
template<typename T>
__global__ void argmax_sample_kernel(T *input, int32_t *output, int hidden_dim) {
float maxValue = -1e20;
int argmax = 0;
int argmax = 0;
for (int i = threadIdx.x; i < hidden_dim; i += blockDim.x) {
float data = (float)input[blockIdx.x * hidden_dim + i];
if (data > maxValue) {
......@@ -105,7 +114,7 @@ __global__ void argmax_sample_kernel(T *input, int32_t *output, int hidden_dim)
argmax = i;
}
}
// blockAllReduceMax seems to be broken when T=half
// blockAllReduceMax seems to be broken when T=half
float maxValueBlock = vllm::blockAllReduceMax(maxValue);
if (maxValue == maxValueBlock) {
output[blockIdx.x] = argmax;
......@@ -127,14 +136,14 @@ __global__ void splitqkv_kernel(T *qkv, T *q, T *k, T *v, int q_size, int kv_siz
}
}
template <typename T, int unroll>
__global__ void quant_kernel_static(const T * input, int8_t * output, T scale, size_t length) {
template<typename T, int unroll>
__global__ void quant_kernel_static(const T *input, int8_t *output, T scale, size_t length) {
int i = (blockIdx.x * blockDim.x + threadIdx.x) * unroll;
if (i >= length) {
return;
}
using Tvec = nunchaku::kernels::Tvec<T, unroll>;
using Tvec = nunchaku::kernels::Tvec<T, unroll>;
using I8vec = nunchaku::kernels::Tvec<int8_t, unroll>;
Tvec rinput = *reinterpret_cast<const Tvec *>(&input[i]);
......@@ -149,14 +158,14 @@ __global__ void quant_kernel_static(const T * input, int8_t * output, T scale, s
*reinterpret_cast<I8vec *>(&output[i]) = routput;
}
template <typename T, int unroll>
__global__ void quant_kernel_static_fuse_gelu(const T * input, int8_t * output, T scale, size_t length) {
template<typename T, int unroll>
__global__ void quant_kernel_static_fuse_gelu(const T *input, int8_t *output, T scale, size_t length) {
int i = (blockIdx.x * blockDim.x + threadIdx.x) * unroll;
if (i >= length) {
return;
}
using Tvec = nunchaku::kernels::Tvec<T, unroll>;
using Tvec = nunchaku::kernels::Tvec<T, unroll>;
using I8vec = nunchaku::kernels::Tvec<int8_t, unroll>;
Tvec rinput = *reinterpret_cast<const Tvec *>(&input[i]);
......@@ -175,10 +184,10 @@ template<typename Tin, typename Tout, int unroll>
__global__ void cast_kernel(const Tin *input, Tout *output, size_t length) {
const int i = (blockIdx.x * blockDim.x + threadIdx.x) * unroll;
using Tvec_in = nunchaku::kernels::Tvec<Tin, unroll>;
using Tvec_in = nunchaku::kernels::Tvec<Tin, unroll>;
using Tvec_out = nunchaku::kernels::Tvec<Tout, unroll>;
Tvec_in rinput = *reinterpret_cast<const Tvec_in *>(&input[i]);
Tvec_in rinput = *reinterpret_cast<const Tvec_in *>(&input[i]);
Tvec_out routput;
#pragma unroll
......@@ -196,16 +205,15 @@ __global__ void cast_kernel(const Tin *input, Tout *output, size_t length) {
// input: [..., N]
// output: [..., K] of index in reverse order
template<typename T, int K>
__global__
void topk_kernel(const T *input, int *output, int N, int strideInput, int numRows) {
const int row = blockIdx.x * blockDim.x + threadIdx.x;
__global__ void topk_kernel(const T *input, int *output, int N, int strideInput, int numRows) {
const int row = blockIdx.x * blockDim.x + threadIdx.x;
const int offset = row * strideInput;
if (row >= numRows) {
return;
}
T val[K];
T val[K];
int16_t idx[K];
#pragma unroll
......@@ -224,7 +232,7 @@ void topk_kernel(const T *input, int *output, int N, int strideInput, int numRow
for (int i = K; i < N; i++) {
T newval = input[offset + i];
T minval = val[0];
T minval = val[0];
int minpos = 0;
#pragma unroll
for (int j = 1; j < K; j++) {
......@@ -259,4 +267,4 @@ void topk_kernel(const T *input, int *output, int N, int strideInput, int numRow
}
}
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/reduction_utils.cuh
View file @ 37a27712
/*
* Adapted from https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/reduce_kernel_utils.cuh
* Adapted from
* https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/reduce_kernel_utils.cuh
* Copyright (c) 2023, The vLLM team.
* Copyright (c) 2020-2023, NVIDIA CORPORATION. All rights reserved.
*
......@@ -18,86 +19,80 @@
#pragma once
#define FINAL_MASK 0xffffffff
namespace vllm {
template<typename T>
__inline__ __device__ T warpReduceSum(T val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1)
val += __shfl_xor_sync(0xffffffff, val, mask, 32);
return val;
for (int mask = 16; mask > 0; mask >>= 1)
val += __shfl_xor_sync(0xffffffff, val, mask, 32);
return val;
}
template <typename T, int NUM>
__inline__ __device__ T warpReduceSumV2(T* val)
{
template<typename T, int NUM>
__inline__ __device__ T warpReduceSumV2(T *val) {
#pragma unroll
for (int i = 0; i < NUM; i++)
{
for (int i = 0; i < NUM; i++) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1)
val[i] += __shfl_xor_sync(FINAL_MASK, val[i], mask, 32);
}
return (T) (0.0f);
return (T)(0.0f);
}
/* Calculate the sum of all elements in a block */
template<typename T>
__inline__ __device__ T blockReduceSum(T val) {
static __shared__ T shared[32];
int lane = threadIdx.x & 0x1f;
int wid = threadIdx.x >> 5;
static __shared__ T shared[32];
int lane = threadIdx.x & 0x1f;
int wid = threadIdx.x >> 5;
val = warpReduceSum<T>(val);
val = warpReduceSum<T>(val);
if (lane == 0)
shared[wid] = val;
if (lane == 0)
shared[wid] = val;
__syncthreads();
__syncthreads();
// Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
// blockDim.x is not divided by 32
val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
val = warpReduceSum<T>(val);
return val;
// Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
// blockDim.x is not divided by 32
val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
val = warpReduceSum<T>(val);
return val;
}
/* Calculate the sum of all elements in a block */
template<typename T>
__inline__ __device__ T blockAllReduceSum(T val) {
static __shared__ T shared[32];
int lane = threadIdx.x & 0x1f;
int wid = threadIdx.x >> 5;
static __shared__ T shared[32];
int lane = threadIdx.x & 0x1f;
int wid = threadIdx.x >> 5;
val = warpReduceSum<T>(val);
val = warpReduceSum<T>(val);
if (lane == 0)
shared[wid] = val;
if (lane == 0)
shared[wid] = val;
__syncthreads();
__syncthreads();
// Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
// blockDim.x is not divided by 32
val = (lane < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
val = warpReduceSum<T>(val);
return val;
// Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
// blockDim.x is not divided by 32
val = (lane < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
val = warpReduceSum<T>(val);
return val;
}
template <typename T, int NUM>
__inline__ __device__ T blockReduceSumV2(T* val)
{
template<typename T, int NUM>
__inline__ __device__ T blockReduceSumV2(T *val) {
static __shared__ T shared[NUM][33];
int lane = threadIdx.x & 0x1f;
int wid = threadIdx.x >> 5;
int wid = threadIdx.x >> 5;
warpReduceSumV2<T, NUM>(val);
if (lane == 0)
{
if (lane == 0) {
#pragma unroll
for (int i = 0; i < NUM; i++)
{
for (int i = 0; i < NUM; i++) {
shared[i][wid] = val[i];
}
}
......@@ -106,17 +101,15 @@ __inline__ __device__ T blockReduceSumV2(T* val)
bool is_mask = threadIdx.x < (blockDim.x / 32.f);
#pragma unroll
for (int i = 0; i < NUM; i++)
{
val[i] = is_mask ? shared[i][lane] : (T) (0.0f);
for (int i = 0; i < NUM; i++) {
val[i] = is_mask ? shared[i][lane] : (T)(0.0f);
}
warpReduceSumV2<T, NUM>(val);
return (T) 0.0f;
return (T)0.0f;
}
template<typename T>
__inline__ __device__ T warpReduceMax(T val)
{
__inline__ __device__ T warpReduceMax(T val) {
#pragma unroll
for (int mask = 16; mask > 0; mask >>= 1)
val = max(val, __shfl_xor_sync(0xffffffff, val, mask, 32));
......@@ -124,13 +117,12 @@ __inline__ __device__ T warpReduceMax(T val)
}
/* Calculate the maximum of all elements in a block */
template<typename T>
__inline__ __device__ T blockReduceMax(T val)
{
__inline__ __device__ T blockReduceMax(T val) {
static __shared__ T shared[32];
int lane = threadIdx.x & 0x1f; // in-warp idx
int wid = threadIdx.x >> 5; // warp idx
val = warpReduceMax(val); // get maxx in each warp
if (lane == 0) // record in-warp maxx by warp Idx
int lane = threadIdx.x & 0x1f; // in-warp idx
int wid = threadIdx.x >> 5; // warp idx
val = warpReduceMax(val); // get maxx in each warp
if (lane == 0) // record in-warp maxx by warp Idx
shared[wid] = val;
__syncthreads();
// Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
......@@ -141,16 +133,15 @@ __inline__ __device__ T blockReduceMax(T val)
}
/* Calculate the maximum of all elements in a block */
template <typename T>
__inline__ __device__ T blockAllReduceMax(T val)
{
template<typename T>
__inline__ __device__ T blockAllReduceMax(T val) {
static __shared__ T shared[32];
int lane = threadIdx.x & 0x1f; // in-warp idx
int wid = threadIdx.x >> 5; // warp idx
int wid = threadIdx.x >> 5; // warp idx
val = warpReduceMax(val); // get maxx in each warp
val = warpReduceMax(val); // get maxx in each warp
if (lane == 0) // record in-warp maxx by warp Idx
if (lane == 0) // record in-warp maxx by warp Idx
shared[wid] = val;
__syncthreads();
......@@ -163,8 +154,4 @@ __inline__ __device__ T blockAllReduceMax(T val)
return val;
}
} // namespace vllm
src/kernels/utils.cuh
View file @ 37a27712
// Adated from FasterTransformer, https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
// Adated from FasterTransformer,
// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/kernels/decoder_masked_multihead_attention/decoder_masked_multihead_attention_template.hpp
#pragma once
#include <cassert>
......@@ -14,8 +15,7 @@
#include <cuda_bf16.h>
#endif
__device__ __forceinline__
static void trap_unsupported_arch() {
__device__ __forceinline__ static void trap_unsupported_arch() {
if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0) {
printf("This kernel is not supported on your GPU\n");
}
......@@ -25,64 +25,143 @@ static void trap_unsupported_arch() {
}
#if defined(ENABLE_BF16) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
__device__ __forceinline__
static __nv_bfloat162 __hfma2(const __nv_bfloat162 a, const __nv_bfloat162 b, const __nv_bfloat162 c) {
__device__ __forceinline__ static __nv_bfloat162
__hfma2(const __nv_bfloat162 a, const __nv_bfloat162 b, const __nv_bfloat162 c) {
trap_unsupported_arch();
return __nv_bfloat162(0.0f, 0.0f);
}
#endif
template<typename T> struct num_elems;
template <> struct num_elems<float> { static constexpr int value = 1; };
template <> struct num_elems<float2> { static constexpr int value = 2; };
template <> struct num_elems<float4> { static constexpr int value = 4; };
template <> struct num_elems<half> { static constexpr int value = 1; };
template <> struct num_elems<half2> { static constexpr int value = 2; };
template<typename T>
struct num_elems;
template<>
struct num_elems<float> {
static constexpr int value = 1;
};
template<>
struct num_elems<float2> {
static constexpr int value = 2;
};
template<>
struct num_elems<float4> {
static constexpr int value = 4;
};
template<>
struct num_elems<half> {
static constexpr int value = 1;
};
template<>
struct num_elems<half2> {
static constexpr int value = 2;
};
#ifdef ENABLE_BF16
template <> struct num_elems<__nv_bfloat16> { static constexpr int value = 1; };
template <> struct num_elems<__nv_bfloat162> { static constexpr int value = 2; };
template<>
struct num_elems<__nv_bfloat16> {
static constexpr int value = 1;
};
template<>
struct num_elems<__nv_bfloat162> {
static constexpr int value = 2;
};
#endif
#ifdef ENABLE_FP8
template <> struct num_elems<__nv_fp8_e4m3> { static constexpr int value = 1; };
template <> struct num_elems<__nv_fp8x2_e4m3> { static constexpr int value = 2; };
template<>
struct num_elems<__nv_fp8_e4m3> {
static constexpr int value = 1;
};
template<>
struct num_elems<__nv_fp8x2_e4m3> {
static constexpr int value = 2;
};
#endif
template<typename T, int num> struct packed_as;
template<typename T> struct packed_as<T, 1> { using type = T; };
template<> struct packed_as<half, 2> { using type = half2; };
template<> struct packed_as<float, 2> { using type = float2; };
template<> struct packed_as<int8_t, 2> { using type = int16_t; };
template<> struct packed_as<int32_t, 2> { using type = int2; };
template<> struct packed_as<half2, 1> { using type = half; };
template<> struct packed_as<float2, 1> { using type = float; };
template<typename T, int num>
struct packed_as;
template<typename T>
struct packed_as<T, 1> {
using type = T;
};
template<>
struct packed_as<half, 2> {
using type = half2;
};
template<>
struct packed_as<float, 2> {
using type = float2;
};
template<>
struct packed_as<int8_t, 2> {
using type = int16_t;
};
template<>
struct packed_as<int32_t, 2> {
using type = int2;
};
template<>
struct packed_as<half2, 1> {
using type = half;
};
template<>
struct packed_as<float2, 1> {
using type = float;
};
#ifdef ENABLE_BF16
template<> struct packed_as<__nv_bfloat16, 2> { using type = __nv_bfloat162; };
template<> struct packed_as<__nv_bfloat162, 1> { using type = __nv_bfloat16; };
template<>
struct packed_as<__nv_bfloat16, 2> {
using type = __nv_bfloat162;
};
template<>
struct packed_as<__nv_bfloat162, 1> {
using type = __nv_bfloat16;
};
#endif
#ifdef ENABLE_FP8
template<> struct packed_as<__nv_fp8_e4m3, 2> { using type = __nv_fp8x2_e4m3; };
template<> struct packed_as<__nv_fp8x2_e4m3, 1> { using type = __nv_fp8_e4m3; };
template<> struct packed_as<__nv_fp8_e5m2, 2> { using type = __nv_fp8x2_e5m2; };
template<> struct packed_as<__nv_fp8x2_e5m2, 1> { using type = __nv_fp8_e5m2; };
template<>
struct packed_as<__nv_fp8_e4m3, 2> {
using type = __nv_fp8x2_e4m3;
};
template<>
struct packed_as<__nv_fp8x2_e4m3, 1> {
using type = __nv_fp8_e4m3;
};
template<>
struct packed_as<__nv_fp8_e5m2, 2> {
using type = __nv_fp8x2_e5m2;
};
template<>
struct packed_as<__nv_fp8x2_e5m2, 1> {
using type = __nv_fp8_e5m2;
};
#endif
inline __device__ float2 operator*(float2 a, float2 b) { return make_float2(a.x * b.x, a.y * b.y); }
inline __device__ float2 operator+(float2 a, float2 b) { return make_float2(a.x + b.x, a.y + b.y); }
inline __device__ float2 operator-(float2 a, float2 b) { return make_float2(a.x - b.x, a.y - b.y); }
inline __device__ float2 operator*(float2 a, float2 b) {
return make_float2(a.x * b.x, a.y * b.y);
}
inline __device__ float2 operator+(float2 a, float2 b) {
return make_float2(a.x + b.x, a.y + b.y);
}
inline __device__ float2 operator-(float2 a, float2 b) {
return make_float2(a.x - b.x, a.y - b.y);
}
inline __device__ float2 operator*(float2 a, float b) { return make_float2(a.x * b, a.y * b); }
inline __device__ float2 operator+(float2 a, float b) { return make_float2(a.x + b, a.y + b); }
inline __device__ float2 operator-(float2 a, float b) { return make_float2(a.x - b, a.y - b); }
inline __device__ float2 operator*(float2 a, float b) {
return make_float2(a.x * b, a.y * b);
}
inline __device__ float2 operator+(float2 a, float b) {
return make_float2(a.x + b, a.y + b);
}
inline __device__ float2 operator-(float2 a, float b) {
return make_float2(a.x - b, a.y - b);
}
static inline __device__ int8_t float_to_int8_rn(float x)
{
static inline __device__ int8_t float_to_int8_rn(float x) {
uint32_t dst;
asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
return reinterpret_cast<const int8_t&>(dst);
return reinterpret_cast<const int8_t &>(dst);
}
template<typename T>
inline __device__ T ldg(const T* val) {
inline __device__ T ldg(const T *val) {
return __ldg(val);
}
......@@ -90,15 +169,13 @@ inline __device__ T ldg(const T* val) {
#define bf1622float2 __bfloat1622float2
#define float22bf162 __float22bfloat162_rn
#define bf162bf162 __bfloat162bfloat162
inline __device__ int16_t bf1622int16(__nv_bfloat162 val)
{
inline __device__ int16_t bf1622int16(__nv_bfloat162 val) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
float2 f_val;
f_val.x = max(min(__low2float(val), 127.f), -128.f);
f_val.y = max(min(__high2float(val), 127.f), -128.f);
union
{
union {
int8_t int8[2];
int16_t int16;
};
......@@ -110,8 +187,7 @@ inline __device__ int16_t bf1622int16(__nv_bfloat162 val)
val = __hmin2(val, make_bfloat162(127., 127.));
val = __hmax2(val, make_bfloat162(-128., -128.));
union
{
union {
int8_t int8[2];
int16_t int16;
};
......@@ -125,7 +201,7 @@ inline __device__ int16_t bf1622int16(__nv_bfloat162 val)
#if ENABLE_BF16
template<>
inline __device__ __nv_bfloat162 ldg(const __nv_bfloat162* val) {
inline __device__ __nv_bfloat162 ldg(const __nv_bfloat162 *val) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
return val[0];
#else
......@@ -134,7 +210,7 @@ inline __device__ __nv_bfloat162 ldg(const __nv_bfloat162* val) {
}
template<>
inline __device__ __nv_bfloat16 ldg(const __nv_bfloat16* val) {
inline __device__ __nv_bfloat16 ldg(const __nv_bfloat16 *val) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
return val[0];
#else
......@@ -143,59 +219,49 @@ inline __device__ __nv_bfloat16 ldg(const __nv_bfloat16* val) {
}
#endif // ENABLE_BF16
template <typename T_OUT, typename T_IN>
__device__ inline T_OUT cuda_cast(T_IN val)
{
template<typename T_OUT, typename T_IN>
__device__ inline T_OUT cuda_cast(T_IN val) {
return val;
}
template <>
__device__ inline float2 cuda_cast<float2, int2>(int2 val)
{
template<>
__device__ inline float2 cuda_cast<float2, int2>(int2 val) {
return make_float2(val.x, val.y);
}
template <>
__device__ inline float2 cuda_cast<float2, float>(float val)
{
template<>
__device__ inline float2 cuda_cast<float2, float>(float val) {
return make_float2(val, val);
}
template <>
__device__ inline float2 cuda_cast<float2, half2>(half2 val)
{
template<>
__device__ inline float2 cuda_cast<float2, half2>(half2 val) {
return __half22float2(val);
}
template <>
__device__ inline half2 cuda_cast<half2, float2>(float2 val)
{
template<>
__device__ inline half2 cuda_cast<half2, float2>(float2 val) {
return __float22half2_rn(val);
}
template <>
__device__ inline half2 cuda_cast<half2, float>(float val)
{
template<>
__device__ inline half2 cuda_cast<half2, float>(float val) {
return __float2half2_rn(val);
}
template <>
__device__ inline half2 cuda_cast<half2, half>(half val)
{
template<>
__device__ inline half2 cuda_cast<half2, half>(half val) {
return __half2half2(val);
}
template <>
__device__ inline int8_t cuda_cast<int8_t, half>(half val)
{
union
{
template<>
__device__ inline int8_t cuda_cast<int8_t, half>(half val) {
union {
int8_t int8[2];
int16_t int16;
};
union
{
union {
half fp16;
int16_t int16_in;
};
......@@ -205,11 +271,9 @@ __device__ inline int8_t cuda_cast<int8_t, half>(half val)
return int8[0];
}
template <>
__device__ inline int16_t cuda_cast<int16_t, half2>(half2 val)
{
union
{
template<>
__device__ inline int16_t cuda_cast<int16_t, half2>(half2 val) {
union {
int8_t int8[2];
int16_t int16;
};
......@@ -219,11 +283,9 @@ __device__ inline int16_t cuda_cast<int16_t, half2>(half2 val)
return int16;
}
template <>
__device__ inline int8_t cuda_cast<int8_t, float>(float val)
{
union
{
template<>
__device__ inline int8_t cuda_cast<int8_t, float>(float val) {
union {
int8_t int8[2];
int16_t int16;
};
......@@ -232,11 +294,9 @@ __device__ inline int8_t cuda_cast<int8_t, float>(float val)
return int8[0];
}
template <>
__device__ inline int16_t cuda_cast<int16_t, float2>(float2 val)
{
union
{
template<>
__device__ inline int16_t cuda_cast<int16_t, float2>(float2 val) {
union {
int8_t int8[2];
int16_t int16;
};
......@@ -246,11 +306,9 @@ __device__ inline int16_t cuda_cast<int16_t, float2>(float2 val)
return int16;
}
template <>
__device__ inline half2 cuda_cast<half2, int16_t>(int16_t val)
{
union
{
template<>
__device__ inline half2 cuda_cast<half2, int16_t>(int16_t val) {
union {
int8_t int8[2];
int16_t int16;
};
......@@ -259,11 +317,9 @@ __device__ inline half2 cuda_cast<half2, int16_t>(int16_t val)
return make_half2(int8[0], int8[1]);
}
template <>
__device__ inline float2 cuda_cast<float2, int16_t>(int16_t val)
{
union
{
template<>
__device__ inline float2 cuda_cast<float2, int16_t>(int16_t val) {
union {
int8_t int8[2];
int16_t int16;
};
......@@ -273,83 +329,69 @@ __device__ inline float2 cuda_cast<float2, int16_t>(int16_t val)
}
#ifdef ENABLE_BF16
template <>
__device__ inline __nv_bfloat16 cuda_cast(int32_t val)
{
template<>
__device__ inline __nv_bfloat16 cuda_cast(int32_t val) {
return static_cast<float>(val);
}
template <>
__device__ inline __nv_bfloat16 cuda_cast(int8_t val)
{
template<>
__device__ inline __nv_bfloat16 cuda_cast(int8_t val) {
return static_cast<float>(val);
}
template <>
__device__ inline int8_t cuda_cast(__nv_bfloat16 val)
{
template<>
__device__ inline int8_t cuda_cast(__nv_bfloat16 val) {
return static_cast<float>(val);
}
template <>
__device__ inline float cuda_cast<float, __nv_bfloat16>(__nv_bfloat16 val)
{
template<>
__device__ inline float cuda_cast<float, __nv_bfloat16>(__nv_bfloat16 val) {
return __bfloat162float(val);
}
template <>
__device__ inline float2 cuda_cast<float2, __nv_bfloat162>(__nv_bfloat162 val)
{
template<>
__device__ inline float2 cuda_cast<float2, __nv_bfloat162>(__nv_bfloat162 val) {
return bf1622float2(val);
}
template <>
__device__ inline half cuda_cast<half, __nv_bfloat16>(__nv_bfloat16 val)
{
template<>
__device__ inline half cuda_cast<half, __nv_bfloat16>(__nv_bfloat16 val) {
return __float2half(__bfloat162float(val));
}
template <>
__device__ inline int16_t cuda_cast<int16_t, __nv_bfloat162>(__nv_bfloat162 val)
{
template<>
__device__ inline int16_t cuda_cast<int16_t, __nv_bfloat162>(__nv_bfloat162 val) {
return bf1622int16(val);
}
template <>
__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val)
{
template<>
__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, float>(float val) {
return __float2bfloat16(val);
}
template <>
__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val)
{
template<>
__device__ inline __nv_bfloat16 cuda_cast<__nv_bfloat16, half>(half val) {
return __float2bfloat16(__half2float(val));
}
template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val)
{
template<>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, __nv_bfloat16>(__nv_bfloat16 val) {
return bf162bf162(val);
}
template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val)
{
template<>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float>(float val) {
return __float2bfloat162_rn(val);
}
template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val)
{
template<>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, float2>(float2 val) {
return float22bf162(val);
}
template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val)
{
union
{
template<>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val) {
union {
int8_t int8[2];
int16_t int16;
};
......@@ -361,72 +403,57 @@ __device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, int16_t>(int16_t val)
return res;
}
template <>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val)
{
template<>
__device__ inline __nv_bfloat162 cuda_cast<__nv_bfloat162, half2>(half2 val) {
return float22bf162(__half22float2(val));
}
#endif // ENABLE BF16
template <typename f16_t>
__device__ __forceinline__
packed_as<f16_t, 2>::type
f162f162(f16_t x);
template<typename f16_t>
__device__ __forceinline__ packed_as<f16_t, 2>::type f162f162(f16_t x);
template <>
__device__ __forceinline__
packed_as<half, 2>::type
f162f162<half>(half x)
{
return __half2half2(x);
template<>
__device__ __forceinline__ packed_as<half, 2>::type f162f162<half>(half x) {
return __half2half2(x);
}
#ifdef ENABLE_BF16
template <>
__device__ __forceinline__
packed_as<__nv_bfloat16, 2>::type
f162f162<__nv_bfloat16>(__nv_bfloat16 x)
{
return __bfloat162bfloat162(x);
}
# endif
template <typename To, typename Ti>
__device__ inline To cuda_sum(Ti val)
{
template<>
__device__ __forceinline__ packed_as<__nv_bfloat16, 2>::type f162f162<__nv_bfloat16>(__nv_bfloat16 x) {
return __bfloat162bfloat162(x);
}
#endif
template<typename To, typename Ti>
__device__ inline To cuda_sum(Ti val) {
return cuda_cast<To>(val);
};
template <typename To>
__device__ inline To cuda_sum(float2 val)
{
template<typename To>
__device__ inline To cuda_sum(float2 val) {
return cuda_cast<To>(val.x + val.y);
};
// Unary maximum: compute the max of a vector type
template <typename To, typename Ti>
__device__ inline To cuda_max(Ti val)
{
template<typename To, typename Ti>
__device__ inline To cuda_max(Ti val) {
return cuda_cast<To>(val);
};
template <>
__device__ inline float cuda_max(float2 val)
{
template<>
__device__ inline float cuda_max(float2 val) {
return fmaxf(val.x, val.y);
}
template <>
__device__ inline half cuda_max(half2 val)
{
template<>
__device__ inline half cuda_max(half2 val) {
return __hmax(val.x, val.y);
}
#ifdef ENABLE_BF16
template <>
__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val)
{
template<>
__device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val) {
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800))
return __hmax(val.x, val.y);
#else
......@@ -437,57 +464,49 @@ __device__ inline __nv_bfloat16 cuda_max(__nv_bfloat162 val)
#endif
// Binary maximum: compute the max of two scalar types
template <typename T>
__device__ inline T cuda_max(T val1, T val2)
{
template<typename T>
__device__ inline T cuda_max(T val1, T val2) {
return (val1 > val2) ? val1 : val2;
}
template <typename T>
__device__ inline T cuda_abs(T val)
{
template<typename T>
__device__ inline T cuda_abs(T val) {
assert(false);
return {};
}
template <>
__device__ inline float cuda_abs(float val)
{
template<>
__device__ inline float cuda_abs(float val) {
return fabs(val);
}
template <>
__device__ inline float2 cuda_abs(float2 val)
{
template<>
__device__ inline float2 cuda_abs(float2 val) {
return make_float2(fabs(val.x), fabs(val.y));
}
template <>
__device__ inline half cuda_abs(half val)
{
template<>
__device__ inline half cuda_abs(half val) {
return __habs(val);
}
template <>
__device__ inline half2 cuda_abs(half2 val)
{
template<>
__device__ inline half2 cuda_abs(half2 val) {
return __habs2(val);
}
#ifdef ENABLE_BF16
#if __CUDA_ARCH__ >= 800 || !defined(__CUDA_ARCH__)
template <>
__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val)
{
template<>
__device__ inline __nv_bfloat16 cuda_abs(__nv_bfloat16 val) {
return __habs(val);
}
template <>
__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val)
{
template<>
__device__ inline __nv_bfloat162 cuda_abs(__nv_bfloat162 val) {
return __habs2(val);
}
#endif
#endif // ENABLE_FP16
\ No newline at end of file
#endif // ENABLE_FP16
src/kernels/zgemm/attention.cu
View file @ 37a27712
......@@ -7,17 +7,15 @@
namespace nunchaku::kernels {
void attention_fp16(
Tensor q, // packed [Batch, Head, TokensQ, HEAD_DIM]
Tensor k, // packed [Batch, Head, TokensKV, HEAD_DIM]
Tensor v, // packed [Batch, Head, TokensKV, HEAD_DIM]
Tensor o, // linear [Batch, TokensQ, Head * HEAD_DIM]
float scale
) {
int sizeBatch = q.shape[0];
int numHeads = q.shape[1];
int numTokensQ = q.shape[2];
int headDim = q.shape[3];
void attention_fp16(Tensor q, // packed [Batch, Head, TokensQ, HEAD_DIM]
Tensor k, // packed [Batch, Head, TokensKV, HEAD_DIM]
Tensor v, // packed [Batch, Head, TokensKV, HEAD_DIM]
Tensor o, // linear [Batch, TokensQ, Head * HEAD_DIM]
float scale) {
int sizeBatch = q.shape[0];
int numHeads = q.shape[1];
int numTokensQ = q.shape[2];
int headDim = q.shape[3];
int numTokensKV = k.shape[2];
assert(o.ndims() == 3);
......@@ -55,20 +53,20 @@ void attention_fp16(
assert(headDim == Attention::HEAD_DIM);
auto launch = [&]<typename Epilogue>(Epilogue::Arguments args) {
dim3 grid(numTokensQ / Attention::BLOCK_M, numHeads, sizeBatch);
using packed_q_t = typename Attention::packed_q_t;
using packed_k_t = typename Attention::packed_k_t;
using packed_v_t = typename Attention::packed_v_t;
auto func = invoke_kernel<typename Attention::attention_fp16_kernel<Epilogue>,
const packed_q_t *,
const packed_k_t *,
const packed_v_t *,
float,
int, int,
typename Epilogue::Arguments,
bool>;
auto func = invoke_kernel<typename Attention::attention_fp16_kernel<Epilogue>,
const packed_q_t *,
const packed_k_t *,
const packed_v_t *,
float,
int,
int,
typename Epilogue::Arguments,
bool>;
shmem = std::max(shmem, Attention::template attention_fp16_kernel<Epilogue>::SHMEM_SIZE);
......@@ -76,26 +74,23 @@ void attention_fp16(
checkCUDA(cudaFuncSetAttribute(func, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem));
}
func<<<grid, GEMM::WARP_SIZE * GEMM::NUM_WARPS, shmem, getCurrentCUDAStream()>>>(
q.data_ptr<packed_q_t>(),
k.data_ptr<packed_k_t>(),
v.data_ptr<packed_v_t>(),
scale,
numTokensQ, numTokensKV,
args,
false
);
func<<<grid, GEMM::WARP_SIZE * GEMM::NUM_WARPS, shmem, getCurrentCUDAStream()>>>(q.data_ptr<packed_q_t>(),
k.data_ptr<packed_k_t>(),
v.data_ptr<packed_v_t>(),
scale,
numTokensQ,
numTokensKV,
args,
false);
checkCUDA(cudaGetLastError());
};
launch.template operator()<typename GEMM::EpilogueDefault>(typename GEMM::EpilogueDefault::Arguments{
.out = o.data_ptr<typename GEMM::half_t>(),
.out = o.data_ptr<typename GEMM::half_t>(),
.actualM = sizeBatch * numTokensQ,
.actualN = numHeads * headDim,
});
});
}
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/zgemm/attention.cuh
View file @ 37a27712
......@@ -4,7 +4,7 @@
namespace nunchaku::kernels {
// M: Q tokens
// M: Q tokens
// N: V HEAD_DIM
// K: K tokens
// D: QK HEAD_DIM
......@@ -12,21 +12,21 @@ template<bool bf16out>
struct AttentionFP16Config {
static constexpr int HEAD_DIM = 128;
static constexpr int BLOCK_M = 128;
static constexpr int BLOCK_M = 128;
static constexpr int WARP_SIZE = 32;
static constexpr int NUM_WARPS = 8;
static constexpr int WARP_K = 32;
static constexpr int INSN_M = 16;
static constexpr int INSN_N = 16;
static constexpr int INSN_M = 16;
static constexpr int INSN_N = 16;
static constexpr int INSN_K_QK = 16;
static constexpr int INSN_K_PV = 16;
using half_t = half;
using half2_t = half2;
using epilogue_half_t = typename std::conditional_t<bf16out, __nv_bfloat16, half>;
using epilogue_half_t = typename std::conditional_t<bf16out, __nv_bfloat16, half>;
using epilogue_half2_t = typename std::conditional_t<bf16out, __nv_bfloat162, half2>;
};
......@@ -67,8 +67,8 @@ public:
#endif
struct GEMMConfig {
static constexpr int BLOCK_M = AttentionConfig::BLOCK_M;
static constexpr int BLOCK_N = AttentionConfig::HEAD_DIM;
static constexpr int BLOCK_M = AttentionConfig::BLOCK_M;
static constexpr int BLOCK_N = AttentionConfig::HEAD_DIM;
static constexpr int WARP_SIZE = AttentionConfig::WARP_SIZE;
static constexpr int NUM_WARPS = AttentionConfig::NUM_WARPS;
......@@ -86,23 +86,24 @@ public:
static constexpr int WARP_N = HEAD_DIM;
static constexpr int WARP_D = HEAD_DIM;
static constexpr int WARP_M_TILES = WARP_M / INSN_M;
static constexpr int WARP_N_TILES = WARP_N / INSN_N;
static constexpr int WARP_K_TILES_QK = WARP_K / INSN_N; // when multiplying Q*K, K is on dimension of N in MMA instruction
static constexpr int WARP_M_TILES = WARP_M / INSN_M;
static constexpr int WARP_N_TILES = WARP_N / INSN_N;
static constexpr int WARP_K_TILES_QK =
WARP_K / INSN_N; // when multiplying Q*K, K is on dimension of N in MMA instruction
static constexpr int WARP_K_TILES_PV = WARP_K / INSN_K_PV;
static constexpr int WARP_D_TILES = WARP_D / INSN_K_QK;
using packed_q_t = uint4;
using packed_k_t = uint4;
using packed_v_t = uint4;
using q_warp = std::array<packed_q_t, WARP_M_TILES * WARP_D_TILES>;
using k_warp = std::array<packed_k_t, WARP_K_TILES_QK * WARP_D_TILES>;
using v_warp = std::array<packed_v_t, WARP_K_TILES_PV * WARP_N_TILES>;
using q_warp = std::array<packed_q_t, WARP_M_TILES * WARP_D_TILES>;
using k_warp = std::array<packed_k_t, WARP_K_TILES_QK * WARP_D_TILES>;
using v_warp = std::array<packed_v_t, WARP_K_TILES_PV * WARP_N_TILES>;
using packed_p_t = uint4;
using p_warp = std::array<packed_v_t, WARP_M_TILES * WARP_K_TILES_PV>;
using p_warp = std::array<packed_v_t, WARP_M_TILES * WARP_K_TILES_PV>;
using packed_fpsum_t = uint4;
using packed_fpsum_t = uint4;
using packed_f32psum_t = typename GEMM::packed_f32psum_t;
using qk_warp = std::array<packed_f32psum_t, WARP_M_TILES * WARP_K_TILES_QK>;
......@@ -112,16 +113,15 @@ public:
using rowval_warp = std::array<float2, WARP_M_TILES>;
struct BlockInfo {
int bm; // M: Q tokens, bm: block id of M
int head; // H: head
int batch; // B: batch
int numBlocksM;
int bm; // M: Q tokens, bm: block id of M
int head; // H: head
int batch; // B: batch
int numBlocksM;
int numHeads;
int numBatch;
};
__device__ __forceinline__
static packed_fpsum_t packed_fp32_to_fp16(packed_f32psum_t input) {
__device__ __forceinline__ static packed_fpsum_t packed_fp32_to_fp16(packed_f32psum_t input) {
std::array<half2_t, 4> results;
for (int i = 0; i < 4; i++) {
results[i] = float22half2<half2_t>(float2(input.data[i * 2], input.data[i * 2 + 1]));
......@@ -129,21 +129,19 @@ public:
return kernels::bit_cast<packed_fpsum_t>(results);
}
__device__ __forceinline__
static packed_f32psum_t packed_fp16_to_fp32(packed_fpsum_t input) {
__device__ __forceinline__ static packed_f32psum_t packed_fp16_to_fp32(packed_fpsum_t input) {
auto arr = kernels::bit_cast<std::array<half2_t, 4>>(input);
packed_f32psum_t results;
for (int i = 0; i < 4; i++) {
float2 tmp = half22float2(arr[i]);
results.data[i * 2] = tmp.x;
float2 tmp = half22float2(arr[i]);
results.data[i * 2] = tmp.x;
results.data[i * 2 + 1] = tmp.y;
}
return results;
}
// q: [batch, head, bm, NUM_WARPS, WARP_M_TILES, WARP_D_TILES, WARP_SIZE] of packed_q_t
__device__ __forceinline__
static void load_q(const packed_q_t *ptr, q_warp &out, bool pred) {
__device__ __forceinline__ static void load_q(const packed_q_t *ptr, q_warp &out, bool pred) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -157,8 +155,7 @@ public:
}
// k: [batch, head, ktile, WARP_K_TILES_QK, WARP_D_TILES, WARP_SIZE] of packed_k_t
__device__ __forceinline__
static void load_k(const packed_k_t *ptr, int ktile, k_warp &out, bool pred) {
__device__ __forceinline__ static void load_k(const packed_k_t *ptr, int ktile, k_warp &out, bool pred) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -172,8 +169,7 @@ public:
}
// v: [batch, head, ktile, WARP_K_TILES_PV, WARP_N_TILES, WARP_SIZE] of packed_v_t
__device__ __forceinline__
static void load_v(const packed_v_t *ptr, int ktile, v_warp &out, bool pred) {
__device__ __forceinline__ static void load_v(const packed_v_t *ptr, int ktile, v_warp &out, bool pred) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -186,34 +182,31 @@ public:
});
}
__device__ __forceinline__
static packed_fpsum_t mma_f16xf16_f16(packed_fpsum_t a, packed_fpsum_t b, packed_fpsum_t psum) {
__device__ __forceinline__ static packed_fpsum_t
mma_f16xf16_f16(packed_fpsum_t a, packed_fpsum_t b, packed_fpsum_t psum) {
uint2 out1 = mma_m16n8k16_f16f16f16f16(a, uint2(b.x, b.y), uint2(psum.x, psum.y));
uint2 out2 = mma_m16n8k16_f16f16f16f16(a, uint2(b.z, b.w), uint2(psum.z, psum.w));
return packed_fpsum_t{out1.x, out1.y, out2.x, out2.y};
}
// set nan values to -inf
__device__ __forceinline__
static half2_t fix_nan(half2_t input) {
__device__ __forceinline__ static half2_t fix_nan(half2_t input) {
static constexpr float neginf = -std::numeric_limits<float>::infinity();
/**
* In accordance to the IEEE-754R standard,
* if one of the input parameters to fminf(), fmin(), fmaxf(), or fmax() is NaN,
* but not the other,
* In accordance to the IEEE-754R standard,
* if one of the input parameters to fminf(), fmin(), fmaxf(), or fmax() is NaN,
* but not the other,
* the result is the non-NaN parameter.
*/
return __hmax2(input, half2_t(neginf, neginf));
}
__device__ __forceinline__
static float fix_nan(float input) {
__device__ __forceinline__ static float fix_nan(float input) {
static constexpr float neginf = -std::numeric_limits<float>::infinity();
return fmaxf(input, neginf);
}
__device__ __forceinline__
static packed_fpsum_t fix_nan(packed_fpsum_t input) {
__device__ __forceinline__ static packed_fpsum_t fix_nan(packed_fpsum_t input) {
input.x = kernels::bit_cast<int>(fix_nan(kernels::bit_cast<half2_t>(input.x)));
input.y = kernels::bit_cast<int>(fix_nan(kernels::bit_cast<half2_t>(input.y)));
input.z = kernels::bit_cast<int>(fix_nan(kernels::bit_cast<half2_t>(input.z)));
......@@ -221,30 +214,28 @@ public:
return input;
}
__device__ __forceinline__
static packed_f32psum_t fix_nan(packed_f32psum_t input) {
#pragma unroll
__device__ __forceinline__ static packed_f32psum_t fix_nan(packed_f32psum_t input) {
#pragma unroll
for (int i = 0; i < 8; i++) {
input.data[i] = fix_nan(input.data[i]);
}
return input;
}
__device__ __forceinline__
static qk_warp compute_qk(q_warp Q, k_warp K) {
__device__ __forceinline__ static qk_warp compute_qk(q_warp Q, k_warp K) {
qk_warp QK;
#pragma unroll
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
#pragma unroll
#pragma unroll
for (int k = 0; k < WARP_K_TILES_QK; k++) {
#if 0
#pragma unroll
#pragma unroll
for (int d = 0; d < WARP_D_TILES; d++) {
packed_fpsum_t psum = make_uint4(0, 0, 0, 0);
psum = mma_f16xf16_f16(Q[m * WARP_D_TILES + d], K[k * WARP_D_TILES + d], psum);
auto f32psum = packed_fp16_to_fp32(psum);
#pragma unroll
#pragma unroll
for (int i = 0; i < 8; i++) {
QK[m * WARP_K_TILES_QK + k].data[i] += f32psum.data[i];
}
......@@ -252,34 +243,32 @@ public:
#else
packed_fpsum_t psum = make_uint4(0, 0, 0, 0);
#pragma unroll
#pragma unroll
for (int d = 0; d < WARP_D_TILES; d++) {
psum = mma_f16xf16_f16(Q[m * WARP_D_TILES + d], K[k * WARP_D_TILES + d], psum);
}
if constexpr (IS_SM80) {
psum = fix_nan(psum);
psum = fix_nan(psum);
QK[m * WARP_K_TILES_QK + k] = packed_fp16_to_fp32(psum);
} else {
QK[m * WARP_K_TILES_QK + k] = fix_nan(packed_fp16_to_fp32(psum));
}
#endif
}
}
return QK;
}
__device__ __forceinline__
static rowval_warp compute_rowmax(qk_warp QK, rowval_warp rowmax, float scale) {
#pragma unroll
__device__ __forceinline__ static rowval_warp compute_rowmax(qk_warp QK, rowval_warp rowmax, float scale) {
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
float2 maxv;
#pragma unroll
#pragma unroll
for (int k = 0; k < WARP_K_TILES_QK; k++) {
packed_f32psum_t &val = QK[m * WARP_K_TILES_QK + k];
float x = fmaxf(fmaxf(val.data[0], val.data[1]), fmaxf(val.data[4], val.data[5]));
float y = fmaxf(fmaxf(val.data[2], val.data[3]), fmaxf(val.data[6], val.data[7]));
float x = fmaxf(fmaxf(val.data[0], val.data[1]), fmaxf(val.data[4], val.data[5]));
float y = fmaxf(fmaxf(val.data[2], val.data[3]), fmaxf(val.data[6], val.data[7]));
if (k == 0) {
maxv = make_float2(x, y);
} else {
......@@ -287,7 +276,7 @@ public:
maxv.y = fmaxf(maxv.y, y);
}
}
#pragma unroll
#pragma unroll
for (int mask = 1; mask <= 2; mask *= 2) {
maxv.x = fmaxf(maxv.x, __shfl_xor_sync(~0, maxv.x, mask));
maxv.y = fmaxf(maxv.y, __shfl_xor_sync(~0, maxv.y, mask));
......@@ -298,40 +287,38 @@ public:
return rowmax;
}
__device__ __forceinline__
static qk_warp softmax(qk_warp QK, rowval_warp rowmax_scaled, float scale) {
#pragma unroll
__device__ __forceinline__ static qk_warp softmax(qk_warp QK, rowval_warp rowmax_scaled, float scale) {
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
float2 shift = rowmax_scaled[m];
#pragma unroll
#pragma unroll
for (int k = 0; k < WARP_K_TILES_QK; k++) {
packed_f32psum_t &val = QK[m * WARP_K_TILES_QK + k];
val.data[0] = cuda_exp2(fmaf(val.data[0], scale, -shift.x));
val.data[1] = cuda_exp2(fmaf(val.data[1], scale, -shift.x));
val.data[4] = cuda_exp2(fmaf(val.data[4], scale, -shift.x));
val.data[5] = cuda_exp2(fmaf(val.data[5], scale, -shift.x));
val.data[2] = cuda_exp2(fmaf(val.data[2], scale, -shift.y));
val.data[3] = cuda_exp2(fmaf(val.data[3], scale, -shift.y));
val.data[6] = cuda_exp2(fmaf(val.data[6], scale, -shift.y));
val.data[7] = cuda_exp2(fmaf(val.data[7], scale, -shift.y));
val.data[0] = cuda_exp2(fmaf(val.data[0], scale, -shift.x));
val.data[1] = cuda_exp2(fmaf(val.data[1], scale, -shift.x));
val.data[4] = cuda_exp2(fmaf(val.data[4], scale, -shift.x));
val.data[5] = cuda_exp2(fmaf(val.data[5], scale, -shift.x));
val.data[2] = cuda_exp2(fmaf(val.data[2], scale, -shift.y));
val.data[3] = cuda_exp2(fmaf(val.data[3], scale, -shift.y));
val.data[6] = cuda_exp2(fmaf(val.data[6], scale, -shift.y));
val.data[7] = cuda_exp2(fmaf(val.data[7], scale, -shift.y));
}
}
return QK;
}
__device__ __forceinline__
static rowval_warp compute_rowsum(qk_warp QK) {
__device__ __forceinline__ static rowval_warp compute_rowsum(qk_warp QK) {
rowval_warp rowsum;
#pragma unroll
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
float2 sumv = make_float2(0.0f, 0.0f);
#pragma unroll
#pragma unroll
for (int k = 0; k < WARP_K_TILES_QK; k++) {
packed_f32psum_t &val = QK[m * WARP_K_TILES_QK + k];
sumv.x += val.data[0] + val.data[1] + val.data[4] + val.data[5];
sumv.y += val.data[2] + val.data[3] + val.data[6] + val.data[7];
}
#pragma unroll
#pragma unroll
for (int mask = 1; mask <= 2; mask *= 2) {
sumv.x += __shfl_xor_sync(~0, sumv.x, mask);
sumv.y += __shfl_xor_sync(~0, sumv.y, mask);
......@@ -341,10 +328,9 @@ public:
return rowsum;
}
__device__ __forceinline__
static rowval_warp compute_rescale(rowval_warp rowmax0, rowval_warp rowmax1) {
__device__ __forceinline__ static rowval_warp compute_rescale(rowval_warp rowmax0, rowval_warp rowmax1) {
rowval_warp rescale;
#pragma unroll
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
rescale[m].x = cuda_exp2(rowmax0[m].x - rowmax1[m].x);
rescale[m].y = cuda_exp2(rowmax0[m].y - rowmax1[m].y);
......@@ -352,36 +338,34 @@ public:
return rescale;
}
__device__ __forceinline__
static o_warp compute_pv(p_warp P, v_warp V, o_warp O, rowval_warp rescale) {
#pragma unroll
__device__ __forceinline__ static o_warp compute_pv(p_warp P, v_warp V, o_warp O, rowval_warp rescale) {
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
#pragma unroll
#pragma unroll
for (int n = 0; n < WARP_N_TILES; n++) {
packed_fpsum_t psum = make_uint4(0, 0, 0, 0);
#pragma unroll
#pragma unroll
for (int k = 0; k < WARP_K_TILES_PV; k++) {
psum = mma_f16xf16_f16(P[m * WARP_K_TILES_PV + k], V[n * WARP_K_TILES_PV + k], psum);
}
packed_f32psum_t pv = packed_fp16_to_fp32(psum);
packed_f32psum_t pv = packed_fp16_to_fp32(psum);
packed_f32psum_t &oval = O[m * WARP_N_TILES + n];
oval.data[0] = oval.data[0] * rescale[m].x + pv.data[0];
oval.data[1] = oval.data[1] * rescale[m].x + pv.data[1];
oval.data[4] = oval.data[4] * rescale[m].x + pv.data[4];
oval.data[5] = oval.data[5] * rescale[m].x + pv.data[5];
oval.data[2] = oval.data[2] * rescale[m].y + pv.data[2];
oval.data[3] = oval.data[3] * rescale[m].y + pv.data[3];
oval.data[6] = oval.data[6] * rescale[m].y + pv.data[6];
oval.data[7] = oval.data[7] * rescale[m].y + pv.data[7];
oval.data[0] = oval.data[0] * rescale[m].x + pv.data[0];
oval.data[1] = oval.data[1] * rescale[m].x + pv.data[1];
oval.data[4] = oval.data[4] * rescale[m].x + pv.data[4];
oval.data[5] = oval.data[5] * rescale[m].x + pv.data[5];
oval.data[2] = oval.data[2] * rescale[m].y + pv.data[2];
oval.data[3] = oval.data[3] * rescale[m].y + pv.data[3];
oval.data[6] = oval.data[6] * rescale[m].y + pv.data[6];
oval.data[7] = oval.data[7] * rescale[m].y + pv.data[7];
}
}
return O;
}
__device__ __forceinline__
static rowval_warp compute_l(rowval_warp L, rowval_warp rescale, rowval_warp rowsum) {
#pragma unroll
__device__ __forceinline__ static rowval_warp compute_l(rowval_warp L, rowval_warp rescale, rowval_warp rowsum) {
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
L[m].x = fmaf(L[m].x, rescale[m].x, rowsum[m].x);
L[m].y = fmaf(L[m].y, rescale[m].y, rowsum[m].y);
......@@ -389,13 +373,12 @@ public:
return L;
}
__device__ __forceinline__
static p_warp qk_to_p(qk_warp QK) {
__device__ __forceinline__ static p_warp qk_to_p(qk_warp QK) {
static_assert(WARP_K_TILES_QK == WARP_K_TILES_PV);
p_warp P;
#pragma unroll
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
#pragma unroll
#pragma unroll
for (int k = 0; k < WARP_K_TILES_PV; k++) {
P[m * WARP_K_TILES_PV + k] = packed_fp32_to_fp16(QK[m * WARP_K_TILES_QK + k]);
}
......@@ -416,43 +399,41 @@ public:
// O = compute_pv(P, V, O, rescale);
// }
__device__ __forceinline__
static std::tuple<p_warp, rowval_warp> compute(q_warp Q, k_warp K, rowval_warp &M, rowval_warp &L, float scale) {
qk_warp qk = compute_qk(Q, K);
rowval_warp M1 = compute_rowmax(qk, M, scale);
qk = softmax(qk, M1, scale);
rowval_warp rowsum = compute_rowsum(qk);
p_warp P = qk_to_p(qk);
__device__ __forceinline__ static std::tuple<p_warp, rowval_warp>
compute(q_warp Q, k_warp K, rowval_warp &M, rowval_warp &L, float scale) {
qk_warp qk = compute_qk(Q, K);
rowval_warp M1 = compute_rowmax(qk, M, scale);
qk = softmax(qk, M1, scale);
rowval_warp rowsum = compute_rowsum(qk);
p_warp P = qk_to_p(qk);
rowval_warp rescale = compute_rescale(M, M1);
M = M1;
L = compute_l(L, rescale, rowsum);
M = M1;
L = compute_l(L, rescale, rowsum);
return {P, rescale};
}
__device__ __forceinline__
static o_warp compute_o(o_warp O, rowval_warp L) {
#pragma unroll
__device__ __forceinline__ static o_warp compute_o(o_warp O, rowval_warp L) {
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
float2 inv;
inv.x = cuda_frcp(L[m].x);
inv.y = cuda_frcp(L[m].y);
#pragma unroll
#pragma unroll
for (int n = 0; n < WARP_N_TILES; n++) {
packed_f32psum_t &oval = O[m * WARP_N_TILES + n];
oval.data[0] = oval.data[0] * inv.x;
oval.data[1] = oval.data[1] * inv.x;
oval.data[4] = oval.data[4] * inv.x;
oval.data[5] = oval.data[5] * inv.x;
oval.data[2] = oval.data[2] * inv.y;
oval.data[3] = oval.data[3] * inv.y;
oval.data[6] = oval.data[6] * inv.y;
oval.data[7] = oval.data[7] * inv.y;
oval.data[0] = oval.data[0] * inv.x;
oval.data[1] = oval.data[1] * inv.x;
oval.data[4] = oval.data[4] * inv.x;
oval.data[5] = oval.data[5] * inv.x;
oval.data[2] = oval.data[2] * inv.y;
oval.data[3] = oval.data[3] * inv.y;
oval.data[6] = oval.data[6] * inv.y;
oval.data[7] = oval.data[7] * inv.y;
}
}
return O;
}
#if 0
template<typename Epilogue>
__device__ __forceinline__
......@@ -462,10 +443,10 @@ public:
const packed_k_t *ptr_k,
const packed_v_t *ptr_v,
float scale,
int ntokens_q,
int ntokens_q,
int ntokens_kv,
Epilogue::Arguments epilogueArgs,
bool alwaysfalse)
bool alwaysfalse)
{
constexpr int NUM_STAGES = 2;
......@@ -485,9 +466,9 @@ public:
load_v(ptr_v, k, V[k], true);
}
#pragma unroll
#pragma unroll
for (auto &pack : O) {
#pragma unroll
#pragma unroll
for (int i = 0; i < 8; i++) {
pack.data[i] = 0;
}
......@@ -498,7 +479,7 @@ public:
M.fill(make_float2(neginf, neginf));
__shared__ q_warp Q_shmem[NUM_WARPS];
#pragma unroll
#pragma unroll
for (int i = 0; i < Q.size(); i++) {
store<true>(&Q_shmem[warpId][i], Q[i]);
}
......@@ -507,9 +488,9 @@ public:
// TODO: mask tokens in last block
for (int k1 = 0; k1 < ntokens_kv / WARP_K; k1 += NUM_STAGES) {
#pragma unroll
#pragma unroll
for (int k2 = 0; k2 < NUM_STAGES; k2++) {
#pragma unroll
#pragma unroll
for (int i = 0; i < Q.size(); i++) {
Q[i] = load<true>(&Q_shmem[warpId][i]);
}
......@@ -519,7 +500,7 @@ public:
bool pred = nextk < ntokens_kv / WARP_K;
load_k(ptr_k, nextk, K[idx], pred);
load_v(ptr_v, nextk, V[idx], pred);
// __syncthreads();
// if (alwaysfalse) {
// dummy = clock();
......@@ -549,42 +530,40 @@ public:
}
#else
template<typename Epilogue>
__device__ __forceinline__
static void attention_fp16_block(
const BlockInfo binfo,
const packed_q_t *ptr_q,
const packed_k_t *ptr_k,
const packed_v_t *ptr_v,
float scale,
int ntokens_q,
int ntokens_kv,
Epilogue::Arguments epilogueArgs,
bool alwaysfalse)
{
__device__ __forceinline__ static void attention_fp16_block(const BlockInfo binfo,
const packed_q_t *ptr_q,
const packed_k_t *ptr_k,
const packed_v_t *ptr_v,
float scale,
int ntokens_q,
int ntokens_kv,
Epilogue::Arguments epilogueArgs,
bool alwaysfalse) {
// constexpr int NUM_STAGES = 2;
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
q_warp Q; // 32
k_warp K; // 64
v_warp V; // 64
o_warp O; // 64
rowval_warp L; // 2
rowval_warp M; // 2
q_warp Q; // 32
k_warp K; // 64
v_warp V; // 64
o_warp O; // 64
rowval_warp L; // 2
rowval_warp M; // 2
load_q(ptr_q, Q, true);
load_k(ptr_k, 0, K, true);
#pragma unroll
#pragma unroll
for (auto &pack : O) {
#pragma unroll
#pragma unroll
for (int i = 0; i < 8; i++) {
pack.data[i] = 0;
}
}
static constexpr float neginf = -std::numeric_limits<float>::max(); // not real inf, to prevent nan during computation
static constexpr float neginf =
-std::numeric_limits<float>::max(); // not real inf, to prevent nan during computation
L.fill(make_float2(0.0f, 0.0f));
M.fill(make_float2(neginf, neginf));
......@@ -593,7 +572,7 @@ public:
using q_shmem_t = packed_q_t[NUM_WARPS][SHMEM_TILES][WARP_SIZE];
__shared__ q_shmem_t Q_shmem;
#pragma unroll
#pragma unroll
for (int i = 0; i < SHMEM_TILES; i++) {
store<true>(&Q_shmem[warpId][i][laneId], Q[Q.size() - 1 - i]);
}
......@@ -602,12 +581,12 @@ public:
int dummy = 0;
// TODO: mask tokens in last block
for (int k1 = 0; k1 < ntokens_kv / WARP_K; k1 ++) {
for (int k1 = 0; k1 < ntokens_kv / WARP_K; k1++) {
if (alwaysfalse) {
ptr_v += K[0].x;
}
#pragma unroll
#pragma unroll
for (int i = 0; i < SHMEM_TILES; i++) {
Q[Q.size() - 1 - i] = load<true>(&Q_shmem[warpId][i][laneId]);
}
......@@ -628,8 +607,6 @@ public:
dummy = clock();
}
auto [P, rescale] = compute(Q, K, M, L, scale);
if (alwaysfalse) {
......@@ -644,9 +621,7 @@ public:
// dummy = clock();
// }
load_k(ptr_k, k1+1, K, k1+1 < ntokens_kv / WARP_K);
load_k(ptr_k, k1 + 1, K, k1 + 1 < ntokens_kv / WARP_K);
// if (alwaysfalse) {
// dummy = clock();
......@@ -665,38 +640,41 @@ public:
auto f16psum = GEMM::packed_fp32_to_fp16(O);
Epilogue()(typename GEMM::BlockInfo{
.bm = binfo.batch * binfo.numBlocksM + binfo.bm,
.bn = binfo.head,
.numBlocksM = binfo.numBatch * binfo.numBlocksM,
.numBlocksN = binfo.numHeads,
}, f16psum, binfo.numBatch * binfo.numBlocksM * BLOCK_M, binfo.numHeads * HEAD_DIM, 0, epilogueArgs);
Epilogue()(
typename GEMM::BlockInfo{
.bm = binfo.batch * binfo.numBlocksM + binfo.bm,
.bn = binfo.head,
.numBlocksM = binfo.numBatch * binfo.numBlocksM,
.numBlocksN = binfo.numHeads,
},
f16psum,
binfo.numBatch * binfo.numBlocksM * BLOCK_M,
binfo.numHeads * HEAD_DIM,
0,
epilogueArgs);
}
#endif
template<typename Epilogue>
struct attention_fp16_kernel {
static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
static constexpr int SHMEM_SIZE = 0; // sizeof(q_shmem_t);
__device__
void operator()(
const packed_q_t *ptr_q,
const packed_k_t *ptr_k,
const packed_v_t *ptr_v,
float scale,
int ntokens_q,
int ntokens_kv,
Epilogue::Arguments epilogueArgs,
bool alwaysfalse)
{
__device__ void operator()(const packed_q_t *ptr_q,
const packed_k_t *ptr_k,
const packed_v_t *ptr_v,
float scale,
int ntokens_q,
int ntokens_kv,
Epilogue::Arguments epilogueArgs,
bool alwaysfalse) {
BlockInfo binfo = {
.bm = (int)blockIdx.x,
.head = (int)blockIdx.y,
.batch = (int)blockIdx.z,
.bm = (int)blockIdx.x,
.head = (int)blockIdx.y,
.batch = (int)blockIdx.z,
.numBlocksM = (int)gridDim.x,
.numHeads = (int)gridDim.y,
.numBatch = (int)gridDim.z,
.numHeads = (int)gridDim.y,
.numBatch = (int)gridDim.z,
};
// extern __shared__ uint8_t shmem[];
......@@ -706,21 +684,20 @@ public:
attention_fp16_block<Epilogue>(
binfo,
ptr_q + ((binfo.batch * binfo.numHeads + binfo.head) * binfo.numBlocksM + binfo.bm) * NUM_WARPS * WARP_M_TILES * WARP_D_TILES * WARP_SIZE,
ptr_k + (binfo.batch * binfo.numHeads + binfo.head) * ktiles * WARP_K_TILES_QK * WARP_D_TILES * WARP_SIZE,
ptr_v + (binfo.batch * binfo.numHeads + binfo.head) * ktiles * WARP_K_TILES_PV * WARP_N_TILES * WARP_SIZE,
ptr_q + ((binfo.batch * binfo.numHeads + binfo.head) * binfo.numBlocksM + binfo.bm) * NUM_WARPS *
WARP_M_TILES * WARP_D_TILES * WARP_SIZE,
ptr_k +
(binfo.batch * binfo.numHeads + binfo.head) * ktiles * WARP_K_TILES_QK * WARP_D_TILES * WARP_SIZE,
ptr_v +
(binfo.batch * binfo.numHeads + binfo.head) * ktiles * WARP_K_TILES_PV * WARP_N_TILES * WARP_SIZE,
scale,
ntokens_q,
ntokens_kv,
// *Q_shmem,
epilogueArgs,
alwaysfalse
);
alwaysfalse);
}
};
};
}; // namespace nunchaku::kernels
}; // namespace nunchaku::kernels
src/kernels/zgemm/epilogues.cuh
View file @ 37a27712
......@@ -19,22 +19,23 @@ public:
IMPORT_GEMM_BASE(Config);
public:
struct EpilogueGelu {
struct Arguments { size_t unused; };
struct Arguments {
size_t unused;
};
// static constexpr float SHIFT_VALUE = 0.171875f;
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
#pragma unroll
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
#pragma unroll
for (int i = 0; i < WARP_M_TILES; i++) {
#pragma unroll
#pragma unroll
for (int j = 0; j < WARP_N_TILES; j++) {
#pragma unroll
#pragma unroll
for (int k = 0; k < 4; k++) {
half2_t &data = fpsum[i * WARP_N_TILES + j].data[k];
data = gelu_half2(data);
data = gelu_half2(data);
// data = __hadd2(data, half2_t(SHIFT_VALUE, SHIFT_VALUE));
}
}
......@@ -48,33 +49,41 @@ public:
half_t *out;
int actualM, actualN;
half_t *pool_out; // [M / PoolSize, N]
half_t *pool_out; // [M / PoolSize, N]
const float *rotary_emb; // [M, HEAD_DIM / 2, ROTARY_EMB_NUM_ELEMENTS]
const half_t *rmsnorm_weight_q; // [HEAD_DIM]
const half_t *rmsnorm_weight_k; // [HEAD_DIM]
float epsilon;
};
static constexpr int HEAD_DIM = 128;
static constexpr int HEAD_DIM = 128;
static constexpr int NUM_HEADS_PER_WARP = WARP_N / HEAD_DIM;
static constexpr int PoolSize = 128;
static constexpr int NUM_WARPS_PER_POOL = PoolSize / WARP_M;
static constexpr int PoolSize = 128;
static constexpr int NUM_WARPS_PER_POOL = PoolSize / WARP_M;
static constexpr int NUM_POOLS_PER_BLOCK = BLOCK_M / PoolSize;
static constexpr int ROTARY_EMB_NUM_ELEMENTS = 2; // 1 for theta, 2 for {sin, cos} pair
__device__ __forceinline__
static void apply(fpsum_warp fpsum, half_t *out, int M, int N, int K, half_t *pool_out, const float *rotary_emb, const half_t *rmsnorm_weight, float epsilon, int maxRows) {
static constexpr int ROTARY_EMB_NUM_ELEMENTS = 2; // 1 for theta, 2 for {sin, cos} pair
__device__ __forceinline__ static void apply(fpsum_warp fpsum,
half_t *out,
int M,
int N,
int K,
half_t *pool_out,
const float *rotary_emb,
const half_t *rmsnorm_weight,
float epsilon,
int maxRows) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
__shared__ alignas(128) uint8_t shmem[NUM_WARPS][ceilDiv(unpack_fpsum::SHMEM_SIZE, 128) * 128];
constexpr int PACK_SIZE = unpack_fpsum::PACK_SIZE;
using pack_t = unpack_fpsum::pack_t;
using pack_t = unpack_fpsum::pack_t;
using pack_rope_t = std::array<float, PACK_SIZE / 2 * ROTARY_EMB_NUM_ELEMENTS>;
using pack_rope_t = std::array<float, PACK_SIZE / 2 * ROTARY_EMB_NUM_ELEMENTS>;
constexpr int LANES_PER_HEAD = HEAD_DIM / PACK_SIZE;
pack_t reduce_tmp;
......@@ -91,98 +100,107 @@ public:
}
}
const float *rotary_emb_base_addr = &rotary_emb[(warpId * WARP_M) * HEAD_DIM / 2 * ROTARY_EMB_NUM_ELEMENTS + laneId * PACK_SIZE / 2 * ROTARY_EMB_NUM_ELEMENTS];
const float *rotary_emb_base_addr = &rotary_emb[(warpId * WARP_M) * HEAD_DIM / 2 * ROTARY_EMB_NUM_ELEMENTS +
laneId * PACK_SIZE / 2 * ROTARY_EMB_NUM_ELEMENTS];
CHECK_NAN(fpsum, "fpsum");
unpack_fpsum()(fpsum, out + warpId * WARP_M * N, N, maxRows - warpId * WARP_M, INT_MAX, shmem[warpId], [&](int rowId, pack_t &pack) ALWAYSINLINE {
// load rope
pack_rope_t rope;
if (laneId < LANES_PER_HEAD) {
// freq = load(reinterpret_cast<pack_freq_t *>(&freqs_cis[(warpId * WARP_M + rowId) * HEAD_DIM * 2 + laneId * PACK_SIZE * 2]));
rope = load(reinterpret_cast<const pack_rope_t *>(&rotary_emb_base_addr[rowId * HEAD_DIM / 2 * ROTARY_EMB_NUM_ELEMENTS]));
}
if constexpr (LANES_PER_HEAD < WARP_SIZE) {
for (int i = 0; i < rope.size(); i++) {
rope[i] = __shfl_sync(~0, rope[i], laneId % LANES_PER_HEAD);
}
}
// rmsnorm
float sqrsum = 0.0f;
for (int i = 0; i < PACK_SIZE; i++) {
sqrsum += float(pack[i]) * float(pack[i]);
CHECK_NAN(sqrsum, "sqrsum");
}
#pragma unroll
for (int mask = LANES_PER_HEAD / 2; mask > 0; mask /= 2) {
sqrsum += __shfl_xor_sync(~0, sqrsum, mask);
}
sqrsum /= HEAD_DIM;
float coef = cuda_frsqrt(sqrsum + epsilon);
CHECK_NAN(coef, "coef");
for (int i = 0; i < PACK_SIZE; i++) {
pack[i] *= coef * float(rms[i]);
CHECK_NAN(rms[i], "rms.wgt");
CHECK_NAN(pack[i], "rms.out");
}
#if 1
// rope
for (int i = 0; i < PACK_SIZE; i += 2) {
float2 pack2 = half22float2(half2_t(pack[i], pack[i+1]));
CHECK_NAN(freq[i].x, "rope.freq");
CHECK_NAN(freq[i].y, "rope.freq");
CHECK_NAN(freq[i+1].x, "rope.freq");
CHECK_NAN(freq[i+1].y, "rope.freq");
// half2_t tmp = __hmul2(freq[i], pack2);
// tmp = __hfma2(freq[i+1], pack2, tmp);
// pack[i] = tmp.x;
// pack[i+1] = tmp.y;
// printf("block.x=%d block.y=%d warpId=%d rowId=%d (%d) freqs = %f %f %f %f\n",
// blockIdx.x, blockIdx.y, warpId, rowId,
// blockIdx.x * BLOCK_M + warpId * WARP_M + rowId,
// (float)freq[i].x, (float)freq[i].y, (float)freq[i+1].x, (float)freq[i+1].y
// );
// __trap();
// half2_t tmp = __hmul2(half2_t(pack2.x, pack2.x), freq[i]);
// tmp = __hfma2(half2_t(pack2.y, pack2.y), freq[i+1], tmp);
// pack[i] = tmp.x;
// pack[i+1] = tmp.y;
float sin, cos;
if constexpr (ROTARY_EMB_NUM_ELEMENTS == 1) {
sin = cuda_sin(rope[i / 2]);
cos = cuda_cos(rope[i / 2]);
}
if constexpr (ROTARY_EMB_NUM_ELEMENTS == 2) {
sin = rope[i];
cos = rope[i+1];
}
// pack[i] = pack2.x * freq[i].x + pack2.y * freq[i].y;
// pack[i+1] = pack2.x * freq[i+1].x + pack2.y * freq[i+1].y;
pack[i] = half_t(pack2.x * cos - pack2.y * sin);
pack[i+1] = half_t(pack2.x * sin + pack2.y * cos);
CHECK_NAN(pack[i], "rope.out");
CHECK_NAN(pack[i+1], "rope.out");
}
#endif
unpack_fpsum()(fpsum,
out + warpId * WARP_M * N,
N,
maxRows - warpId * WARP_M,
INT_MAX,
shmem[warpId],
[&](int rowId, pack_t &pack) ALWAYSINLINE {
// load rope
pack_rope_t rope;
if (laneId < LANES_PER_HEAD) {
// freq = load(reinterpret_cast<pack_freq_t *>(&freqs_cis[(warpId * WARP_M + rowId) *
// HEAD_DIM * 2 + laneId * PACK_SIZE * 2]));
rope = load(reinterpret_cast<const pack_rope_t *>(
&rotary_emb_base_addr[rowId * HEAD_DIM / 2 * ROTARY_EMB_NUM_ELEMENTS]));
}
if constexpr (LANES_PER_HEAD < WARP_SIZE) {
for (int i = 0; i < rope.size(); i++) {
rope[i] = __shfl_sync(~0, rope[i], laneId % LANES_PER_HEAD);
}
}
// rmsnorm
float sqrsum = 0.0f;
for (int i = 0; i < PACK_SIZE; i++) {
sqrsum += float(pack[i]) * float(pack[i]);
CHECK_NAN(sqrsum, "sqrsum");
}
#pragma unroll
for (int mask = LANES_PER_HEAD / 2; mask > 0; mask /= 2) {
sqrsum += __shfl_xor_sync(~0, sqrsum, mask);
}
sqrsum /= HEAD_DIM;
float coef = cuda_frsqrt(sqrsum + epsilon);
CHECK_NAN(coef, "coef");
for (int i = 0; i < PACK_SIZE; i++) {
pack[i] *= coef * float(rms[i]);
CHECK_NAN(rms[i], "rms.wgt");
CHECK_NAN(pack[i], "rms.out");
}
#if 1
// rope
for (int i = 0; i < PACK_SIZE; i += 2) {
float2 pack2 = half22float2(half2_t(pack[i], pack[i + 1]));
CHECK_NAN(freq[i].x, "rope.freq");
CHECK_NAN(freq[i].y, "rope.freq");
CHECK_NAN(freq[i + 1].x, "rope.freq");
CHECK_NAN(freq[i + 1].y, "rope.freq");
// half2_t tmp = __hmul2(freq[i], pack2);
// tmp = __hfma2(freq[i+1], pack2, tmp);
// pack[i] = tmp.x;
// pack[i+1] = tmp.y;
// printf("block.x=%d block.y=%d warpId=%d rowId=%d (%d) freqs = %f %f %f %f\n",
// blockIdx.x, blockIdx.y, warpId, rowId,
// blockIdx.x * BLOCK_M + warpId * WARP_M + rowId,
// (float)freq[i].x, (float)freq[i].y, (float)freq[i+1].x, (float)freq[i+1].y
// );
// __trap();
// half2_t tmp = __hmul2(half2_t(pack2.x, pack2.x), freq[i]);
// tmp = __hfma2(half2_t(pack2.y, pack2.y), freq[i+1], tmp);
// pack[i] = tmp.x;
// pack[i+1] = tmp.y;
float sin, cos;
if constexpr (ROTARY_EMB_NUM_ELEMENTS == 1) {
sin = cuda_sin(rope[i / 2]);
cos = cuda_cos(rope[i / 2]);
}
if constexpr (ROTARY_EMB_NUM_ELEMENTS == 2) {
sin = rope[i];
cos = rope[i + 1];
}
// pack[i] = pack2.x * freq[i].x + pack2.y * freq[i].y;
// pack[i+1] = pack2.x * freq[i+1].x + pack2.y * freq[i+1].y;
pack[i] = half_t(pack2.x * cos - pack2.y * sin);
pack[i + 1] = half_t(pack2.x * sin + pack2.y * cos);
CHECK_NAN(pack[i], "rope.out");
CHECK_NAN(pack[i + 1], "rope.out");
}
#endif
// mean pool
for (int i = 0; i < PACK_SIZE; i++) {
reduce_tmp[i] += pack[i];
}
});
// mean pool
for (int i = 0; i < PACK_SIZE; i++) {
reduce_tmp[i] += pack[i];
}
});
if (!pool_out) {
return;
......@@ -193,7 +211,7 @@ public:
if (warpId < NUM_POOLS_PER_BLOCK) {
const int row = warpId * NUM_WARPS_PER_POOL;
reduce_tmp = load<true>(&pool[row]);
reduce_tmp = load<true>(&pool[row]);
for (int i = 1; i < NUM_WARPS_PER_POOL; i++) {
pack_t pack = load<true>(&pool[row + i]);
......@@ -210,8 +228,8 @@ public:
__syncthreads();
}
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
const int bm = binfo.bm;
const int bn = binfo.bn;
......@@ -223,34 +241,39 @@ public:
assert(args.actualN == N);
if (is_q || is_k) {
apply(
fpsum,
args.out + bm * BLOCK_M * args.actualN + bn * BLOCK_N,
M, N, K,
args.pool_out ? args.pool_out + bm * BLOCK_M / PoolSize * N : nullptr,
args.rotary_emb + bm * BLOCK_M * (HEAD_DIM / 2 * ROTARY_EMB_NUM_ELEMENTS),
is_q ? args.rmsnorm_weight_q : args.rmsnorm_weight_k,
args.epsilon,
args.actualM - bm * BLOCK_M
);
apply(fpsum,
args.out + bm * BLOCK_M * args.actualN + bn * BLOCK_N,
M,
N,
K,
args.pool_out ? args.pool_out + bm * BLOCK_M / PoolSize * N : nullptr,
args.rotary_emb + bm * BLOCK_M * (HEAD_DIM / 2 * ROTARY_EMB_NUM_ELEMENTS),
is_q ? args.rmsnorm_weight_q : args.rmsnorm_weight_k,
args.epsilon,
args.actualM - bm * BLOCK_M);
} else {
EpilogueDefault()(binfo, fpsum, M, N, K, typename EpilogueDefault::Arguments{
.out = args.out,
.actualM = args.actualM,
.actualN = args.actualN,
});
EpilogueDefault()(binfo,
fpsum,
M,
N,
K,
typename EpilogueDefault::Arguments{
.out = args.out,
.actualM = args.actualM,
.actualN = args.actualN,
});
}
}
};
struct EpilogueRMSNormRope {
static constexpr int HEAD_DIM = 128;
static constexpr int NUM_HEADS_PER_WARP = WARP_N / HEAD_DIM;
static constexpr int HEAD_DIM = 128;
static constexpr int NUM_HEADS_PER_WARP = WARP_N / HEAD_DIM;
static constexpr int WARP_N_TILES_PER_HEAD = WARP_N_TILES / NUM_HEADS_PER_WARP;
static constexpr int ROTARY_EMB_NUM_ELEMENTS = 2;
using packed_rotemb_t = float4;
using packed_rotemb_t = float4;
static constexpr int WARP_N_ROTEMB_TILES = WARP_N_TILES / NUM_HEADS_PER_WARP * 2;
using rotemb_warp = std::array<packed_rotemb_t, WARP_M_TILES * WARP_N_ROTEMB_TILES>; // 128 regs
......@@ -263,17 +286,17 @@ public:
float epsilon;
};
__device__ __forceinline__
static rotemb_warp load_rotemb(const packed_rotemb_t *ptr_rotemb) {
__device__ __forceinline__ static rotemb_warp load_rotemb(const packed_rotemb_t *ptr_rotemb) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
rotemb_warp rotemb;
const packed_rotemb_t *ptrlane = &ptr_rotemb[warpId * WARP_M_TILES * WARP_N_ROTEMB_TILES * WARP_SIZE + laneId];
const packed_rotemb_t *ptrlane =
&ptr_rotemb[warpId * WARP_M_TILES * WARP_N_ROTEMB_TILES * WARP_SIZE + laneId];
unrolled_loop<WARP_M_TILES>([&]<int i>() {
unrolled_loop<WARP_N_ROTEMB_TILES>([&]<int j>() {
constexpr int offset = (i * WARP_N_ROTEMB_TILES + j) * WARP_SIZE;
constexpr int offset = (i * WARP_N_ROTEMB_TILES + j) * WARP_SIZE;
rotemb[i * WARP_N_ROTEMB_TILES + j] = load(&ptrlane[offset]);
});
});
......@@ -281,28 +304,26 @@ public:
return rotemb;
}
__device__ __forceinline__
static void load_rmsnorm(const half_t *ptr_rmsnorm_weight, half_t *shmem) {
__device__ __forceinline__ static void load_rmsnorm(const half_t *ptr_rmsnorm_weight, half_t *shmem) {
const int laneId = threadIdx.x % WARP_SIZE;
static constexpr int PACK_SIZE = HEAD_DIM / WARP_SIZE;
using packed_t = std::array<half_t, PACK_SIZE>;
using packed_t = std::array<half_t, PACK_SIZE>;
packed_t pack = load(reinterpret_cast<const packed_t *>(ptr_rmsnorm_weight + laneId * PACK_SIZE));
store<true>(reinterpret_cast<packed_t *>(shmem + laneId * PACK_SIZE), pack);
}
__device__ __forceinline__
static packed_fpsum_t load_rmsnorm_from_shmem(half_t *shmem, int n) {
__device__ __forceinline__ static packed_fpsum_t load_rmsnorm_from_shmem(half_t *shmem, int n) {
const int laneId = threadIdx.x % WARP_SIZE;
const int col = n * INSN_N + laneId / 16 * 8; // lane 0-15: n*16+0, lane 16-31: n*16+8
const int col = n * INSN_N + laneId / 16 * 8; // lane 0-15: n*16+0, lane 16-31: n*16+8
uint4 tmp;
ldmatrix(shmem + col, tmp);
return kernels::bit_cast<packed_fpsum_t>(tmp);
}
__device__ __forceinline__
static void apply(fpsum_warp &fpsum, const packed_rotemb_t *ptr_rotemb, const half_t *ptr_rmsnorm_weight, float epsilon) {
__device__ __forceinline__ static void
apply(fpsum_warp &fpsum, const packed_rotemb_t *ptr_rotemb, const half_t *ptr_rmsnorm_weight, float epsilon) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -319,21 +340,21 @@ public:
return fval.x * fval.x + fval.y * fval.y;
};
#pragma unroll
#pragma unroll
for (int head = 0; head < NUM_HEADS_PER_WARP; head++) {
const int n_offset = head * WARP_N_TILES_PER_HEAD;
#pragma unroll
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
float sqrsum[2] = {0.0f, 0.0f};
#pragma unroll
#pragma unroll
for (int n = 0; n < WARP_N_TILES_PER_HEAD; n++) {
sqrsum[0] += sqr(fpsum[m * WARP_N_TILES + n + n_offset].data[0]);
sqrsum[1] += sqr(fpsum[m * WARP_N_TILES + n + n_offset].data[1]);
sqrsum[0] += sqr(fpsum[m * WARP_N_TILES + n + n_offset].data[2]);
sqrsum[1] += sqr(fpsum[m * WARP_N_TILES + n + n_offset].data[3]);
}
#pragma unroll
#pragma unroll
for (int mask = 1; mask <= 2; mask *= 2) {
sqrsum[0] += __shfl_xor_sync(~0, sqrsum[0], mask);
sqrsum[1] += __shfl_xor_sync(~0, sqrsum[1], mask);
......@@ -343,14 +364,14 @@ public:
}
}
#pragma unroll
#pragma unroll
for (int head = 0; head < NUM_HEADS_PER_WARP; head++) {
const int n_offset = head * WARP_N_TILES_PER_HEAD;
#pragma unroll
#pragma unroll
for (int n = 0; n < WARP_N_TILES_PER_HEAD; n++) {
packed_f32psum_t rms = packed_fp16_to_fp32(load_rmsnorm_from_shmem(&shmem_rmsnorm[warpId][0], n));
#pragma unroll
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
packed_f32psum_t pack = packed_fp16_to_fp32(fpsum[m * WARP_N_TILES + n + n_offset]);
pack.data[0] *= rmsnorm_coef[head][m][0] * rms.data[0];
......@@ -385,8 +406,8 @@ public:
}
}
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
const int bm = binfo.bm;
const int bn = binfo.bn;
......@@ -395,22 +416,20 @@ public:
const bool is_k = !is_q && bn < binfo.numBlocksN / 3 * 2;
if (is_q || is_k) {
apply(
fpsum,
args.rotary_emb + bm * NUM_WARPS * WARP_M_TILES * WARP_N_ROTEMB_TILES * WARP_SIZE,
is_q ? args.rmsnorm_weight_q : args.rmsnorm_weight_k,
args.epsilon
);
apply(fpsum,
args.rotary_emb + bm * NUM_WARPS * WARP_M_TILES * WARP_N_ROTEMB_TILES * WARP_SIZE,
is_q ? args.rmsnorm_weight_q : args.rmsnorm_weight_k,
args.epsilon);
}
}
};
struct EpiloguePackQKV {
using attn_half_t = half;
using attn_half_t = half;
using attn_half2_t = half2;
using packed_qkv_t = uint4;
static constexpr int HEAD_DIM = 128;
static constexpr int HEAD_DIM = 128;
static constexpr int INSN_K_QK = 16;
static constexpr int INSN_K_PV = 16;
......@@ -424,8 +443,7 @@ public:
int strideHead_v;
};
__device__ __forceinline__
static attn_half2_t convert_half2(half2_t input) {
__device__ __forceinline__ static attn_half2_t convert_half2(half2_t input) {
if constexpr (std::is_same_v<half2_t, attn_half2_t>) {
return input;
} else {
......@@ -434,8 +452,7 @@ public:
}
}
__device__ __forceinline__
static packed_qkv_t pack_q(packed_fpsum_t input) {
__device__ __forceinline__ static packed_qkv_t pack_q(packed_fpsum_t input) {
packed_qkv_t output;
output.x = kernels::bit_cast<int>(convert_half2(input.data[0]));
output.y = kernels::bit_cast<int>(convert_half2(input.data[1]));
......@@ -444,8 +461,7 @@ public:
return output;
}
__device__ __forceinline__
static packed_qkv_t pack_k(packed_fpsum_t input) {
__device__ __forceinline__ static packed_qkv_t pack_k(packed_fpsum_t input) {
packed_qkv_t output;
output.x = kernels::bit_cast<int>(convert_half2(input.data[0]));
output.y = kernels::bit_cast<int>(convert_half2(input.data[2]));
......@@ -454,8 +470,7 @@ public:
return output;
}
__device__ __forceinline__
static packed_qkv_t pack_v(packed_fpsum_t input) {
__device__ __forceinline__ static packed_qkv_t pack_v(packed_fpsum_t input) {
packed_qkv_t output;
output.x = kernels::bit_cast<int>(convert_half2(movmatrix(input.data[0])));
output.y = kernels::bit_cast<int>(convert_half2(movmatrix(input.data[1])));
......@@ -464,8 +479,7 @@ public:
return output;
}
__device__ __forceinline__
static void mask(packed_qkv_t &pack, uint32_t maskVal, int m, int maxRows) {
__device__ __forceinline__ static void mask(packed_qkv_t &pack, uint32_t maskVal, int m, int maxRows) {
const int laneId = threadIdx.x % WARP_SIZE;
if (m * INSN_M + laneId / 4 >= maxRows) {
pack.x = maskVal;
......@@ -479,15 +493,15 @@ public:
// qkv: [batch, head, bm, NUM_WARPS, WARP_M_TILES, WARP_N_TILES, WARP_SIZE] of packed_qkv_t
template<typename F>
__device__ __forceinline__
static void apply(fpsum_warp &fpsum, packed_qkv_t *ptr_output, int maxRows, F &&funcPack, attn_half2_t maskVal) {
__device__ __forceinline__ static void
apply(fpsum_warp &fpsum, packed_qkv_t *ptr_output, int maxRows, F &&funcPack, attn_half2_t maskVal) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
static_assert(HEAD_DIM == WARP_N);
packed_qkv_t *ptrlane = &ptr_output[((warpId * WARP_M_TILES + 0) * WARP_N_TILES + 0) * WARP_SIZE + laneId];
unrolled_loop<WARP_M_TILES>([&]<int m>() ALWAYSINLINE {
unrolled_loop<WARP_N_TILES>([&]<int n>() ALWAYSINLINE {
packed_qkv_t pack = funcPack(fpsum[m * WARP_N_TILES + n]);
......@@ -497,35 +511,34 @@ public:
});
}
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
const int bm = binfo.bm;
const int bn = binfo.bn;
assert(binfo.numBlocksN % 3 == 0);
const int numBlocksQ = binfo.numBlocksN / 3;
const bool is_q = bn < numBlocksQ;
const bool is_k = !is_q && bn < numBlocksQ * 2;
const bool is_q = bn < numBlocksQ;
const bool is_k = !is_q && bn < numBlocksQ * 2;
// bn is head_id (assume HEAD_DIM == WARP_N)
int head_id, strideHead;
if (is_q) {
head_id = bn;
head_id = bn;
strideHead = args.strideHead_q;
} else if (is_k) {
head_id = bn - numBlocksQ;
head_id = bn - numBlocksQ;
strideHead = args.strideHead_k;
} else {
head_id = bn - numBlocksQ * 2;
head_id = bn - numBlocksQ * 2;
strideHead = args.strideHead_v;
}
int block_offset = head_id * strideHead + bm * NUM_WARPS * WARP_M_TILES * WARP_N_TILES * WARP_SIZE;
int maxRows = args.actualM - bm * BLOCK_M;
int maxRows = args.actualM - bm * BLOCK_M;
// static constexpr float neginf = -std::numeric_limits<float>::infinity();
if (is_q) {
apply(fpsum, args.out_q + block_offset, maxRows, pack_q, attn_half2_t(0.0f, 0.0f));
} else if (is_k) {
......@@ -537,9 +550,9 @@ public:
};
struct EpilogueLiteLA {
__device__ __forceinline__
static packed_f32psum_t mma_litela(packed_fpsum_t k, packed_fpsum_t v, packed_f32psum_t psum) {
__device__ __forceinline__ static packed_f32psum_t
mma_litela(packed_fpsum_t k, packed_fpsum_t v, packed_f32psum_t psum) {
for (int i = 0; i < 4; i++) {
k.data[i] = movmatrix(k.data[i]);
v.data[i] = movmatrix(v.data[i]);
......@@ -549,20 +562,20 @@ public:
}
static constexpr int LITELA_HEAD_DIM = 32;
static constexpr int LITELA_K_TILES = LITELA_HEAD_DIM / 16;
static constexpr int LITELA_V_TILES = LITELA_HEAD_DIM / 16;
static constexpr int LITELA_K_TILES = LITELA_HEAD_DIM / 16;
static constexpr int LITELA_V_TILES = LITELA_HEAD_DIM / 16;
static constexpr int SHMEM_SIZE = NUM_WARPS * (LITELA_HEAD_DIM + 1) * (LITELA_HEAD_DIM + 8) * sizeof(float);
// out_vk: [batch_size, num_heads, head_dim + 1, head_dim]
__device__ __forceinline__
static void apply_litela(const BlockInfo binfo, fpsum_warp fpsum, float *out_vk, int num_blocks_per_batch) {
__device__ __forceinline__ static void
apply_litela(const BlockInfo binfo, fpsum_warp fpsum, float *out_vk, int num_blocks_per_batch) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
using vk_t = float[NUM_WARPS][LITELA_HEAD_DIM + 1][LITELA_HEAD_DIM + 8];
extern __shared__ uint8_t shmem[];
vk_t &shmem_vk = *reinterpret_cast<vk_t *>(shmem);
static_assert(sizeof(vk_t) == SHMEM_SIZE);
......@@ -570,11 +583,13 @@ public:
assert(binfo.numBlocksN % 3 == 0);
const int num_heads = binfo.numBlocksN / 3 * 2 * (WARP_N / (LITELA_HEAD_DIM * 2));
const int batch_id = binfo.bm / num_blocks_per_batch;
const int batch_id = binfo.bm / num_blocks_per_batch;
for (int head_id = 0; head_id < WARP_N / (LITELA_HEAD_DIM * 2); head_id++) {
const int global_head_id = (binfo.bn - binfo.numBlocksN / 3) * (WARP_N / (LITELA_HEAD_DIM * 2)) + head_id;
float *out_vk_current_head = out_vk + (batch_id * num_heads + global_head_id) * (LITELA_HEAD_DIM + 1) * LITELA_HEAD_DIM;
const int global_head_id =
(binfo.bn - binfo.numBlocksN / 3) * (WARP_N / (LITELA_HEAD_DIM * 2)) + head_id;
float *out_vk_current_head =
out_vk + (batch_id * num_heads + global_head_id) * (LITELA_HEAD_DIM + 1) * LITELA_HEAD_DIM;
for (int i = laneId; i < sizeof(shmem_vk) / sizeof(float) / NUM_WARPS; i += WARP_SIZE) {
*((&shmem_vk[warpId][0][0]) + i) = 0;
......@@ -583,12 +598,13 @@ public:
for (int tile_v = 0; tile_v < LITELA_V_TILES; tile_v++) {
for (int tile_k = 0; tile_k < LITELA_K_TILES; tile_k++) {
packed_f32psum_t attn_sum = { 0 };
packed_f32psum_t attn_sum = {0};
for (int i = 0; i < WARP_M_TILES; i++) {
packed_fpsum_t k = fpsum[i * WARP_N_TILES + head_id * (LITELA_HEAD_DIM * 2) / 16 + tile_k];
packed_fpsum_t v = fpsum[i * WARP_N_TILES + head_id * (LITELA_HEAD_DIM * 2) / 16 + LITELA_HEAD_DIM / 16 + tile_v];
packed_fpsum_t v = fpsum[i * WARP_N_TILES + head_id * (LITELA_HEAD_DIM * 2) / 16 +
LITELA_HEAD_DIM / 16 + tile_v];
for (int j = 0; j < 4; j++) {
k.data[j] = __hmax2(k.data[j], half2_t(0, 0)); // relu
k.data[j] = __hmax2(k.data[j], half2_t(0, 0)); // relu
}
attn_sum = mma_litela(k, v, attn_sum);
}
......@@ -607,14 +623,14 @@ public:
}
}
for (int tile_k = 0; tile_k < LITELA_K_TILES; tile_k++) {
packed_f32psum_t attn_sum = { 0 };
packed_f32psum_t attn_sum = {0};
for (int i = 0; i < WARP_M_TILES; i++) {
packed_fpsum_t k = fpsum[i * WARP_N_TILES + head_id * (LITELA_HEAD_DIM * 2) / 16 + tile_k];
packed_fpsum_t v = {};
for (int j = 0; j < 4; j++) {
k.data[j] = __hmax2(k.data[j], half2_t(0, 0)); // relu
k.data[j] = __hmax2(k.data[j], half2_t(0, 0)); // relu
}
#pragma unroll
#pragma unroll
for (int i = 0; i < 4; i++) {
v.data[i] = half2_t(1, 1);
}
......@@ -660,21 +676,23 @@ public:
int actualM;
};
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
const int bm = binfo.bm;
const int bn = binfo.bn;
if (bn < binfo.numBlocksN / 3) {
fpsum = apply_act(fpsum, [](half_t x) { return __hmax(x, 0); }); // relu
return EpilogueDefault()(
binfo,
fpsum,
M, N / 3, K, typename EpilogueDefault::Arguments{
.out = args.out_q,
.actualM = args.actualM,
.actualN = N / 3,
});
fpsum = apply_act(fpsum, [](half_t x) { return __hmax(x, 0); }); // relu
return EpilogueDefault()(binfo,
fpsum,
M,
N / 3,
K,
typename EpilogueDefault::Arguments{
.out = args.out_q,
.actualM = args.actualM,
.actualN = N / 3,
});
}
return apply_litela(binfo, fpsum, args.out_vk, args.num_blocks_per_batch);
......@@ -686,46 +704,49 @@ public:
struct vk_mul_q_kernel {
static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
// FIXME FIXME FIXME
__device__
void operator()(half_t *q, const float *vk, float eps, int num_tokens) {
__device__ void operator()(half_t *q, const float *vk, float eps, int num_tokens) {
const int block_id = blockIdx.x;
const int head_id = blockIdx.y;
const int batch_id = blockIdx.z;
const int num_blocks = gridDim.x;
const int num_heads = gridDim.y;
const int num_heads = gridDim.y;
const int block_size = blockDim.x;
bool pred = block_id * block_size + threadIdx.x < num_tokens;
half_t *localq = &q[(((batch_id * num_blocks + block_id) * block_size + threadIdx.x) * num_heads + head_id) * LITELA_HEAD_DIM];
half_t *localq =
&q[(((batch_id * num_blocks + block_id) * block_size + threadIdx.x) * num_heads + head_id) *
LITELA_HEAD_DIM];
const float *localvk = &vk[(batch_id * num_heads + head_id) * (LITELA_HEAD_DIM + 1) * LITELA_HEAD_DIM];
// half_t *localout = &out[(((batch_id * num_blocks + block_id) * block_size + threadIdx.x) * num_heads + head_id) * LITELA_HEAD_DIM];
// half_t *localout = &out[(((batch_id * num_blocks + block_id) * block_size + threadIdx.x) * num_heads
// + head_id) * LITELA_HEAD_DIM];
using packed_q = std::array<half_t, 8>;
using packed_q = std::array<half_t, 8>;
using packed_vk = std::array<float, 4>;
half_t qblock[LITELA_HEAD_DIM];
for (int i = 0; i < LITELA_HEAD_DIM; i += sizeof(packed_q) / sizeof(half_t)) {
if (pred) {
*reinterpret_cast<packed_q *>(&qblock[i]) = load(reinterpret_cast<const packed_q *>(&localq[i]));
*reinterpret_cast<packed_q *>(&qblock[i]) =
load(reinterpret_cast<const packed_q *>(&localq[i]));
}
}
float outblock[LITELA_HEAD_DIM + 1];
#pragma unroll
#pragma unroll
for (int j = 0; j < LITELA_HEAD_DIM + 1; j++) {
outblock[j] = 0;
#pragma unroll
#pragma unroll
for (int i = 0; i < LITELA_HEAD_DIM; i += sizeof(packed_vk) / sizeof(float)) {
packed_vk vkpack = load(reinterpret_cast<const packed_vk *>(&localvk[j * LITELA_HEAD_DIM + i]));
#pragma unroll
#pragma unroll
for (int k = 0; k < vkpack.size(); k++) {
outblock[j] += (float)qblock[i + k] * vkpack[k];
}
}
}
for (int i = 0; i < LITELA_HEAD_DIM; i += sizeof(packed_q) / sizeof(half_t)) {
packed_q opack;
for (int k = 0; k < opack.size(); k++) {
......@@ -739,11 +760,11 @@ public:
};
};
template<typename Epilogue>
struct test_epilogue_kernel {
static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
static constexpr size_t SHMEM_PER_WARP = ceilDiv<size_t>(Base::template load_act_to_fpsum<false>::SHMEM_SIZE, 128) * 128;
static constexpr size_t SHMEM_PER_WARP =
ceilDiv<size_t>(Base::template load_act_to_fpsum<false>::SHMEM_SIZE, 128) * 128;
static constexpr size_t SHMEM_SIZE = SHMEM_PER_WARP * NUM_WARPS;
struct Arguments {
......@@ -757,18 +778,16 @@ public:
typename Epilogue::Arguments argsEpilogue;
};
__device__ __forceinline__
void operator()(Arguments args)
{
__device__ __forceinline__ void operator()(Arguments args) {
const BlockInfo binfo = {
.bm = (int)blockIdx.x,
.bn = (int)blockIdx.y,
.bm = (int)blockIdx.x,
.bn = (int)blockIdx.y,
.numBlocksM = (int)gridDim.x,
.numBlocksN = (int)gridDim.y,
};
const int bm = binfo.bm;
const int bn = binfo.bn;
const int bm = binfo.bm;
const int bn = binfo.bn;
const int warpId = threadIdx.x / WARP_SIZE;
const int m_offset = bm * BLOCK_M + warpId * WARP_M;
......@@ -778,26 +797,27 @@ public:
fpsum_warp fpsum;
Base::template load_act_to_fpsum<false>()(
args.input + m_offset * args.actualN + n_offset,
args.actualN,
args.actualM - m_offset,
args.actualN - n_offset,
fpsum,
shmem + warpId * SHMEM_PER_WARP
);
Base::template load_act_to_fpsum<false>()(args.input + m_offset * args.actualN + n_offset,
args.actualN,
args.actualM - m_offset,
args.actualN - n_offset,
fpsum,
shmem + warpId * SHMEM_PER_WARP);
Epilogue()(binfo, fpsum, args.M, args.N, 0, args.argsEpilogue);
EpilogueDefault()(binfo, fpsum, args.M, args.N, 0, typename EpilogueDefault::Arguments{
.out = args.output,
.actualM = args.actualM,
.actualN = args.actualN,
});
EpilogueDefault()(binfo,
fpsum,
args.M,
args.N,
0,
typename EpilogueDefault::Arguments{
.out = args.output,
.actualM = args.actualM,
.actualN = args.actualN,
});
}
};
};
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/zgemm/gemm_base.cuh
View file @ 37a27712
......@@ -7,7 +7,6 @@
#include "gemm_utils.cuh"
#include "mma_earlycuda.cuh"
#pragma nv_diag_suppress 177
#ifdef _MSC_VER
......@@ -32,8 +31,8 @@ class GEMMConfig_W4A4 {
public:
// BE CAREFUL: weights need to be repacked when the tiling size changes
static constexpr int BLOCK_M = 256;
static constexpr int BLOCK_N = 128;
static constexpr int BLOCK_M = 256;
static constexpr int BLOCK_N = 128;
static constexpr int WARP_SIZE = 32;
static constexpr int NUM_WARPS = 8;
......@@ -45,18 +44,18 @@ public:
// may generate incorrect results in certain circumstances
static constexpr bool FASTER_I2F = faster_i2f;
using half_t = typename std::conditional_t<bf16, __nv_bfloat16, half>;
using half_t = typename std::conditional_t<bf16, __nv_bfloat16, half>;
using half2_t = typename std::conditional_t<bf16, __nv_bfloat162, half2>;
};
using GEMMConfig_W4A4_FP16 = GEMMConfig_W4A4<false>;
using GEMMConfig_W4A4_BF16 = GEMMConfig_W4A4<true>;
using GEMMConfig_W4A4_FP16 = GEMMConfig_W4A4<false>;
using GEMMConfig_W4A4_BF16 = GEMMConfig_W4A4<true>;
using GEMMConfig_W4A4_FP16_FasterI2F = GEMMConfig_W4A4<false, true>;
class GEMMConfig_W8A8 {
public:
static constexpr int BLOCK_M = 256;
static constexpr int BLOCK_N = 128;
static constexpr int BLOCK_M = 256;
static constexpr int BLOCK_N = 128;
static constexpr int WARP_SIZE = 32;
static constexpr int NUM_WARPS = 8;
......@@ -97,13 +96,13 @@ public:
/**
* refer to https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#mma-16864-c
*
*
* wscales store order: (pack = 4)
* 0 1 8 9 <-- load by lane 0, broadcast to lane {0, 4, 8, ..., 28} (8x)
* 2 3 10 11 <-- load by lane 1, broadcast to lane {1, 5, 9, ..., 29} (8x)
* 4 5 12 13 <-- load by lane 2, broadcast to lane {2, 6, 10, ..., 30} (8x)
* 6 7 14 15 <-- load by lane 3, broadcast to lane {3, 7, 11, ..., 31} (8x)
*
*
* 16 17 24 25 <-- load by lane 4, broadcast to lane {0, 4, 8, ..., 28} (8x)
* ...
* 22 23 30 31 <-- load by lane 7, broadcast to lane {3, 7, 11, ..., 31} (8x)
......@@ -111,24 +110,25 @@ public:
* 112 113 120 121 <-- load by lane 28, broadcast to lane {0, 4, 8, ..., 28} (8x)
* ...
* 118 119 126 127 <-- load by lane 31, broadcast to lane {3, 7, 11, ..., 31} (8x)
*
*
* wscales store order: (pack = 8)
* 0 1 8 9 16 17 24 25 <-- load by lane 0, broadcast to lane {0, 4, 8, ..., 28} (8x)
* 2 3 10 11 18 19 26 27 <-- load by lane 1, broadcast to lane {1, 5, 9, ..., 29} (8x)
* 4 5 12 13 20 21 28 29 <-- load by lane 2, broadcast to lane {2, 6, 10, ..., 30} (8x)
* 6 7 14 15 22 23 30 31 <-- load by lane 3, broadcast to lane {3, 7, 11, ..., 31} (8x)
*
*
* 224 225 232 233 240 241 248 249 <-- load by lane 28, broadcast to lane {0, 4, 8, ..., 28} (8x)
* ...
* 230 231 238 239 246 247 254 255 <-- load by lane 31, broadcast to lane {3, 7, 11, ..., 31} (8x)
*
* {k}-th wscale used by lane {i} => {k // (WSCALES_PACK_SIZE * WARP_SIZE)}-th pack, in lane {4*(k // WSCALES_PACK_SIZE) + i % 4}, element {k % WSCALES_PACK_SIZE}
*
*
* {k}-th wscale used by lane {i} => {k // (WSCALES_PACK_SIZE * WARP_SIZE)}-th pack, in lane {4*(k //
* WSCALES_PACK_SIZE) + i % 4}, element {k % WSCALES_PACK_SIZE}
*
* max pack size set to 8 since max load size is 16 bytes / lane
* min pack size set to 2 since shuffle granularity is 32b 2*half
* */
static constexpr int WSCALES_PACK_SIZE = clamp(WARP_N / WARP_SIZE, 4 / sizeof(half), 16 / sizeof(half));
static constexpr int WSCALES_NUM_PACKS = ceilDiv(WARP_N, (WSCALES_PACK_SIZE * WARP_SIZE));
* */
static constexpr int WSCALES_PACK_SIZE = clamp(WARP_N / WARP_SIZE, 4 / sizeof(half), 16 / sizeof(half));
static constexpr int WSCALES_NUM_PACKS = ceilDiv(WARP_N, (WSCALES_PACK_SIZE * WARP_SIZE));
static constexpr int WSCALES_VALID_LANES = std::min(WARP_SIZE, WARP_N / WSCALES_PACK_SIZE);
/**
......@@ -145,16 +145,17 @@ public:
* ...
* 54 62
* 55 63 <-- load by lane 31, broadcast to lane {28, 29, 30, 31} (4x)
*
* {k}-th wscale used by lane {i} => {k // (ASCALES_PACK_SIZE * WARP_SIZE)}-th pack, in lane {8*(k // ASCALES_PACK_SIZE) + i // 4}, element {k % ASCALES_PACK_SIZE}
*
* {k}-th wscale used by lane {i} => {k // (ASCALES_PACK_SIZE * WARP_SIZE)}-th pack, in lane {8*(k //
* ASCALES_PACK_SIZE) + i // 4}, element {k % ASCALES_PACK_SIZE}
*/
static constexpr int ASCALES_PACK_SIZE = clamp(WARP_M / WARP_SIZE, 4 / sizeof(half), 16 / sizeof(half));
static constexpr int ASCALES_NUM_PACKS = ceilDiv(WARP_M, (ASCALES_PACK_SIZE * WARP_SIZE));
static constexpr int ASCALES_PACK_SIZE = clamp(WARP_M / WARP_SIZE, 4 / sizeof(half), 16 / sizeof(half));
static constexpr int ASCALES_NUM_PACKS = ceilDiv(WARP_M, (ASCALES_PACK_SIZE * WARP_SIZE));
static constexpr int ASCALES_VALID_LANES = std::min(WARP_SIZE, WARP_M / ASCALES_PACK_SIZE);
using packed_act_t = uint4;
using packed_wgt_t = uint4;
struct alignas(32) packed_psum_t {
int data[8];
};
......@@ -184,15 +185,14 @@ public:
half2_t data[ASCALES_PACK_SIZE / 2];
};
using act_warp = std::array<packed_act_t, WARP_M_TILES>;
using wgt_warp = std::array<packed_wgt_t, WARP_N_TILES>;
using ascale_warp = std::array<packed_ascale_t, ASCALES_NUM_PACKS>;
using wscale_warp = std::array<packed_wscale_t, WSCALES_NUM_PACKS>;
using fpsum_warp = std::array<packed_fpsum_t, WARP_M_TILES * WARP_N_TILES>;
using f32psum_warp = std::array<packed_f32psum_t, WARP_M_TILES * WARP_N_TILES>;
using act_warp = std::array<packed_act_t, WARP_M_TILES>;
using wgt_warp = std::array<packed_wgt_t, WARP_N_TILES>;
using ascale_warp = std::array<packed_ascale_t, ASCALES_NUM_PACKS>;
using wscale_warp = std::array<packed_wscale_t, WSCALES_NUM_PACKS>;
using fpsum_warp = std::array<packed_fpsum_t, WARP_M_TILES * WARP_N_TILES>;
using f32psum_warp = std::array<packed_f32psum_t, WARP_M_TILES * WARP_N_TILES>;
using gated_fpsum_warp = std::array<packed_gated_fpsum_t, WARP_M_TILES * WARP_N_TILES>;
struct BlockInfo {
int bm;
int bn;
......@@ -200,19 +200,19 @@ public:
int numBlocksN;
};
__device__ __forceinline__
static packed_f32psum_t mma_f16xf16_f32(packed_fpsum_t a, packed_fpsum_t b, packed_f32psum_t psum) {
__device__ __forceinline__ static packed_f32psum_t
mma_f16xf16_f32(packed_fpsum_t a, packed_fpsum_t b, packed_f32psum_t psum) {
static_assert(std::is_same_v<half_t, half> || std::is_same_v<half_t, __nv_bfloat16>);
static constexpr bool is_bf16 = std::is_same_v<half_t, __nv_bfloat16>;
uint4 out1 = mma_m16n8k16_f32f16f16f32<is_bf16>(
kernels::bit_cast<uint4>(a),
kernels::bit_cast<uint2>(std::array<half2_t, 2>(b.data[0], b.data[1])),
kernels::bit_cast<uint4>(a),
kernels::bit_cast<uint2>(std::array<half2_t, 2>(b.data[0], b.data[1])),
kernels::bit_cast<uint4>(float4(psum.data[0], psum.data[1], psum.data[2], psum.data[3])));
uint4 out2 = mma_m16n8k16_f32f16f16f32<is_bf16>(
kernels::bit_cast<uint4>(a),
kernels::bit_cast<uint2>(std::array<half2_t, 2>(b.data[2], b.data[3])),
kernels::bit_cast<uint4>(a),
kernels::bit_cast<uint2>(std::array<half2_t, 2>(b.data[2], b.data[3])),
kernels::bit_cast<uint4>(float4(psum.data[4], psum.data[5], psum.data[6], psum.data[7])));
psum.data[0] = kernels::bit_cast<float>(out1.x);
psum.data[1] = kernels::bit_cast<float>(out1.y);
......@@ -222,12 +222,11 @@ public:
psum.data[5] = kernels::bit_cast<float>(out2.y);
psum.data[6] = kernels::bit_cast<float>(out2.z);
psum.data[7] = kernels::bit_cast<float>(out2.w);
return psum;
}
__device__ __forceinline__
static packed_fpsum_t packed_fp32_to_fp16(packed_f32psum_t input) {
__device__ __forceinline__ static packed_fpsum_t packed_fp32_to_fp16(packed_f32psum_t input) {
packed_fpsum_t results;
for (int i = 0; i < 4; i++) {
results.data[i] = float22half2<half2_t>(float2(input.data[i * 2], input.data[i * 2 + 1]));
......@@ -235,31 +234,28 @@ public:
return results;
}
__device__ __forceinline__
static packed_f32psum_t packed_fp16_to_fp32(packed_fpsum_t input) {
__device__ __forceinline__ static packed_f32psum_t packed_fp16_to_fp32(packed_fpsum_t input) {
packed_f32psum_t results;
for (int i = 0; i < 4; i++) {
float2 tmp = half22float2(input.data[i]);
results.data[i * 2] = tmp.x;
float2 tmp = half22float2(input.data[i]);
results.data[i * 2] = tmp.x;
results.data[i * 2 + 1] = tmp.y;
}
return results;
}
__device__ __forceinline__
static fpsum_warp packed_fp32_to_fp16(f32psum_warp input) {
__device__ __forceinline__ static fpsum_warp packed_fp32_to_fp16(f32psum_warp input) {
fpsum_warp results;
#pragma unroll
#pragma unroll
for (int i = 0; i < results.size(); i++) {
results[i] = packed_fp32_to_fp16(input[i]);
}
return results;
}
__device__ __forceinline__
static f32psum_warp packed_fp16_to_fp32(fpsum_warp input) {
__device__ __forceinline__ static f32psum_warp packed_fp16_to_fp32(fpsum_warp input) {
f32psum_warp results;
#pragma unroll
#pragma unroll
for (int i = 0; i < results.size(); i++) {
results[i] = packed_fp16_to_fp32(input[i]);
}
......@@ -267,108 +263,110 @@ public:
}
// activation: row major, [M / BLOCK_M, K / WARP_K, NUM_WARPS, WARP_M_TILES, WARP_SIZE] of packed_act_t
__device__ __forceinline__
static void load_act(const packed_act_t *act, int k, int K, act_warp &out, bool pred) {
__device__ __forceinline__ static void load_act(const packed_act_t *act, int k, int K, act_warp &out, bool pred) {
int laneId = threadIdx.x % WARP_SIZE;
int warpId = threadIdx.x / WARP_SIZE;
#pragma unroll
#pragma unroll
for (int i = 0; i < WARP_M_TILES; i++) {
//if (pred) {
// out[i] = load(&act[((warpId * WARP_M_TILES + i) * K / WARP_K + k) * WARP_SIZE + laneId]);
out[i] = load_pred(&act[((k * NUM_WARPS + warpId) * WARP_M_TILES + i) * WARP_SIZE + laneId], pred);
// if (pred) {
// out[i] = load(&act[((warpId * WARP_M_TILES + i) * K / WARP_K + k) * WARP_SIZE + laneId]);
out[i] = load_pred(&act[((k * NUM_WARPS + warpId) * WARP_M_TILES + i) * WARP_SIZE + laneId], pred);
//}
}
}
// weight: column major: [N / BLOCK_N, 1, K / WARP_K, WARP_N_TILES, WARP_SIZE] of packed_wgt_t
__device__ __forceinline__
static void load_wgt(const packed_wgt_t *wgt, int k, int K, wgt_warp &out, bool pred) {
__device__ __forceinline__ static void load_wgt(const packed_wgt_t *wgt, int k, int K, wgt_warp &out, bool pred) {
int laneId = threadIdx.x % WARP_SIZE;
// const packed_wgt_t *ptr = &wgt[(0 * K / WARP_K + k) * WARP_SIZE + laneId];
const packed_wgt_t *ptr = &wgt[(0 + k * WARP_N_TILES) * WARP_SIZE + laneId];
// int offset = K / WARP_K * WARP_SIZE;
#pragma unroll
#pragma unroll
for (int i = 0; i < WARP_N_TILES; i++) {
//if (pred) {
// out[i] = load(&wgt[(i * K / WARP_K + k) * WARP_SIZE + laneId]);
// out[i] = load(&wgt[(i + k * WARP_N_TILES) * WARP_SIZE + laneId]);
out[i] = load_pred(&ptr[i * WARP_SIZE], pred);
// ptr += offset;
// if (pred) {
// out[i] = load(&wgt[(i * K / WARP_K + k) * WARP_SIZE + laneId]);
// out[i] = load(&wgt[(i + k * WARP_N_TILES) * WARP_SIZE + laneId]);
out[i] = load_pred(&ptr[i * WARP_SIZE], pred);
// ptr += offset;
//}
}
}
// ascales: row major [M / BLOCK_M, K / group size, NUM_WARPS, ASCALES_NUM_PACKS, ASCALES_VALID_LANES] of packed_ascale_t
__device__ __forceinline__
static void load_ascale(const packed_ascale_t *ascales, int group, int M, ascale_warp &out, bool pred) {
// ascales: row major [M / BLOCK_M, K / group size, NUM_WARPS, ASCALES_NUM_PACKS, ASCALES_VALID_LANES] of
// packed_ascale_t
__device__ __forceinline__ static void
load_ascale(const packed_ascale_t *ascales, int group, int M, ascale_warp &out, bool pred) {
int laneId = threadIdx.x % WARP_SIZE;
int warpId = threadIdx.x / WARP_SIZE;
#pragma unroll
#pragma unroll
for (int i = 0; i < ASCALES_NUM_PACKS; i++) {
// if (pred && laneId < ASCALES_VALID_LANES) {
// out[i] = ascales[(group * M / WARP_M + warpId) * ASCALES_VALID_LANES * ASCALES_NUM_PACKS + i * ASCALES_VALID_LANES + laneId];
out[i] = load_pred(&ascales[(group * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES + i * ASCALES_VALID_LANES + laneId], pred && laneId < ASCALES_VALID_LANES);
// out[i] = ascales[(group * M / WARP_M + warpId) * ASCALES_VALID_LANES * ASCALES_NUM_PACKS + i *
// ASCALES_VALID_LANES + laneId];
out[i] = load_pred(&ascales[(group * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES +
i * ASCALES_VALID_LANES + laneId],
pred && laneId < ASCALES_VALID_LANES);
// }
}
}
// wscales: column major [N / BLOCK_N, K / group size, 1, WSCALES_NUM_PACKS, WSCALES_VALID_LANES] of packed_wscale_t </del>
__device__ __forceinline__
static void load_wscale(const packed_wscale_t *wscales, int group, int N, wscale_warp &out, bool pred) {
// wscales: column major [N / BLOCK_N, K / group size, 1, WSCALES_NUM_PACKS, WSCALES_VALID_LANES] of packed_wscale_t
// </del>
__device__ __forceinline__ static void
load_wscale(const packed_wscale_t *wscales, int group, int N, wscale_warp &out, bool pred) {
int laneId = threadIdx.x % WARP_SIZE;
// static_assert(WSCALES_NUM_PACKS * WSCALES_VALID_LANES == 32);
// static_assert(sizeof(packed_wscale_t) == 8);
// const packed_wscale_t *ptr = &wscales[(group * WSCALES_NUM_PACKS + 0) * WSCALES_VALID_LANES + laneId];
// // const packed_wscale_t *ptr = (const packed_wscale_t *)((const char *)wscales) + ((group * WSCALES_NUM_PACKS + 0) * WSCALES_VALID_LANES + laneId) * sizeof(packed_wscale_t);
// // const packed_wscale_t *ptr = (const packed_wscale_t *)((const char *)wscales) + ((group *
// WSCALES_NUM_PACKS + 0) * WSCALES_VALID_LANES + laneId) * sizeof(packed_wscale_t);
#pragma unroll
#pragma unroll
for (int i = 0; i < WSCALES_NUM_PACKS; i++) {
// if (pred && laneId < WSCALES_VALID_LANES) {
// out[i] = wscales[group * N / WARP_N * WSCALES_VALID_LANES * WSCALES_NUM_PACKS + i * WSCALES_VALID_LANES + laneId];
// out[i] = load(&wscales[group * N / WARP_N * WSCALES_VALID_LANES * WSCALES_NUM_PACKS + i * WSCALES_VALID_LANES + laneId]);
out[i] = load_pred(&wscales[(group * WSCALES_NUM_PACKS + i) * WSCALES_VALID_LANES + laneId], pred && laneId < WSCALES_VALID_LANES);
// out[i] = load(&ptr[i * WSCALES_VALID_LANES]);
// out[i] = wscales[group * N / WARP_N * WSCALES_VALID_LANES * WSCALES_NUM_PACKS + i * WSCALES_VALID_LANES +
// laneId]; out[i] = load(&wscales[group * N / WARP_N * WSCALES_VALID_LANES * WSCALES_NUM_PACKS + i *
// WSCALES_VALID_LANES + laneId]);
out[i] = load_pred(&wscales[(group * WSCALES_NUM_PACKS + i) * WSCALES_VALID_LANES + laneId],
pred && laneId < WSCALES_VALID_LANES);
// out[i] = load(&ptr[i * WSCALES_VALID_LANES]);
// }
}
}
// get {k}-th and {k+1}-th wscale from the block, k must be multiples of 2, k must be uniform across all lanes
__device__ __forceinline__
static half2_t broadcast_wscale(wscale_warp block, int k, int laneId) {
const int packIdx = k / (WSCALES_PACK_SIZE * WARP_SIZE);
const int srcLane = 4 * (k / WSCALES_PACK_SIZE) + laneId % 4;
__device__ __forceinline__ static half2_t broadcast_wscale(wscale_warp block, int k, int laneId) {
const int packIdx = k / (WSCALES_PACK_SIZE * WARP_SIZE);
const int srcLane = 4 * (k / WSCALES_PACK_SIZE) + laneId % 4;
const int elementIdx = k % WSCALES_PACK_SIZE / 2;
return __shfl_sync(~0, block[packIdx].data[elementIdx], srcLane);
}
// get {k}-th and {k+1}-th ascale from the block, k must be multiples of 2, k must be uniform across all lanes
__device__ __forceinline__
static half2_t broadcast_ascale(ascale_warp block, int k, int laneId) {
const int packIdx = k / (ASCALES_PACK_SIZE * WARP_SIZE);
const int srcLane = 8 * (k / ASCALES_PACK_SIZE) + laneId / 4;
__device__ __forceinline__ static half2_t broadcast_ascale(ascale_warp block, int k, int laneId) {
const int packIdx = k / (ASCALES_PACK_SIZE * WARP_SIZE);
const int srcLane = 8 * (k / ASCALES_PACK_SIZE) + laneId / 4;
const int elementIdx = k % ASCALES_PACK_SIZE / 2;
return __shfl_sync(~0, block[packIdx].data[elementIdx], srcLane);
}
struct i2f_normal {
__device__ __forceinline__
static float2 int2float2(int x, int y) {
__device__ __forceinline__ static float2 int2float2(int x, int y) {
return make_float2(__int2float_rn(x), __int2float_rn(y));
}
__device__ __forceinline__
static half2_t int2half2(int x, int y) {
__device__ __forceinline__ static half2_t int2half2(int x, int y) {
return float22half2<half2_t>(int2float2(x, y));
}
};
template<typename i2f = i2f_normal, typename F>
__device__ __forceinline__
static void apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, fpsum_warp &fpsum) {
__device__ __forceinline__ static void
apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, fpsum_warp &fpsum) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -377,8 +375,8 @@ public:
for (int i = 0; i < WARP_M_TILES; i++) {
half2_t as = broadcast_ascale(ascale, i * 2, laneId);
asx[i] = half2_t(as.x, as.x);
asy[i] = half2_t(as.y, as.y);
asx[i] = half2_t(as.x, as.x);
asy[i] = half2_t(as.y, as.y);
}
for (int j = 0; j < WARP_N_TILES; j++) {
......@@ -392,8 +390,9 @@ public:
// constexpr int target = 0;
// if (threadIdx.x == 3 && j == 1 && i == 0) {
// printf("before ws2 = %f %f fsum.data[%d] = %f %f\n", (float)ws2.x, (float)ws2.y, target, (float)fsum.data[target].x, (float)fsum.data[target].y);
// printf("before ws2 = %f %f fsum.data[%d] = %f %f\n", (float)ws2.x, (float)ws2.y, target,
// (float)fsum.data[target].x, (float)fsum.data[target].y);
// }
fsum.data[0] = __hfma2(i2f::int2half2(psum.data[0], psum.data[1]), __hmul2(asx[i], ws1), fsum.data[0]);
......@@ -402,15 +401,16 @@ public:
fsum.data[3] = __hfma2(i2f::int2half2(psum.data[6], psum.data[7]), __hmul2(asy[i], ws2), fsum.data[3]);
// if (threadIdx.x == 3 && j == 1 && i == 0) {
// printf("before ws2 = %f %f fsum.data[%d] = %f %f\n", (float)ws2.x, (float)ws2.y, target, (float)fsum.data[target].x, (float)fsum.data[target].y);
// printf("before ws2 = %f %f fsum.data[%d] = %f %f\n", (float)ws2.x, (float)ws2.y, target,
// (float)fsum.data[target].x, (float)fsum.data[target].y);
// }
}
}
}
template<typename i2f = i2f_normal, typename F>
__device__ __forceinline__
static void apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, f32psum_warp &fpsum) {
__device__ __forceinline__ static void
apply_scales(F &&getpsum, ascale_warp ascale, wscale_warp wscale, f32psum_warp &fpsum) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -419,8 +419,8 @@ public:
for (int i = 0; i < WARP_M_TILES; i++) {
half2_t as = broadcast_ascale(ascale, i * 2, laneId);
asx[i] = half22float2(half2_t(as.x, as.x));
asy[i] = half22float2(half2_t(as.y, as.y));
asx[i] = half22float2(half2_t(as.x, as.x));
asy[i] = half22float2(half2_t(as.y, as.y));
}
auto fma2 = [](float2 a, float2 b, float &cx, float &cy) ALWAYSINLINE {
......@@ -449,17 +449,18 @@ public:
* input: WARP_M of half (in shared memory, per-warp)
* output: [..., ASCALES_NUM_PACKS, ASCALES_VALID_LANES] in global memory, per-warp
*/
__device__ __forceinline__
static void pack_ascales(const half_t *input, packed_ascale_t *output) {
__device__ __forceinline__ static void pack_ascales(const half_t *input, packed_ascale_t *output) {
const int laneId = threadIdx.x % WARP_SIZE;
#pragma unroll
#pragma unroll
for (int j = 0; j < ASCALES_NUM_PACKS; j++) {
if (laneId < ASCALES_VALID_LANES) {
packed_ascale_t tmp;
#pragma unroll
#pragma unroll
for (int i = 0; i < ASCALES_PACK_SIZE; i += 2) {
tmp.data[i / 2].x = input[j * ASCALES_PACK_SIZE * WARP_SIZE + laneId / 8 * 8 * ASCALES_PACK_SIZE + laneId % 8 + i * 8];
tmp.data[i / 2].y = input[j * ASCALES_PACK_SIZE * WARP_SIZE + laneId / 8 * 8 * ASCALES_PACK_SIZE + laneId % 8 + (i + 1) * 8];
tmp.data[i / 2].x = input[j * ASCALES_PACK_SIZE * WARP_SIZE + laneId / 8 * 8 * ASCALES_PACK_SIZE +
laneId % 8 + i * 8];
tmp.data[i / 2].y = input[j * ASCALES_PACK_SIZE * WARP_SIZE + laneId / 8 * 8 * ASCALES_PACK_SIZE +
laneId % 8 + (i + 1) * 8];
}
output[j * ASCALES_VALID_LANES + laneId] = tmp;
}
......@@ -470,16 +471,17 @@ public:
* input: WARP_N of half (in shared memory, per-warp)
* output: [..., WSCALES_NUM_PACKS, WSCALES_VALID_LANES] in global memory, per-warp
*/
__device__ __forceinline__
static void pack_wscales(const half_t *input, packed_wscale_t *output) {
__device__ __forceinline__ static void pack_wscales(const half_t *input, packed_wscale_t *output) {
const int laneId = threadIdx.x % WARP_SIZE;
#pragma unroll
#pragma unroll
for (int j = 0; j < WSCALES_NUM_PACKS; j++) {
if (laneId < WSCALES_VALID_LANES) {
packed_wscale_t tmp;
#pragma unroll
#pragma unroll
for (int i = 0; i < WSCALES_PACK_SIZE; i += 2) {
tmp.data[i / 2] = *reinterpret_cast<const half2_t *>(&input[j * WSCALES_PACK_SIZE * WARP_SIZE + laneId / 4 * 4 * WSCALES_PACK_SIZE + laneId % 4 * 2 + i * 4]);
tmp.data[i / 2] = *reinterpret_cast<const half2_t *>(
&input[j * WSCALES_PACK_SIZE * WARP_SIZE + laneId / 4 * 4 * WSCALES_PACK_SIZE + laneId % 4 * 2 +
i * 4]);
}
store(&output[j * WSCALES_VALID_LANES + laneId], tmp);
}
......@@ -491,17 +493,16 @@ public:
using matrix_t = half_t[8][WARP_N + 8];
static constexpr int SHMEM_SIZE = sizeof(matrix_t);
static constexpr int PACK_SIZE = WARP_N / WARP_SIZE;
using pack_t = std::array<half_t, PACK_SIZE>;
static constexpr int PACK_SIZE = WARP_N / WARP_SIZE;
using pack_t = std::array<half_t, PACK_SIZE>;
// F (int rowId, pack_t &pack)
template<typename ...F>
__device__ __forceinline__
void operator()(fpsum_warp fpsum, half_t *output, int stride, int maxRows, int maxCols, void *shmem, F &&...plugins) {
template<typename... F>
__device__ __forceinline__ void operator()(
fpsum_warp fpsum, half_t *output, int stride, int maxRows, int maxCols, void *shmem, F &&...plugins) {
const int laneId = threadIdx.x % WARP_SIZE;
matrix_t &mat = *reinterpret_cast<matrix_t *>(shmem);
matrix_t &mat = *reinterpret_cast<matrix_t *>(shmem);
// pack_t reduce_tmp;
// constexpr bool enableReduce = !std::is_void_v<FuncReduce>;
......@@ -518,19 +519,19 @@ public:
// }
// };
#pragma unroll
#pragma unroll
for (int i = 0; i < WARP_M_TILES; i++) {
#pragma unroll
#pragma unroll
for (int j = 0; j < WARP_N_TILES; j++) {
packed_fpsum_t &fsum = fpsum[i * WARP_N_TILES + j];
int row = laneId / 4;
int col = laneId % 4 * 2 + j * INSN_N;
packed_fpsum_t &fsum = fpsum[i * WARP_N_TILES + j];
int row = laneId / 4;
int col = laneId % 4 * 2 + j * INSN_N;
*reinterpret_cast<half2_t *>(&mat[row][col + 0]) = fsum.data[0];
*reinterpret_cast<half2_t *>(&mat[row][col + 8]) = fsum.data[2];
}
__syncwarp();
#pragma unroll
#pragma unroll
for (int row = 0; row < 8; row++) {
pack_t pack = *reinterpret_cast<pack_t *>(&mat[row][laneId * PACK_SIZE]);
......@@ -542,22 +543,24 @@ public:
bool pred = i * INSN_M + row < maxRows && laneId * PACK_SIZE < maxCols;
// if (pred) {
store_pred(reinterpret_cast<pack_t *>(&output[(i * INSN_M + row) * stride + laneId * PACK_SIZE]), pack, pred);
store_pred(reinterpret_cast<pack_t *>(&output[(i * INSN_M + row) * stride + laneId * PACK_SIZE]),
pack,
pred);
// }
}
__syncwarp();
#pragma unroll
#pragma unroll
for (int j = 0; j < WARP_N_TILES; j++) {
packed_fpsum_t &fsum = fpsum[i * WARP_N_TILES + j];
int row = laneId / 4;
int col = laneId % 4 * 2 + j * INSN_N;
packed_fpsum_t &fsum = fpsum[i * WARP_N_TILES + j];
int row = laneId / 4;
int col = laneId % 4 * 2 + j * INSN_N;
*reinterpret_cast<half2_t *>(&mat[row][col + 0]) = fsum.data[1];
*reinterpret_cast<half2_t *>(&mat[row][col + 8]) = fsum.data[3];
}
__syncwarp();
#pragma unroll
#pragma unroll
for (int row = 0; row < 8; row++) {
pack_t pack = *reinterpret_cast<pack_t *>(&mat[row][laneId * PACK_SIZE]);
......@@ -569,7 +572,10 @@ public:
bool pred = i * INSN_M + 8 + row < maxRows && laneId * PACK_SIZE < maxCols;
// if (pred) {
store_pred(reinterpret_cast<pack_t *>(&output[(i * INSN_M + 8 + row) * stride + laneId * PACK_SIZE]), pack, pred);
store_pred(
reinterpret_cast<pack_t *>(&output[(i * INSN_M + 8 + row) * stride + laneId * PACK_SIZE]),
pack,
pred);
// }
}
__syncwarp();
......@@ -584,22 +590,22 @@ public:
// [WARP_M, WARP_N * 2] when fuse_glu
template<bool fuse_glu>
struct load_act_to_fpsum {
using matrix_t = half_t[INSN_M][WARP_N + 8];
using matrix_t = half_t[INSN_M][WARP_N + 8];
static constexpr size_t SHMEM_SIZE = sizeof(matrix_t);
__device__ __forceinline__
void operator()(const half_t *input, int stride, int maxRows, int maxCols, fpsum_warp &out, void *shmem) {
__device__ __forceinline__ void
operator()(const half_t *input, int stride, int maxRows, int maxCols, fpsum_warp &out, void *shmem) {
const int laneId = threadIdx.x % WARP_SIZE;
matrix_t &mat = *reinterpret_cast<matrix_t *>(shmem);
constexpr int PACK_SIZE = WARP_N / WARP_SIZE;
using packed_input = std::array<half_t, PACK_SIZE>;
using packed_raw_input = std::array<half2_t, PACK_SIZE>;
using packed_input = std::array<half_t, PACK_SIZE>;
using packed_raw_input = std::array<half2_t, PACK_SIZE>;
#pragma unroll
#pragma unroll
for (int m = 0; m < WARP_M_TILES; m++) {
#pragma unroll
#pragma unroll
for (int row = 0; row < INSN_M; row++) {
packed_input pack;
// TODO: numCols not multiples of PACK_SIZE
......@@ -608,9 +614,10 @@ public:
raw.fill(half2_t(0, 0));
bool pred = (m * INSN_M + row) < maxRows && laneId * PACK_SIZE * 2 < maxCols;
if (pred) {
raw = load(reinterpret_cast<const packed_raw_input *>(input + (m * INSN_M + row) * stride + laneId * PACK_SIZE * 2));
raw = load(reinterpret_cast<const packed_raw_input *>(input + (m * INSN_M + row) * stride +
laneId * PACK_SIZE * 2));
}
#pragma unroll
#pragma unroll
for (int j = 0; j < PACK_SIZE; j++) {
pack[j] = raw[j].x * silu(raw[j].y);
}
......@@ -618,7 +625,8 @@ public:
pack.fill(half_t(0));
bool pred = (m * INSN_M + row) < maxRows && laneId * PACK_SIZE < maxCols;
if (pred) {
pack = load(reinterpret_cast<const packed_input *>(input + (m * INSN_M + row) * stride + laneId * PACK_SIZE));
pack = load(reinterpret_cast<const packed_input *>(input + (m * INSN_M + row) * stride +
laneId * PACK_SIZE));
}
}
store<true>(reinterpret_cast<packed_input *>(&mat[row][laneId * PACK_SIZE]), pack);
......@@ -637,27 +645,24 @@ public:
}
};
template<typename F>
__device__ __forceinline__
static fpsum_warp apply_act(fpsum_warp fpsum, F func) {
__device__ __forceinline__ static fpsum_warp apply_act(fpsum_warp fpsum, F func) {
fpsum_warp result;
#pragma unroll
#pragma unroll
for (int i = 0; i < WARP_M_TILES; i++) {
#pragma unroll
#pragma unroll
for (int j = 0; j < WARP_N_TILES; j++) {
#pragma unroll
#pragma unroll
for (int k = 0; k < 4; k++) {
half2_t &dst = result[i * WARP_N_TILES + j].data[k];
half2_t src = fpsum[i * WARP_N_TILES + j].data[k];
dst.x = func(src.x);
dst.y = func(src.y);
half2_t src = fpsum[i * WARP_N_TILES + j].data[k];
dst.x = func(src.x);
dst.y = func(src.y);
}
}
}
return result;
}
struct EpilogueDefault {
struct Arguments {
......@@ -665,57 +670,57 @@ public:
int actualM, actualN;
};
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
const int warpId = threadIdx.x / WARP_SIZE;
__shared__ alignas(128) uint8_t shmem[NUM_WARPS][ceilDiv(unpack_fpsum::SHMEM_SIZE, 128) * 128];
const int m_offset = binfo.bm * BLOCK_M + warpId * WARP_M;
const int n_offset = binfo.bn * BLOCK_N;
unpack_fpsum()(
fpsum,
args.out + m_offset * args.actualN + n_offset,
args.actualN,
args.actualM - m_offset,
args.actualN - n_offset,
shmem[warpId],
[&](int rowId, unpack_fpsum::pack_t &pack) ALWAYSINLINE {
if constexpr (std::is_same_v<half_t, half>) {
#pragma unroll
for (int i = 0; i < pack.size(); i++) {
pack[i] = __hmin(pack[i], (half)65504);
pack[i] = __hmax(pack[i], (half)-65504);
}
}
}
);
unpack_fpsum()(fpsum,
args.out + m_offset * args.actualN + n_offset,
args.actualN,
args.actualM - m_offset,
args.actualN - n_offset,
shmem[warpId],
[&](int rowId, unpack_fpsum::pack_t &pack) ALWAYSINLINE {
if constexpr (std::is_same_v<half_t, half>) {
#pragma unroll
for (int i = 0; i < pack.size(); i++) {
pack[i] = __hmin(pack[i], (half)65504);
pack[i] = __hmax(pack[i], (half)-65504);
}
}
});
}
};
struct EpilogueNop {
// workaround for layout mismatch between host and device code
struct Arguments { size_t unused; };
struct Arguments {
size_t unused;
};
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
}
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {}
};
template<bool USE_BIAS = true, bool USE_SCALE = false>
struct EpilogueBias {
struct Arguments {
const packed_wscale_t *bias; // [N / BLOCK_N, WSCALES_NUM_PACKS, WSCALES_VALID_LANES] of packed_wscale_t
const packed_wscale_t *bias; // [N / BLOCK_N, WSCALES_NUM_PACKS, WSCALES_VALID_LANES] of packed_wscale_t
const packed_wscale_t *scale;
};
__device__ __forceinline__
void apply_bias(fpsum_warp &fpsum, int M, int N, int K, const packed_wscale_t *bias, const packed_wscale_t *scale) {
__device__ __forceinline__ void
apply_bias(fpsum_warp &fpsum, int M, int N, int K, const packed_wscale_t *bias, const packed_wscale_t *scale) {
const int laneId = threadIdx.x % WARP_SIZE;
// if (laneId == 0) {
// printf("block.x=%d block.y=%d warpId=%d bias=%p\n", blockIdx.x, blockIdx.y, threadIdx.x / WARP_SIZE, bias);
// printf("block.x=%d block.y=%d warpId=%d bias=%p\n", blockIdx.x, blockIdx.y, threadIdx.x / WARP_SIZE,
// bias);
// }
wscale_warp b, s;
......@@ -737,7 +742,6 @@ public:
s1 = broadcast_wscale(s, j * 4, laneId);
s2 = broadcast_wscale(s, j * 4 + 2, laneId);
}
for (int i = 0; i < WARP_M_TILES; i++) {
auto &fsum = fpsum[i * WARP_N_TILES + j];
......@@ -762,114 +766,114 @@ public:
}
}
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
const int bn = binfo.bn;
if constexpr (USE_BIAS || USE_SCALE) {
apply_bias(
fpsum, M, N, K,
args.bias + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES,
args.scale + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES
);
apply_bias(fpsum,
M,
N,
K,
args.bias + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES,
args.scale + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES);
}
}
};
struct EpilogueSilu {
struct Arguments { size_t unused; };
struct Arguments {
size_t unused;
};
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
fpsum = apply_act(fpsum, [](half_t x) { return silu(x); });
}
};
template<typename ...Epilogues>
template<typename... Epilogues>
struct EpilogueCombination {
using Arguments = std::tuple<typename Epilogues::Arguments...>;
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp &fpsum, int M, int N, int K, const Arguments &args) {
// this function makes intellisense crashes :(
#if __INTELLISENSE__
__trap(); // should not happen when actually compiling
#else
#if __INTELLISENSE__
__trap(); // should not happen when actually compiling
#else
std::tuple<Epilogues...> epilogues;
auto run = [&]<size_t idx>() {
std::get<idx>(epilogues).operator()(binfo, fpsum, M, N, K, std::get<idx>(args));
};
auto foreach = [&]<size_t ...Is>(std::index_sequence<Is...>) {
(run.template operator()<Is>(), ...);
};
foreach(std::make_index_sequence<sizeof...(Epilogues)>());
#endif
auto foreach = [&]<size_t... Is>(std::index_sequence<Is...>) { (run.template operator()<Is>(), ...); };
foreach (std::make_index_sequence<sizeof...(Epilogues)>())
;
#endif
}
};
};
#define IMPORT_GEMM_BASE(config) \
using Base = GEMMBase<config>; \
using Base::BLOCK_M; \
using Base::BLOCK_N; \
using Base::WARP_SIZE; \
using Base::NUM_WARPS; \
using Base::INSN_M; \
using Base::INSN_N; \
using Base::INSN_K; \
using typename Base::half_t; \
using typename Base::half2_t; \
using Base::WARP_M; \
using Base::WARP_N; \
using Base::WARP_K; \
using Base::WARP_M_TILES; \
using Base::WARP_N_TILES; \
using Base::WARP_K_TILES; \
using Base::WSCALES_PACK_SIZE; \
using Base::WSCALES_NUM_PACKS; \
using Base::WSCALES_VALID_LANES; \
using Base::ASCALES_PACK_SIZE; \
using Base::ASCALES_NUM_PACKS; \
using Base::ASCALES_VALID_LANES; \
using typename Base::packed_act_t; \
using typename Base::packed_wgt_t; \
using typename Base::packed_psum_t; \
using typename Base::packed_fpsum_t; \
using typename Base::packed_gated_fpsum_t; \
using typename Base::packed_f32psum_t; \
using typename Base::packed_wscale_t; \
using typename Base::packed_ascale_t; \
using typename Base::act_warp; \
using typename Base::wgt_warp; \
using typename Base::ascale_warp; \
using typename Base::wscale_warp; \
using typename Base::fpsum_warp; \
using typename Base::f32psum_warp; \
using typename Base::gated_fpsum_warp; \
using typename Base::BlockInfo; \
using typename Base::unpack_fpsum; \
using typename Base::EpilogueDefault; \
using typename Base::EpilogueNop; \
template<bool USE_BIAS, bool USE_SCALE> \
using EpilogueBias = typename Base::EpilogueBias<USE_BIAS, USE_SCALE>; \
using Base::mma_f16xf16_f32; \
using Base::packed_fp32_to_fp16; \
using Base::packed_fp16_to_fp32; \
using Base::load_act; \
using Base::load_wgt; \
using Base::load_ascale; \
using Base::load_wscale; \
using Base::broadcast_wscale; \
using Base::broadcast_ascale; \
using Base::apply_scales; \
using Base::pack_ascales; \
using Base::pack_wscales; \
#define IMPORT_GEMM_BASE(config) \
using Base = GEMMBase<config>; \
using Base::BLOCK_M; \
using Base::BLOCK_N; \
using Base::WARP_SIZE; \
using Base::NUM_WARPS; \
using Base::INSN_M; \
using Base::INSN_N; \
using Base::INSN_K; \
using typename Base::half_t; \
using typename Base::half2_t; \
using Base::WARP_M; \
using Base::WARP_N; \
using Base::WARP_K; \
using Base::WARP_M_TILES; \
using Base::WARP_N_TILES; \
using Base::WARP_K_TILES; \
using Base::WSCALES_PACK_SIZE; \
using Base::WSCALES_NUM_PACKS; \
using Base::WSCALES_VALID_LANES; \
using Base::ASCALES_PACK_SIZE; \
using Base::ASCALES_NUM_PACKS; \
using Base::ASCALES_VALID_LANES; \
using typename Base::packed_act_t; \
using typename Base::packed_wgt_t; \
using typename Base::packed_psum_t; \
using typename Base::packed_fpsum_t; \
using typename Base::packed_gated_fpsum_t; \
using typename Base::packed_f32psum_t; \
using typename Base::packed_wscale_t; \
using typename Base::packed_ascale_t; \
using typename Base::act_warp; \
using typename Base::wgt_warp; \
using typename Base::ascale_warp; \
using typename Base::wscale_warp; \
using typename Base::fpsum_warp; \
using typename Base::f32psum_warp; \
using typename Base::gated_fpsum_warp; \
using typename Base::BlockInfo; \
using typename Base::unpack_fpsum; \
using typename Base::EpilogueDefault; \
using typename Base::EpilogueNop; \
template<bool USE_BIAS, bool USE_SCALE> \
using EpilogueBias = typename Base::EpilogueBias<USE_BIAS, USE_SCALE>; \
using Base::mma_f16xf16_f32; \
using Base::packed_fp32_to_fp16; \
using Base::packed_fp16_to_fp32; \
using Base::load_act; \
using Base::load_wgt; \
using Base::load_ascale; \
using Base::load_wscale; \
using Base::broadcast_wscale; \
using Base::broadcast_ascale; \
using Base::apply_scales; \
using Base::pack_ascales; \
using Base::pack_wscales; \
using Base::apply_act;
template<typename kernel>
constexpr int min_arch() {
if constexpr (requires {kernel::MIN_ARCH;}) {
if constexpr (requires { kernel::MIN_ARCH; }) {
return kernel::MIN_ARCH;
} else {
return 0;
......@@ -877,16 +881,15 @@ constexpr int min_arch() {
}
template<typename kernel>
constexpr int max_arch() {
if constexpr (requires {kernel::MAX_ARCH;}) {
if constexpr (requires { kernel::MAX_ARCH; }) {
return kernel::MAX_ARCH;
} else {
return INT_MAX;
}
}
template<typename kernel, typename ...T>
__global__
static void invoke_kernel(T ...args) {
template<typename kernel, typename... T>
__global__ static void invoke_kernel(T... args) {
#ifdef __CUDA_ARCH__
if constexpr (__CUDA_ARCH__ >= min_arch<kernel>() && __CUDA_ARCH__ <= max_arch<kernel>()) {
kernel()(args...);
......@@ -900,8 +903,7 @@ static void invoke_kernel(T ...args) {
}
template<typename T>
__global__
static void test_sizeof_device() {
__global__ static void test_sizeof_device() {
printf("sizeof on device = %d\n", (int)sizeof(T));
}
......@@ -918,4 +920,4 @@ static void test_sizeof() {
checkCUDA(cudaDeviceSynchronize());
}
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/zgemm/gemm_utils.cuh
View file @ 37a27712
......@@ -7,7 +7,7 @@
namespace nunchaku::kernels {
static constexpr int clamp(int val, int min, int max) {
if (val < min)
if (val < min)
return min;
if (val > max)
return max;
......@@ -15,17 +15,20 @@ static constexpr int clamp(int val, int min, int max) {
}
template<bool shmem = false, typename T>
__device__ __forceinline__
static T load(const T *addr) {
__device__ __forceinline__ static T load(const T *addr) {
if constexpr (shmem) {
if constexpr (sizeof(T) == 8) {
uint2 data;
asm volatile ("ld.shared.v2.b32 {%0, %1}, [%2];" : "=r"(data.x), "=r"(data.y) : "l"(__cvta_generic_to_shared(addr)));
asm volatile("ld.shared.v2.b32 {%0, %1}, [%2];"
: "=r"(data.x), "=r"(data.y)
: "l"(__cvta_generic_to_shared(addr)));
return *reinterpret_cast<T *>(&data);
}
if constexpr (sizeof(T) == 16) {
uint4 data;
asm volatile ("ld.shared.v4.b32 {%0, %1, %2, %3}, [%4];" : "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w) : "l"(__cvta_generic_to_shared(addr)));
asm volatile("ld.shared.v4.b32 {%0, %1, %2, %3}, [%4];"
: "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w)
: "l"(__cvta_generic_to_shared(addr)));
return *reinterpret_cast<T *>(&data);
}
return *addr;
......@@ -44,30 +47,32 @@ static T load(const T *addr) {
}
template<typename T>
__device__ __forceinline__
static T load_pred(const T *addr, bool pred) {
__device__ __forceinline__ static T load_pred(const T *addr, bool pred) {
if constexpr (sizeof(T) == 4) {
uint32_t data;
asm volatile (
"{ .reg .pred loadpred; setp.ne.b32 loadpred, %2, 0;"
"@loadpred ld.global.nc.b32 %0, [%1];"
"}" : "=r"(data) : "l"(addr), "r"((int)pred));
asm volatile("{ .reg .pred loadpred; setp.ne.b32 loadpred, %2, 0;"
"@loadpred ld.global.nc.b32 %0, [%1];"
"}"
: "=r"(data)
: "l"(addr), "r"((int)pred));
return *reinterpret_cast<T *>(&data);
}
if constexpr (sizeof(T) == 8) {
uint2 data;
asm volatile (
"{ .reg .pred loadpred; setp.ne.b32 loadpred, %3, 0;"
"@loadpred ld.global.nc.v2.b32 {%0, %1}, [%2];"
"}" : "=r"(data.x), "=r"(data.y) : "l"(addr), "r"((int)pred));
asm volatile("{ .reg .pred loadpred; setp.ne.b32 loadpred, %3, 0;"
"@loadpred ld.global.nc.v2.b32 {%0, %1}, [%2];"
"}"
: "=r"(data.x), "=r"(data.y)
: "l"(addr), "r"((int)pred));
return *reinterpret_cast<T *>(&data);
}
if constexpr (sizeof(T) == 16) {
uint4 data;
asm volatile (
"{ .reg .pred loadpred; setp.ne.b32 loadpred, %5, 0;"
"@loadpred ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];"
"}" : "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w) : "l"(addr), "r"((int)pred));
asm volatile("{ .reg .pred loadpred; setp.ne.b32 loadpred, %5, 0;"
"@loadpred ld.global.nc.v4.b32 {%0, %1, %2, %3}, [%4];"
"}"
: "=r"(data.x), "=r"(data.y), "=r"(data.z), "=r"(data.w)
: "l"(addr), "r"((int)pred));
return *reinterpret_cast<T *>(&data);
}
......@@ -79,17 +84,21 @@ static T load_pred(const T *addr, bool pred) {
}
template<bool shmem = false, typename T>
__device__ __forceinline__
static void store(T *addr, T val) {
__device__ __forceinline__ static void store(T *addr, T val) {
if constexpr (shmem) {
if constexpr (sizeof(T) == 8) {
uint2 data = *reinterpret_cast<uint2 *>(&val);
asm volatile ("st.shared.v2.b32 [%0], {%1, %2};" :: "l"(__cvta_generic_to_shared(addr)), "r"(data.x), "r"(data.y));
asm volatile(
"st.shared.v2.b32 [%0], {%1, %2};" ::"l"(__cvta_generic_to_shared(addr)), "r"(data.x), "r"(data.y));
return;
}
if constexpr (sizeof(T) == 16) {
uint4 data = *reinterpret_cast<uint4 *>(&val);
asm volatile ("st.shared.v4.b32 [%0], {%1, %2, %3, %4};" :: "l"(__cvta_generic_to_shared(addr)), "r"(data.x), "r"(data.y), "r"(data.z), "r"(data.w));
asm volatile("st.shared.v4.b32 [%0], {%1, %2, %3, %4};" ::"l"(__cvta_generic_to_shared(addr)),
"r"(data.x),
"r"(data.y),
"r"(data.z),
"r"(data.w));
return;
}
*addr = val;
......@@ -107,35 +116,41 @@ static void store(T *addr, T val) {
if constexpr (sizeof(T) == 16) {
__stcg(reinterpret_cast<uint4 *>(addr), *reinterpret_cast<uint4 *>(&val));
return;
}
}
*addr = val;
}
template<typename T>
__device__ __forceinline__
static void store_pred(T *addr, T val, bool pred) {
__device__ __forceinline__ static void store_pred(T *addr, T val, bool pred) {
if constexpr (sizeof(T) == 4) {
uint32_t data = *reinterpret_cast<uint32_t *>(&val);
asm volatile (
"{ .reg .pred storepred; setp.ne.b32 storepred, %0, 0;"
"@storepred st.global.cg.b32 [%1], %2;"
"}" :: "r"((int)pred), "l"(addr), "r"(data));
asm volatile("{ .reg .pred storepred; setp.ne.b32 storepred, %0, 0;"
"@storepred st.global.cg.b32 [%1], %2;"
"}" ::"r"((int)pred),
"l"(addr),
"r"(data));
return;
}
if constexpr (sizeof(T) == 8) {
uint2 data = *reinterpret_cast<uint2 *>(&val);
asm volatile (
"{ .reg .pred storepred; setp.ne.b32 storepred, %0, 0;"
"@storepred st.global.cg.v2.b32 [%1], {%2, %3};"
"}" :: "r"((int)pred), "l"(addr), "r"(data.x), "r"(data.y));
asm volatile("{ .reg .pred storepred; setp.ne.b32 storepred, %0, 0;"
"@storepred st.global.cg.v2.b32 [%1], {%2, %3};"
"}" ::"r"((int)pred),
"l"(addr),
"r"(data.x),
"r"(data.y));
return;
}
if constexpr (sizeof(T) == 16) {
uint4 data = *reinterpret_cast<uint4 *>(&val);
asm volatile (
"{ .reg .pred storepred; setp.ne.b32 storepred, %0, 0;"
"@storepred st.global.cg.v4.b32 [%1], {%2, %3, %4, %5};"
"}" :: "r"((int)pred), "l"(addr), "r"(data.x), "r"(data.y), "r"(data.z), "r"(data.w));
asm volatile("{ .reg .pred storepred; setp.ne.b32 storepred, %0, 0;"
"@storepred st.global.cg.v4.b32 [%1], {%2, %3, %4, %5};"
"}" ::"r"((int)pred),
"l"(addr),
"r"(data.x),
"r"(data.y),
"r"(data.z),
"r"(data.w));
return;
}
......@@ -144,198 +159,174 @@ static void store_pred(T *addr, T val, bool pred) {
}
}
__device__ __forceinline__
static float2 half22float2(half2 val) {
__device__ __forceinline__ static float2 half22float2(half2 val) {
return __half22float2(val);
}
__device__ __forceinline__
static float2 half22float2(__nv_bfloat162 val) {
__device__ __forceinline__ static float2 half22float2(__nv_bfloat162 val) {
return __bfloat1622float2(val);
}
template<typename T>
__device__ __forceinline__
static T float22half2(float2 val) = delete;
__device__ __forceinline__ static T float22half2(float2 val) = delete;
template<>
__device__ __forceinline__
half2 float22half2<half2>(float2 val) {
__device__ __forceinline__ half2 float22half2<half2>(float2 val) {
return __float22half2_rn(val);
}
template<>
__device__ __forceinline__
__nv_bfloat162 float22half2<__nv_bfloat162>(float2 val) {
__device__ __forceinline__ __nv_bfloat162 float22half2<__nv_bfloat162>(float2 val) {
return __float22bfloat162_rn(val);
}
template<typename T>
__device__ __forceinline__
static void unused_var(T &val, bool alwaysfalse) {
__device__ __forceinline__ static void unused_var(T &val, bool alwaysfalse) {
volatile T *ptr = nullptr;
if (alwaysfalse) {
*ptr = val;
}
}
__device__ __forceinline__
static void ldmatrix(const void *ptr, uint4 &out) {
asm volatile(
"ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
: "=r"(out.x), "=r"(out.y), "=r"(out.z), "=r"(out.w)
: "l"(__cvta_generic_to_shared(ptr))
);
__device__ __forceinline__ static void ldmatrix(const void *ptr, uint4 &out) {
asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
: "=r"(out.x), "=r"(out.y), "=r"(out.z), "=r"(out.w)
: "l"(__cvta_generic_to_shared(ptr)));
}
template<typename T>
__device__ __forceinline__
static T movmatrix(T x) {
asm volatile ("movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;" : "=r"(*reinterpret_cast<uint32_t *>(&x)) : "r"(*reinterpret_cast<uint32_t *>(&x)));
__device__ __forceinline__ static T movmatrix(T x) {
asm volatile("movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;"
: "=r"(*reinterpret_cast<uint32_t *>(&x))
: "r"(*reinterpret_cast<uint32_t *>(&x)));
return x;
}
// x in low bit, y in high bit
template<int bitwidth, bool use_unsigned>
__device__ __forceinline__
uint32_t quantize_float2(float2 value) = delete;
__device__ __forceinline__ uint32_t quantize_float2(float2 value) = delete;
template<>
__device__ __forceinline__
uint32_t quantize_float2<4, false>(float2 value) {
__device__ __forceinline__ uint32_t quantize_float2<4, false>(float2 value) {
int v1, v2;
uint32_t result;
asm volatile ("cvt.rni.s32.f32 %0, %1;" : "=r"(v1) : "f"(value.x));
asm volatile ("cvt.rni.s32.f32 %0, %1;" : "=r"(v2) : "f"(value.y));
asm volatile ("cvt.pack.sat.s4.s32.b32 %0, %1, %2, 0;" : "=r"(result) : "r"(v2), "r"(v1));
asm volatile("cvt.rni.s32.f32 %0, %1;" : "=r"(v1) : "f"(value.x));
asm volatile("cvt.rni.s32.f32 %0, %1;" : "=r"(v2) : "f"(value.y));
asm volatile("cvt.pack.sat.s4.s32.b32 %0, %1, %2, 0;" : "=r"(result) : "r"(v2), "r"(v1));
return result;
}
template<>
__device__ __forceinline__
uint32_t quantize_float2<4, true>(float2 value) {
__device__ __forceinline__ uint32_t quantize_float2<4, true>(float2 value) {
int v1, v2;
uint32_t result;
asm volatile ("cvt.rni.s32.f32 %0, %1;" : "=r"(v1) : "f"(value.x));
asm volatile ("cvt.rni.s32.f32 %0, %1;" : "=r"(v2) : "f"(value.y));
asm volatile ("cvt.pack.sat.u4.s32.b32 %0, %1, %2, 0;" : "=r"(result) : "r"(v2), "r"(v1));
asm volatile("cvt.rni.s32.f32 %0, %1;" : "=r"(v1) : "f"(value.x));
asm volatile("cvt.rni.s32.f32 %0, %1;" : "=r"(v2) : "f"(value.y));
asm volatile("cvt.pack.sat.u4.s32.b32 %0, %1, %2, 0;" : "=r"(result) : "r"(v2), "r"(v1));
return result;
}
template<>
__device__ __forceinline__
uint32_t quantize_float2<8, false>(float2 value) {
__device__ __forceinline__ uint32_t quantize_float2<8, false>(float2 value) {
int v1, v2;
uint32_t result;
asm volatile ("cvt.rni.s32.f32 %0, %1;" : "=r"(v1) : "f"(value.x));
asm volatile ("cvt.rni.s32.f32 %0, %1;" : "=r"(v2) : "f"(value.y));
asm volatile ("cvt.pack.sat.s8.s32.b32 %0, %1, %2, 0;" : "=r"(result) : "r"(v2), "r"(v1));
asm volatile("cvt.rni.s32.f32 %0, %1;" : "=r"(v1) : "f"(value.x));
asm volatile("cvt.rni.s32.f32 %0, %1;" : "=r"(v2) : "f"(value.y));
asm volatile("cvt.pack.sat.s8.s32.b32 %0, %1, %2, 0;" : "=r"(result) : "r"(v2), "r"(v1));
return result;
}
__device__ __forceinline__
uint32_t quantize_float2_fp4(float2 value) {
__device__ __forceinline__ uint32_t quantize_float2_fp4(float2 value) {
uint32_t result;
asm volatile ("{ .reg .b8 tmp; cvt.rn.satfinite.e2m1x2.f32 tmp, %1, %2; cvt.u32.u8 %0, tmp; }" : "=r"(result) : "f"(value.y), "f"(value.x));
asm volatile("{ .reg .b8 tmp; cvt.rn.satfinite.e2m1x2.f32 tmp, %1, %2; cvt.u32.u8 %0, tmp; }"
: "=r"(result)
: "f"(value.y), "f"(value.x));
return result;
}
__device__ __forceinline__
uint32_t quantize_float4_fp8(float4 value) {
__device__ __forceinline__ uint32_t quantize_float4_fp8(float4 value) {
uint16_t lo, hi;
asm volatile ("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(lo) : "f"(value.y), "f"(value.x));
asm volatile ("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(hi) : "f"(value.w), "f"(value.z));
asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(lo) : "f"(value.y), "f"(value.x));
asm volatile("cvt.rn.satfinite.e4m3x2.f32 %0, %1, %2;" : "=h"(hi) : "f"(value.w), "f"(value.z));
return uint32_t(lo) | (uint32_t(hi) << 16);
}
__device__ __forceinline__
static float cuda_tanhf(float x) {
__device__ __forceinline__ static float cuda_tanhf(float x) {
float result;
asm ("tanh.approx.f32 %0, %1;" : "=f"(result) : "f"(x));
asm("tanh.approx.f32 %0, %1;" : "=f"(result) : "f"(x));
return result;
}
__device__ __forceinline__
static float cuda_frcp(float x) {
__device__ __forceinline__ static float cuda_frcp(float x) {
float result;
asm ("rcp.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
asm("rcp.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
return result;
}
__device__ __forceinline__
static float cuda_frsqrt(float x) {
__device__ __forceinline__ static float cuda_frsqrt(float x) {
float result;
asm ("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
asm("rsqrt.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
return result;
}
__device__ __forceinline__
static float cuda_sin(float x) {
__device__ __forceinline__ static float cuda_sin(float x) {
float result;
asm ("sin.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
asm("sin.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
return result;
}
__device__ __forceinline__
static float cuda_cos(float x) {
__device__ __forceinline__ static float cuda_cos(float x) {
float result;
asm ("cos.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
asm("cos.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
return result;
}
__device__ __forceinline__
static float cuda_exp2(float x) {
__device__ __forceinline__ static float cuda_exp2(float x) {
float result;
asm ("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(result) : "f"(x));
return result;
}
// https://forums.developer.nvidia.com/t/hardware-accelerated-computation-of-the-sigmoid-logistic-function/266206
__forceinline__ __device__
static float cuda_sigmoidf (float a)
{
__forceinline__ __device__ static float cuda_sigmoidf(float a) {
#if USE_TANH
return fmaf (0.5, __tanhf (0.5f * a), 0.5f);
#else // USE_TANH
return fmaf(0.5, __tanhf(0.5f * a), 0.5f);
#else // USE_TANH
const float L2E = 1.442695041f; // log2(exp(1))
float t, d, e, r;
t = -L2E * a;
asm ("ex2.approx.ftz.f32 %0,%1;\n\t" : "=f"(e) : "f"(t));
asm("ex2.approx.ftz.f32 %0,%1;\n\t" : "=f"(e) : "f"(t));
d = e + 1.0f;
asm ("rcp.approx.ftz.f32 %0,%1;\n\t" : "=f"(r) : "f"(d));
asm("rcp.approx.ftz.f32 %0,%1;\n\t" : "=f"(r) : "f"(d));
return r;
#endif // USE_TANH
}
template<typename T>
__device__ __forceinline__
static T gelu_half2(T x) {
__device__ __forceinline__ static T gelu_half2(T x) {
float2 xf = half22float2(x);
float2 x3f = xf * xf * xf;
float t1 = 0.5f + 0.5f * cuda_tanhf(0.79788456f * (xf.x + (0.044715f * x3f.x)));
float t2 = 0.5f + 0.5f * cuda_tanhf(0.79788456f * (xf.y + (0.044715f * x3f.y)));
float t1 = 0.5f + 0.5f * cuda_tanhf(0.79788456f * (xf.x + (0.044715f * x3f.x)));
float t2 = 0.5f + 0.5f * cuda_tanhf(0.79788456f * (xf.y + (0.044715f * x3f.y)));
return float22half2<T>(xf * make_float2(t1, t2));
}
template<typename T>
__device__ __forceinline__
static T gelu_half(T x) {
__device__ __forceinline__ static T gelu_half(T x) {
float xf = float(x);
float x3f = xf * xf * xf;
float t = 0.5f + 0.5f * cuda_tanhf(0.79788456f * (xf + (0.044715f * x3f)));
float t = 0.5f + 0.5f * cuda_tanhf(0.79788456f * (xf + (0.044715f * x3f)));
return (T)(xf * t);
}
template <typename T>
__device__ __forceinline__
static T silu(const T &x) {
// x * sigmoid(x)
return (T)((float)x * cuda_sigmoidf((float)x));
// return (T)__fdividef((float)x, 1.0f + __expf((float)-x));
template<typename T>
__device__ __forceinline__ static T silu(const T &x) {
// x * sigmoid(x)
return (T)((float)x * cuda_sigmoidf((float)x));
// return (T)__fdividef((float)x, 1.0f + __expf((float)-x));
}
__device__ __forceinline__
static half2 h2div(half2 a, half2 b) {
__device__ __forceinline__ static half2 h2div(half2 a, half2 b) {
float2 af = half22float2(a);
float2 bf = half22float2(b);
float2 of;
......@@ -343,8 +334,7 @@ static half2 h2div(half2 a, half2 b) {
of.y = __fdividef(af.y, bf.y);
return float22half2<half2>(of);
};
__device__ __forceinline__
static __nv_bfloat162 h2div(__nv_bfloat162 a, __nv_bfloat162 b) {
__device__ __forceinline__ static __nv_bfloat162 h2div(__nv_bfloat162 a, __nv_bfloat162 b) {
float2 af = half22float2(a);
float2 bf = half22float2(b);
float2 of;
......@@ -353,41 +343,37 @@ static __nv_bfloat162 h2div(__nv_bfloat162 a, __nv_bfloat162 b) {
return float22half2<__nv_bfloat162>(of);
};
__device__ __forceinline__
static void reduce_add(float *addr, float val) {
asm volatile ("red.relaxed.gpu.global.add.f32 [%0], %1;" :: "l"(addr), "f"(val));
__device__ __forceinline__ static void reduce_add(float *addr, float val) {
asm volatile("red.relaxed.gpu.global.add.f32 [%0], %1;" ::"l"(addr), "f"(val));
}
__device__ __forceinline__
static void reduce_add_pred(float *addr, float val, bool pred) {
asm volatile (
"{ .reg .pred storepred; setp.ne.b32 storepred, %0, 0;"
"@storepred red.relaxed.gpu.global.add.f32 [%1], %2;"
"}" :: "r"((int)pred), "l"(addr), "f"(val));
__device__ __forceinline__ static void reduce_add_pred(float *addr, float val, bool pred) {
asm volatile("{ .reg .pred storepred; setp.ne.b32 storepred, %0, 0;"
"@storepred red.relaxed.gpu.global.add.f32 [%1], %2;"
"}" ::"r"((int)pred),
"l"(addr),
"f"(val));
}
template<int cnt, typename F>
__device__ __forceinline__
static void unrolled_loop(F &&lambda) {
auto call = [&]<int ...Is>(std::integer_sequence<int, Is...>) {
(lambda.template operator()<Is>(), ...);
};
__device__ __forceinline__ static void unrolled_loop(F &&lambda) {
auto call = [&]<int... Is>(std::integer_sequence<int, Is...>) { (lambda.template operator()<Is>(), ...); };
call(std::make_integer_sequence<int, cnt>());
}
// int2float is slow on sm_80 and before
// val in [-4194304, 4194303]
__device__ __forceinline__
static float int2float_fast(int val) {
__device__ __forceinline__ static float int2float_fast(int val) {
float fval;
// fval = (val & 0x7FFFFF) ^ 0x4B400000
asm volatile ("lop3.b32 %0, %1, %2, %3, %4;" : "=f"(fval) : "r"(val), "n"(0x7FFFFF), "n"(0x4B400000), "n"((0xF0 & 0xCC) ^ 0xAA));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;"
: "=f"(fval)
: "r"(val), "n"(0x7FFFFF), "n"(0x4B400000), "n"((0xF0 & 0xCC) ^ 0xAA));
return fval - 12582912.0f;
}
template<typename To, typename From>
__device__ __forceinline__
static To bit_cast(const From &input) {
__device__ __forceinline__ static To bit_cast(const From &input) {
static_assert(sizeof(To) == sizeof(From));
// not safe but anyway
return *reinterpret_cast<const To *>(&input);
......@@ -395,20 +381,20 @@ static To bit_cast(const From &input) {
// both int2float and float2half are slow on sm_75 and before
// val in [-8192, 8191], steps of 16, round to negative inf
__device__ __forceinline__
static half2 int2half2_fast_8192(int x, int y) {
__device__ __forceinline__ static half2 int2half2_fast_8192(int x, int y) {
uint32_t ival;
uint32_t hval;
// ival.lo = x.lo; ival.hi = y.lo;
asm volatile ("prmt.b32 %0, %1, %2, %3;" : "=r"(ival) : "r"(x), "r"(y), "n"(0x5410));
asm volatile("prmt.b32 %0, %1, %2, %3;" : "=r"(ival) : "r"(x), "r"(y), "n"(0x5410));
ival = ival >> 4;
// (val & 0x03FF03FF) ^ 0x76007600
asm volatile ("lop3.b32 %0, %1, %2, %3, %4;" : "=r"(hval) : "r"(ival), "n"(0x03FF03FF), "n"(0x76007600), "n"((0xF0 & 0xCC) ^ 0xAA));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;"
: "=r"(hval)
: "r"(ival), "n"(0x03FF03FF), "n"(0x76007600), "n"((0xF0 & 0xCC) ^ 0xAA));
return __hadd2(kernels::bit_cast<half2>(hval), half2(-24576.0f, -24576.0f));
}
// val in [-4096, 4095], steps of 8, round to nearest
__device__ __forceinline__
static half2 int2half2_fast_4096_rn(int x, int y) {
__device__ __forceinline__ static half2 int2half2_fast_4096_rn(int x, int y) {
// x = max(min(x, 4095), -4096);
// y = max(min(y, 4095), -4096);
// TODO: round to even?
......@@ -416,24 +402,27 @@ static half2 int2half2_fast_4096_rn(int x, int y) {
y = y * 8192 + 32768;
uint32_t ival;
uint32_t hval;
// ival.lo = x.hi; ival.hi = y.hi;
// ival.lo = x.hi; ival.hi = y.hi;
// <=> divide x and y by 65536 and pack them
asm volatile ("prmt.b32 %0, %1, %2, %3;" : "=r"(ival) : "r"(x), "r"(y), "n"(0x7632));
asm volatile("prmt.b32 %0, %1, %2, %3;" : "=r"(ival) : "r"(x), "r"(y), "n"(0x7632));
// (val & 0x03FF03FF) ^ 0x72007200
asm volatile ("lop3.b32 %0, %1, %2, %3, %4;" : "=r"(hval) : "r"(ival), "n"(0x03FF03FF), "n"(0x72007200), "n"((0xF0 & 0xCC) ^ 0xAA));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;"
: "=r"(hval)
: "r"(ival), "n"(0x03FF03FF), "n"(0x72007200), "n"((0xF0 & 0xCC) ^ 0xAA));
return __hadd2(kernels::bit_cast<half2>(hval), half2(-12288.0f, -12288.0f));
}
// val in [-512, 511]
__device__ __forceinline__
static half2 int2half2_fast_512(int x, int y) {
__device__ __forceinline__ static half2 int2half2_fast_512(int x, int y) {
uint32_t ival;
uint32_t hval;
// ival.lo = x.lo; ival.hi = y.lo;
// ival.lo = x.lo; ival.hi = y.lo;
// <=> divide x and y by 65536 and pack them
asm volatile ("prmt.b32 %0, %1, %2, %3;" : "=r"(ival) : "r"(x), "r"(y), "n"(0x5410));
asm volatile("prmt.b32 %0, %1, %2, %3;" : "=r"(ival) : "r"(x), "r"(y), "n"(0x5410));
// (val & 0x03FF03FF) ^ 0x66006600
asm volatile ("lop3.b32 %0, %1, %2, %3, %4;" : "=r"(hval) : "r"(ival), "n"(0x03FF03FF), "n"(0x66006600), "n"((0xF0 & 0xCC) ^ 0xAA));
asm volatile("lop3.b32 %0, %1, %2, %3, %4;"
: "=r"(hval)
: "r"(ival), "n"(0x03FF03FF), "n"(0x66006600), "n"((0xF0 & 0xCC) ^ 0xAA));
return __hadd2(kernels::bit_cast<half2>(hval), half2(-1536.0f, -1536.0f));
}
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/zgemm/gemm_w4a4.cu
View file @ 37a27712
......@@ -14,7 +14,6 @@ struct FasterI2FMode {
static bool check(bool act_unsigned);
};
template<typename F>
static void invoke_launch(Tensor::ScalarType dtype, bool use_fp4, bool fasterI2F, F &&launch) {
if (fasterI2F && dtype == Tensor::FP16) {
......@@ -32,37 +31,35 @@ static void invoke_launch(Tensor::ScalarType dtype, bool use_fp4, bool fasterI2F
}
}
void gemm_w4a4(
Tensor act, // packed act [M, K / 2]
Tensor wgt, // packed act [N, K / 2]
Tensor out, // linear [M, N]
Tensor qout, // packed act [M, N / 2]
Tensor ascales, // packed as [K / 64, M]
Tensor wscales, // packed ws [K / 64, N]
Tensor oscales, // packed as [N / 64, M]
Tensor poolout, // linear [M / PoolSize, N]
Tensor lora_act_in, // packed lora_act [M, R]
Tensor lora_up, // packed lora_wgt [N, R]
Tensor lora_down, // packed lora_wgt [N, R]
Tensor lora_act_out, // packed lora_act [M, R]
Tensor norm_q, // linear [HEAD_DIM]
Tensor norm_k, // linear [HEAD_DIM]
Tensor rotary_emb, // linear [M, HEAD_DIM / 2, 2, 2]
Tensor bias, // packed ws [N]
Tensor smooth_factor, // packed ws [N], for quantization of the next layer
Tensor out_vk, // linear [B, num_heads, head_dim + 1, head_dim]
Tensor out_linearattn,// linear [B, (M), N / 3]
bool act_unsigned,
std::vector<float> lora_scales, // [R / 16]
bool fuse_silu,
bool fp4,
float alpha,
Tensor wcscales,
Tensor out_q, // packed attention [B, H, M, D]
Tensor out_k, // packed attention [B, H, M, D]
Tensor out_v, // packed attention [B, H, M, D]
int attn_tokens
) {
void gemm_w4a4(Tensor act, // packed act [M, K / 2]
Tensor wgt, // packed act [N, K / 2]
Tensor out, // linear [M, N]
Tensor qout, // packed act [M, N / 2]
Tensor ascales, // packed as [K / 64, M]
Tensor wscales, // packed ws [K / 64, N]
Tensor oscales, // packed as [N / 64, M]
Tensor poolout, // linear [M / PoolSize, N]
Tensor lora_act_in, // packed lora_act [M, R]
Tensor lora_up, // packed lora_wgt [N, R]
Tensor lora_down, // packed lora_wgt [N, R]
Tensor lora_act_out, // packed lora_act [M, R]
Tensor norm_q, // linear [HEAD_DIM]
Tensor norm_k, // linear [HEAD_DIM]
Tensor rotary_emb, // linear [M, HEAD_DIM / 2, 2, 2]
Tensor bias, // packed ws [N]
Tensor smooth_factor, // packed ws [N], for quantization of the next layer
Tensor out_vk, // linear [B, num_heads, head_dim + 1, head_dim]
Tensor out_linearattn, // linear [B, (M), N / 3]
bool act_unsigned,
std::vector<float> lora_scales, // [R / 16]
bool fuse_silu,
bool fp4,
float alpha,
Tensor wcscales,
Tensor out_q, // packed attention [B, H, M, D]
Tensor out_k, // packed attention [B, H, M, D]
Tensor out_v, // packed attention [B, H, M, D]
int attn_tokens) {
Tensor::ScalarType dtype = Tensor::INVALID_SCALAR_TYPE;
if (!fp4) {
dtype = ascales.dtype();
......@@ -75,37 +72,35 @@ void gemm_w4a4(
}
}
invoke_launch(dtype, fp4, FasterI2FMode::check(act_unsigned), [&]<typename Config, bool USE_FP4>() {
GEMM_W4A4_Launch<Config, USE_FP4>::gemm_w4a4(
act,
wgt,
out,
qout,
ascales,
wscales,
oscales,
poolout,
lora_act_in,
lora_up,
lora_down,
lora_act_out,
norm_q,
norm_k,
rotary_emb,
bias,
smooth_factor,
out_vk,
out_linearattn,
act_unsigned,
lora_scales,
fuse_silu,
fp4,
alpha,
wcscales,
out_q,
out_k,
out_v,
attn_tokens
);
GEMM_W4A4_Launch<Config, USE_FP4>::gemm_w4a4(act,
wgt,
out,
qout,
ascales,
wscales,
oscales,
poolout,
lora_act_in,
lora_up,
lora_down,
lora_act_out,
norm_q,
norm_k,
rotary_emb,
bias,
smooth_factor,
out_vk,
out_linearattn,
act_unsigned,
lora_scales,
fuse_silu,
fp4,
alpha,
wcscales,
out_q,
out_k,
out_v,
attn_tokens);
});
}
......@@ -115,26 +110,28 @@ void linearattn_vk_mul_q(Tensor q, Tensor vk) {
});
}
void quantize_w4a4_act_fuse_lora(Tensor input, Tensor output, Tensor oscales, Tensor lora_down, Tensor lora_act_out, Tensor smooth, bool fuse_glu, bool fp4) {
void quantize_w4a4_act_fuse_lora(Tensor input,
Tensor output,
Tensor oscales,
Tensor lora_down,
Tensor lora_act_out,
Tensor smooth,
bool fuse_glu,
bool fp4) {
invoke_launch(input.dtype(), fp4, false, [&]<typename Config, bool USE_FP4>() {
GEMM_W4A4_Launch<Config, USE_FP4>::quantize_w4a4_act_fuse_lora(
input, output, oscales, lora_down, lora_act_out, smooth, fuse_glu, fp4
);
input, output, oscales, lora_down, lora_act_out, smooth, fuse_glu, fp4);
});
}
void quantize_w4a4_act(Tensor input, Tensor output, Tensor oscales) {
invoke_launch(input.dtype(), false, false, [&]<typename Config, bool USE_FP4>() {
GEMM_W4A4_Launch<Config, false>::quantize_w4a4_act(
input, output, oscales
);
GEMM_W4A4_Launch<Config, false>::quantize_w4a4_act(input, output, oscales);
});
}
void quantize_w4a4_wgt(Tensor input, Tensor output, Tensor oscales) {
invoke_launch(input.dtype(), false, false, [&]<typename Config, bool USE_FP4>() {
GEMM_W4A4_Launch<Config, false>::quantize_w4a4_wgt(
input, output, oscales
);
GEMM_W4A4_Launch<Config, false>::quantize_w4a4_wgt(input, output, oscales);
});
}
......@@ -143,7 +140,7 @@ bool FasterI2FMode::check(bool act_unsigned) {
if (prop->major != 7 || prop->minor != 5) {
return false;
}
if (mode == Always) {
return true;
} else if (mode == Enabled && !act_unsigned) {
......@@ -162,4 +159,4 @@ void set_faster_i2f_mode(std::string mode) {
FasterI2FMode::mode = mapping.at(mode);
}
};
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/zgemm/gemm_w4a4.cuh
View file @ 37a27712
......@@ -31,8 +31,7 @@ public:
static constexpr bool FP4_AVAILABLE = false;
#endif
__device__ __forceinline__
static void trap_no_fp4() {
__device__ __forceinline__ static void trap_no_fp4() {
if (blockIdx.x == 0 && blockIdx.y == 0 && threadIdx.x == 0) {
printf("FP4 is not available on this device\n");
}
......@@ -44,12 +43,12 @@ public:
static_assert(WARP_N % 32 == 0);
static_assert(WARP_M % 32 == 0);
static constexpr int WMSCALES_PACK_SIZE = clamp(WARP_N / 32, 1, 4);
static constexpr int WMSCALES_NUM_PACKS = ceilDiv(WARP_N / 32, WMSCALES_PACK_SIZE);
static constexpr int WMSCALES_PACK_SIZE = clamp(WARP_N / 32, 1, 4);
static constexpr int WMSCALES_NUM_PACKS = ceilDiv(WARP_N / 32, WMSCALES_PACK_SIZE);
static constexpr int WMSCALES_VALID_LANES = WARP_SIZE;
static constexpr int AMSCALES_PACK_SIZE = clamp(WARP_M / 32, 1, 4);
static constexpr int AMSCALES_NUM_PACKS = ceilDiv(WARP_M / 32, AMSCALES_PACK_SIZE);
static constexpr int AMSCALES_PACK_SIZE = clamp(WARP_M / 32, 1, 4);
static constexpr int AMSCALES_NUM_PACKS = ceilDiv(WARP_M / 32, AMSCALES_PACK_SIZE);
static constexpr int AMSCALES_VALID_LANES = WARP_SIZE;
struct packed_wmscale_t {
......@@ -62,48 +61,50 @@ public:
using wmscale_warp = std::array<packed_wmscale_t, WMSCALES_NUM_PACKS>;
// amscales: [M / BLOCK_M, K / group size, NUM_WARPS, AMSCALES_NUM_PACKS, WARP_SIZE] of packed_amscale_t
__device__ __forceinline__
static void load_amscale(const packed_amscale_t *ptr, int group, amscale_warp &out, bool pred) {
__device__ __forceinline__ static void
load_amscale(const packed_amscale_t *ptr, int group, amscale_warp &out, bool pred) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
#pragma unroll
#pragma unroll
for (int i = 0; i < AMSCALES_NUM_PACKS; i++) {
out[i] = load_pred(&ptr[(group * NUM_WARPS + warpId) * AMSCALES_NUM_PACKS * AMSCALES_VALID_LANES + i * AMSCALES_VALID_LANES + laneId], pred);
out[i] = load_pred(&ptr[(group * NUM_WARPS + warpId) * AMSCALES_NUM_PACKS * AMSCALES_VALID_LANES +
i * AMSCALES_VALID_LANES + laneId],
pred);
}
}
// wmscales: [N / BLOCK_N, 1, K / group size, WMSCALES_NUM_PACKS, WMSCALES_VALID_LANES] of packed_wmscale_t
__device__ __forceinline__
static void load_wmscale(const packed_wmscale_t *ptr, int group, wmscale_warp &out, bool pred) {
__device__ __forceinline__ static void
load_wmscale(const packed_wmscale_t *ptr, int group, wmscale_warp &out, bool pred) {
const int laneId = threadIdx.x % WARP_SIZE;
#pragma unroll
#pragma unroll
for (int i = 0; i < WMSCALES_NUM_PACKS; i++) {
out[i] = load_pred(&ptr[(group * WMSCALES_NUM_PACKS + i) * WMSCALES_VALID_LANES + laneId], pred);
}
}
__device__ __forceinline__
static void quantize_w4a4_fp4_from_fpsum_warp(const packed_fpsum_t (&fpsum)[INSN_K / INSN_N], packed_act_t &output, uint32_t &output_scale, int ida) {
__device__ __forceinline__ static void quantize_w4a4_fp4_from_fpsum_warp(
const packed_fpsum_t (&fpsum)[INSN_K / INSN_N], packed_act_t &output, uint32_t &output_scale, int ida) {
constexpr int NUM_GROUPS = 4;
static_assert(NUM_GROUPS == INSN_K / INSN_N);
constexpr float QVALUE_MAX = 6.0f;
constexpr float QVALUE_MAX = 6.0f;
constexpr float RECPI_QVALUE_MAX = 1 / QVALUE_MAX;
constexpr float MSCALE_MAX = 448.0f;
constexpr float MSCALE_MAX = 448.0f;
const int laneId = threadIdx.x % WARP_SIZE;
// 0 for row 0-7; 1 for row 8-15
// each half2_t represents a 8*8 matrix
half2_t input[2][INSN_K / INSN_N * 2];
#pragma unroll
#pragma unroll
for (int i = 0; i < INSN_K / INSN_N; i++) {
input[0][i * 2 + 0] = fpsum[i].data[0];
input[0][i * 2 + 1] = fpsum[i].data[2];
input[1][i * 2 + 0] = fpsum[i].data[1];
input[1][i * 2 + 1] = fpsum[i].data[3];
}
auto maxabs = [](half2_t val) ALWAYSINLINE {
val = __habs2(val);
return __hmax(val.x, val.y);
......@@ -111,14 +112,14 @@ public:
// each half_t represents maxvalue in a 8*16 matrix
half_t maxvalue[2][NUM_GROUPS];
#pragma unroll
#pragma unroll
for (int i = 0; i < NUM_GROUPS; i++) {
maxvalue[0][i] = __hmax(maxabs(input[0][i * 2]), maxabs(input[0][i * 2 + 1]));
maxvalue[1][i] = __hmax(maxabs(input[1][i * 2]), maxabs(input[1][i * 2 + 1]));
}
#pragma unroll
#pragma unroll
for (int mask = 2; mask > 0; mask /= 2) {
#pragma unroll
#pragma unroll
for (int i = 0; i < NUM_GROUPS; i++) {
maxvalue[0][i] = __hmax(maxvalue[0][i], __shfl_xor_sync(~0, maxvalue[0][i], mask));
maxvalue[1][i] = __hmax(maxvalue[1][i], __shfl_xor_sync(~0, maxvalue[1][i], mask));
......@@ -128,10 +129,10 @@ public:
float scale[2][NUM_GROUPS];
float rscale[2][NUM_GROUPS];
#pragma unroll
#pragma unroll
for (int i = 0; i < NUM_GROUPS; i++) {
scale[0][i] = fminf(float(maxvalue[0][i]) * RECPI_QVALUE_MAX, MSCALE_MAX);
scale[1][i] = fminf(float(maxvalue[1][i]) * RECPI_QVALUE_MAX, MSCALE_MAX);
scale[0][i] = fminf(float(maxvalue[0][i]) * RECPI_QVALUE_MAX, MSCALE_MAX);
scale[1][i] = fminf(float(maxvalue[1][i]) * RECPI_QVALUE_MAX, MSCALE_MAX);
// TODO: check whether (1 / scale) or (1 / fp8scale) is better
rscale[0][i] = cuda_frcp(scale[0][i]);
rscale[1][i] = cuda_frcp(scale[1][i]);
......@@ -152,30 +153,29 @@ public:
if (laneId % 4 / 2 == ida) {
output_scale = (laneId % 2 == 0) ? fp8scale[0] : fp8scale[1];
}
uint32_t qpacks[2][INSN_K / INSN_M * 2];
#pragma unroll
#pragma unroll
for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
#pragma unroll
#pragma unroll
for (int j = 0; j < 2; j++) {
float2 fval = half22float2(input[j][i]) * make_float2(rscale[j][i / 2], rscale[j][i / 2]);
float2 fval = half22float2(input[j][i]) * make_float2(rscale[j][i / 2], rscale[j][i / 2]);
qpacks[j][i] = quantize_float2_fp4(fval) << (laneId % 4 * 8);
}
}
#pragma unroll
#pragma unroll
for (int mask = 1; mask <= 2; mask *= 2) {
#pragma unroll
#pragma unroll
for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
#pragma unroll
#pragma unroll
for (int j = 0; j < 2; j++) {
qpacks[j][i] |= __shfl_xor_sync(~0, qpacks[j][i], mask);
}
}
}
// lane 0,1,2,3 / 4,5,6,7 / ... should have identical qpacks now
#pragma unroll
#pragma unroll
for (int i = 0; i < 4; i++) {
if (laneId % 4 == i) {
output.x = qpacks[0][0 + i];
......@@ -188,88 +188,110 @@ public:
// m16n16k64 MMA
// ida, idb in {0, 1}
__device__ __forceinline__
static packed_f32psum_t mma_fp4(packed_act_t act, packed_wgt_t wgt, packed_f32psum_t psum, uint32_t amscale, uint32_t wmscale, int ida, int idb) {
__device__ __forceinline__ static packed_f32psum_t mma_fp4(packed_act_t act,
packed_wgt_t wgt,
packed_f32psum_t psum,
uint32_t amscale,
uint32_t wmscale,
int ida,
int idb) {
packed_f32psum_t out;
asm volatile (
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue4m3 "
"{%0, %1, %2, %3}, "
"{%4, %5, %6, %7}, "
"{%8, %9}, "
"{%10, %11, %12, %13}, "
"{%14}, {%15, %16}, "
"{%17}, {%18, %19};"
:
"=f"(out.data[0]), "=f"(out.data[1]), "=f"(out.data[2]), "=f"(out.data[3])
:
"r"(act.x), "r"(act.y), "r"(act.z), "r"(act.w),
"r"(wgt.x), "r"(wgt.y),
"f"(psum.data[0]), "f"(psum.data[1]), "f"(psum.data[2]), "f"(psum.data[3]),
"r"(amscale), "n"(0), "h"((short)ida),
"r"(wmscale), "n"(0), "h"((short)(idb * 2))
);
asm volatile (
"{%17}, {%18, %19};"
: "=f"(out.data[0]), "=f"(out.data[1]), "=f"(out.data[2]), "=f"(out.data[3])
: "r"(act.x),
"r"(act.y),
"r"(act.z),
"r"(act.w),
"r"(wgt.x),
"r"(wgt.y),
"f"(psum.data[0]),
"f"(psum.data[1]),
"f"(psum.data[2]),
"f"(psum.data[3]),
"r"(amscale),
"n"(0),
"h"((short)ida),
"r"(wmscale),
"n"(0),
"h"((short)(idb * 2)));
asm volatile(
"mma.sync.aligned.m16n8k64.row.col.kind::mxf4nvf4.block_scale.scale_vec::4X.f32.e2m1.e2m1.f32.ue4m3 "
"{%0, %1, %2, %3}, "
"{%4, %5, %6, %7}, "
"{%8, %9}, "
"{%10, %11, %12, %13}, "
"{%14}, {%15, %16}, "
"{%17}, {%18, %19};"
:
"=f"(out.data[4]), "=f"(out.data[5]), "=f"(out.data[6]), "=f"(out.data[7])
:
"r"(act.x), "r"(act.y), "r"(act.z), "r"(act.w),
"r"(wgt.z), "r"(wgt.w),
"f"(psum.data[4]), "f"(psum.data[5]), "f"(psum.data[6]), "f"(psum.data[7]),
"r"(amscale), "n"(0), "h"((short)ida),
"r"(wmscale), "n"(0), "h"((short)(idb * 2 + 1))
);
"{%17}, {%18, %19};"
: "=f"(out.data[4]), "=f"(out.data[5]), "=f"(out.data[6]), "=f"(out.data[7])
: "r"(act.x),
"r"(act.y),
"r"(act.z),
"r"(act.w),
"r"(wgt.z),
"r"(wgt.w),
"f"(psum.data[4]),
"f"(psum.data[5]),
"f"(psum.data[6]),
"f"(psum.data[7]),
"r"(amscale),
"n"(0),
"h"((short)ida),
"r"(wmscale),
"n"(0),
"h"((short)(idb * 2 + 1)));
return out;
}
__device__ __forceinline__
static void compute_fp4(act_warp A, wgt_warp W, amscale_warp amscale, wmscale_warp wmscale, f32psum_warp &psum) {
__device__ __forceinline__ static void
compute_fp4(act_warp A, wgt_warp W, amscale_warp amscale, wmscale_warp wmscale, f32psum_warp &psum) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
#pragma unroll
#pragma unroll
for (int j = 0; j < WARP_N_TILES; j++) {
#pragma unroll
#pragma unroll
for (int i = 0; i < WARP_M_TILES; i++) {
psum[i * WARP_N_TILES + j] = mma_fp4(
A[i], W[j], psum[i * WARP_N_TILES + j],
amscale[i / 2 / AMSCALES_PACK_SIZE].data[i / 2 % AMSCALES_PACK_SIZE],
wmscale[j / 2 / WMSCALES_PACK_SIZE].data[j / 2 % WMSCALES_PACK_SIZE],
i % 2, j % 2
);
psum[i * WARP_N_TILES + j] =
mma_fp4(A[i],
W[j],
psum[i * WARP_N_TILES + j],
amscale[i / 2 / AMSCALES_PACK_SIZE].data[i / 2 % AMSCALES_PACK_SIZE],
wmscale[j / 2 / WMSCALES_PACK_SIZE].data[j / 2 % WMSCALES_PACK_SIZE],
i % 2,
j % 2);
}
}
}
template<typename Epilogue, bool USE_ALPHA>
__device__ __forceinline__
static void gemm_w4a4_fp4_block(
const BlockInfo binfo,
const packed_act_t *act,
const packed_wgt_t *wgt,
const packed_amscale_t *ascales,
const packed_wmscale_t *wscales,
float alpha, // per-tensor scale of weight
int M, int N, int K,
const Epilogue::Arguments &epilogueArgs,
bool alwaysfalse)
{
__device__ __forceinline__ static void gemm_w4a4_fp4_block(const BlockInfo binfo,
const packed_act_t *act,
const packed_wgt_t *wgt,
const packed_amscale_t *ascales,
const packed_wmscale_t *wscales,
float alpha, // per-tensor scale of weight
int M,
int N,
int K,
const Epilogue::Arguments &epilogueArgs,
bool alwaysfalse) {
constexpr int NUM_STAGES = 2;
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
act_warp A[NUM_STAGES]; // 8 * 2
wgt_warp W[NUM_STAGES]; // 32 * 2
amscale_warp amscale[NUM_STAGES]; // 1 * 2
wmscale_warp wmscale[NUM_STAGES]; // 4 * 2
f32psum_warp fpsum; // 128
act_warp A[NUM_STAGES]; // 8 * 2
wgt_warp W[NUM_STAGES]; // 32 * 2
amscale_warp amscale[NUM_STAGES]; // 1 * 2
wmscale_warp wmscale[NUM_STAGES]; // 4 * 2
f32psum_warp fpsum; // 128
for (int k = 0; k < NUM_STAGES - 1; k++) {
load_act(act, k, K, A[k], true);
......@@ -278,21 +300,21 @@ public:
load_wmscale(wscales, k, wmscale[k], true);
}
#pragma unroll
#pragma unroll
for (auto &pack : fpsum) {
#pragma unroll
#pragma unroll
for (int i = 0; i < 8; i++) {
pack.data[i] = 0;
}
}
int dummy = 0;
for (int k1 = 0; k1 < K / WARP_K; k1 += NUM_STAGES) {
#pragma unroll
#pragma unroll
for (int k2 = 0; k2 < NUM_STAGES; k2++) {
int nextk = k1 + k2 + NUM_STAGES - 1;
int idx = (k2 + NUM_STAGES - 1) % NUM_STAGES;
int idx = (k2 + NUM_STAGES - 1) % NUM_STAGES;
bool pred = nextk < K / WARP_K;
load_act(act, nextk, K, A[idx], pred);
load_wgt(wgt, nextk, K, W[idx], pred);
......@@ -317,15 +339,14 @@ public:
unused_var(dummy, alwaysfalse);
if constexpr (USE_ALPHA) {
#pragma unroll
#pragma unroll
for (auto &pack : fpsum) {
#pragma unroll
#pragma unroll
for (int i = 0; i < 8; i++) {
pack.data[i] *= alpha;
}
}
}
auto f16psum = packed_fp32_to_fp16(fpsum);
......@@ -337,21 +358,20 @@ public:
template<typename Epilogue, bool USE_ALPHA>
struct gemm_w4a4_fp4_kernel {
static constexpr int MIN_ARCH = 1200;
__device__
void operator()(
const packed_act_t *act,
const packed_wgt_t *wgt,
const packed_amscale_t *ascales,
const packed_wmscale_t *wscales,
float alpha,
int M, int N, int K,
Epilogue::Arguments epilogueArgs,
bool swapBlockXY,
bool alwaysfalse)
{
__device__ void operator()(const packed_act_t *act,
const packed_wgt_t *wgt,
const packed_amscale_t *ascales,
const packed_wmscale_t *wscales,
float alpha,
int M,
int N,
int K,
Epilogue::Arguments epilogueArgs,
bool swapBlockXY,
bool alwaysfalse) {
BlockInfo binfo = {
.bm = (int)blockIdx.x,
.bn = (int)blockIdx.y,
.bm = (int)blockIdx.x,
.bn = (int)blockIdx.y,
.numBlocksM = (int)gridDim.x,
.numBlocksN = (int)gridDim.y,
};
......@@ -372,27 +392,27 @@ public:
ascales + bm * (K / WARP_K) * NUM_WARPS * AMSCALES_NUM_PACKS * AMSCALES_VALID_LANES,
wscales + bn * (K / WARP_K) * WMSCALES_NUM_PACKS * WMSCALES_VALID_LANES,
alpha,
M, N, K,
M,
N,
K,
epilogueArgs,
alwaysfalse
);
alwaysfalse);
} else {
trap_no_fp4();
}
}
};
public:
template<bool ACT_UNSIGNED>
__device__ __forceinline__
static packed_psum_t mma(packed_act_t act, packed_wgt_t wgt) {
__device__ __forceinline__ static packed_psum_t mma(packed_act_t act, packed_wgt_t wgt) {
packed_psum_t psum;
uint4 out1 = mma_m16n8kx_s32common<mma_helper::s4u4<ACT_UNSIGNED>, mma_helper::s4>(act, uint2(wgt.x, wgt.y), uint4(0, 0, 0, 0));
uint4 out2 = mma_m16n8kx_s32common<mma_helper::s4u4<ACT_UNSIGNED>, mma_helper::s4>(act, uint2(wgt.z, wgt.w), uint4(0, 0, 0, 0));
uint4 out1 = mma_m16n8kx_s32common<mma_helper::s4u4<ACT_UNSIGNED>, mma_helper::s4>(
act, uint2(wgt.x, wgt.y), uint4(0, 0, 0, 0));
uint4 out2 = mma_m16n8kx_s32common<mma_helper::s4u4<ACT_UNSIGNED>, mma_helper::s4>(
act, uint2(wgt.z, wgt.w), uint4(0, 0, 0, 0));
psum.data[0] = out1.x;
psum.data[1] = out1.y;
psum.data[2] = out1.z;
......@@ -401,29 +421,30 @@ public:
psum.data[5] = out2.y;
psum.data[6] = out2.z;
psum.data[7] = out2.w;
return psum;
}
// template<bool si>
template<bool use_unsigned>
__device__ __forceinline__
static void quantize_w4a4_from_fpsum_warp(const packed_fpsum_t (&fpsum)[INSN_K / INSN_N], packed_act_t &output, half_t *output_scale) {
__device__ __forceinline__ static void quantize_w4a4_from_fpsum_warp(const packed_fpsum_t (&fpsum)[INSN_K / INSN_N],
packed_act_t &output,
half_t *output_scale) {
const int laneId = threadIdx.x % WARP_SIZE;
constexpr float QVALUE_MAX_SIGNED = 7.0f;
constexpr float QVALUE_MAX_UNSIGNED = 15.0f;
constexpr float RECPI_QVALUE_MAX_SIGNED = 1 / QVALUE_MAX_SIGNED;
constexpr float QVALUE_MAX_SIGNED = 7.0f;
constexpr float QVALUE_MAX_UNSIGNED = 15.0f;
constexpr float RECPI_QVALUE_MAX_SIGNED = 1 / QVALUE_MAX_SIGNED;
constexpr float RECPI_QVALUE_MAX_UNSIGNED = 1 / QVALUE_MAX_UNSIGNED;
constexpr float QVALUE_MAX = use_unsigned ? QVALUE_MAX_UNSIGNED : QVALUE_MAX_SIGNED;
constexpr float QVALUE_MAX = use_unsigned ? QVALUE_MAX_UNSIGNED : QVALUE_MAX_SIGNED;
constexpr float RECPI_QVALUE_MAX = use_unsigned ? RECPI_QVALUE_MAX_UNSIGNED : RECPI_QVALUE_MAX_SIGNED;
// constexpr int QUANTIZE_BITMASK = 0xf;
// 0 for row 0-7; 1 for row 8-15
half2_t input[2][INSN_K / INSN_N * 2];
#pragma unroll
#pragma unroll
for (int i = 0; i < INSN_K / INSN_N; i++) {
input[0][i * 2 + 0] = fpsum[i].data[0];
input[0][i * 2 + 1] = fpsum[i].data[2];
......@@ -434,14 +455,14 @@ public:
half_t maxvalue[2];
maxvalue[0] = 0;
maxvalue[1] = 0;
#pragma unroll
#pragma unroll
for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
half2_t abs0 = __habs2(input[0][i]);
half2_t abs1 = __habs2(input[1][i]);
maxvalue[0] = __hmax(maxvalue[0], __hmax(abs0.x, abs0.y));
maxvalue[1] = __hmax(maxvalue[1], __hmax(abs1.x, abs1.y));
maxvalue[0] = __hmax(maxvalue[0], __hmax(abs0.x, abs0.y));
maxvalue[1] = __hmax(maxvalue[1], __hmax(abs1.x, abs1.y));
}
#pragma unroll
#pragma unroll
for (int mask = 2; mask > 0; mask /= 2) {
maxvalue[0] = __hmax(maxvalue[0], __shfl_xor_sync(~0, maxvalue[0], mask));
maxvalue[1] = __hmax(maxvalue[1], __shfl_xor_sync(~0, maxvalue[1], mask));
......@@ -455,7 +476,7 @@ public:
scale[0] = float(maxvalue[0]) * RECPI_QVALUE_MAX;
scale[1] = float(maxvalue[1]) * RECPI_QVALUE_MAX;
if (laneId % 4 == 0) {
output_scale[laneId / 4] = half_t(scale[0]);
output_scale[laneId / 4] = half_t(scale[0]);
output_scale[laneId / 4 + 8] = half_t(scale[1]);
}
......@@ -466,23 +487,23 @@ public:
rscale[1] = cuda_frcp(scale[1]);
uint32_t qpacks[2][INSN_K / INSN_M * 2];
#pragma unroll
#pragma unroll
for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
#pragma unroll
#pragma unroll
for (int j = 0; j < 2; j++) {
// half2_t hval = __hmul2(input[j][i], half2_t(rscale[j], rscale[j]));
// float2 fval = half22float2(hval);
float2 fval = half22float2(input[j][i]) * make_float2(rscale[j], rscale[j]);
float2 fval = half22float2(input[j][i]) * make_float2(rscale[j], rscale[j]);
qpacks[j][i] = quantize_float2<4, use_unsigned>(fval) << (laneId % 4 * 8);
}
}
// 2 * 8 * 2 = 32 instructions => 256 cycles
#pragma unroll
#pragma unroll
for (int mask = 1; mask <= 2; mask *= 2) {
#pragma unroll
#pragma unroll
for (int i = 0; i < INSN_K / INSN_M * 2; i++) {
#pragma unroll
#pragma unroll
for (int j = 0; j < 2; j++) {
qpacks[j][i] |= __shfl_xor_sync(~0, qpacks[j][i], mask);
}
......@@ -490,7 +511,7 @@ public:
}
// lane 0,1,2,3 / 4,5,6,7 / ... should have identical qpacks now
#pragma unroll
#pragma unroll
for (int i = 0; i < 4; i++) {
if (laneId % 4 == i) {
output.x = qpacks[0][0 + i];
......@@ -501,73 +522,74 @@ public:
}
}
/**
* each warp quantizes a INSN_M * INSN_K (16 * 64) matrix
* input is per-warp (in global memory)
* output is per-thread (in regs)
* output_scale is per-warp (in shared memory)
* shmem must be at least INSN_M * INSN_K * sizeof(element) (16 * 64 * 0.5 = 512 Bytes)
* default to quantize activation, if quantize weight, input should be column-majored and output should be transposed ({x, y, z, w} = {x, z, y, w})
* default to quantize activation, if quantize weight, input should be column-majored and output should be
* transposed ({x, y, z, w} = {x, z, y, w})
*/
__device__ __forceinline__
static void quantize_w4a4_warp(const half_t *input, int stride, packed_act_t &output, half_t *output_scale, void *shmem) {
__device__ __forceinline__ static void
quantize_w4a4_warp(const half_t *input, int stride, packed_act_t &output, half_t *output_scale, void *shmem) {
const int laneId = threadIdx.x % WARP_SIZE;
constexpr int QUANTIZE_BITWIDTH = 4;
constexpr int QVALUE_MAX = 7; // 4 bit => [-8, 7]
constexpr int QVALUE_MAX = 7; // 4 bit => [-8, 7]
// 1 lane = 1 pack
// 1 warp = 32 lanes = 32 packs = 1 packwarp
// a pack is {a0, ..., a7} in figure https://docs.nvidia.com/cuda/parallel-thread-execution/index.html?highlight=ex2#mma-16864-a
// PACK_SIZE * 4 = INSN_K / 2
constexpr int PACK_SIZE = INSN_K / 8; // = 8 for 4bit
constexpr int NUM_PACKS_PER_ROW = INSN_K / PACK_SIZE;
// a pack is {a0, ..., a7} in figure
// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html?highlight=ex2#mma-16864-a PACK_SIZE * 4 =
// INSN_K / 2
constexpr int PACK_SIZE = INSN_K / 8; // = 8 for 4bit
constexpr int NUM_PACKS_PER_ROW = INSN_K / PACK_SIZE;
constexpr int NUM_ROWS_PER_PACKWARP = PACK_SIZE * WARP_SIZE / INSN_K;
constexpr int NUM_PACKWARPS = INSN_M / NUM_ROWS_PER_PACKWARP;
using packed_input = std::array<half_t, PACK_SIZE>;
constexpr int NUM_PACKWARPS = INSN_M / NUM_ROWS_PER_PACKWARP;
using packed_input = std::array<half_t, PACK_SIZE>;
packed_input packs[NUM_PACKWARPS];
// load
#pragma unroll
#pragma unroll
for (int i = 0; i < NUM_PACKWARPS; i++) {
int rowId = i * NUM_ROWS_PER_PACKWARP + laneId / NUM_PACKS_PER_ROW;
int colId = laneId % NUM_PACKS_PER_ROW * PACK_SIZE;
packs[i] = load(reinterpret_cast<const packed_input *>(input + rowId * stride + colId));
packs[i] = load(reinterpret_cast<const packed_input *>(input + rowId * stride + colId));
}
// find max
half_t maxvalue[NUM_PACKWARPS];
#pragma unroll
#pragma unroll
for (int i = 0; i < NUM_PACKWARPS; i++) {
maxvalue[i] = __habs(packs[i][0]);
#pragma unroll
maxvalue[i] = __habs(packs[i][0]);
#pragma unroll
for (int j = 1; j < PACK_SIZE; j++) {
maxvalue[i] = __hmax(maxvalue[i], __habs(packs[i][j]));
}
}
// warp reduce (max)
#pragma unroll
#pragma unroll
for (int mask = NUM_PACKS_PER_ROW / 2; mask > 0; mask /= 2) {
#pragma unroll
#pragma unroll
for (int i = 0; i < NUM_PACKWARPS; i++) {
maxvalue[i] = __hmax(maxvalue[i], __shfl_xor_sync(~0, maxvalue[i], mask));
}
}
// broadcast (max)
#pragma unroll
#pragma unroll
for (int i = 0; i < NUM_PACKWARPS; i++) {
maxvalue[i] = __shfl_sync(~0, maxvalue[i], laneId / NUM_PACKS_PER_ROW * NUM_PACKS_PER_ROW);
}
// quantize
using matrix_t = uint32_t[INSN_M][NUM_PACKS_PER_ROW];
matrix_t &mat = *reinterpret_cast<matrix_t *>(shmem);
#pragma unroll
matrix_t &mat = *reinterpret_cast<matrix_t *>(shmem);
#pragma unroll
for (int i = 0; i < NUM_PACKWARPS; i++) {
half_t scale = maxvalue[i] / half_t(QVALUE_MAX);
half_t rscale = half_t(QVALUE_MAX) / maxvalue[i];
......@@ -576,13 +598,13 @@ public:
}
uint32_t qpack = 0;
// #pragma unroll
// for (int j = 0; j < PACK_SIZE; j++) {
// int intvalue = __half2int_rn(packs[i][j] / scale);
// intvalue = clamp(intvalue, -QVALUE_MAX, QVALUE_MAX);
// qpack |= (intvalue & QUANTIZE_BITMASK) << (QUANTIZE_BITWIDTH * j);
// }
#pragma unroll
// #pragma unroll
// for (int j = 0; j < PACK_SIZE; j++) {
// int intvalue = __half2int_rn(packs[i][j] / scale);
// intvalue = clamp(intvalue, -QVALUE_MAX, QVALUE_MAX);
// qpack |= (intvalue & QUANTIZE_BITMASK) << (QUANTIZE_BITWIDTH * j);
// }
#pragma unroll
for (int j = 0; j < PACK_SIZE; j += 2) {
half2_t hval = __hmul2(half2_t(rscale, rscale), half2_t(packs[i][j], packs[i][j + 1]));
qpack |= quantize_float2<QUANTIZE_BITWIDTH, false>(half22float2(hval)) << (j * QUANTIZE_BITWIDTH);
......@@ -590,7 +612,7 @@ public:
mat[i * NUM_ROWS_PER_PACKWARP + laneId / NUM_PACKS_PER_ROW][laneId % NUM_PACKS_PER_ROW] = qpack;
}
__syncwarp();
// convert to imma format
int row = laneId % 16;
int col = laneId / 16 * 4;
......@@ -602,12 +624,11 @@ public:
// each thread block (1 warp) quantize WARP_M * WARP_K tile (32 * 64)
struct quantize_w4a4_act_kernel {
static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
__device__
void operator()(const half_t *input, packed_act_t *output, packed_ascale_t *oscales, int K) {
__device__ void operator()(const half_t *input, packed_act_t *output, packed_ascale_t *oscales, int K) {
const int laneId = threadIdx.x % WARP_SIZE;
const int bm = blockIdx.x / (BLOCK_M / WARP_M);
const int bk = blockIdx.y;
const int bm = blockIdx.x / (BLOCK_M / WARP_M);
const int bk = blockIdx.y;
const int warpId = blockIdx.x % (BLOCK_M / WARP_M);
const int row = blockIdx.x * WARP_M;
......@@ -620,28 +641,27 @@ public:
packed_act_t tmpout;
quantize_w4a4_warp(
input + (row + tileId * INSN_M) * K + col,
K,
tmpout,
oscale_shmem + tileId * INSN_M,
tmp_shmem
);
input + (row + tileId * INSN_M) * K + col, K, tmpout, oscale_shmem + tileId * INSN_M, tmp_shmem);
store(&output[(((bm * K / WARP_K + bk) * NUM_WARPS + warpId) * WARP_M_TILES + tileId) * WARP_SIZE + laneId], tmpout);
store(&output[(((bm * K / WARP_K + bk) * NUM_WARPS + warpId) * WARP_M_TILES + tileId) * WARP_SIZE +
laneId],
tmpout);
}
// if (threadIdx.x == 0) {
// printf("Block (%d, %d) => offset = %d\n", blockIdx.x, blockIdx.y, (bm * K / WARP_K + bk) * NUM_WARPS + warpId);
// printf("Block (%d, %d) => offset = %d\n", blockIdx.x, blockIdx.y, (bm * K / WARP_K + bk) * NUM_WARPS
// + warpId);
// }
pack_ascales(oscale_shmem, &oscales[((bm * K / WARP_K + bk) * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES]);
pack_ascales(
oscale_shmem,
&oscales[((bm * K / WARP_K + bk) * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES]);
}
};
// each thread block (1 warp) quantize WARP_N * WARP_K tile (128 * 64)
struct quantize_w4a4_wgt_kernel {
static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
__device__
void operator()(const half_t *input, packed_wgt_t *output, packed_wscale_t *oscales, int K) {
__device__ void operator()(const half_t *input, packed_wgt_t *output, packed_wscale_t *oscales, int K) {
const int laneId = threadIdx.x % WARP_SIZE;
const int bn = blockIdx.x / (BLOCK_N / WARP_N);
......@@ -657,12 +677,7 @@ public:
packed_wgt_t tmpout;
quantize_w4a4_warp(
input + (col + tileId * INSN_N) * K + row,
K,
tmpout,
oscale_shmem + tileId * INSN_N,
tmp_shmem
);
input + (col + tileId * INSN_N) * K + row, K, tmpout, oscale_shmem + tileId * INSN_N, tmp_shmem);
std::swap(tmpout.y, tmpout.z);
......@@ -674,59 +689,52 @@ public:
};
struct i2f_sm80 {
__device__ __forceinline__
static float2 int2float2(int x, int y) {
__device__ __forceinline__ static float2 int2float2(int x, int y) {
return make_float2(int2float_fast(x), int2float_fast(y));
}
__device__ __forceinline__
static half2_t int2half2(int x, int y) {
__device__ __forceinline__ static half2_t int2half2(int x, int y) {
return float22half2<half2_t>(int2float2(x, y));
}
};
struct i2f_sm75 {
__device__ __forceinline__
static float2 int2float2(int x, int y) {
__device__ __forceinline__ static float2 int2float2(int x, int y) {
return make_float2(int2float_fast(x), int2float_fast(y));
}
__device__ __forceinline__
static half2_t int2half2(int x, int y) {
__device__ __forceinline__ static half2_t int2half2(int x, int y) {
return half2(__int2half_rn(x), __int2half_rn(y));
}
};
struct i2f_sm75_fast {
__device__ __forceinline__
static float2 int2float2(int x, int y) {
__device__ __forceinline__ static float2 int2float2(int x, int y) {
return make_float2(int2float_fast(x), int2float_fast(y));
}
__device__ __forceinline__
static half2_t int2half2(int x, int y) {
__device__ __forceinline__ static half2_t int2half2(int x, int y) {
return int2half2_fast_512(x, y);
}
};
template<bool ACT_UNSIGNED, typename T>
__device__ __forceinline__
static void compute(act_warp A, wgt_warp W, ascale_warp ascale, wscale_warp wscale, T &fpsum) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 800
__device__ __forceinline__ static void
compute(act_warp A, wgt_warp W, ascale_warp ascale, wscale_warp wscale, T &fpsum) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 800
using int2half2 = i2f_sm80;
#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
using int2half2 = std::conditional_t<Config::FASTER_I2F, i2f_sm75_fast, i2f_sm75>;;
#else
using int2half2 = Base::i2f_normal;
#endif
Base::template apply_scales<int2half2>([&](int i, int j) {
return mma<ACT_UNSIGNED>(A[i], W[j]);
}, ascale, wscale, fpsum);
#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750
using int2half2 = std::conditional_t<Config::FASTER_I2F, i2f_sm75_fast, i2f_sm75>;
;
#else
using int2half2 = Base::i2f_normal;
#endif
Base::template apply_scales<int2half2>(
[&](int i, int j) { return mma<ACT_UNSIGNED>(A[i], W[j]); }, ascale, wscale, fpsum);
}
__device__ __forceinline__
static void checkNan(fpsum_warp fpsum, const char *info = "") {
__device__ __forceinline__ static void checkNan(fpsum_warp fpsum, const char *info = "") {
#if ENABLE_NAN_CHECK
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -735,14 +743,17 @@ public:
for (int j = 0; j < 4; j++) {
bool abnormal = !isfinite((float)fpsum[i].data[j].x) || !isfinite((float)fpsum[i].data[j].y);
if (abnormal) {
printf("abnormal value detected at block.x=%d block.y=%d warpId=%d laneId=%d fpsum_warp (%s) i=%d j=%d data.x=%f data.y=%f\n",
blockIdx.x, blockIdx.y,
warpId, laneId,
info,
i, j,
(float)fpsum[i].data[j].x,
(float)fpsum[i].data[j].y
);
printf("abnormal value detected at block.x=%d block.y=%d warpId=%d laneId=%d fpsum_warp (%s) i=%d "
"j=%d data.x=%f data.y=%f\n",
blockIdx.x,
blockIdx.y,
warpId,
laneId,
info,
i,
j,
(float)fpsum[i].data[j].x,
(float)fpsum[i].data[j].y);
__trap();
}
}
......@@ -750,8 +761,7 @@ public:
#endif
}
__device__ __forceinline__
static void checkNan(packed_f32psum_t fpsum, const char *info = "") {
__device__ __forceinline__ static void checkNan(packed_f32psum_t fpsum, const char *info = "") {
#if ENABLE_NAN_CHECK
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -759,21 +769,22 @@ public:
for (int j = 0; j < 8; j++) {
bool abnormal = !isfinite(fpsum.data[j]);
if (abnormal) {
printf("abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_f32psum_t (%s) j=%d data=%f\n",
blockIdx.x, blockIdx.y,
warpId, laneId,
printf(
"abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_f32psum_t (%s) j=%d data=%f\n",
blockIdx.x,
blockIdx.y,
warpId,
laneId,
info,
j,
fpsum.data[j]
);
fpsum.data[j]);
__trap();
}
}
#endif
}
__device__ __forceinline__
static void checkNan(packed_fpsum_t fpsum, const char *info = "") {
__device__ __forceinline__ static void checkNan(packed_fpsum_t fpsum, const char *info = "") {
#if ENABLE_NAN_CHECK
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -781,34 +792,36 @@ public:
for (int j = 0; j < 4; j++) {
bool abnormal = !isfinite((float)fpsum.data[j].x) || !isfinite((float)fpsum.data[j].y);
if (abnormal) {
printf("abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_fpsum_t (%s) j=%d data.x=%f data.y=%f\n",
blockIdx.x, blockIdx.y,
warpId, laneId,
info,
j,
(float)fpsum.data[j].x,
(float)fpsum.data[j].y
);
printf("abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_fpsum_t (%s) j=%d data.x=%f "
"data.y=%f\n",
blockIdx.x,
blockIdx.y,
warpId,
laneId,
info,
j,
(float)fpsum.data[j].x,
(float)fpsum.data[j].y);
__trap();
}
}
#endif
}
__device__ __forceinline__
static void checkNan(float data, const char *info = "") {
__device__ __forceinline__ static void checkNan(float data, const char *info = "") {
#if ENABLE_NAN_CHECK
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
bool abnormal = !isfinite(data);
if (abnormal) {
printf("abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_fpsum_t (%s) data=%f\n",
blockIdx.x, blockIdx.y,
warpId, laneId,
info,
data
);
printf("abnormal value detected at bm=%d bn=%d warpId=%d laneId=%d packed_fpsum_t (%s) data=%f\n",
blockIdx.x,
blockIdx.y,
warpId,
laneId,
info,
data);
__trap();
}
#endif
......@@ -816,19 +829,18 @@ public:
// out: [M / BLOCK_M, N / BLOCK_N, NUM_WARPS, 1, NUM_M_TILES, NUM_N_TILES, WARP_SIZE] of fpsum_warp
template<typename Epilogue, bool ACT_UNSIGNED, bool USE_FP32_ACCUM>
__device__ __forceinline__
static void gemm_w4a4_block(
const BlockInfo binfo,
const packed_act_t *act,
const packed_wgt_t *wgt,
const packed_ascale_t *ascales,
const packed_wscale_t *wscales,
// const packed_wscale_t *bias_ptr,
// half_t *out,
int M, int N, int K,
const Epilogue::Arguments &epilogueArgs,
bool alwaysfalse)
{
__device__ __forceinline__ static void gemm_w4a4_block(const BlockInfo binfo,
const packed_act_t *act,
const packed_wgt_t *wgt,
const packed_ascale_t *ascales,
const packed_wscale_t *wscales,
// const packed_wscale_t *bias_ptr,
// half_t *out,
int M,
int N,
int K,
const Epilogue::Arguments &epilogueArgs,
bool alwaysfalse) {
constexpr int NUM_STAGES = 2;
const int laneId = threadIdx.x % WARP_SIZE;
......@@ -838,11 +850,11 @@ public:
fpsum_warp fpsum;
GEMM_W4A4_Block<Config>()(act, wgt, ascales, wscales, K, fpsum, alwaysfalse);
#else
act_warp A[NUM_STAGES]; // 8
wgt_warp W[NUM_STAGES]; // 32
ascale_warp ascale[NUM_STAGES]; // 1
wscale_warp wscale[NUM_STAGES]; // 2
std::conditional_t<USE_FP32_ACCUM, f32psum_warp, fpsum_warp> fpsum; // 64
act_warp A[NUM_STAGES]; // 8
wgt_warp W[NUM_STAGES]; // 32
ascale_warp ascale[NUM_STAGES]; // 1
wscale_warp wscale[NUM_STAGES]; // 2
std::conditional_t<USE_FP32_ACCUM, f32psum_warp, fpsum_warp> fpsum; // 64
// load_wscale<true>(wscales, wscale[0], true);
// load_wscale<false>(wscales, wscale[1], true);
......@@ -867,14 +879,14 @@ public:
}
}
}
int dummy = 0;
for (int k1 = 0; k1 < K / WARP_K; k1 += NUM_STAGES) {
#pragma unroll
#pragma unroll
for (int k2 = 0; k2 < NUM_STAGES; k2++) {
int nextk = k1 + k2 + NUM_STAGES - 1;
int idx = (k2 + NUM_STAGES - 1) % NUM_STAGES;
int idx = (k2 + NUM_STAGES - 1) % NUM_STAGES;
bool pred = nextk < K / WARP_K;
load_act(act, nextk, K, A[idx], pred);
load_wgt(wgt, nextk, K, W[idx], pred);
......@@ -889,11 +901,11 @@ public:
compute<ACT_UNSIGNED>(A[k2], W[k2], ascale[k2], wscale[k2], fpsum);
//#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
// #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
if (alwaysfalse) {
dummy = clock();
}
//#endif
// #endif
// asm volatile ("membar.cta;");
}
......@@ -927,11 +939,17 @@ public:
const packed_wscale_t *smooth_factor;
};
static constexpr int NUM_PACKS = INSN_K / INSN_N;
static constexpr int NUM_PACKS = INSN_K / INSN_N;
static constexpr int NUM_GROUPS = WARP_N_TILES / NUM_PACKS;
__device__ __forceinline__
void apply_quantize(fpsum_warp fpsum, int M, int N, int K, packed_act_t *qout, oscales_t *oscales, half_t shift_value, const packed_wscale_t *smooth_factor) {
__device__ __forceinline__ void apply_quantize(fpsum_warp fpsum,
int M,
int N,
int K,
packed_act_t *qout,
oscales_t *oscales,
half_t shift_value,
const packed_wscale_t *smooth_factor) {
const int laneId = threadIdx.x % WARP_SIZE;
const int warpId = threadIdx.x / WARP_SIZE;
......@@ -940,21 +958,21 @@ public:
wscale_warp smooth;
load_wscale(smooth_factor, 0, N, smooth, true);
#pragma unroll
#pragma unroll
for (int group = 0; group < NUM_GROUPS; group++) {
amscale_warp omscale;
#pragma unroll
#pragma unroll
for (int i = 0; i < WARP_M_TILES; i++) {
packed_fpsum_t tmp[NUM_PACKS];
#pragma unroll
#pragma unroll
for (int j = 0; j < NUM_PACKS; j++) {
half2_t ws1 = broadcast_wscale(smooth, (group * NUM_PACKS + j) * 4, laneId);
half2_t ws2 = broadcast_wscale(smooth, (group * NUM_PACKS + j) * 4 + 2, laneId);
#pragma unroll
#pragma unroll
for (int k = 0; k < 4; k++) {
half2_t src = fpsum[i * WARP_N_TILES + group * NUM_PACKS + j].data[k];
half2_t src = fpsum[i * WARP_N_TILES + group * NUM_PACKS + j].data[k];
half2_t &dst = tmp[j].data[k];
// dst.x = gelu(src.x);
......@@ -977,7 +995,8 @@ public:
packed_act_t qresult;
if constexpr (USE_FP4) {
quantize_w4a4_fp4_from_fpsum_warp(tmp, qresult, omscale[i / 2 / AMSCALES_PACK_SIZE].data[i / 2 % AMSCALES_PACK_SIZE], i % 2);
quantize_w4a4_fp4_from_fpsum_warp(
tmp, qresult, omscale[i / 2 / AMSCALES_PACK_SIZE].data[i / 2 % AMSCALES_PACK_SIZE], i % 2);
} else {
quantize_w4a4_from_fpsum_warp<USE_UNSIGNED>(tmp, qresult, &oscale_shmem[warpId][i * INSN_M]);
}
......@@ -985,34 +1004,38 @@ public:
}
if constexpr (USE_FP4) {
#pragma unroll
#pragma unroll
for (int k = 0; k < AMSCALES_NUM_PACKS; k++) {
store(&oscales[((group * NUM_WARPS + warpId) * AMSCALES_NUM_PACKS + k) * AMSCALES_VALID_LANES + laneId], omscale[k]);
store(&oscales[((group * NUM_WARPS + warpId) * AMSCALES_NUM_PACKS + k) * AMSCALES_VALID_LANES +
laneId],
omscale[k]);
}
}
if constexpr (!USE_FP4) {
__syncwarp();
pack_ascales(&oscale_shmem[warpId][0], &oscales[(group * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES]);
pack_ascales(&oscale_shmem[warpId][0],
&oscales[(group * NUM_WARPS + warpId) * ASCALES_NUM_PACKS * ASCALES_VALID_LANES]);
__syncwarp();
}
}
}
__device__ __forceinline__
void operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
__device__ __forceinline__ void
operator()(const BlockInfo binfo, fpsum_warp fpsum, int M, int N, int K, const Arguments &args) {
const int bm = binfo.bm;
const int bn = binfo.bn;
if constexpr (!USE_FP4 || FP4_AVAILABLE) {
apply_quantize(
fpsum, M, N, K,
args.qout + (bm * N / WARP_K + bn * NUM_GROUPS) * NUM_WARPS * WARP_M_TILES * WARP_SIZE,
args.oscales + (bm * N / WARP_K + bn * NUM_GROUPS) * NUM_WARPS *
(USE_FP4 ? AMSCALES_NUM_PACKS * AMSCALES_VALID_LANES : ASCALES_NUM_PACKS * ASCALES_VALID_LANES),
args.shift_value,
args.smooth_factor + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES
);
apply_quantize(fpsum,
M,
N,
K,
args.qout + (bm * N / WARP_K + bn * NUM_GROUPS) * NUM_WARPS * WARP_M_TILES * WARP_SIZE,
args.oscales + (bm * N / WARP_K + bn * NUM_GROUPS) * NUM_WARPS *
(USE_FP4 ? AMSCALES_NUM_PACKS * AMSCALES_VALID_LANES
: ASCALES_NUM_PACKS * ASCALES_VALID_LANES),
args.shift_value,
args.smooth_factor + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES);
} else {
trap_no_fp4();
}
......@@ -1025,22 +1048,21 @@ public:
static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
static constexpr int MAX_ARCH = Config::FASTER_I2F ? 750 : INT_MAX; // FASTER_I2F is only needed on sm_75
__device__
void operator()(
const packed_act_t *act,
const packed_wgt_t *wgt,
const packed_ascale_t *ascales,
const packed_wscale_t *wscales,
int M, int N, int K,
Epilogue::Arguments epilogueArgs,
bool swapBlockXY,
bool alwaysfalse)
{
__device__ void operator()(const packed_act_t *act,
const packed_wgt_t *wgt,
const packed_ascale_t *ascales,
const packed_wscale_t *wscales,
int M,
int N,
int K,
Epilogue::Arguments epilogueArgs,
bool swapBlockXY,
bool alwaysfalse) {
// printf("Device sizeof(args) = %d", (int)sizeof(epilogueArgs));
BlockInfo binfo = {
.bm = (int)blockIdx.x,
.bn = (int)blockIdx.y,
.bm = (int)blockIdx.x,
.bn = (int)blockIdx.y,
.numBlocksM = (int)gridDim.x,
.numBlocksN = (int)gridDim.y,
};
......@@ -1064,20 +1086,21 @@ public:
// bias ? bias + bn * WSCALES_NUM_PACKS * WSCALES_VALID_LANES : nullptr,
// out + (bm * BLOCK_M * N) + bn * BLOCK_N,
// out + (bm * N / BLOCK_N + bn) * NUM_WARPS * WARP_M_TILES * WARP_N_TILES * WARP_SIZE,
M, N, K,
M,
N,
K,
epilogueArgs,
alwaysfalse
);
alwaysfalse);
}
};
template<bool fuse_glu, bool use_fp4>
struct quantize_w4a4_fuse_lora_kernel {
using oscales_t = typename std::conditional_t<use_fp4, packed_amscale_t, packed_ascale_t>;
static constexpr int MIN_ARCH = std::is_same_v<half_t, __nv_bfloat16> ? 800 : 750;
static constexpr size_t SHMEM_PER_WARP = ceilDiv<size_t>(Base::template load_act_to_fpsum<fuse_glu>::SHMEM_SIZE, 128) * 128;
static constexpr size_t SHMEM_PER_WARP =
ceilDiv<size_t>(Base::template load_act_to_fpsum<fuse_glu>::SHMEM_SIZE, 128) * 128;
static constexpr size_t SHMEM_SIZE = SHMEM_PER_WARP * NUM_WARPS;
struct Arguments {
......@@ -1091,25 +1114,23 @@ public:
int lora_rank;
// aligned to BLOCK_M and BLOCK_N
int M, N; // N should be the actual K in the next GEMM (needs /2 if fuse_glu)
int M, N; // N should be the actual K in the next GEMM (needs /2 if fuse_glu)
// the actual M and N (no need to /2 if fuse_glu)
int actualM, actualN;
bool alwaysfalse;
};
__device__ __forceinline__
void operator()(Arguments args)
{
__device__ __forceinline__ void operator()(Arguments args) {
const BlockInfo binfo = {
.bm = (int)blockIdx.x,
.bn = (int)blockIdx.y,
.bm = (int)blockIdx.x,
.bn = (int)blockIdx.y,
.numBlocksM = (int)gridDim.x,
.numBlocksN = (int)gridDim.y,
};
const int bm = binfo.bm;
const int bn = binfo.bn;
const int bm = binfo.bm;
const int bn = binfo.bn;
const int warpId = threadIdx.x / WARP_SIZE;
const int m_offset = bm * BLOCK_M + warpId * WARP_M;
......@@ -1119,42 +1140,48 @@ public:
fpsum_warp fpsum;
Base::template load_act_to_fpsum<fuse_glu>()(
args.input + m_offset * args.actualN + n_offset,
args.actualN,
args.actualM - m_offset,
args.actualN - n_offset,
fpsum,
shmem + warpId * SHMEM_PER_WARP
// args.smooth_factor ? args.smooth_factor + n_offset : nullptr
Base::template load_act_to_fpsum<fuse_glu>()(args.input + m_offset * args.actualN + n_offset,
args.actualN,
args.actualM - m_offset,
args.actualN - n_offset,
fpsum,
shmem + warpId * SHMEM_PER_WARP
// args.smooth_factor ? args.smooth_factor + n_offset : nullptr
);
CHECK_NAN(fpsum, "fpsum");
// for (int i = 0; i < 16; i++) {
// printf("bm=%d bn=%d warp=%d lane=%d fpsum[%d][0:1]=%f %f\n",
// printf("bm=%d bn=%d warp=%d lane=%d fpsum[%d][0:1]=%f %f\n",
// bm, bn, warpId, threadIdx.x % WARP_SIZE, i,
// (float)fpsum[i].data[0].x, (float)fpsum[i].data[0].y);
// }
using EpilogueLoraDown = typename Lora<Config>::EpilogueLoraDown;
EpilogueLoraDown()(binfo, fpsum, args.M, args.N, 0, typename EpilogueLoraDown::Arguments{
.lora_wgt_down = args.lora_wgt_down,
.lora_act = args.lora_act,
.rank = args.lora_rank,
.alwaysfalse = args.alwaysfalse,
});
EpilogueQuantize<false, false, use_fp4>()(binfo, fpsum, args.M, args.N, 0, typename EpilogueQuantize<false, false, use_fp4>::Arguments{
.qout = args.output,
.oscales = args.oscales,
.shift_value = 0,
.smooth_factor = args.smooth_factor
});
EpilogueLoraDown()(binfo,
fpsum,
args.M,
args.N,
0,
typename EpilogueLoraDown::Arguments{
.lora_wgt_down = args.lora_wgt_down,
.lora_act = args.lora_act,
.rank = args.lora_rank,
.alwaysfalse = args.alwaysfalse,
});
EpilogueQuantize<false, false, use_fp4>()(
binfo,
fpsum,
args.M,
args.N,
0,
typename EpilogueQuantize<false, false, use_fp4>::Arguments{.qout = args.output,
.oscales = args.oscales,
.shift_value = 0,
.smooth_factor = args.smooth_factor});
}
};
};
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
src/kernels/zgemm/gemm_w4a4_launch.cuh
View file @ 37a27712
......@@ -5,57 +5,61 @@ namespace nunchaku::kernels {
template<typename Config, bool USE_FP4>
class GEMM_W4A4_Launch {
using GEMM = GEMM_W4A4<Config>;
using GEMM = GEMM_W4A4<Config>;
using Epilogues = Epilogues<Config>;
using Lora = Lora<Config>;
using Lora = Lora<Config>;
using packed_act_t = typename GEMM::packed_act_t;
using packed_wgt_t = typename GEMM::packed_wgt_t;
using packed_ascale_t = typename GEMM::packed_ascale_t;
using packed_wscale_t = typename GEMM::packed_wscale_t;
using packed_act_t = typename GEMM::packed_act_t;
using packed_wgt_t = typename GEMM::packed_wgt_t;
using packed_ascale_t = typename GEMM::packed_ascale_t;
using packed_wscale_t = typename GEMM::packed_wscale_t;
using packed_amscale_t = typename GEMM::packed_amscale_t;
using packed_wmscale_t = typename GEMM::packed_wmscale_t;
using packed_fpsum_t = typename GEMM::packed_fpsum_t;
using half_t = typename GEMM::half_t;
using packed_fpsum_t = typename GEMM::packed_fpsum_t;
using half_t = typename GEMM::half_t;
public:
static void gemm_w4a4(
Tensor act, // packed act [M, K / 2]
Tensor wgt, // packed act [N, K / 2]
Tensor out, // linear [M, N]
Tensor qout, // packed act [M, N / 2]
Tensor ascales, // packed as [K / 64, M]
Tensor wscales, // packed ws [K / 64, N]
Tensor oscales, // packed as [N / 64, M]
Tensor poolout, // linear [M / PoolSize, N]
Tensor lora_act_in, // packed lora_act [M, R]
Tensor lora_up, // packed lora_wgt [N, R]
Tensor lora_down, // packed lora_wgt [N, R]
Tensor lora_act_out, // packed lora_act [M, R]
Tensor norm_q, // linear [HEAD_DIM]
Tensor norm_k, // linear [HEAD_DIM]
Tensor rotary_emb, // linear [M, HEAD_DIM / 2, 2, 2]
Tensor bias, // packed ws [N]
Tensor smooth_factor, // packed ws [N], for quantization of the next layer
Tensor out_vk, // linear [B, num_heads, head_dim + 1, head_dim]
Tensor out_linearattn,// linear [B, (M), N / 3]
bool act_unsigned,
std::vector<float> lora_scales, // [R / 16]
bool fuse_silu,
bool fp4,
float alpha,
Tensor wcscales, // packed ws [N]
Tensor out_q, // packed attention [B, H, M, D]
Tensor out_k, // packed attention [B, H, M, D]
Tensor out_v, // packed attention [B, H, M, D]
int attn_tokens
);
static void quantize_w4a4_act_fuse_lora(Tensor input, Tensor output, Tensor oscales, Tensor lora_down, Tensor lora_act_out, Tensor smooth, bool fuse_glu, bool fp4);
static void gemm_w4a4(Tensor act, // packed act [M, K / 2]
Tensor wgt, // packed act [N, K / 2]
Tensor out, // linear [M, N]
Tensor qout, // packed act [M, N / 2]
Tensor ascales, // packed as [K / 64, M]
Tensor wscales, // packed ws [K / 64, N]
Tensor oscales, // packed as [N / 64, M]
Tensor poolout, // linear [M / PoolSize, N]
Tensor lora_act_in, // packed lora_act [M, R]
Tensor lora_up, // packed lora_wgt [N, R]
Tensor lora_down, // packed lora_wgt [N, R]
Tensor lora_act_out, // packed lora_act [M, R]
Tensor norm_q, // linear [HEAD_DIM]
Tensor norm_k, // linear [HEAD_DIM]
Tensor rotary_emb, // linear [M, HEAD_DIM / 2, 2, 2]
Tensor bias, // packed ws [N]
Tensor smooth_factor, // packed ws [N], for quantization of the next layer
Tensor out_vk, // linear [B, num_heads, head_dim + 1, head_dim]
Tensor out_linearattn, // linear [B, (M), N / 3]
bool act_unsigned,
std::vector<float> lora_scales, // [R / 16]
bool fuse_silu,
bool fp4,
float alpha,
Tensor wcscales, // packed ws [N]
Tensor out_q, // packed attention [B, H, M, D]
Tensor out_k, // packed attention [B, H, M, D]
Tensor out_v, // packed attention [B, H, M, D]
int attn_tokens);
static void quantize_w4a4_act_fuse_lora(Tensor input,
Tensor output,
Tensor oscales,
Tensor lora_down,
Tensor lora_act_out,
Tensor smooth,
bool fuse_glu,
bool fp4);
static void quantize_w4a4_act(Tensor input, Tensor output, Tensor oscales);
static void quantize_w4a4_wgt(Tensor input, Tensor output, Tensor oscales);
static void linearattn_vk_mul_q(Tensor q, Tensor vk);
};
}; // namespace nunchaku::kernels
\ No newline at end of file
}; // namespace nunchaku::kernels
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