| 1 | // see csrc/quantization/w8a8/cutlass/moe/get_group_starts.cuh |
| 2 | #pragma once |
| 3 | |
| 4 | #include <cuda.h> |
| 5 | #include <torch/csrc/stable/tensor.h> |
| 6 | #include "libtorch_stable/torch_utils.h" |
| 7 | |
| 8 | #include "cutlass/bfloat16.h" |
| 9 | #include "cutlass/float8.h" |
| 10 | |
| 11 | // ElementB is int32 (packed int4) |
| 12 | // ElementGroupScale is cutlass::Array<cutlass::float_e4m3_t, 8> (packed fp8) |
| 13 | template <typename ElementA, typename ElementB, typename ElementC, |
| 14 | typename ElementAccumulator, typename ElementGroupScale> |
| 15 | __global__ void get_group_gemm_starts( |
| 16 | int64_t* expert_offsets, ElementA** a_offsets, ElementB** b_offsets, |
| 17 | ElementC** out_offsets, ElementAccumulator** a_scales_offsets, |
| 18 | ElementAccumulator** b_scales_offsets, |
| 19 | ElementGroupScale** b_group_scales_offsets, ElementA* a_base_as_int, |
| 20 | ElementB* b_base_as_int, ElementC* out_base_as_int, |
| 21 | ElementAccumulator* a_scales_base_as_int, |
| 22 | ElementAccumulator* b_scales_base_as_int, |
| 23 | ElementGroupScale* b_group_scales_base_as_int, int64_t n, int64_t k, |
| 24 | int64_t scale_k) { |
| 25 | int expert_id = threadIdx.x; |
| 26 | |
| 27 | int64_t expert_offset = expert_offsets[expert_id]; |
| 28 | |
| 29 | // same as w8a8 |
| 30 | a_offsets[expert_id] = a_base_as_int + expert_offset * k; |
| 31 | out_offsets[expert_id] = out_base_as_int + expert_offset * n; |
| 32 | a_scales_offsets[expert_id] = a_scales_base_as_int + expert_offset; |
| 33 | b_scales_offsets[expert_id] = b_scales_base_as_int + (n * expert_id); |
| 34 | |
| 35 | // w4a8 specific |
| 36 | constexpr int pack_factor = 8; // pack 8 int4 into int32 |
| 37 | b_offsets[expert_id] = b_base_as_int + (expert_id * k * n / pack_factor); |
| 38 | b_group_scales_offsets[expert_id] = |
| 39 | b_group_scales_base_as_int + (expert_id * scale_k * n); |
| 40 | } |
| 41 | |
| 42 | #define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE) \ |
| 43 | else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \ |
| 44 | get_group_gemm_starts<cutlass::float_e4m3_t, int32_t, C_TYPE, float, \ |
| 45 | cutlass::Array<cutlass::float_e4m3_t, 8>> \ |
| 46 | <<<1, num_experts, 0, stream>>>( \ |
| 47 | static_cast<int64_t*>(expert_offsets.data_ptr()), \ |
| 48 | static_cast<cutlass::float_e4m3_t**>(a_ptrs.data_ptr()), \ |
| 49 | static_cast<int32_t**>(b_ptrs.data_ptr()), \ |
| 50 | static_cast<C_TYPE**>(out_ptrs.data_ptr()), \ |
| 51 | static_cast<float**>(a_scales_ptrs.data_ptr()), \ |
| 52 | static_cast<float**>(b_scales_ptrs.data_ptr()), \ |
| 53 | static_cast<cutlass::Array<cutlass::float_e4m3_t, 8>**>( \ |
| 54 | b_group_scales_ptrs.data_ptr()), \ |
| 55 | static_cast<cutlass::float_e4m3_t*>(a_tensors.data_ptr()), \ |
| 56 | static_cast<int32_t*>(b_tensors.data_ptr()), \ |
| 57 | static_cast<C_TYPE*>(out_tensors.data_ptr()), \ |
| 58 | static_cast<float*>(a_scales.data_ptr()), \ |
| 59 | static_cast<float*>(b_scales.data_ptr()), \ |
| 60 | static_cast<cutlass::Array<cutlass::float_e4m3_t, 8>*>( \ |
| 61 | b_group_scales.data_ptr()), \ |
| 62 | n, k, scale_k); \ |
| 63 | } |
| 64 | |
| 65 | namespace { |
| 66 | |
| 67 | void run_get_group_gemm_starts( |
| 68 | torch::stable::Tensor const& expert_offsets, torch::stable::Tensor& a_ptrs, |
| 69 | torch::stable::Tensor& b_ptrs, torch::stable::Tensor& out_ptrs, |
| 70 | torch::stable::Tensor& a_scales_ptrs, torch::stable::Tensor& b_scales_ptrs, |
| 71 | torch::stable::Tensor& b_group_scales_ptrs, |
| 72 | torch::stable::Tensor const& a_tensors, |
| 73 | torch::stable::Tensor const& b_tensors, torch::stable::Tensor& out_tensors, |
| 74 | torch::stable::Tensor const& a_scales, |
| 75 | torch::stable::Tensor const& b_scales, |
| 76 | torch::stable::Tensor const& b_group_scales, const int64_t b_group_size) { |
| 77 | STD_TORCH_CHECK(a_tensors.scalar_type() == |
| 78 | torch::headeronly::ScalarType::Float8_e4m3fn); |
| 79 | STD_TORCH_CHECK( |
| 80 | b_tensors.scalar_type() == |
| 81 | torch::headeronly::ScalarType::Int); // int4 8x packed into int32 |
| 82 | STD_TORCH_CHECK(a_scales.scalar_type() == |
| 83 | torch::headeronly::ScalarType::Float); |
| 84 | STD_TORCH_CHECK(b_scales.scalar_type() == |
| 85 | torch::headeronly::ScalarType::Float); |
| 86 | STD_TORCH_CHECK( |
| 87 | b_group_scales.scalar_type() == |
| 88 | torch::headeronly::ScalarType::Float8_e4m3fn); // the underlying torch |
| 89 | // type is e4m3 |
| 90 | STD_TORCH_CHECK( |
| 91 | out_tensors.scalar_type() == |
| 92 | torch::headeronly::ScalarType::BFloat16); // only support bf16 for now |
| 93 | // expect int64_t to avoid overflow during offset calculations |
| 94 | STD_TORCH_CHECK(expert_offsets.scalar_type() == |
| 95 | torch::headeronly::ScalarType::Long); |
| 96 | |
| 97 | int num_experts = static_cast<int>(expert_offsets.size(0)); |
| 98 | // logical k, n |
| 99 | int64_t n = out_tensors.size(1); |
| 100 | int64_t k = a_tensors.size(1); |
| 101 | int64_t scale_k = cutlass::ceil_div(k, b_group_size); |
| 102 | |
| 103 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 104 | a_tensors.get_device_index()); |
| 105 | auto stream = get_current_cuda_stream(a_tensors.get_device_index()); |
| 106 | |
| 107 | if (false) { |
| 108 | } |
| 109 | __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::BFloat16, |
| 110 | cutlass::bfloat16_t) |
| 111 | __CALL_GET_STARTS_KERNEL(torch::headeronly::ScalarType::Half, half) |
| 112 | else { |
| 113 | STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)"); |
| 114 | } |
| 115 | } |
| 116 | |
| 117 | } // namespace |