huangyanpei/VLA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Init

Browse Source
This commit is contained in:
Starrick
2025-09-07 14:59:17 +08:00
commit 24dbdbd24b
40 changed files with 10754 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/README.md
+23
View File
@@ -0,0 +1,23 @@
# Fusion Operators (CSRC)
High-performance CUDA kernels for accelerating model training.
## Operators
### Asymmetric Dual Expert GEMM
- `asym_dual_gmm`: Simultaneous matrix multiplication for two experts
- Supports all transpose combinations (NN, TN, NT, TT)
### Token Permutation
- `permute`: Token permutation for MoE routing
- `unpermute`: Token recovery after expert computation
- `unpermute_bwd`: Backward pass for token recovery
### Multimodal RoPE
- `rope`: Rotary Position Embedding forward pass
- `rope_bwd`: RoPE backward pass
- Support for multimodal inputs with configurable sections
## Acknowledgments
The `permute` and `unpermute` operators are adapted from [fanshiqing/grouped_gemm](https://github.com/fanshiqing/grouped_gemm). Thanks for their open-source contributions.
csrc/dual_asym_grouped_gemm.cu
+366
View File
@@ -0,0 +1,366 @@
#include "dual_asym_grouped_gemm.h"
#include <ATen/cuda/CUDAContext.h>
#include <c10/util/BFloat16.h>
#include <c10/cuda/CUDAStream.h>
#include <torch/extension.h>
#include "cutlass/bfloat16.h"
#include "cutlass/complex.h"
#include "cutlass/gemm/kernel/gemm_grouped.h"
#include "cutlass/gemm/kernel/default_gemm_grouped.h"
#include "cutlass/gemm/device/gemm_grouped.h"
#define NUM_STREAM 4
#define CUDA_CALL(code) \
do \
{ \
cudaError_t status = code; \
std::string err = cudaGetErrorString(status); \
TORCH_CHECK(status == cudaSuccess, err); \
} while (0)
#define CUBLAS_CALL(code) \
do \
{ \
cublasStatus_t status = code; \
TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, "CuBLAS Error"); \
} while (0)
#define GROUPED_GEMM_STRINGIFY_HELPER(x) #x
#define GROUPED_GEMM_STRINGIFY(x) \
GROUPED_GEMM_STRINGIFY_HELPER(x)
template <typename T>
torch::Tensor CopyToDevice(const std::vector<T> &x, const torch::Device &device)
{
size_t bytes = x.size() * sizeof(T);
auto options = torch::TensorOptions().dtype(torch::kInt8).device(device);
torch::Tensor out = torch::empty(bytes, options);
CUDA_CALL(cudaMemcpyAsync(out.data_ptr(),
x.data(), bytes,
cudaMemcpyHostToDevice,
c10::cuda::getCurrentCUDAStream()));
return out;
}
using DualExpertGemmKernelNN = typename cutlass::gemm::kernel::DefaultGemmGrouped<
::cutlass::bfloat16_t,
::cutlass::layout::RowMajor,
::cutlass::ComplexTransform::kNone,
8,
::cutlass::bfloat16_t,
::cutlass::layout::RowMajor,
::cutlass::ComplexTransform::kNone,
8,
::cutlass::bfloat16_t,
::cutlass::layout::RowMajor,
float,
::cutlass::arch::OpClassTensorOp,
::cutlass::arch::Sm80,
::cutlass::gemm::GemmShape<128, 128, 32>,
::cutlass::gemm::GemmShape<64, 64, 32>,
::cutlass::gemm::GemmShape<16, 8, 16>,
::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>,
::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
4>::GemmKernel;
using DualExpertGemmKernelTN = typename cutlass::gemm::kernel::DefaultGemmGrouped<
::cutlass::bfloat16_t,
::cutlass::layout::ColumnMajor,
::cutlass::ComplexTransform::kNone,
8,
::cutlass::bfloat16_t,
::cutlass::layout::RowMajor,
::cutlass::ComplexTransform::kNone,
8,
::cutlass::bfloat16_t,
::cutlass::layout::RowMajor,
float,
::cutlass::arch::OpClassTensorOp,
::cutlass::arch::Sm80,
::cutlass::gemm::GemmShape<128, 128, 32>,
::cutlass::gemm::GemmShape<64, 64, 32>,
::cutlass::gemm::GemmShape<16, 8, 16>,
::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>,
::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
4>::GemmKernel;
using DualExpertGemmKernelNT = typename cutlass::gemm::kernel::DefaultGemmGrouped<
::cutlass::bfloat16_t,
::cutlass::layout::RowMajor,
::cutlass::ComplexTransform::kNone,
8,
::cutlass::bfloat16_t,
::cutlass::layout::ColumnMajor,
::cutlass::ComplexTransform::kNone,
8,
::cutlass::bfloat16_t,
::cutlass::layout::RowMajor,
float,
::cutlass::arch::OpClassTensorOp,
::cutlass::arch::Sm80,
::cutlass::gemm::GemmShape<128, 128, 32>,
::cutlass::gemm::GemmShape<64, 64, 32>,
::cutlass::gemm::GemmShape<16, 8, 16>,
::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>,
::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
4>::GemmKernel;
using DualExpertGemmKernelTT = typename cutlass::gemm::kernel::DefaultGemmGrouped<
::cutlass::bfloat16_t,
::cutlass::layout::ColumnMajor,
::cutlass::ComplexTransform::kNone,
8,
::cutlass::bfloat16_t,
::cutlass::layout::ColumnMajor,
::cutlass::ComplexTransform::kNone,
8,
::cutlass::bfloat16_t,
::cutlass::layout::RowMajor,
float,
::cutlass::arch::OpClassTensorOp,
::cutlass::arch::Sm80,
::cutlass::gemm::GemmShape<128, 128, 32>,
::cutlass::gemm::GemmShape<64, 64, 32>,
::cutlass::gemm::GemmShape<16, 8, 16>,
::cutlass::epilogue::thread::LinearCombination<::cutlass::bfloat16_t, 8, float, float>,
::cutlass::gemm::threadblock::GemmBatchedIdentityThreadblockSwizzle,
4>::GemmKernel;
using DualExpertGemmNN = ::cutlass::gemm::device::GemmGrouped<DualExpertGemmKernelNN>;
using DualExpertGemmTN = ::cutlass::gemm::device::GemmGrouped<DualExpertGemmKernelTN>;
using DualExpertGemmNT = ::cutlass::gemm::device::GemmGrouped<DualExpertGemmKernelNT>;
using DualExpertGemmTT = ::cutlass::gemm::device::GemmGrouped<DualExpertGemmKernelTT>;
template <typename Gemm>
typename Gemm::Arguments MakeAsymmetricArgumentsSeparated(
torch::Tensor input_expert0,
torch::Tensor input_expert1,
torch::Tensor weight_expert0,
torch::Tensor weight_expert1,
torch::Tensor output_expert0,
torch::Tensor output_expert1)
{
TORCH_CHECK(input_expert0.dim() == 2 && input_expert1.dim() == 2,
"Input tensors must be 2D");
TORCH_CHECK(weight_expert0.dim() == 2 && weight_expert1.dim() == 2,
"Weight tensors must be 2D");
TORCH_CHECK(output_expert0.dim() == 2 && output_expert1.dim() == 2,
"Output tensors must be 2D");
using LayoutA = typename Gemm::LayoutA;
using LayoutB = typename Gemm::LayoutB;
using LayoutC = typename Gemm::LayoutC;
bool a_is_column_major = std::is_same_v<LayoutA, cutlass::layout::ColumnMajor>;
bool b_is_column_major = std::is_same_v<LayoutB, cutlass::layout::ColumnMajor>;
int64_t m0, k0, n0;
if (a_is_column_major)
{
m0 = input_expert0.size(1);
k0 = input_expert0.size(0);
}
else
{
m0 = input_expert0.size(0);
k0 = input_expert0.size(1);
}
if (b_is_column_major)
{
n0 = weight_expert0.size(0);
TORCH_CHECK(weight_expert0.size(1) == k0, "Expert 0: k dimensions must match");
}
else
{
n0 = weight_expert0.size(1);
TORCH_CHECK(weight_expert0.size(0) == k0, "Expert 0: k dimensions must match");
}
int64_t m1, k1, n1;
if (a_is_column_major)
{
m1 = input_expert1.size(1);
k1 = input_expert1.size(0);
}
else
{
m1 = input_expert1.size(0);
k1 = input_expert1.size(1);
}
if (b_is_column_major)
{
n1 = weight_expert1.size(0);
TORCH_CHECK(weight_expert1.size(1) == k1, "Expert 1: k dimensions must match");
}
else
{
n1 = weight_expert1.size(1);
TORCH_CHECK(weight_expert1.size(0) == k1, "Expert 1: k dimensions must match");
}
std::vector<cutlass::gemm::GemmCoord> problem_sizes_host(2);
problem_sizes_host[0] = cutlass::gemm::GemmCoord(m0, n0, k0);
problem_sizes_host[1] = cutlass::gemm::GemmCoord(m1, n1, k1);
int64_t num_experts = 2;
int threadblock_count = Gemm::sufficient(problem_sizes_host.data(), num_experts);
if (!threadblock_count)
{
TORCH_CHECK(false, "Dual Expert Grouped GEMM execution not possible with HW");
}
std::vector<int64_t> lda_host(num_experts);
std::vector<int64_t> ldb_host(num_experts);
std::vector<int64_t> ldc_host(num_experts);
using ElementA = typename Gemm::ElementA;
using ElementB = typename Gemm::ElementB;
using ElementC = typename Gemm::ElementC;
std::vector<ElementA *> ptr_a_host(num_experts);
std::vector<ElementB *> ptr_b_host(num_experts);
std::vector<ElementC *> ptr_c_host(num_experts);
auto problem_0 = problem_sizes_host[0];
lda_host[0] = LayoutA::packed({problem_0.m(), problem_0.k()}).stride(0);
ldb_host[0] = LayoutB::packed({problem_0.k(), problem_0.n()}).stride(0);
ldc_host[0] = LayoutC::packed({problem_0.m(), problem_0.n()}).stride(0);
ptr_a_host[0] = (ElementA *)input_expert0.data_ptr();
ptr_b_host[0] = (ElementB *)weight_expert0.data_ptr();
ptr_c_host[0] = (ElementC *)output_expert0.data_ptr();
auto problem_1 = problem_sizes_host[1];
lda_host[1] = LayoutA::packed({problem_1.m(), problem_1.k()}).stride(0);
ldb_host[1] = LayoutB::packed({problem_1.k(), problem_1.n()}).stride(0);
ldc_host[1] = LayoutC::packed({problem_1.m(), problem_1.n()}).stride(0);
ptr_a_host[1] = (ElementA *)input_expert1.data_ptr();
ptr_b_host[1] = (ElementB *)weight_expert1.data_ptr();
ptr_c_host[1] = (ElementC *)output_expert1.data_ptr();
torch::Tensor lda = CopyToDevice(lda_host, input_expert0.device());
torch::Tensor ldb = CopyToDevice(ldb_host, input_expert0.device());
torch::Tensor ldc = CopyToDevice(ldc_host, input_expert0.device());
torch::Tensor ptr_a = CopyToDevice(ptr_a_host, input_expert0.device());
torch::Tensor ptr_b = CopyToDevice(ptr_b_host, input_expert0.device());
torch::Tensor ptr_c = CopyToDevice(ptr_c_host, input_expert0.device());
torch::Tensor problem_sizes = CopyToDevice(problem_sizes_host, input_expert0.device());
typename Gemm::EpilogueOutputOp::Params epilogue_op(/*alpha=*/1.0f, /*beta=*/0.0f);
typename Gemm::Arguments arguments(
(cutlass::gemm::GemmCoord *)problem_sizes.data_ptr(),
(int)num_experts,
(int)threadblock_count,
epilogue_op,
(ElementA **)ptr_a.data_ptr(),
(ElementB **)ptr_b.data_ptr(),
(ElementC **)ptr_c.data_ptr(),
(ElementC **)ptr_c.data_ptr(),
(int64_t *)lda.data_ptr(),
(int64_t *)ldb.data_ptr(),
(int64_t *)ldc.data_ptr(),
(int64_t *)ldc.data_ptr(),
(cutlass::gemm::GemmCoord *)problem_sizes_host.data());
return arguments;
}
template <typename Gemm>
void executeDualExpertGemm(
torch::Tensor input_expert0,
torch::Tensor input_expert1,
torch::Tensor weight_expert0,
torch::Tensor weight_expert1,
torch::Tensor output_expert0,
torch::Tensor output_expert1)
{
Gemm gemm;
auto arguments = MakeAsymmetricArgumentsSeparated<Gemm>(
input_expert0, input_expert1,
weight_expert0, weight_expert1,
output_expert0, output_expert1);
int64_t workspace_size = gemm.get_workspace_size(arguments);
auto options = torch::TensorOptions().dtype(torch::kInt8).device(input_expert0.device());
torch::Tensor workspace = torch::empty(workspace_size, options);
if (gemm.initialize(arguments, workspace.data_ptr()) != cutlass::Status::kSuccess)
{
TORCH_CHECK(false, "Failed to initialize CUTLASS Asymmetric Dual Expert GEMM");
}
if (gemm.run(c10::cuda::getCurrentCUDAStream()) != cutlass::Status::kSuccess)
{
TORCH_CHECK(false, "Failed to run CUTLASS Asymmetric Dual Expert GEMM");
}
}
void AsymmetricDualExpertGemm(
torch::Tensor input_expert0,
torch::Tensor input_expert1,
torch::Tensor weight_expert0,
torch::Tensor weight_expert1,
torch::Tensor output_expert0,
torch::Tensor output_expert1,
bool trans_a, bool trans_b)
{
TORCH_CHECK(input_expert0.device() == input_expert1.device() &&
input_expert0.device() == weight_expert0.device() &&
input_expert0.device() == weight_expert1.device() &&
input_expert0.device() == output_expert0.device() &&
input_expert0.device() == output_expert1.device(),
"All tensors must be on the same device");
TORCH_CHECK(input_expert0.device().is_cuda(),
"All tensors must be on CUDA device for CUTLASS GEMM");
TORCH_CHECK(input_expert0.dtype() == torch::kBFloat16,
"All tensors must be BFloat16 for this kernel");
torch::Tensor input0_contiguous = input_expert0.contiguous();
torch::Tensor input1_contiguous = input_expert1.contiguous();
torch::Tensor weight0_contiguous = weight_expert0.contiguous();
torch::Tensor weight1_contiguous = weight_expert1.contiguous();
torch::Tensor output0_contiguous = output_expert0.contiguous();
torch::Tensor output1_contiguous = output_expert1.contiguous();
if (!trans_a && !trans_b)
{
using Gemm = DualExpertGemmNN;
executeDualExpertGemm<Gemm>(input0_contiguous, input1_contiguous,
weight0_contiguous, weight1_contiguous,
output0_contiguous, output1_contiguous);
}
else if (trans_a && !trans_b)
{
using Gemm = DualExpertGemmTN;
executeDualExpertGemm<Gemm>(input0_contiguous, input1_contiguous,
weight0_contiguous, weight1_contiguous,
output0_contiguous, output1_contiguous);
}
else if (!trans_a && trans_b)
{
using Gemm = DualExpertGemmNT;
executeDualExpertGemm<Gemm>(input0_contiguous, input1_contiguous,
weight0_contiguous, weight1_contiguous,
output0_contiguous, output1_contiguous);
}
else
{
using Gemm = DualExpertGemmTT;
executeDualExpertGemm<Gemm>(input0_contiguous, input1_contiguous,
weight0_contiguous, weight1_contiguous,
output0_contiguous, output1_contiguous);
}
}
csrc/dual_asym_grouped_gemm.h
+10
View File
@@ -0,0 +1,10 @@
#include <torch/extension.h>
void AsymmetricDualExpertGemm(
torch::Tensor input_expert0,
torch::Tensor input_expert1,
torch::Tensor weight_expert0,
torch::Tensor weight_expert1,
torch::Tensor output_expert0,
torch::Tensor output_expert1,
bool trans_a, bool trans_b);
csrc/ops.cu
+16
View File
@@ -0,0 +1,16 @@
#include "dual_asym_grouped_gemm.h"
#include "permute.h"
#include "rope.h"
#include <torch/extension.h>
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("asym_dual_gmm", &AsymmetricDualExpertGemm, "Asymmetric Dual Expert Grouped GEMM.");
m.def("permute", &moe_permute_topK_op, "Token permutation kernel");
m.def("unpermute", &moe_recover_topK_op, "Token un-permutation kernel");
m.def("unpermute_bwd", &moe_recover_topK_bwd_op, "Token un-permutation backward kernel");
m.def("rope", &launch_multimodal_rope_forward, "Multimodal RoPE forward kernel");
m.def("rope_bwd", &launch_multimodal_rope_backward, "Multimodal RoPE backward kernel");
}
csrc/permute.cu
+942
View File
@@ -0,0 +1,942 @@
/*************************************************************************
* Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
*
* See LICENSE for license information.
************************************************************************/
#include "permute.h"
#include <torch/torch.h>
#include <cub/cub.cuh>
#include <cuda_bf16.h>
#include "cuda_runtime.h"
#include "device_launch_parameters.h"
#include "ATen/cuda/CUDAContext.h"
#include "cutlass/arch/memory.h"
#include "cutlass/arch/cache_operation.h"
#include "cutlass/array.h"
#include "cutlass/numeric_conversion.h"
using torch::Tensor;
template <typename T>
inline T *get_ptr(torch::Tensor &t)
{
return reinterpret_cast<T *>(t.data_ptr());
}
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Top K
//
/////////////////////////////////////////////////////////////////////////////////////////////////
static __global__ void moe_permute_topK_row_map(
const int *sorted_row_id,
int *row_id_map,
const int num_rows,
const int num_topK,
const int num_out_tokens)
{
// Each block corresponds to one source token
// row_id_map[num_topK][num_rows]
const int bid = blockIdx.x;
const int tid = threadIdx.x;
const int idx = bid * blockDim.x + tid;
if (idx >= num_rows * num_topK)
return;
int source_row = sorted_row_id[idx];
int source_token_id = source_row / num_topK;
int source_topK_id = source_row % num_topK;
if (idx >= num_out_tokens)
{
row_id_map[source_topK_id * num_rows + source_token_id] = -1;
}
else
{
row_id_map[source_topK_id * num_rows + source_token_id] = idx;
}
}
template <typename T, typename TCompute, int kElementsPerAccess, bool hasProb>
__global__ void moe_recover_topK_kernel(const T *input,
T *unpermuted_output,
const int *row_id_map,
const float *prob,
const int num_rows,
const int num_topK,
const int num_cols)
{
extern __shared__ int8_t s_mem[];
TCompute *s_prob = reinterpret_cast<TCompute *>(s_mem);
using FragmentLoadStore = cutlass::Array<T, kElementsPerAccess>;
using FragmentCompute = cutlass::Array<TCompute, kElementsPerAccess>;
cutlass::NumericArrayConverter<TCompute, T, kElementsPerAccess> src_converter;
cutlass::NumericArrayConverter<T, TCompute, kElementsPerAccess> dst_converter;
// each block corresponds to one source token
const int source_token = blockIdx.x;
const int tid = threadIdx.x;
if (hasProb)
{
for (int i = tid; i < num_topK; i += blockDim.x * blockDim.y)
{
s_prob[i] = TCompute(prob[source_token * num_topK + i]);
}
__syncthreads();
}
for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess)
{
FragmentLoadStore frag_load_store;
FragmentCompute frag_elem;
FragmentCompute frag_sum;
int source_row = row_id_map[source_token];
if (source_row != -1)
{
const T *source_row_ptr = input + source_row * num_cols;
cutlass::arch::global_load<FragmentLoadStore, sizeof(FragmentLoadStore), cutlass::arch::CacheOperation::LastUse>(
frag_load_store, (source_row_ptr + i), true);
frag_sum = src_converter(frag_load_store);
if (hasProb)
{
frag_sum = frag_sum * s_prob[0];
}
}
else
{
frag_sum.clear();
}
for (int k = 1; k < num_topK; k++)
{
source_row = row_id_map[k * num_rows + source_token];
if (source_row == -1)
continue;
const T *source_row_ptr = input + source_row * num_cols;
cutlass::arch::global_load<FragmentLoadStore, sizeof(FragmentLoadStore), cutlass::arch::CacheOperation::LastUse>(
frag_load_store, (source_row_ptr + i), true);
frag_elem = src_converter(frag_load_store);
if (hasProb)
{
frag_elem = frag_elem * s_prob[k];
}
for (int e = 0; e < kElementsPerAccess; e++)
{
frag_sum.at(e) = frag_sum.at(e) + frag_elem.at(e);
}
}
T *dest_row_ptr = unpermuted_output + source_token * num_cols;
frag_load_store = dst_converter(frag_sum);
*(float4 *)(dest_row_ptr + i) = *(float4 *)(frag_load_store.data());
}
}
template <typename T,
typename TCompute,
int kElementsPerAccess,
int topKTile,
bool hasProb>
__global__ void moe_permute_topK_kernel(const T *input_bwd,
const T *input_fwd,
T *act_grad,
const float *prob,
float *prob_grad,
const int *row_id_map,
const int num_rows,
const int num_topK,
const int num_cols)
{
extern __shared__ int8_t s_mem[];
TCompute *s_prob = reinterpret_cast<TCompute *>(s_mem);
using FragmentLoadStore = cutlass::Array<T, kElementsPerAccess>;
using FragmentCompute = cutlass::Array<TCompute, kElementsPerAccess>;
cutlass::NumericArrayConverter<TCompute, T, kElementsPerAccess> src_converter;
cutlass::NumericArrayConverter<T, TCompute, kElementsPerAccess> dst_converter;
const int source_token = blockIdx.x;
const int tid = threadIdx.x;
if (hasProb)
{
for (int i = tid; i < num_topK; i += blockDim.x)
{
s_prob[i] = TCompute(prob[source_token * num_topK + i]);
}
__syncthreads();
}
float accum[topKTile] = {0.0f};
FragmentLoadStore frag_load_store;
const T *source_row_ptr = input_bwd + source_token * num_cols;
for (int i = tid * kElementsPerAccess; i < num_cols; i += blockDim.x * kElementsPerAccess)
{
cutlass::arch::global_load<FragmentLoadStore, sizeof(FragmentLoadStore), cutlass::arch::CacheOperation::LastUse>(
frag_load_store, (source_row_ptr + i), true);
FragmentCompute frag_src = src_converter(frag_load_store);
int index = source_token;
for (int k = 0; k < topKTile; k++)
{
if (k == num_topK) break;
int dest_row = row_id_map[index];
index += num_rows;
if (dest_row == -1)
continue;
if (hasProb)
{
frag_load_store = dst_converter(frag_src * s_prob[k]);
}
else
{
frag_load_store = dst_converter(frag_src);
}
T *dest_row_ptr = act_grad + dest_row * num_cols;
*(float4 *)(dest_row_ptr + i) = *(float4 *)(frag_load_store.data());
if (hasProb)
{
const T *input_fwd_ptr = input_fwd + dest_row * num_cols;
cutlass::arch::global_load<FragmentLoadStore, sizeof(FragmentLoadStore), cutlass::arch::CacheOperation::LastUse>(
frag_load_store, (input_fwd_ptr + i), true);
FragmentCompute frag_input_fwd = src_converter(frag_load_store);
for (int e = 0; e < kElementsPerAccess; e++)
{
accum[k] += float(frag_src.at(e) * frag_input_fwd.at(e));
}
}
}
}
if (hasProb)
{
for (int k = 0; k < topKTile; k++)
{
if (k == num_topK) break;
for (int mask = 16; mask > 0; mask /= 2)
{
accum[k] = accum[k] + __shfl_xor_sync(0xffffffff, accum[k], mask, 32);
}
}
if (tid == 0)
{
for (int k = 0; k < topKTile; k++)
{
if (k == num_topK) break;
prob_grad[source_token * num_topK + k] = accum[k];
}
}
}
}
template <typename T, typename TCompute, bool FWD, int kElementsPerAccess>
void moe_permute_topK_kernel_launcher(
const T *input,
T *output,
const int *sorted_row_id,
int *row_id_map,
const float *prob,
const int num_rows,
const int num_topK,
const int num_cols,
const int num_out_tokens,
cudaStream_t stream,
float *prob_grad = nullptr,
const T *input_fwd = nullptr)
{
if (FWD)
{
if (prob_grad == nullptr)
{
// permute_topK fwd
int threads = 64;
int blocks = (num_rows * num_topK + threads - 1) / threads;
moe_permute_topK_row_map<<<blocks, threads, 0, stream>>>(
sorted_row_id,
row_id_map,
num_rows,
num_topK,
num_out_tokens);
blocks = num_rows;
threads = std::min(num_cols / kElementsPerAccess, 1024);
moe_permute_topK_kernel<T, T, kElementsPerAccess, 128, false><<<blocks, threads, 0, stream>>>(
input,
nullptr,
output,
nullptr,
nullptr,
row_id_map,
num_rows,
num_topK,
num_cols);
}
else
{
// unpermute_topK bwd
int blocks = num_rows;
int threads = 32;
size_t smem_bytes = num_topK * sizeof(TCompute);
if (num_topK == 1)
{
moe_permute_topK_kernel<T, T, kElementsPerAccess, 1, false><<<blocks, threads, 0, stream>>>(
input,
input_fwd,
output,
prob,
prob_grad,
row_id_map,
num_rows,
num_topK,
num_cols);
}
else if (num_topK <= 8)
{
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 8, true><<<blocks, threads, smem_bytes, stream>>>(
input,
input_fwd,
output,
prob,
prob_grad,
row_id_map,
num_rows,
num_topK,
num_cols);
}
else if (num_topK <= 16)
{
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 16, true><<<blocks, threads, smem_bytes, stream>>>(
input,
input_fwd,
output,
prob,
prob_grad,
row_id_map,
num_rows,
num_topK,
num_cols);
}
else if (num_topK <= 32)
{
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 32, true><<<blocks, threads, smem_bytes, stream>>>(
input,
input_fwd,
output,
prob,
prob_grad,
row_id_map,
num_rows,
num_topK,
num_cols);
}
else if (num_topK <= 64)
{
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 64, true><<<blocks, threads, smem_bytes, stream>>>(
input,
input_fwd,
output,
prob,
prob_grad,
row_id_map,
num_rows,
num_topK,
num_cols);
}
else if (num_topK <= 128)
{
moe_permute_topK_kernel<T, TCompute, kElementsPerAccess, 128, true><<<blocks, threads, smem_bytes, stream>>>(
input,
input_fwd,
output,
prob,
prob_grad,
row_id_map,
num_rows,
num_topK,
num_cols);
}
else
{
throw std::runtime_error("num_topK cannot exceed 128.");
}
}
}
else
{
int blocks = num_rows;
int threads = std::min(num_cols / kElementsPerAccess, 1024);
size_t smem_bytes = num_topK * sizeof(TCompute);
if (num_topK == 1)
{
// permute_topK bwd with topK==1
moe_recover_topK_kernel<T, T, kElementsPerAccess, false><<<blocks, threads, smem_bytes, stream>>>(
input,
output,
row_id_map,
prob,
num_rows,
num_topK,
num_cols);
}
else if (prob == nullptr)
{
// permute_topK bwd
moe_recover_topK_kernel<T, TCompute, kElementsPerAccess, false><<<blocks, threads, smem_bytes, stream>>>(
input,
output,
row_id_map,
prob,
num_rows,
num_topK,
num_cols);
}
else
{
// unpermute_topK fwd
moe_recover_topK_kernel<T, TCompute, kElementsPerAccess, true><<<blocks, threads, smem_bytes, stream>>>(
input,
output,
row_id_map,
prob,
num_rows,
num_topK,
num_cols);
}
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Permute_topK OP
//
/////////////////////////////////////////////////////////////////////////////////////////////////
std::tuple<Tensor, Tensor, std::vector<Tensor>> moe_permute_topK_op(
Tensor input,
Tensor indices,
int64_t num_out_tokens,
std::vector<Tensor> workspace,
int64_t max_expanded_token_num)
{
const int num_tokens = input.size(0);
const int num_cols = input.size(1);
const int num_topK = indices.size(1);
// initialize the workspace on the first run
if (workspace.empty()) {
auto options = torch::TensorOptions().dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false);
Tensor sorted_indices = torch::empty(max_expanded_token_num, options);
Tensor row_id = torch::range(0, max_expanded_token_num - 1, 1, options);
Tensor sorted_row_id =
torch::empty(max_expanded_token_num, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
size_t temp_storage_bytes = 0;
int *temp_ptr = nullptr;
cub::DeviceRadixSort::SortPairs(nullptr, temp_storage_bytes,
temp_ptr, temp_ptr,
temp_ptr, temp_ptr, max_expanded_token_num);
Tensor temp_storage =
torch::empty(temp_storage_bytes, torch::dtype(torch::kInt8).device(torch::kCUDA).requires_grad(false));
workspace.push_back(sorted_indices);
workspace.push_back(row_id);
workspace.push_back(sorted_row_id);
workspace.push_back(temp_storage);
}
int *indices_ptr = get_ptr<int>(indices);
int *sorted_indices_ptr = get_ptr<int>(workspace[0]);
int *row_id_ptr = get_ptr<int>(workspace[1]);
int *sorted_row_id_ptr = get_ptr<int>(workspace[2]);
void *d_temp_storage = get_ptr<void>(workspace[3]);
size_t temp_storage_bytes = std::numeric_limits<size_t>::max();
cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes,
indices_ptr, sorted_indices_ptr,
row_id_ptr, sorted_row_id_ptr, num_tokens * num_topK);
// activations type
const at::ScalarType _st = input.scalar_type();
// Output buffer alloc
num_out_tokens = (num_out_tokens > 0) ? num_out_tokens : num_tokens * num_topK;
Tensor permuted_output =
torch::empty({num_out_tokens, num_cols}, torch::dtype(_st).device(torch::kCUDA).requires_grad(false));
Tensor row_id_map =
torch::empty({num_tokens * num_topK}, torch::dtype(torch::kInt32).device(torch::kCUDA).requires_grad(false));
int *row_id_map_ptr = get_ptr<int>(row_id_map);
auto stream = at::cuda::getCurrentCUDAStream().stream();
switch (_st)
{
case at::ScalarType::Float:
{
using dType = float;
using dTypeCompute = float;
dType *input_ptr = get_ptr<dType>(input);
dType *permuted_output_ptr = get_ptr<dType>(permuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 4>(
input_ptr,
permuted_output_ptr,
sorted_row_id_ptr,
row_id_map_ptr,
nullptr,
num_tokens,
num_topK,
num_cols,
num_out_tokens,
stream);
break;
}
case at::ScalarType::Half:
{
using dType = cutlass::half_t;
using dTypeCompute = cutlass::half_t;
dType *input_ptr = get_ptr<dType>(input);
dType *permuted_output_ptr = get_ptr<dType>(permuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 8>(
input_ptr,
permuted_output_ptr,
sorted_row_id_ptr,
row_id_map_ptr,
nullptr,
num_tokens,
num_topK,
num_cols,
num_out_tokens,
stream);
break;
}
#ifdef ENABLE_BF16
case at::ScalarType::BFloat16:
{
using dType = cutlass::bfloat16_t;
using dTypeCompute = cutlass::bfloat16_t;
dType *input_ptr = get_ptr<dType>(input);
dType *permuted_output_ptr = get_ptr<dType>(permuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 8>(
input_ptr,
permuted_output_ptr,
sorted_row_id_ptr,
row_id_map_ptr,
nullptr,
num_tokens,
num_topK,
num_cols,
num_out_tokens,
stream);
break;
}
#endif
#ifdef ENABLE_FP8
case at::ScalarType::Float8_e5m2:
{
using dType = cutlass::float_e5m2_t;
using dTypeCompute = cutlass::half_t;
dType *input_ptr = get_ptr<dType>(input);
dType *permuted_output_ptr = get_ptr<dType>(permuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 16>(
input_ptr,
permuted_output_ptr,
sorted_row_id_ptr,
row_id_map_ptr,
nullptr,
num_tokens,
num_topK,
num_cols,
num_out_tokens,
stream);
break;
}
case at::ScalarType::Float8_e4m3fn:
{
using dType = cutlass::float_e4m3_t;
using dTypeCompute = cutlass::half_t;
dType *input_ptr = get_ptr<dType>(input);
dType *permuted_output_ptr = get_ptr<dType>(permuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 16>(
input_ptr,
permuted_output_ptr,
sorted_row_id_ptr,
row_id_map_ptr,
nullptr,
num_tokens,
num_topK,
num_cols,
num_out_tokens,
stream);
break;
}
#endif
default:
throw std::runtime_error("Wrong activation tensor type.");
}
return std::make_tuple(permuted_output, row_id_map, workspace);
}
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// Unpermute_topK OP
//
/////////////////////////////////////////////////////////////////////////////////////////////////
Tensor moe_recover_topK_op(
Tensor input,
Tensor row_id_map,
Tensor prob,
int64_t num_tokens,
int64_t num_topK)
{
const int num_cols = input.size(1);
// activations type
const at::ScalarType _st = input.scalar_type();
// Output buffer alloc
Tensor unpermuted_output =
torch::empty({num_tokens, num_cols}, torch::dtype(_st).device(torch::kCUDA).requires_grad(false));
int *row_id_map_ptr = get_ptr<int>(row_id_map);
float *prob_ptr = (prob.defined()) ? get_ptr<float>(prob) : nullptr;
auto stream = at::cuda::getCurrentCUDAStream().stream();
switch (_st)
{
case at::ScalarType::Float:
{
using dType = float;
using dTypeCompute = float;
dType *input_ptr = get_ptr<dType>(input);
dType *unpermuted_output_ptr = get_ptr<dType>(unpermuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, false, 4>(
input_ptr,
unpermuted_output_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream);
break;
}
case at::ScalarType::Half:
{
using dType = cutlass::half_t;
using dTypeCompute = cutlass::half_t;
dType *input_ptr = get_ptr<dType>(input);
dType *unpermuted_output_ptr = get_ptr<dType>(unpermuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, false, 8>(
input_ptr,
unpermuted_output_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream);
break;
}
#ifdef ENABLE_BF16
case at::ScalarType::BFloat16:
{
using dType = cutlass::bfloat16_t;
using dTypeCompute = cutlass::bfloat16_t;
dType *input_ptr = get_ptr<dType>(input);
dType *unpermuted_output_ptr = get_ptr<dType>(unpermuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, false, 8>(
input_ptr,
unpermuted_output_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream);
break;
}
#endif
#ifdef ENABLE_FP8
case at::ScalarType::Float8_e5m2:
{
using dType = cutlass::float_e5m2_t;
using dTypeCompute = cutlass::half_t;
dType *input_ptr = get_ptr<dType>(input);
dType *unpermuted_output_ptr = get_ptr<dType>(unpermuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, false, 16>(
input_ptr,
unpermuted_output_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream);
break;
}
case at::ScalarType::Float8_e4m3fn:
{
using dType = cutlass::float_e4m3_t;
using dTypeCompute = cutlass::half_t;
dType *input_ptr = get_ptr<dType>(input);
dType *unpermuted_output_ptr = get_ptr<dType>(unpermuted_output);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, false, 16>(
input_ptr,
unpermuted_output_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream);
break;
}
#endif
default:
throw std::runtime_error("Wrong activation tensor type.");
}
return unpermuted_output;
}
std::tuple<Tensor, Tensor> moe_recover_topK_bwd_op(
Tensor input_bwd,
Tensor input_fwd,
Tensor row_id_map,
Tensor prob)
{
const int num_tokens = prob.size(0);
const int num_topK = prob.size(1);
const int num_cols = input_bwd.size(1);
int *row_id_map_ptr = get_ptr<int>(row_id_map);
float *prob_ptr = get_ptr<float>(prob);
// activations type
const at::ScalarType _st = input_bwd.scalar_type();
// Output buffer alloc
Tensor act_grad =
torch::empty({input_fwd.size(0), num_cols}, torch::dtype(_st).device(torch::kCUDA).requires_grad(false));
Tensor prob_grad =
torch::empty({num_tokens, num_topK}, torch::dtype(torch::kFloat32).device(torch::kCUDA).requires_grad(false));
float *prob_grad_ptr = get_ptr<float>(prob_grad);
auto stream = at::cuda::getCurrentCUDAStream().stream();
switch (_st)
{
case at::ScalarType::Float:
{
using dType = float;
using dTypeCompute = float;
dType *input_bwd_ptr = get_ptr<dType>(input_bwd);
dType *input_fwd_ptr = get_ptr<dType>(input_fwd);
dType *act_grad_ptr = get_ptr<dType>(act_grad);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 4>(
input_bwd_ptr,
act_grad_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream,
prob_grad_ptr,
input_fwd_ptr);
break;
}
case at::ScalarType::Half:
{
using dType = cutlass::half_t;
using dTypeCompute = cutlass::half_t;
dType *input_bwd_ptr = get_ptr<dType>(input_bwd);
dType *input_fwd_ptr = get_ptr<dType>(input_fwd);
dType *act_grad_ptr = get_ptr<dType>(act_grad);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 8>(
input_bwd_ptr,
act_grad_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream,
prob_grad_ptr,
input_fwd_ptr);
break;
}
#ifdef ENABLE_BF16
case at::ScalarType::BFloat16:
{
using dType = cutlass::bfloat16_t;
using dTypeCompute = cutlass::bfloat16_t;
dType *input_bwd_ptr = get_ptr<dType>(input_bwd);
dType *input_fwd_ptr = get_ptr<dType>(input_fwd);
dType *act_grad_ptr = get_ptr<dType>(act_grad);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 8>(
input_bwd_ptr,
act_grad_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream,
prob_grad_ptr,
input_fwd_ptr);
break;
}
#endif
#ifdef ENABLE_FP8
case at::ScalarType::Float8_e5m2:
{
using dType = cutlass::float_e5m2_t;
using dTypeCompute = cutlass::half_t;
dType *input_bwd_ptr = get_ptr<dType>(input_bwd);
dType *input_fwd_ptr = get_ptr<dType>(input_fwd);
dType *act_grad_ptr = get_ptr<dType>(act_grad);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 16>(
input_bwd_ptr,
act_grad_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream,
prob_grad_ptr,
input_fwd_ptr);
break;
}
case at::ScalarType::Float8_e4m3fn:
{
using dType = cutlass::float_e4m3_t;
using dTypeCompute = cutlass::half_t;
dType *input_bwd_ptr = get_ptr<dType>(input_bwd);
dType *input_fwd_ptr = get_ptr<dType>(input_fwd);
dType *act_grad_ptr = get_ptr<dType>(act_grad);
moe_permute_topK_kernel_launcher<dType, dTypeCompute, true, 16>(
input_bwd_ptr,
act_grad_ptr,
nullptr,
row_id_map_ptr,
prob_ptr,
num_tokens,
num_topK,
num_cols,
0,
stream,
prob_grad_ptr,
input_fwd_ptr);
break;
}
#endif
default:
throw std::runtime_error("Wrong activation tensor type.");
}
return std::make_tuple(act_grad, prob_grad);
}
csrc/permute.h
+31
View File
@@ -0,0 +1,31 @@
/*************************************************************************
* Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
*
* See LICENSE for license information.
************************************************************************/
#pragma once
#include <torch/extension.h>
using torch::Tensor;
std::tuple<Tensor, Tensor, std::vector<Tensor>> moe_permute_topK_op(
Tensor input,
Tensor indices,
int64_t num_out_tokens,
std::vector<Tensor> workspace,
int64_t max_expanded_token_num);
torch::Tensor moe_recover_topK_op(
torch::Tensor input,
torch::Tensor row_id_map,
torch::Tensor prob_opt,
int64_t num_tokens,
int64_t num_topK);
std::tuple<torch::Tensor, torch::Tensor> moe_recover_topK_bwd_op(
Tensor input_bwd,
Tensor input_fwd,
Tensor row_id_map,
Tensor prob);
csrc/rope.cu
+552
View File
@@ -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;
}
}
csrc/rope.h
+14
View File
@@ -0,0 +1,14 @@
#include <torch/extension.h>
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
);
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
);