| 1 | /* |
| 2 | * SPDX-License-Identifier: Apache-2.0 |
| 3 | * SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 4 | * |
| 5 | * MXFP4 x MXFP4 block-scaled grouped GEMM kernel for MoE on SM100. |
| 6 | * Uses Cutlass mx_float4_t operands, E8M0 block scales, and 32-element groups. |
| 7 | */ |
| 8 | |
| 9 | #include <torch/csrc/stable/library.h> |
| 10 | #include <torch/csrc/stable/tensor.h> |
| 11 | #include "libtorch_stable/torch_utils.h" |
| 12 | |
| 13 | #include <cutlass/arch/arch.h> |
| 14 | |
| 15 | #include "libtorch_stable/cutlass_extensions/common.hpp" |
| 16 | |
| 17 | #include "cute/tensor.hpp" |
| 18 | #include "cutlass/tensor_ref.h" |
| 19 | #include "cutlass/epilogue/collective/default_epilogue.hpp" |
| 20 | #include "cutlass/epilogue/thread/linear_combination.h" |
| 21 | #include "cutlass/gemm/dispatch_policy.hpp" |
| 22 | #include "cutlass/gemm/group_array_problem_shape.hpp" |
| 23 | #include "cutlass/gemm/collective/collective_builder.hpp" |
| 24 | #include "cutlass/epilogue/collective/collective_builder.hpp" |
| 25 | #include "cutlass/gemm/device/gemm_universal_adapter.h" |
| 26 | #include "cutlass/gemm/kernel/gemm_universal.hpp" |
| 27 | |
| 28 | #include "cutlass/util/packed_stride.hpp" |
| 29 | #include <cassert> |
| 30 | |
| 31 | using namespace cute; |
| 32 | |
| 33 | // Offset-computation kernel for MXFP4 grouped GEMM (group size 32). |
| 34 | template <typename ElementAB, typename ElementC, typename ElementSF, |
| 35 | typename LayoutSFA, typename LayoutSFB, typename ScaleConfig> |
| 36 | __global__ void __mxfp4_get_group_gemm_starts( |
| 37 | ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets, |
| 38 | ElementSF** a_scales_offsets, ElementSF** b_scales_offsets, |
| 39 | LayoutSFA* layout_sfa_base_as_int, LayoutSFB* layout_sfb_base_as_int, |
| 40 | ElementAB* a_base_as_int, ElementAB* b_base_as_int, |
| 41 | ElementC* out_base_as_int, ElementSF* a_scales_base_as_int, |
| 42 | ElementSF* b_scales_base_as_int, const int32_t* expert_offsets, |
| 43 | const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes, |
| 44 | int64_t* a_strides, int64_t* b_strides, int64_t* c_strides, |
| 45 | const int64_t a_stride_val, const int64_t b_stride_val, |
| 46 | const int64_t c_stride_val, const int K, const int N) { |
| 47 | int64_t expert_id = threadIdx.x; |
| 48 | if (expert_id >= gridDim.x * blockDim.x) { |
| 49 | return; |
| 50 | } |
| 51 | int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]); |
| 52 | int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]); |
| 53 | int64_t group_size = 32; |
| 54 | int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]); |
| 55 | int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]); |
| 56 | int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]); |
| 57 | assert((m >= 0 && n == N && k == K && k % 2 == 0) && |
| 58 | "unexpected problem sizes"); |
| 59 | |
| 60 | int64_t half_k = static_cast<int64_t>(k / 2); |
| 61 | int64_t group_k = static_cast<int64_t>(k / group_size); |
| 62 | // Shape of A as uint8/byte = [M, K // 2] |
| 63 | a_offsets[expert_id] = a_base_as_int + expert_offset * half_k; |
| 64 | // Shape of B as uint8/byte = [E, N, K // 2] |
| 65 | b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k; |
| 66 | // Shape of C = [M, N] |
| 67 | out_offsets[expert_id] = out_base_as_int + expert_offset * n; |
| 68 | // Shape of a_scale = [sum(sf_sizes), K // group_size] |
| 69 | a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k; |
| 70 | |
| 71 | assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) == |
| 72 | 0 && |
| 73 | "TMA requires 128-byte alignment"); |
| 74 | |
| 75 | // Shape of B scale = [E, N, K // group_size] |
| 76 | b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k; |
| 77 | assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) == |
| 78 | 0 && |
| 79 | "TMA requires 128-byte alignment"); |
| 80 | |
| 81 | // Initialize strides |
| 82 | a_strides[expert_id] = a_stride_val; |
| 83 | b_strides[expert_id] = b_stride_val; |
| 84 | c_strides[expert_id] = c_stride_val; |
| 85 | |
| 86 | LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id; |
| 87 | LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id; |
| 88 | |
| 89 | *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape( |
| 90 | static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1)); |
| 91 | *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape( |
| 92 | static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1)); |
| 93 | } |
| 94 | |
| 95 | #define __CALL_MXFP4_GET_STARTS_KERNEL(ELEMENT_AB_TYPE, SF_TYPE, \ |
| 96 | TENSOR_C_TYPE, C_TYPE, LayoutSFA, \ |
| 97 | LayoutSFB, ScaleConfig) \ |
| 98 | else if (out_tensors.scalar_type() == TENSOR_C_TYPE) { \ |
| 99 | __mxfp4_get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, LayoutSFA, \ |
| 100 | LayoutSFB, ScaleConfig> \ |
| 101 | <<<1, num_experts, 0, stream>>>( \ |
| 102 | static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \ |
| 103 | static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()), \ |
| 104 | static_cast<C_TYPE**>(out_starts.data_ptr()), \ |
| 105 | static_cast<SF_TYPE**>(a_scales_starts.data_ptr()), \ |
| 106 | static_cast<SF_TYPE**>(b_scales_starts.data_ptr()), \ |
| 107 | reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \ |
| 108 | reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \ |
| 109 | static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()), \ |
| 110 | static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()), \ |
| 111 | static_cast<C_TYPE*>(out_tensors.data_ptr()), \ |
| 112 | static_cast<SF_TYPE*>(a_scales.data_ptr()), \ |
| 113 | static_cast<SF_TYPE*>(b_scales.data_ptr()), \ |
| 114 | static_cast<int32_t*>(expert_offsets.data_ptr()), \ |
| 115 | static_cast<int32_t*>(sf_offsets.data_ptr()), \ |
| 116 | static_cast<int32_t*>(problem_sizes.data_ptr()), \ |
| 117 | static_cast<int64_t*>(a_strides.data_ptr()), \ |
| 118 | static_cast<int64_t*>(b_strides.data_ptr()), \ |
| 119 | static_cast<int64_t*>(c_strides.data_ptr()), a_stride_val, \ |
| 120 | b_stride_val, c_stride_val, K, N); \ |
| 121 | } |
| 122 | |
| 123 | template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig> |
| 124 | void mxfp4_run_get_group_gemm_starts( |
| 125 | const torch::stable::Tensor& a_starts, |
| 126 | const torch::stable::Tensor& b_starts, |
| 127 | const torch::stable::Tensor& out_starts, |
| 128 | const torch::stable::Tensor& a_scales_starts, |
| 129 | const torch::stable::Tensor& b_scales_starts, |
| 130 | const torch::stable::Tensor& layout_sfa, |
| 131 | const torch::stable::Tensor& layout_sfb, |
| 132 | const torch::stable::Tensor& a_strides, |
| 133 | const torch::stable::Tensor& b_strides, |
| 134 | const torch::stable::Tensor& c_strides, int64_t a_stride_val, |
| 135 | int64_t b_stride_val, int64_t c_stride_val, |
| 136 | torch::stable::Tensor const& a_tensors, |
| 137 | torch::stable::Tensor const& b_tensors, |
| 138 | torch::stable::Tensor const& out_tensors, |
| 139 | torch::stable::Tensor const& a_scales, |
| 140 | torch::stable::Tensor const& b_scales, |
| 141 | torch::stable::Tensor const& expert_offsets, |
| 142 | torch::stable::Tensor const& sf_offsets, |
| 143 | torch::stable::Tensor const& problem_sizes, int M, int N, int K) { |
| 144 | int num_experts = (int)expert_offsets.size(0); |
| 145 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 146 | a_tensors.get_device_index()); |
| 147 | auto stream = get_current_cuda_stream(a_tensors.get_device_index()); |
| 148 | |
| 149 | STD_TORCH_CHECK(out_tensors.size(1) == N, |
| 150 | "Output tensor shape doesn't match expected shape"); |
| 151 | STD_TORCH_CHECK(K / 2 == b_tensors.size(2), |
| 152 | "b_tensors(dim = 2) and a_tensors(dim = 1) trailing" |
| 153 | " dimension must match"); |
| 154 | if (false) { |
| 155 | } |
| 156 | // MXFP4 uses E8M0 (float_ue8m0_t) scale factors |
| 157 | __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t, |
| 158 | torch::headeronly::ScalarType::BFloat16, |
| 159 | cutlass::bfloat16_t, LayoutSFA, LayoutSFB, |
| 160 | ScaleConfig) |
| 161 | __CALL_MXFP4_GET_STARTS_KERNEL(cutlass::float_e2m1_t, cutlass::float_ue8m0_t, |
| 162 | torch::headeronly::ScalarType::Half, half, |
| 163 | LayoutSFA, LayoutSFB, ScaleConfig) |
| 164 | else { |
| 165 | STD_TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)"); |
| 166 | } |
| 167 | } |
| 168 | |
| 169 | template <typename OutType> |
| 170 | void run_mxfp4_blockwise_scaled_group_mm_sm100( |
| 171 | torch::stable::Tensor& output, const torch::stable::Tensor& a, |
| 172 | const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale, |
| 173 | const torch::stable::Tensor& b_blockscales, |
| 174 | const torch::stable::Tensor& problem_sizes, |
| 175 | const torch::stable::Tensor& expert_offsets, |
| 176 | const torch::stable::Tensor& sf_offsets, int M, int N, int K) { |
| 177 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 178 | a.get_device_index()); |
| 179 | using ProblemShape = |
| 180 | cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>; |
| 181 | using ElementType = cutlass::float_e2m1_t; |
| 182 | using ElementSFType = cutlass::float_ue8m0_t; |
| 183 | using ElementA = cutlass::mx_float4_t<cutlass::float_e2m1_t>; |
| 184 | using ElementB = cutlass::mx_float4_t<cutlass::float_e2m1_t>; |
| 185 | |
| 186 | using ElementC = OutType; |
| 187 | using ElementD = ElementC; |
| 188 | using ElementAccumulator = float; |
| 189 | // Layout definitions |
| 190 | using LayoutA = cutlass::layout::RowMajor; |
| 191 | using LayoutB = cutlass::layout::ColumnMajor; |
| 192 | using LayoutC = cutlass::layout::RowMajor; |
| 193 | using LayoutD = LayoutC; |
| 194 | |
| 195 | static constexpr int AlignmentA = 32; |
| 196 | static constexpr int AlignmentB = 32; |
| 197 | static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value; |
| 198 | static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value; |
| 199 | |
| 200 | // Architecture definitions |
| 201 | using ArchTag = cutlass::arch::Sm100; |
| 202 | using EpilogueOperatorClass = cutlass::arch::OpClassTensorOp; |
| 203 | using MainloopOperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; |
| 204 | using StageCountType = cutlass::gemm::collective::StageCountAuto; |
| 205 | |
| 206 | using ClusterShape = Shape<_1, _1, _1>; |
| 207 | struct MMA1SMConfig { |
| 208 | using MmaTileShape = Shape<_128, _128, _128>; |
| 209 | using KernelSchedule = |
| 210 | cutlass::gemm::KernelPtrArrayTmaWarpSpecialized1SmMxf4Sm100; |
| 211 | using EpilogueSchedule = cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; |
| 212 | }; |
| 213 | |
| 214 | using CollectiveEpilogue = |
| 215 | typename cutlass::epilogue::collective::CollectiveBuilder< |
| 216 | ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape, |
| 217 | ClusterShape, Shape<_128, _64>, ElementAccumulator, |
| 218 | ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD, |
| 219 | LayoutC*, AlignmentD, |
| 220 | typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp; |
| 221 | |
| 222 | using CollectiveMainloop = |
| 223 | typename cutlass::gemm::collective::CollectiveBuilder< |
| 224 | ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA, |
| 225 | ElementB, LayoutB*, AlignmentB, ElementAccumulator, |
| 226 | typename MMA1SMConfig::MmaTileShape, ClusterShape, |
| 227 | cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>( |
| 228 | sizeof(typename CollectiveEpilogue::SharedStorage))>, |
| 229 | typename MMA1SMConfig::KernelSchedule>::CollectiveOp; |
| 230 | |
| 231 | using GemmKernel = |
| 232 | cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop, |
| 233 | CollectiveEpilogue>; |
| 234 | |
| 235 | using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>; |
| 236 | using StrideA = typename Gemm::GemmKernel::InternalStrideA; |
| 237 | using StrideB = typename Gemm::GemmKernel::InternalStrideB; |
| 238 | using StrideC = typename Gemm::GemmKernel::InternalStrideC; |
| 239 | using StrideD = typename Gemm::GemmKernel::InternalStrideD; |
| 240 | |
| 241 | using LayoutSFA = |
| 242 | typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA; |
| 243 | using LayoutSFB = |
| 244 | typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB; |
| 245 | using ScaleConfig = |
| 246 | typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig; |
| 247 | |
| 248 | using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape; |
| 249 | int num_experts = static_cast<int>(expert_offsets.size(0)); |
| 250 | |
| 251 | torch::stable::Tensor a_ptrs = |
| 252 | torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, |
| 253 | std::nullopt, a.device()); |
| 254 | torch::stable::Tensor b_ptrs = |
| 255 | torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, |
| 256 | std::nullopt, a.device()); |
| 257 | torch::stable::Tensor out_ptrs = |
| 258 | torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, |
| 259 | std::nullopt, a.device()); |
| 260 | torch::stable::Tensor a_scales_ptrs = |
| 261 | torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, |
| 262 | std::nullopt, a.device()); |
| 263 | torch::stable::Tensor b_scales_ptrs = |
| 264 | torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, |
| 265 | std::nullopt, a.device()); |
| 266 | torch::stable::Tensor layout_sfa = torch::stable::empty( |
| 267 | {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt, |
| 268 | a.device()); |
| 269 | torch::stable::Tensor layout_sfb = torch::stable::empty( |
| 270 | {num_experts, 5}, torch::headeronly::ScalarType::Long, std::nullopt, |
| 271 | a.device()); |
| 272 | torch::stable::Tensor a_strides1 = |
| 273 | torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, |
| 274 | std::nullopt, a.device()); |
| 275 | torch::stable::Tensor b_strides1 = |
| 276 | torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, |
| 277 | std::nullopt, a.device()); |
| 278 | torch::stable::Tensor c_strides1 = |
| 279 | torch::stable::empty(num_experts, torch::headeronly::ScalarType::Long, |
| 280 | std::nullopt, a.device()); |
| 281 | |
| 282 | mxfp4_run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>( |
| 283 | a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, layout_sfa, |
| 284 | layout_sfb, a_strides1, b_strides1, c_strides1, a.stride(0) * 2, |
| 285 | b.stride(1) * 2, output.stride(0), a, b, output, a_blockscale, |
| 286 | b_blockscales, expert_offsets, sf_offsets, problem_sizes, M, N, K); |
| 287 | |
| 288 | // Create an instance of the GEMM |
| 289 | Gemm gemm_op; |
| 290 | |
| 291 | UnderlyingProblemShape* problem_sizes_as_shapes = |
| 292 | static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr()); |
| 293 | |
| 294 | // Set the Scheduler info |
| 295 | cutlass::KernelHardwareInfo hw_info; |
| 296 | using RasterOrderOptions = typename cutlass::gemm::kernel::detail:: |
| 297 | PersistentTileSchedulerSm100GroupParams< |
| 298 | typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions; |
| 299 | typename Gemm::GemmKernel::TileSchedulerArguments scheduler; |
| 300 | scheduler.raster_order = RasterOrderOptions::AlongM; |
| 301 | hw_info.device_id = a.get_device_index(); |
| 302 | static std::unordered_map<int, int> cached_sm_counts; |
| 303 | if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) { |
| 304 | cached_sm_counts[hw_info.device_id] = |
| 305 | cutlass::KernelHardwareInfo::query_device_multiprocessor_count( |
| 306 | hw_info.device_id); |
| 307 | } |
| 308 | hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX); |
| 309 | |
| 310 | // Mainloop Arguments |
| 311 | typename GemmKernel::MainloopArguments mainloop_args{ |
| 312 | static_cast<const ElementType**>(a_ptrs.data_ptr()), |
| 313 | static_cast<StrideA*>(a_strides1.data_ptr()), |
| 314 | static_cast<const ElementType**>(b_ptrs.data_ptr()), |
| 315 | static_cast<StrideB*>(b_strides1.data_ptr()), |
| 316 | static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()), |
| 317 | reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), |
| 318 | static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()), |
| 319 | reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())}; |
| 320 | |
| 321 | // Epilogue Arguments |
| 322 | typename GemmKernel::EpilogueArguments epilogue_args{ |
| 323 | {}, // epilogue.thread |
| 324 | nullptr, |
| 325 | static_cast<StrideC*>(c_strides1.data_ptr()), |
| 326 | static_cast<ElementD**>(out_ptrs.data_ptr()), |
| 327 | static_cast<StrideC*>(c_strides1.data_ptr())}; |
| 328 | auto& fusion_args = epilogue_args.thread; |
| 329 | // Scalar epilogue (CUTLASS grouped GEMM): D = 1 * accum + 0 * C |
| 330 | fusion_args.alpha_ptr = nullptr; |
| 331 | fusion_args.beta_ptr = nullptr; |
| 332 | fusion_args.alpha = 1.0f; |
| 333 | fusion_args.alpha_ptr_array = nullptr; |
| 334 | fusion_args.dAlpha = {_0{}, _0{}, 0}; |
| 335 | fusion_args.beta = 0.0f; |
| 336 | fusion_args.beta_ptr_array = nullptr; |
| 337 | fusion_args.dBeta = {_0{}, _0{}, 0}; |
| 338 | |
| 339 | // Gemm Arguments |
| 340 | typename GemmKernel::Arguments args{ |
| 341 | cutlass::gemm::GemmUniversalMode::kGrouped, |
| 342 | {num_experts, problem_sizes_as_shapes, nullptr}, |
| 343 | mainloop_args, |
| 344 | epilogue_args, |
| 345 | hw_info, |
| 346 | scheduler}; |
| 347 | |
| 348 | size_t workspace_size = Gemm::get_workspace_size(args); |
| 349 | auto workspace = |
| 350 | torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte, |
| 351 | std::nullopt, a.device()); |
| 352 | const cudaStream_t stream = get_current_cuda_stream(a.get_device_index()); |
| 353 | |
| 354 | auto can_implement_status = gemm_op.can_implement(args); |
| 355 | STD_TORCH_CHECK( |
| 356 | can_implement_status == cutlass::Status::kSuccess, |
| 357 | "Failed to implement MXFP4 GEMM: status=", (int)can_implement_status); |
| 358 | |
| 359 | // Run the GEMM |
| 360 | auto status = gemm_op.initialize(args, workspace.data_ptr()); |
| 361 | STD_TORCH_CHECK(status == cutlass::Status::kSuccess, |
| 362 | "Failed to initialize MXFP4 GEMM: status=", (int)status, |
| 363 | " workspace_size=", workspace_size, |
| 364 | " num_experts=", num_experts, " M=", M, " N=", N, " K=", K); |
| 365 | |
| 366 | status = gemm_op.run(args, workspace.data_ptr(), stream); |
| 367 | STD_TORCH_CHECK(status == cutlass::Status::kSuccess, |
| 368 | "Failed to run MXFP4 GEMM"); |
| 369 | } |
| 370 | |
| 371 | template <typename OutType> |
| 372 | void run_mxfp4_blockwise_scaled_group_mm( |
| 373 | torch::stable::Tensor& output, const torch::stable::Tensor& a, |
| 374 | const torch::stable::Tensor& b, const torch::stable::Tensor& a_blockscale, |
| 375 | const torch::stable::Tensor& b_blockscales, |
| 376 | const torch::stable::Tensor& problem_sizes, |
| 377 | const torch::stable::Tensor& expert_offsets, |
| 378 | const torch::stable::Tensor& sf_offsets, int M, int N, int K) { |
| 379 | int32_t version_num = get_sm_version_num(); |
| 380 | #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 |
| 381 | if (version_num >= 100 && version_num < 120) { |
| 382 | run_mxfp4_blockwise_scaled_group_mm_sm100<OutType>( |
| 383 | output, a, b, a_blockscale, b_blockscales, problem_sizes, |
| 384 | expert_offsets, sf_offsets, M, N, K); |
| 385 | return; |
| 386 | } |
| 387 | #endif |
| 388 | STD_TORCH_CHECK_NOT_IMPLEMENTED( |
| 389 | false, |
| 390 | "No compiled cutlass_mxfp4_group_mm kernel for CUDA device capability: ", |
| 391 | version_num, ". Required capability: 100"); |
| 392 | } |
| 393 | |
| 394 | #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 |
| 395 | constexpr auto MXFP4_FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte; |
| 396 | // E8M0 scale factors stored as uint8 |
| 397 | constexpr auto MXFP4_SF_DTYPE = torch::headeronly::ScalarType::Byte; |
| 398 | #endif |
| 399 | |
| 400 | #define CHECK_TYPE(x, st, m) \ |
| 401 | STD_TORCH_CHECK(x.scalar_type() == st, \ |
| 402 | ": Inconsistency of torch::stable::Tensor type:", m) |
| 403 | #define CHECK_TH_CUDA(x, m) \ |
| 404 | STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.") |
| 405 | #define CHECK_CONTIGUOUS(x, m) \ |
| 406 | STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.") |
| 407 | #define CHECK_INPUT(x, st, m) \ |
| 408 | CHECK_TH_CUDA(x, m); \ |
| 409 | CHECK_CONTIGUOUS(x, m); \ |
| 410 | CHECK_TYPE(x, st, m) |
| 411 | |
| 412 | void cutlass_mxfp4_group_mm(torch::stable::Tensor& output, |
| 413 | const torch::stable::Tensor& a, |
| 414 | const torch::stable::Tensor& b, |
| 415 | const torch::stable::Tensor& a_blockscale, |
| 416 | const torch::stable::Tensor& b_blockscales, |
| 417 | const torch::stable::Tensor& problem_sizes, |
| 418 | const torch::stable::Tensor& expert_offsets, |
| 419 | const torch::stable::Tensor& sf_offsets) { |
| 420 | #if defined ENABLE_NVFP4_SM100 && ENABLE_NVFP4_SM100 |
| 421 | // Input validation |
| 422 | CHECK_INPUT(a, MXFP4_FLOAT4_E2M1X2, "a"); |
| 423 | CHECK_INPUT(b, MXFP4_FLOAT4_E2M1X2, "b"); |
| 424 | // MXFP4 uses E8M0 scale factors (stored as uint8) |
| 425 | CHECK_INPUT(a_blockscale, MXFP4_SF_DTYPE, "a_blockscale"); |
| 426 | CHECK_INPUT(b_blockscales, MXFP4_SF_DTYPE, "b_blockscales"); |
| 427 | |
| 428 | STD_TORCH_CHECK( |
| 429 | a_blockscale.dim() == 2, |
| 430 | "expected a_blockscale to be of shape [num_experts, rounded_m," |
| 431 | " k // group_size], observed rank: ", |
| 432 | a_blockscale.dim()) |
| 433 | STD_TORCH_CHECK(b_blockscales.dim() == 3, |
| 434 | "expected b_blockscale to be of shape: " |
| 435 | " [num_experts, n, k // group_size], observed rank: ", |
| 436 | b_blockscales.dim()) |
| 437 | STD_TORCH_CHECK(problem_sizes.dim() == 2, |
| 438 | "problem_sizes must be a 2D tensor"); |
| 439 | STD_TORCH_CHECK(problem_sizes.size(1) == 3, |
| 440 | "problem_sizes must have the shape (num_experts, 3)"); |
| 441 | STD_TORCH_CHECK( |
| 442 | problem_sizes.size(0) == expert_offsets.size(0), |
| 443 | "Number of experts in problem_sizes must match expert_offsets"); |
| 444 | STD_TORCH_CHECK( |
| 445 | problem_sizes.scalar_type() == torch::headeronly::ScalarType::Int, |
| 446 | "problem_sizes must be int32."); |
| 447 | |
| 448 | int M = static_cast<int>(a.size(0)); |
| 449 | int N = static_cast<int>(b.size(1)); |
| 450 | int E = static_cast<int>(b.size(0)); |
| 451 | int K = static_cast<int>(2 * b.size(2)); |
| 452 | |
| 453 | if (output.scalar_type() == torch::headeronly::ScalarType::BFloat16) { |
| 454 | run_mxfp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>( |
| 455 | output, a, b, a_blockscale, b_blockscales, problem_sizes, |
| 456 | expert_offsets, sf_offsets, M, N, K); |
| 457 | } else { |
| 458 | run_mxfp4_blockwise_scaled_group_mm<cutlass::half_t>( |
| 459 | output, a, b, a_blockscale, b_blockscales, problem_sizes, |
| 460 | expert_offsets, sf_offsets, M, N, K); |
| 461 | } |
| 462 | #else |
| 463 | STD_TORCH_CHECK_NOT_IMPLEMENTED( |
| 464 | false, |
| 465 | "No compiled cutlass_mxfp4_group_mm kernel; build vLLM with " |
| 466 | "SM100 block-scaled FP4 MoE (ENABLE_NVFP4_SM100) and CUDA 12.8+."); |
| 467 | #endif |
| 468 | } |
| 469 | |
| 470 | STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { |
| 471 | m.impl("cutlass_mxfp4_group_mm", TORCH_BOX(&cutlass_mxfp4_group_mm)); |
| 472 | } |