github-actions[bot]
112 lines
3.5 KiB
C++
112 lines
3.5 KiB
C++
#include <ATen/cuda/CUDAContext.h>
|
|
#include <c10/cuda/CUDAGuard.h>
|
|
#include <torch/extension.h>
|
|
#include <torch/python.h>
|
|
|
|
#include <vector>
|
|
|
|
// Forward declaration of CUDA kernel template
|
|
template<typename out_t>
|
|
void rms_norm_split_rope_cuda(
|
|
void* x,
|
|
void* sin_freqs,
|
|
void* cos_freqs,
|
|
void* weights,
|
|
int b,
|
|
int s,
|
|
int n,
|
|
int h,
|
|
long cos_sb, long cos_sn, long cos_ss,
|
|
long sin_sb, long sin_sn, long sin_ss,
|
|
void* out,
|
|
cudaStream_t stream
|
|
);
|
|
|
|
at::Tensor rms_norm_split_rope(
|
|
at::Tensor &x,
|
|
at::Tensor &sin_freqs,
|
|
at::Tensor &cos_freqs,
|
|
at::Tensor &weights,
|
|
bool out_fp8
|
|
) {
|
|
TORCH_CHECK(x.scalar_type() == at::ScalarType::BFloat16, "Input must be BFloat16");
|
|
TORCH_CHECK(sin_freqs.scalar_type() == at::ScalarType::BFloat16, "sin_freqs must be BFloat16");
|
|
TORCH_CHECK(cos_freqs.scalar_type() == at::ScalarType::BFloat16, "cos_freqs must be BFloat16");
|
|
TORCH_CHECK(x.is_cuda(), "Input must be on CUDA");
|
|
TORCH_CHECK(sin_freqs.is_cuda(), "sin_freqs must be on CUDA");
|
|
TORCH_CHECK(cos_freqs.is_cuda(), "cos_freqs must be on CUDA");
|
|
|
|
// Get dimensions
|
|
// x: [b, s, h]
|
|
// cos, sin: [b, n, s, d] where n*d = h/2
|
|
int b = x.size(0);
|
|
int s = x.size(1);
|
|
int h = x.size(2);
|
|
|
|
TORCH_CHECK(cos_freqs.dim() == 4, "cos_freqs must be 4D");
|
|
TORCH_CHECK(sin_freqs.dim() == 4, "sin_freqs must be 4D");
|
|
|
|
int n = cos_freqs.size(1);
|
|
int d = h / n;
|
|
|
|
|
|
// Require a contiguous innermost (d/2) dim for the vectorized int4 freq load,
|
|
// but keep the outer (b, n, s) strides: apply_split_rotary_emb hands us a
|
|
// swapaxes view (logical [b, n, s, d/2], physical [b, s, n, d/2]) whose inner
|
|
// stride is already 1, so this never copies it. The strides are forwarded to
|
|
// the kernel so the read is correct regardless of the physical layout.
|
|
if (x.stride(-1) != 1) { x = x.contiguous(); }
|
|
if (cos_freqs.stride(-1) != 1) { cos_freqs = cos_freqs.contiguous(); }
|
|
if (sin_freqs.stride(-1) != 1) { sin_freqs = sin_freqs.contiguous(); }
|
|
|
|
long cos_sb = cos_freqs.stride(0), cos_sn = cos_freqs.stride(1), cos_ss = cos_freqs.stride(2);
|
|
long sin_sb = sin_freqs.stride(0), sin_sn = sin_freqs.stride(1), sin_ss = sin_freqs.stride(2);
|
|
|
|
// Create output tensor
|
|
at::Tensor out;
|
|
if (out_fp8) {
|
|
out = torch::empty(x.sizes(), x.options().dtype(torch::kFloat8_e4m3fn));
|
|
} else {
|
|
out = torch::empty(x.sizes(), x.options().dtype(torch::kBFloat16));
|
|
}
|
|
|
|
// Setup CUDA
|
|
at::cuda::CUDAGuard device_guard{(char)x.get_device()};
|
|
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
|
|
|
// Launch kernel
|
|
if (out_fp8) {
|
|
rms_norm_split_rope_cuda<at::Float8_e4m3fn>(
|
|
x.data_ptr(),
|
|
sin_freqs.data_ptr(),
|
|
cos_freqs.data_ptr(),
|
|
weights.data_ptr(), // weights (optional, not used yet)
|
|
b,
|
|
s,
|
|
n,
|
|
h,
|
|
cos_sb, cos_sn, cos_ss,
|
|
sin_sb, sin_sn, sin_ss,
|
|
(void*)out.data_ptr(),
|
|
stream
|
|
);
|
|
} else {
|
|
rms_norm_split_rope_cuda<at::BFloat16>(
|
|
x.data_ptr(),
|
|
sin_freqs.data_ptr(),
|
|
cos_freqs.data_ptr(),
|
|
weights.data_ptr(), // weights (optional, not used yet)
|
|
b,
|
|
s,
|
|
n,
|
|
h,
|
|
cos_sb, cos_sn, cos_ss,
|
|
sin_sb, sin_sn, sin_ss,
|
|
(void*)out.data_ptr(),
|
|
stream
|
|
);
|
|
}
|
|
|
|
return out;
|
|
}
|