| 1 | #include <cuda.h> |
| 2 | #include <cuda_runtime.h> |
| 3 | #include <torch/csrc/stable/accelerator.h> |
| 4 | #include <torch/csrc/stable/ops.h> |
| 5 | #include <torch/csrc/stable/tensor.h> |
| 6 | #include <torch/headeronly/core/ScalarType.h> |
| 7 | #include <torch/headeronly/util/Exception.h> |
| 8 | |
| 9 | #include "core/registration.h" |
| 10 | #include "libtorch_stable/moe/permute_unpermute_kernels/moe_permute_unpermute_kernel.h" |
| 11 | #include "libtorch_stable/torch_utils.h" |
| 12 | |
| 13 | #include <torch/csrc/stable/library.h> |
| 14 | |
| 15 | // moe_permute kernels require at least CUDA 12.0 |
| 16 | #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) |
| 17 | |
| 18 | namespace { |
| 19 | |
| 20 | int64_t product_integers(torch::headeronly::IntHeaderOnlyArrayRef sizes) { |
| 21 | int64_t numel = 1; |
| 22 | for (int64_t s : sizes) { |
| 23 | numel *= s; |
| 24 | } |
| 25 | return numel; |
| 26 | } |
| 27 | |
| 28 | torch::stable::Tensor maybe_allocate_tensor( |
| 29 | const std::optional<torch::stable::Tensor>& maybe_tensor, |
| 30 | torch::headeronly::IntHeaderOnlyArrayRef expected_sizes, |
| 31 | torch::headeronly::ScalarType dtype, torch::stable::Device device, |
| 32 | char const* name) { |
| 33 | auto expected_numel = product_integers(expected_sizes); |
| 34 | if (maybe_tensor.has_value()) { |
| 35 | auto tensor = maybe_tensor.value(); |
| 36 | STD_TORCH_CHECK(tensor.device() == device, name, |
| 37 | " must be on the same device"); |
| 38 | STD_TORCH_CHECK(tensor.scalar_type() == dtype, name, |
| 39 | " has incorrect dtype"); |
| 40 | STD_TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); |
| 41 | STD_TORCH_CHECK(tensor.numel() >= expected_numel, name, |
| 42 | " is too small for the requested shape"); |
| 43 | auto flat_tensor = torch::stable::view(tensor, {tensor.numel()}); |
| 44 | return torch::stable::view( |
| 45 | torch::stable::narrow(flat_tensor, 0, 0, expected_numel), |
| 46 | expected_sizes); |
| 47 | } |
| 48 | return torch::stable::empty(expected_sizes, dtype, std::nullopt, device); |
| 49 | } |
| 50 | |
| 51 | } // namespace |
| 52 | |
| 53 | int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, |
| 54 | int64_t n_expert) { |
| 55 | return static_cast<int64_t>( |
| 56 | CubKeyValueSorter::getWorkspaceSize(num_expanded_rows, n_expert)); |
| 57 | } |
| 58 | |
| 59 | torch::stable::Tensor moe_sort_routing( |
| 60 | const torch::stable::Tensor& topk_ids, // [n_token, topk] |
| 61 | const torch::stable::Tensor& token_expert_indices, // [n_token, topk] |
| 62 | const std::optional<torch::stable::Tensor>& expert_map, // [n_expert] |
| 63 | int64_t n_expert, int64_t n_local_expert, int64_t topk, |
| 64 | torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] |
| 65 | torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] |
| 66 | const std::optional<torch::stable::Tensor>& maybe_sort_workspace, |
| 67 | const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id, |
| 68 | const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx, |
| 69 | const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort, |
| 70 | bool build_inverse) { |
| 71 | STD_TORCH_CHECK(expert_first_token_offset.scalar_type() == |
| 72 | torch::headeronly::ScalarType::Long, |
| 73 | "expert_first_token_offset must be int64"); |
| 74 | STD_TORCH_CHECK(topk_ids.scalar_type() == torch::headeronly::ScalarType::Int, |
| 75 | "topk_ids must be int32"); |
| 76 | STD_TORCH_CHECK( |
| 77 | token_expert_indices.scalar_type() == torch::headeronly::ScalarType::Int, |
| 78 | "token_expert_indices must be int32"); |
| 79 | STD_TORCH_CHECK( |
| 80 | inv_permuted_idx.scalar_type() == torch::headeronly::ScalarType::Int, |
| 81 | "inv_permuted_idx must be int32"); |
| 82 | STD_TORCH_CHECK(expert_first_token_offset.size(0) == n_local_expert + 1, |
| 83 | "expert_first_token_offset shape != n_local_expert+1"); |
| 84 | STD_TORCH_CHECK( |
| 85 | inv_permuted_idx.sizes().equals(token_expert_indices.sizes()), |
| 86 | "token_expert_indices shape must be same as inv_permuted_idx"); |
| 87 | |
| 88 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 89 | topk_ids.get_device_index()); |
| 90 | auto device = topk_ids.device(); |
| 91 | auto n_token = topk_ids.sizes()[0]; |
| 92 | auto expanded_rows = n_token * topk; |
| 93 | auto stream = get_current_cuda_stream(topk_ids.get_device_index()); |
| 94 | |
| 95 | auto sorter_size = moe_permute_sort_workspace_size(expanded_rows, n_expert); |
| 96 | auto sort_workspace = maybe_allocate_tensor( |
| 97 | maybe_sort_workspace, {sorter_size}, torch::headeronly::ScalarType::Char, |
| 98 | device, "sort_workspace"); |
| 99 | auto permuted_experts_id = maybe_allocate_tensor( |
| 100 | maybe_permuted_experts_id, topk_ids.sizes(), |
| 101 | torch::headeronly::ScalarType::Int, device, "permuted_experts_id"); |
| 102 | auto sorted_row_idx = maybe_allocate_tensor( |
| 103 | maybe_sorted_row_idx, inv_permuted_idx.sizes(), |
| 104 | torch::headeronly::ScalarType::Int, device, "sorted_row_idx"); |
| 105 | |
| 106 | CubKeyValueSorter sorter{}; |
| 107 | torch::stable::Tensor topk_ids_for_sort = topk_ids; |
| 108 | |
| 109 | if (expert_map.has_value() || build_inverse) { |
| 110 | const int* expert_map_ptr = |
| 111 | expert_map.has_value() ? get_ptr<int>(expert_map.value()) : nullptr; |
| 112 | topk_ids_for_sort = maybe_allocate_tensor( |
| 113 | maybe_topk_ids_for_sort, topk_ids.sizes(), |
| 114 | torch::headeronly::ScalarType::Int, device, "topk_ids_for_sort"); |
| 115 | torch::stable::copy_(topk_ids_for_sort, topk_ids); |
| 116 | preprocessTopkIdLauncher(get_ptr<int>(topk_ids_for_sort), n_token * topk, |
| 117 | expert_map_ptr, n_expert, stream); |
| 118 | } |
| 119 | |
| 120 | sortAndScanExpert( |
| 121 | get_ptr<const int>(topk_ids_for_sort), get_ptr<int>(token_expert_indices), |
| 122 | get_ptr<int>(permuted_experts_id), get_ptr<int>(sorted_row_idx), |
| 123 | get_ptr<int64_t>(expert_first_token_offset), n_token, n_expert, |
| 124 | n_local_expert, topk, sorter, get_ptr<int>(sort_workspace), stream, |
| 125 | build_inverse ? get_ptr<int>(inv_permuted_idx) : nullptr); |
| 126 | return sorted_row_idx; |
| 127 | } |
| 128 | |
| 129 | void moe_permute_impl( |
| 130 | const torch::stable::Tensor& input, // [n_token, hidden] |
| 131 | const torch::stable::Tensor& topk_ids, // [n_token, topk] |
| 132 | const torch::stable::Tensor& token_expert_indices, // [n_token, topk] |
| 133 | const std::optional<torch::stable::Tensor>& expert_map, // [n_expert] |
| 134 | int64_t n_expert, int64_t n_local_expert, int64_t topk, |
| 135 | torch::stable::Tensor& permuted_input, // [permuted_size, hidden] |
| 136 | torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] |
| 137 | torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] |
| 138 | torch::stable::Tensor& permuted_idx, // [permute_size] |
| 139 | const std::optional<torch::stable::Tensor>& maybe_sort_workspace, |
| 140 | const std::optional<torch::stable::Tensor>& maybe_permuted_experts_id, |
| 141 | const std::optional<torch::stable::Tensor>& maybe_sorted_row_idx, |
| 142 | const std::optional<torch::stable::Tensor>& maybe_topk_ids_for_sort) { |
| 143 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 144 | input.get_device_index()); |
| 145 | auto sorted_row_idx = moe_sort_routing( |
| 146 | topk_ids, token_expert_indices, expert_map, n_expert, n_local_expert, |
| 147 | topk, expert_first_token_offset, inv_permuted_idx, maybe_sort_workspace, |
| 148 | maybe_permuted_experts_id, maybe_sorted_row_idx, maybe_topk_ids_for_sort, |
| 149 | false); |
| 150 | auto stream = get_current_cuda_stream(input.get_device_index()); |
| 151 | auto n_token = input.size(0); |
| 152 | auto n_hidden = input.size(1); |
| 153 | auto valid_num_ptr = |
| 154 | expert_map.has_value() |
| 155 | ? get_ptr<int64_t>(expert_first_token_offset) + n_local_expert |
| 156 | : nullptr; |
| 157 | |
| 158 | MOE_DISPATCH(input.scalar_type(), [&] { |
| 159 | expandInputRowsKernelLauncher<scalar_t>( |
| 160 | get_ptr<scalar_t>(input), get_ptr<scalar_t>(permuted_input), |
| 161 | get_ptr<int>(sorted_row_idx), get_ptr<int>(inv_permuted_idx), |
| 162 | get_ptr<int>(permuted_idx), get_ptr<int64_t>(expert_first_token_offset), |
| 163 | n_token, valid_num_ptr, n_hidden, topk, n_local_expert, stream); |
| 164 | }); |
| 165 | } |
| 166 | |
| 167 | void moe_permute( |
| 168 | const torch::stable::Tensor& input, // [n_token, hidden] |
| 169 | const torch::stable::Tensor& topk_ids, // [n_token, topk] |
| 170 | const torch::stable::Tensor& token_expert_indices, // [n_token, topk] |
| 171 | const std::optional<torch::stable::Tensor>& expert_map, // [n_expert] |
| 172 | int64_t n_expert, int64_t n_local_expert, int64_t topk, |
| 173 | torch::stable::Tensor& permuted_input, // [permuted_size, hidden] |
| 174 | torch::stable::Tensor& expert_first_token_offset, // [n_local_expert + 1] |
| 175 | torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] |
| 176 | torch::stable::Tensor& permuted_idx) { // [permute_size] |
| 177 | moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert, |
| 178 | n_local_expert, topk, permuted_input, |
| 179 | expert_first_token_offset, inv_permuted_idx, permuted_idx, |
| 180 | std::nullopt, std::nullopt, std::nullopt, std::nullopt); |
| 181 | } |
| 182 | |
| 183 | void moe_permute_with_scratch( |
| 184 | const torch::stable::Tensor& input, const torch::stable::Tensor& topk_ids, |
| 185 | const torch::stable::Tensor& token_expert_indices, |
| 186 | const std::optional<torch::stable::Tensor>& expert_map, int64_t n_expert, |
| 187 | int64_t n_local_expert, int64_t topk, torch::stable::Tensor& permuted_input, |
| 188 | torch::stable::Tensor& expert_first_token_offset, |
| 189 | torch::stable::Tensor& inv_permuted_idx, |
| 190 | torch::stable::Tensor& permuted_idx, torch::stable::Tensor& sort_workspace, |
| 191 | torch::stable::Tensor& permuted_experts_id, |
| 192 | torch::stable::Tensor& sorted_row_idx, |
| 193 | torch::stable::Tensor& topk_ids_for_sort) { |
| 194 | moe_permute_impl(input, topk_ids, token_expert_indices, expert_map, n_expert, |
| 195 | n_local_expert, topk, permuted_input, |
| 196 | expert_first_token_offset, inv_permuted_idx, permuted_idx, |
| 197 | sort_workspace, permuted_experts_id, sorted_row_idx, |
| 198 | topk_ids_for_sort); |
| 199 | } |
| 200 | |
| 201 | void moe_prepare_scatter(const torch::stable::Tensor& topk_ids, |
| 202 | const torch::stable::Tensor& token_expert_indices, |
| 203 | const std::optional<torch::stable::Tensor>& expert_map, |
| 204 | int64_t n_expert, int64_t n_local_expert, |
| 205 | torch::stable::Tensor& expert_first_token_offset, |
| 206 | torch::stable::Tensor& scatter_idx, |
| 207 | torch::stable::Tensor& sort_workspace, |
| 208 | torch::stable::Tensor& sorted_experts, |
| 209 | torch::stable::Tensor& sorted_rows, |
| 210 | torch::stable::Tensor& topk_ids_for_sort) { |
| 211 | moe_sort_routing(topk_ids, token_expert_indices, expert_map, n_expert, |
| 212 | n_local_expert, topk_ids.size(1), expert_first_token_offset, |
| 213 | scatter_idx, sort_workspace, sorted_experts, sorted_rows, |
| 214 | topk_ids_for_sort, true); |
| 215 | } |
| 216 | |
| 217 | void moe_unpermute( |
| 218 | const torch::stable::Tensor& |
| 219 | permuted_hidden_states, // [n_token * topk, hidden] |
| 220 | const torch::stable::Tensor& topk_weights, // [n_token, topk] |
| 221 | const torch::stable::Tensor& inv_permuted_idx, // [n_token, topk] |
| 222 | const std::optional<torch::stable::Tensor>& |
| 223 | expert_first_token_offset, // [n_local_expert+1] |
| 224 | int64_t topk, |
| 225 | torch::stable::Tensor& hidden_states) { // [n_token, hidden] |
| 226 | STD_TORCH_CHECK( |
| 227 | permuted_hidden_states.scalar_type() == hidden_states.scalar_type(), |
| 228 | "permuted_hidden_states dtype must be same as hidden_states"); |
| 229 | |
| 230 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 231 | hidden_states.get_device_index()); |
| 232 | auto n_token = hidden_states.size(0); |
| 233 | auto n_hidden = hidden_states.size(1); |
| 234 | auto stream = get_current_cuda_stream(hidden_states.get_device_index()); |
| 235 | |
| 236 | int64_t const* valid_ptr = nullptr; |
| 237 | if (expert_first_token_offset.has_value()) { |
| 238 | int n_local_expert = expert_first_token_offset.value().size(0) - 1; |
| 239 | valid_ptr = |
| 240 | get_ptr<int64_t>(expert_first_token_offset.value()) + n_local_expert; |
| 241 | } |
| 242 | |
| 243 | MOE_DISPATCH(hidden_states.scalar_type(), [&] { |
| 244 | finalizeMoeRoutingKernelLauncher<scalar_t, scalar_t>( |
| 245 | get_ptr<scalar_t>(permuted_hidden_states), |
| 246 | get_ptr<scalar_t>(hidden_states), get_ptr<float>(topk_weights), |
| 247 | get_ptr<int>(inv_permuted_idx), n_token, n_hidden, topk, valid_ptr, |
| 248 | stream); |
| 249 | }); |
| 250 | } |
| 251 | |
| 252 | template <typename T> |
| 253 | __global__ void shuffleInputRowsKernel(const T* input, |
| 254 | const int32_t* dst2src_map, T* output, |
| 255 | int64_t num_src_rows, |
| 256 | int64_t num_dst_rows, int64_t num_cols) { |
| 257 | int64_t dest_row_idx = blockIdx.x; |
| 258 | if (blockIdx.x < num_dst_rows) { |
| 259 | int64_t const source_row_idx = dst2src_map[dest_row_idx]; |
| 260 | // Load 128-bits per thread |
| 261 | constexpr int64_t ELEM_PER_THREAD = 128 / sizeof(T) / 8; |
| 262 | using DataElem = cutlass::Array<T, ELEM_PER_THREAD>; |
| 263 | |
| 264 | auto* dest_row_ptr = |
| 265 | reinterpret_cast<DataElem*>(output + dest_row_idx * num_cols); |
| 266 | int64_t const start_offset = threadIdx.x; |
| 267 | int64_t const stride = blockDim.x; |
| 268 | int64_t const num_elems_in_col = num_cols / ELEM_PER_THREAD; |
| 269 | |
| 270 | if (source_row_idx < 0 || source_row_idx >= num_src_rows) { |
| 271 | DataElem const zero{}; |
| 272 | for (int elem_index = start_offset; elem_index < num_elems_in_col; |
| 273 | elem_index += stride) { |
| 274 | dest_row_ptr[elem_index] = zero; |
| 275 | } |
| 276 | return; |
| 277 | } |
| 278 | |
| 279 | // Duplicate and permute rows |
| 280 | auto const* source_row_ptr = |
| 281 | reinterpret_cast<DataElem const*>(input + source_row_idx * num_cols); |
| 282 | for (int elem_index = start_offset; elem_index < num_elems_in_col; |
| 283 | elem_index += stride) { |
| 284 | dest_row_ptr[elem_index] = source_row_ptr[elem_index]; |
| 285 | } |
| 286 | } |
| 287 | } |
| 288 | |
| 289 | void shuffle_rows(const torch::stable::Tensor& input_tensor, |
| 290 | const torch::stable::Tensor& dst2src_map, |
| 291 | torch::stable::Tensor& output_tensor) { |
| 292 | STD_TORCH_CHECK(input_tensor.scalar_type() == output_tensor.scalar_type(), |
| 293 | "Input and output tensors must have the same data type"); |
| 294 | |
| 295 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 296 | output_tensor.get_device_index()); |
| 297 | auto stream = get_current_cuda_stream(output_tensor.get_device_index()); |
| 298 | const int64_t blocks = output_tensor.size(0); |
| 299 | const int64_t threads = 256; |
| 300 | const int64_t num_dest_rows = output_tensor.size(0); |
| 301 | const int64_t num_src_rows = input_tensor.size(0); |
| 302 | const int64_t num_cols = input_tensor.size(1); |
| 303 | |
| 304 | STD_TORCH_CHECK(!(num_cols % (128 / input_tensor.element_size() / 8)), |
| 305 | "num_cols must be divisible by 128 / " |
| 306 | "input_tensor.element_size() / 8"); |
| 307 | |
| 308 | MOE_DISPATCH(input_tensor.scalar_type(), [&] { |
| 309 | shuffleInputRowsKernel<scalar_t><<<blocks, threads, 0, stream>>>( |
| 310 | reinterpret_cast<const scalar_t*>(input_tensor.const_data_ptr()), |
| 311 | reinterpret_cast<const int32_t*>(dst2src_map.const_data_ptr()), |
| 312 | reinterpret_cast<scalar_t*>(output_tensor.mutable_data_ptr()), |
| 313 | num_src_rows, num_dest_rows, num_cols); |
| 314 | }); |
| 315 | } |
| 316 | |
| 317 | #else |
| 318 | |
| 319 | int64_t moe_permute_sort_workspace_size(int64_t num_expanded_rows, |
| 320 | int64_t n_expert) { |
| 321 | STD_TORCH_CHECK( |
| 322 | false, "moe_permute_sort_workspace_size is not supported on CUDA < 12.0"); |
| 323 | } |
| 324 | |
| 325 | void moe_permute(const torch::stable::Tensor& input, |
| 326 | const torch::stable::Tensor& topk_ids, |
| 327 | const torch::stable::Tensor& token_expert_indices, |
| 328 | const std::optional<torch::stable::Tensor>& expert_map, |
| 329 | int64_t n_expert, int64_t n_local_expert, int64_t topk, |
| 330 | torch::stable::Tensor& permuted_input, |
| 331 | torch::stable::Tensor& expert_first_token_offset, |
| 332 | torch::stable::Tensor& inv_permuted_idx, |
| 333 | torch::stable::Tensor& permuted_idx) { |
| 334 | STD_TORCH_CHECK(false, "moe_permute is not supported on CUDA < 12.0"); |
| 335 | } |
| 336 | |
| 337 | void moe_permute_with_scratch( |
| 338 | const torch::stable::Tensor& input, const torch::stable::Tensor& topk_ids, |
| 339 | const torch::stable::Tensor& token_expert_indices, |
| 340 | const std::optional<torch::stable::Tensor>& expert_map, int64_t n_expert, |
| 341 | int64_t n_local_expert, int64_t topk, torch::stable::Tensor& permuted_input, |
| 342 | torch::stable::Tensor& expert_first_token_offset, |
| 343 | torch::stable::Tensor& inv_permuted_idx, |
| 344 | torch::stable::Tensor& permuted_idx, torch::stable::Tensor& sort_workspace, |
| 345 | torch::stable::Tensor& permuted_experts_id, |
| 346 | torch::stable::Tensor& sorted_row_idx, |
| 347 | torch::stable::Tensor& topk_ids_for_sort) { |
| 348 | STD_TORCH_CHECK(false, |
| 349 | "moe_permute_with_scratch is not supported on CUDA < 12.0"); |
| 350 | } |
| 351 | |
| 352 | void moe_unpermute( |
| 353 | const torch::stable::Tensor& permuted_hidden_states, |
| 354 | const torch::stable::Tensor& topk_weights, |
| 355 | const torch::stable::Tensor& inv_permuted_idx, |
| 356 | const std::optional<torch::stable::Tensor>& expert_first_token_offset, |
| 357 | int64_t topk, torch::stable::Tensor& hidden_states) { |
| 358 | STD_TORCH_CHECK(false, "moe_unpermute is not supported on CUDA < 12.0"); |
| 359 | } |
| 360 | |
| 361 | void moe_prepare_scatter(const torch::stable::Tensor& topk_ids, |
| 362 | const torch::stable::Tensor& token_expert_indices, |
| 363 | const std::optional<torch::stable::Tensor>& expert_map, |
| 364 | int64_t n_expert, int64_t n_local_expert, |
| 365 | torch::stable::Tensor& expert_first_token_offset, |
| 366 | torch::stable::Tensor& scatter_idx, |
| 367 | torch::stable::Tensor& sort_workspace, |
| 368 | torch::stable::Tensor& sorted_experts, |
| 369 | torch::stable::Tensor& sorted_rows, |
| 370 | torch::stable::Tensor& topk_ids_for_sort) { |
| 371 | STD_TORCH_CHECK(false, "moe_prepare_scatter is not supported on CUDA < 12.0"); |
| 372 | } |
| 373 | |
| 374 | #endif |
| 375 | |
| 376 | bool moe_permute_unpermute_supported() { |
| 377 | #if defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) |
| 378 | return true; |
| 379 | #else |
| 380 | return false; |
| 381 | #endif |
| 382 | } |
| 383 | |
| 384 | STABLE_TORCH_LIBRARY_IMPL(_moe_C, CUDA, m) { |
| 385 | m.impl("moe_permute", TORCH_BOX(&moe_permute)); |
| 386 | m.impl("moe_permute_with_scratch", TORCH_BOX(&moe_permute_with_scratch)); |
| 387 | m.impl("moe_prepare_scatter", TORCH_BOX(&moe_prepare_scatter)); |
| 388 | m.impl("moe_unpermute", TORCH_BOX(&moe_unpermute)); |
| 389 | } |