Init
This commit is contained in:
Starrick
+552
@@ -0,0 +1,552 @@
|
||||
#undef __CUDA_NO_HALF_OPERATORS__
|
||||
#undef __CUDA_NO_HALF_CONVERSIONS__
|
||||
#undef __CUDA_NO_BFLOAT16_CONVERSIONS__
|
||||
#undef __CUDA_NO_HALF2_OPERATORS__
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_bf16.h>
|
||||
#include <type_traits>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/util/BFloat16.h>
|
||||
#include <c10/cuda/CUDAStream.h>
|
||||
#include <torch/extension.h>
|
||||
|
||||
|
||||
|
||||
// Type traits for CUDA types
|
||||
template <typename T>
|
||||
struct CudaTypeTraits
|
||||
{
|
||||
static constexpr bool is_supported = false;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct CudaTypeTraits<float>
|
||||
{
|
||||
static constexpr bool is_supported = true;
|
||||
using type = float;
|
||||
static constexpr const char *name = "float";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct CudaTypeTraits<half>
|
||||
{
|
||||
static constexpr bool is_supported = true;
|
||||
using type = half;
|
||||
static constexpr const char *name = "half";
|
||||
};
|
||||
|
||||
template <>
|
||||
struct CudaTypeTraits<__nv_bfloat16>
|
||||
{
|
||||
static constexpr bool is_supported = true;
|
||||
using type = __nv_bfloat16;
|
||||
static constexpr const char *name = "bfloat16";
|
||||
};
|
||||
|
||||
// Optimized forward kernel template
|
||||
template <typename T>
|
||||
__global__ void multimodal_rope_forward_kernel(
|
||||
const T *__restrict__ q, // [batch, q_heads, seq_len, head_dim]
|
||||
const T *__restrict__ k, // [batch, kv_heads, seq_len, head_dim]
|
||||
const T *__restrict__ cos, // [3, batch, seq_len, head_dim]
|
||||
const T *__restrict__ sin, // [3, batch, seq_len, head_dim]
|
||||
T *__restrict__ q_out, // [batch, q_heads, seq_len, head_dim]
|
||||
T *__restrict__ k_out, // [batch, kv_heads, seq_len, head_dim]
|
||||
const int *__restrict__ mrope_section_doubled, // [32, 48, 48]
|
||||
int batch_size,
|
||||
int q_heads,
|
||||
int kv_heads,
|
||||
int seq_len,
|
||||
int head_dim)
|
||||
{
|
||||
static_assert(CudaTypeTraits<T>::is_supported, "Unsupported data type");
|
||||
|
||||
// Block organization: batch_size * (q_heads + kv_heads) blocks
|
||||
int batch_idx = blockIdx.x;
|
||||
int head_idx = blockIdx.y;
|
||||
int seq_paral_size = gridDim.z;
|
||||
int seq_paral_idx = blockIdx.z;
|
||||
|
||||
bool is_q_head = head_idx < q_heads;
|
||||
int actual_head_idx = is_q_head ? head_idx : (head_idx - q_heads);
|
||||
int total_heads = is_q_head ? q_heads : kv_heads;
|
||||
|
||||
// Each warp processes one token (seq position)
|
||||
int warp_id = threadIdx.x / 32;
|
||||
int lane_id = threadIdx.x % 32;
|
||||
int warps_per_block = blockDim.x / 32;
|
||||
|
||||
// Process tokens in batches across warps
|
||||
|
||||
for (int seq_base = warp_id + seq_paral_idx * warps_per_block; seq_base < seq_len; seq_base += seq_paral_size * warps_per_block) {
|
||||
int seq_idx = seq_base;
|
||||
if (seq_idx >= seq_len)
|
||||
break;
|
||||
|
||||
// Each lane processes multiple dimensions (head_dim=128, 32 lanes)
|
||||
constexpr int dims_per_lane = 4;
|
||||
for (int dim_batch = 0; dim_batch < (head_dim + 31) / 32; ++dim_batch)
|
||||
{
|
||||
int dim_start = dim_batch * 32 + lane_id;
|
||||
if (dim_start >= head_dim)
|
||||
break;
|
||||
|
||||
#pragma unroll
|
||||
for (int dim_offset = 0; dim_offset < dims_per_lane; ++dim_offset)
|
||||
{
|
||||
int dim_idx = dim_start + dim_offset * 32;
|
||||
if (dim_idx >= head_dim)
|
||||
break;
|
||||
|
||||
// Determine which section this dimension belongs to
|
||||
int section_idx;
|
||||
int cos_sin_d = dim_idx;
|
||||
if (dim_idx < mrope_section_doubled[0])
|
||||
{
|
||||
section_idx = 0;
|
||||
}
|
||||
else if (dim_idx < mrope_section_doubled[0] + mrope_section_doubled[1])
|
||||
{
|
||||
section_idx = 1;
|
||||
}
|
||||
else
|
||||
{
|
||||
section_idx = 2;
|
||||
}
|
||||
|
||||
// Load cos/sin values (coalesced access)
|
||||
int cos_sin_idx = section_idx * batch_size * seq_len * head_dim +
|
||||
batch_idx * seq_len * head_dim +
|
||||
seq_idx * head_dim + cos_sin_d;
|
||||
|
||||
T cos_val = cos[cos_sin_idx];
|
||||
T sin_val = sin[cos_sin_idx];
|
||||
|
||||
// Calculate tensor indices
|
||||
int tensor_idx = batch_idx * total_heads * seq_len * head_dim +
|
||||
actual_head_idx * seq_len * head_dim +
|
||||
seq_idx * head_dim + dim_idx;
|
||||
|
||||
// Get input value
|
||||
T input_val = is_q_head ? q[tensor_idx] : k[tensor_idx];
|
||||
|
||||
// Calculate rotate_half value
|
||||
int half_dim = head_dim / 2;
|
||||
int rotate_dim = (dim_idx < half_dim) ? (dim_idx + half_dim) : (dim_idx - half_dim);
|
||||
int rotate_tensor_idx = batch_idx * total_heads * seq_len * head_dim +
|
||||
actual_head_idx * seq_len * head_dim +
|
||||
seq_idx * head_dim + rotate_dim;
|
||||
|
||||
T rotate_val = is_q_head ? q[rotate_tensor_idx] : k[rotate_tensor_idx];
|
||||
if (dim_idx < half_dim)
|
||||
{
|
||||
rotate_val = -rotate_val; // First half: negate second half
|
||||
}
|
||||
|
||||
// Apply RoPE: output = input * cos + rotate_half(input) * sin
|
||||
T output_val = input_val * cos_val + rotate_val * sin_val;
|
||||
|
||||
// Store result
|
||||
if (is_q_head)
|
||||
{
|
||||
q_out[tensor_idx] = output_val;
|
||||
}
|
||||
else
|
||||
{
|
||||
k_out[tensor_idx] = output_val;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Optimized backward kernel template
|
||||
template <typename T>
|
||||
__global__ void multimodal_rope_backward_kernel(
|
||||
const T *__restrict__ grad_q_out, // [batch, q_heads, seq_len, head_dim]
|
||||
const T *__restrict__ grad_k_out, // [batch, kv_heads, seq_len, head_dim]
|
||||
const T *__restrict__ q, // [batch, q_heads, seq_len, head_dim]
|
||||
const T *__restrict__ k, // [batch, kv_heads, seq_len, head_dim]
|
||||
const T *__restrict__ cos, // [3, batch, seq_len, head_dim]
|
||||
const T *__restrict__ sin, // [3, batch, seq_len, head_dim]
|
||||
T *__restrict__ grad_q, // [batch, q_heads, seq_len, head_dim]
|
||||
T *__restrict__ grad_k, // [batch, kv_heads, seq_len, head_dim]
|
||||
const int *__restrict__ mrope_section_doubled,
|
||||
int batch_size,
|
||||
int q_heads,
|
||||
int kv_heads,
|
||||
int seq_len,
|
||||
int head_dim)
|
||||
{
|
||||
int batch_idx = blockIdx.x;
|
||||
int head_idx = blockIdx.y;
|
||||
int seq_paral_size = gridDim.z;
|
||||
int seq_paral_idx = blockIdx.z;
|
||||
|
||||
bool is_q_head = head_idx < q_heads;
|
||||
int actual_head_idx = is_q_head ? head_idx : (head_idx - q_heads);
|
||||
int total_heads = is_q_head ? q_heads : kv_heads;
|
||||
|
||||
int warp_id = threadIdx.x / 32;
|
||||
int lane_id = threadIdx.x % 32;
|
||||
int warps_per_block = blockDim.x / 32;
|
||||
|
||||
// Process tokens
|
||||
for (int seq_base = warp_id + seq_paral_idx * warps_per_block; seq_base < seq_len; seq_base += seq_paral_size * warps_per_block)
|
||||
{
|
||||
int seq_idx = seq_base;
|
||||
if (seq_idx >= seq_len)
|
||||
break;
|
||||
|
||||
// Process dimensions in chunks - much simpler now!
|
||||
constexpr int dims_per_lane = 4;
|
||||
for (int dim_batch = 0; dim_batch < (head_dim + 31) / 32; ++dim_batch)
|
||||
{
|
||||
int dim_start = dim_batch * 32 + lane_id;
|
||||
if (dim_start >= head_dim)
|
||||
break;
|
||||
|
||||
#pragma unroll
|
||||
for (int dim_offset = 0; dim_offset < dims_per_lane; ++dim_offset)
|
||||
{
|
||||
int dim_idx = dim_start + dim_offset * 32;
|
||||
if (dim_idx >= head_dim)
|
||||
break;
|
||||
|
||||
// Get section info for cos/sin reconstruction
|
||||
int section_idx;
|
||||
int cos_sin_d = dim_idx;
|
||||
if (dim_idx < mrope_section_doubled[0])
|
||||
{
|
||||
section_idx = 0;
|
||||
}
|
||||
else if (dim_idx < mrope_section_doubled[0] + mrope_section_doubled[1])
|
||||
{
|
||||
section_idx = 1;
|
||||
}
|
||||
else
|
||||
{
|
||||
section_idx = 2;
|
||||
}
|
||||
|
||||
// Global cos/sin index
|
||||
int cos_sin_idx = section_idx * batch_size * seq_len * head_dim +
|
||||
batch_idx * seq_len * head_dim +
|
||||
seq_idx * head_dim + cos_sin_d;
|
||||
|
||||
// Tensor index for current position
|
||||
int tensor_idx = batch_idx * total_heads * seq_len * head_dim +
|
||||
actual_head_idx * seq_len * head_dim +
|
||||
seq_idx * head_dim + dim_idx;
|
||||
|
||||
// Load values
|
||||
T cos_val = cos[cos_sin_idx];
|
||||
T sin_val = sin[cos_sin_idx];
|
||||
T grad_out_val = is_q_head ? grad_q_out[tensor_idx] : grad_k_out[tensor_idx];
|
||||
|
||||
// Calculate paired dimension for rotate_half
|
||||
int half_dim = head_dim / 2;
|
||||
int rotate_dim = (dim_idx < half_dim) ? (dim_idx + half_dim) : (dim_idx - half_dim);
|
||||
int rotate_tensor_idx = batch_idx * total_heads * seq_len * head_dim +
|
||||
actual_head_idx * seq_len * head_dim +
|
||||
seq_idx * head_dim + rotate_dim;
|
||||
|
||||
// Get gradient from the paired dimension
|
||||
T paired_grad_out = is_q_head ? grad_q_out[rotate_tensor_idx] : grad_k_out[rotate_tensor_idx];
|
||||
|
||||
// Get paired sin value
|
||||
int paired_section_idx;
|
||||
int paired_cos_sin_d = rotate_dim;
|
||||
if (rotate_dim < mrope_section_doubled[0])
|
||||
{
|
||||
paired_section_idx = 0;
|
||||
}
|
||||
else if (rotate_dim < mrope_section_doubled[0] + mrope_section_doubled[1])
|
||||
{
|
||||
paired_section_idx = 1;
|
||||
}
|
||||
else
|
||||
{
|
||||
paired_section_idx = 2;
|
||||
}
|
||||
|
||||
int paired_cos_sin_idx = paired_section_idx * batch_size * seq_len * head_dim +
|
||||
batch_idx * seq_len * head_dim +
|
||||
seq_idx * head_dim + paired_cos_sin_d;
|
||||
T paired_sin_val = sin[paired_cos_sin_idx];
|
||||
|
||||
// === Compute input gradients (the only thing we need!) ===
|
||||
|
||||
// Direct term: grad_input = grad_out * cos
|
||||
T grad_input_direct = grad_out_val * cos_val;
|
||||
|
||||
// Cross term from rotate_half
|
||||
T grad_input_from_rotate;
|
||||
if (dim_idx < half_dim)
|
||||
{
|
||||
// First half: gets contribution from second half (positive)
|
||||
grad_input_from_rotate = paired_grad_out * paired_sin_val;
|
||||
}
|
||||
else
|
||||
{
|
||||
// Second half: gets contribution from first half (negative due to rotate_half)
|
||||
grad_input_from_rotate = -paired_grad_out * paired_sin_val;
|
||||
}
|
||||
|
||||
// Total input gradient
|
||||
T total_grad_input = grad_input_direct + grad_input_from_rotate;
|
||||
|
||||
// Store input gradients (no atomic operations needed!)
|
||||
if (is_q_head)
|
||||
{
|
||||
grad_q[tensor_idx] = total_grad_input;
|
||||
}
|
||||
else
|
||||
{
|
||||
grad_k[tensor_idx] = total_grad_input;
|
||||
}
|
||||
|
||||
// No cos/sin gradient computation - they're not learnable parameters!
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Template-based host functions
|
||||
template <typename T>
|
||||
void launch_multimodal_rope_forward_template(
|
||||
const T *q, const T *k, const T *cos, const T *sin,
|
||||
T *q_out, T *k_out,
|
||||
const int *mrope_section_doubled,
|
||||
int batch_size, int q_heads, int kv_heads, int seq_len, int head_dim,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
static_assert(CudaTypeTraits<T>::is_supported, "Unsupported data type for multimodal RoPE");
|
||||
|
||||
// Grid: batch_size x (q_heads + kv_heads)
|
||||
dim3 grid(batch_size, q_heads + kv_heads, 8);
|
||||
|
||||
// Block: enough threads to handle seq_len with multiple warps
|
||||
int threads_per_block = min(512, ((seq_len + 3) / 4) * 32); // 4 warps max
|
||||
threads_per_block = ((threads_per_block + 31) / 32) * 32; // Round to warp size
|
||||
|
||||
multimodal_rope_forward_kernel<T><<<grid, threads_per_block, 0, stream>>>(
|
||||
q, k, cos, sin, q_out, k_out, mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
void launch_multimodal_rope_backward_template(
|
||||
const T *grad_q_out, const T *grad_k_out,
|
||||
const T *q, const T *k, const T *cos, const T *sin,
|
||||
T *grad_q, T *grad_k, const int *mrope_section_doubled,
|
||||
int batch_size, int q_heads, int kv_heads, int seq_len, int head_dim,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
static_assert(CudaTypeTraits<T>::is_supported, "Unsupported data type for multimodal RoPE");
|
||||
|
||||
dim3 grid(batch_size, q_heads + kv_heads, 8);
|
||||
int threads_per_block = min(512, ((seq_len + 3) / 4) * 32);
|
||||
threads_per_block = ((threads_per_block + 31) / 32) * 32;
|
||||
|
||||
multimodal_rope_backward_kernel<T><<<grid, threads_per_block, 0, stream>>>(
|
||||
grad_q_out, grad_k_out, q, k, cos, sin,
|
||||
grad_q, grad_k, mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim);
|
||||
}
|
||||
|
||||
// Enum for data type dispatch
|
||||
enum class DataType
|
||||
{
|
||||
FLOAT32,
|
||||
FLOAT16,
|
||||
BFLOAT16
|
||||
};
|
||||
|
||||
// Template dispatcher
|
||||
template <DataType DT>
|
||||
struct DataTypeDispatcher;
|
||||
|
||||
template <>
|
||||
struct DataTypeDispatcher<DataType::FLOAT32>
|
||||
{
|
||||
using type = float;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeDispatcher<DataType::FLOAT16>
|
||||
{
|
||||
using type = half;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct DataTypeDispatcher<DataType::BFLOAT16>
|
||||
{
|
||||
using type = __nv_bfloat16;
|
||||
};
|
||||
|
||||
// Type-safe host interface
|
||||
template <DataType DT>
|
||||
void launch_multimodal_rope_forward_typed(
|
||||
const void *q, const void *k, const void *cos, const void *sin,
|
||||
void *q_out, void *k_out,
|
||||
const int *mrope_section_doubled,
|
||||
int batch_size, int q_heads, int kv_heads, int seq_len, int head_dim,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
using T = typename DataTypeDispatcher<DT>::type;
|
||||
|
||||
launch_multimodal_rope_forward_template<T>(
|
||||
static_cast<const T *>(q),
|
||||
static_cast<const T *>(k),
|
||||
static_cast<const T *>(cos),
|
||||
static_cast<const T *>(sin),
|
||||
static_cast<T *>(q_out),
|
||||
static_cast<T *>(k_out),
|
||||
mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim,
|
||||
stream);
|
||||
}
|
||||
|
||||
template <DataType DT>
|
||||
void launch_multimodal_rope_backward_typed(
|
||||
const void *grad_q_out, const void *grad_k_out,
|
||||
const void *q, const void *k, const void *cos, const void *sin,
|
||||
void *grad_q, void *grad_k,
|
||||
const int *mrope_section_doubled,
|
||||
int batch_size, int q_heads, int kv_heads, int seq_len, int head_dim,
|
||||
cudaStream_t stream)
|
||||
{
|
||||
using T = typename DataTypeDispatcher<DT>::type;
|
||||
|
||||
launch_multimodal_rope_backward_template<T>(
|
||||
static_cast<const T *>(grad_q_out),
|
||||
static_cast<const T *>(grad_k_out),
|
||||
static_cast<const T *>(q),
|
||||
static_cast<const T *>(k),
|
||||
static_cast<const T *>(cos),
|
||||
static_cast<const T *>(sin),
|
||||
static_cast<T *>(grad_q),
|
||||
static_cast<T *>(grad_k),
|
||||
mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim,
|
||||
stream);
|
||||
}
|
||||
|
||||
void launch_multimodal_rope_forward(
|
||||
torch::Tensor q, torch::Tensor k, torch::Tensor cos, torch::Tensor sin,
|
||||
torch::Tensor q_out, torch::Tensor k_out,
|
||||
std::vector<int> mrope_section_doubled)
|
||||
{
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
|
||||
int batch_size = q.size(0);
|
||||
int q_heads = q.size(1);
|
||||
int seq_len = q.size(2);
|
||||
int head_dim = q.size(3);
|
||||
int kv_heads = k.size(1);
|
||||
|
||||
int data_type;
|
||||
if (q.scalar_type() == torch::kFloat32)
|
||||
{
|
||||
data_type = 0;
|
||||
}
|
||||
else if (q.scalar_type() == torch::kFloat16)
|
||||
{
|
||||
data_type = 1;
|
||||
}
|
||||
else if (q.scalar_type() == torch::kBFloat16)
|
||||
{
|
||||
data_type = 2;
|
||||
}
|
||||
int *d_mrope_section_doubled;
|
||||
cudaMalloc(&d_mrope_section_doubled, 3 * sizeof(int));
|
||||
cudaMemcpyAsync(d_mrope_section_doubled, mrope_section_doubled.data(), 3 * sizeof(int),
|
||||
cudaMemcpyHostToDevice, stream);
|
||||
|
||||
switch (data_type)
|
||||
{
|
||||
case 0: // float32
|
||||
launch_multimodal_rope_forward_typed<DataType::FLOAT32>(
|
||||
q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(),
|
||||
q_out.data_ptr(), k_out.data_ptr(), d_mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim, stream);
|
||||
break;
|
||||
case 1: // float16
|
||||
launch_multimodal_rope_forward_typed<DataType::FLOAT16>(
|
||||
q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(),
|
||||
q_out.data_ptr(), k_out.data_ptr(), d_mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim, stream);
|
||||
break;
|
||||
case 2: // bfloat16
|
||||
launch_multimodal_rope_forward_typed<DataType::BFLOAT16>(
|
||||
q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(),
|
||||
q_out.data_ptr(), k_out.data_ptr(), d_mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim, stream);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
void launch_multimodal_rope_backward(
|
||||
torch::Tensor grad_q_out, torch::Tensor grad_k_out,
|
||||
torch::Tensor q, torch::Tensor k, torch::Tensor cos, torch::Tensor sin,
|
||||
torch::Tensor grad_q, torch::Tensor grad_k,
|
||||
std::vector<int> mrope_section_doubled)
|
||||
{
|
||||
|
||||
cudaStream_t stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
|
||||
int batch_size = q.size(0);
|
||||
int q_heads = q.size(1);
|
||||
int seq_len = q.size(2);
|
||||
int head_dim = q.size(3);
|
||||
int kv_heads = k.size(1);
|
||||
|
||||
int data_type;
|
||||
if (q.scalar_type() == torch::kFloat32)
|
||||
{
|
||||
data_type = 0;
|
||||
}
|
||||
else if (q.scalar_type() == torch::kFloat16)
|
||||
{
|
||||
data_type = 1;
|
||||
}
|
||||
else if (q.scalar_type() == torch::kBFloat16)
|
||||
{
|
||||
data_type = 2;
|
||||
}
|
||||
int *d_mrope_section_doubled;
|
||||
cudaMalloc(&d_mrope_section_doubled, 3 * sizeof(int));
|
||||
cudaMemcpyAsync(d_mrope_section_doubled, mrope_section_doubled.data(), 3 * sizeof(int),
|
||||
cudaMemcpyHostToDevice, stream);
|
||||
|
||||
switch (data_type)
|
||||
{
|
||||
case 0: // float32
|
||||
launch_multimodal_rope_backward_typed<DataType::FLOAT32>(
|
||||
grad_q_out.data_ptr(), grad_k_out.data_ptr(),
|
||||
q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(),
|
||||
grad_q.data_ptr(), grad_k.data_ptr(), d_mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim, stream);
|
||||
break;
|
||||
case 1: // float16
|
||||
launch_multimodal_rope_backward_typed<DataType::FLOAT16>(
|
||||
grad_q_out.data_ptr(), grad_k_out.data_ptr(),
|
||||
q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(),
|
||||
grad_q.data_ptr(), grad_k.data_ptr(), d_mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim, stream);
|
||||
break;
|
||||
case 2: // bfloat16
|
||||
launch_multimodal_rope_backward_typed<DataType::BFLOAT16>(
|
||||
grad_q_out.data_ptr(), grad_k_out.data_ptr(),
|
||||
q.data_ptr(), k.data_ptr(), cos.data_ptr(), sin.data_ptr(),
|
||||
grad_q.data_ptr(), grad_k.data_ptr(), d_mrope_section_doubled,
|
||||
batch_size, q_heads, kv_heads, seq_len, head_dim, stream);
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user