| 1 | /* |
| 2 | * Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. |
| 3 | * |
| 4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | * you may not use this file except in compliance with the License. |
| 6 | * You may obtain a copy of the License at |
| 7 | * |
| 8 | * http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | * |
| 10 | * Unless required by applicable law or agreed to in writing, software |
| 11 | * distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | * See the License for the specific language governing permissions and |
| 14 | * limitations under the License. |
| 15 | */ |
| 16 | |
| 17 | #include <torch/csrc/stable/tensor.h> |
| 18 | |
| 19 | #include "libtorch_stable/torch_utils.h" |
| 20 | |
| 21 | #include "libtorch_stable/cutlass_extensions/common.hpp" |
| 22 | |
| 23 | #include "cutlass/cutlass.h" |
| 24 | |
| 25 | #include <type_traits> |
| 26 | |
| 27 | #include "cutlass/gemm/collective/collective_builder.hpp" |
| 28 | #include "cutlass/epilogue/collective/collective_builder.hpp" |
| 29 | #include "cutlass/gemm/device/gemm_universal_adapter.h" |
| 30 | #include "cutlass/gemm/kernel/gemm_universal.hpp" |
| 31 | |
| 32 | #include "cutlass/util/packed_stride.hpp" |
| 33 | |
| 34 | #include "libtorch_stable/core/math.hpp" |
| 35 | #include "core/batch_invariant.hpp" |
| 36 | |
| 37 | using namespace cute; |
| 38 | |
| 39 | #if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) |
| 40 | |
| 41 | // Configuration for M in (256, inf), also reused for batch-invariant mode |
| 42 | // to keep a fixed large-M tiling across all batch sizes. |
| 43 | // Do not change the tile K or tile scheduler here unless you are also |
| 44 | // updating the batch-invariant behavior; if batch-invariant mode needs a |
| 45 | // different schedule, add a dedicated batch-invariant config/path instead. |
| 46 | struct sm100_fp4_config_default { |
| 47 | using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; |
| 48 | using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; |
| 49 | using TileScheduler = cutlass::gemm::PersistentScheduler; |
| 50 | using TileShape = Shape<_256, _256, _256>; |
| 51 | using ClusterShape = Shape<_2, _1, _1>; |
| 52 | using PerSmTileShape_MNK = Shape<_128, _256, _256>; |
| 53 | }; |
| 54 | |
| 55 | // Configuration for M in (16, 256] |
| 56 | struct sm100_fp4_config_M256 { |
| 57 | using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; |
| 58 | using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; |
| 59 | using TileScheduler = void; |
| 60 | using TileShape = Shape<_256, _128, _256>; |
| 61 | using ClusterShape = Shape<_2, _1, _1>; |
| 62 | using PerSmTileShape_MNK = Shape<_128, _128, _256>; |
| 63 | }; |
| 64 | |
| 65 | // Configuration for M in [1, 16] |
| 66 | struct sm100_fp4_config_M16 { |
| 67 | using KernelSchedule = cutlass::gemm::collective::KernelScheduleAuto; |
| 68 | using EpilogueSchedule = cutlass::epilogue::collective::EpilogueScheduleAuto; |
| 69 | using TileScheduler = void; |
| 70 | using TileShape = Shape<_128, _128, _256>; |
| 71 | using ClusterShape = Shape<_1, _1, _1>; |
| 72 | using PerSmTileShape_MNK = Shape<_128, _128, _256>; |
| 73 | }; |
| 74 | |
| 75 | template <typename Config, typename OutType> |
| 76 | struct Fp4GemmSm100 { |
| 77 | // A matrix configuration |
| 78 | using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>; |
| 79 | using LayoutATag = cutlass::layout::RowMajor; |
| 80 | static constexpr int AlignmentA = 32; |
| 81 | |
| 82 | // B matrix configuration |
| 83 | using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>; |
| 84 | using LayoutBTag = cutlass::layout::ColumnMajor; |
| 85 | static constexpr int AlignmentB = 32; |
| 86 | |
| 87 | // C/D matrix configuration |
| 88 | using ElementD = OutType; |
| 89 | using ElementC = OutType; |
| 90 | using LayoutCTag = cutlass::layout::RowMajor; |
| 91 | using LayoutDTag = cutlass::layout::RowMajor; |
| 92 | static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value; |
| 93 | static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value; |
| 94 | |
| 95 | // Kernel functional config |
| 96 | using ElementAccumulator = float; |
| 97 | using ArchTag = cutlass::arch::Sm100; |
| 98 | using OperatorClass = cutlass::arch::OpClassBlockScaledTensorOp; |
| 99 | |
| 100 | // Use config's tile shapes |
| 101 | using MmaTileShape = typename Config::TileShape; |
| 102 | using ClusterShape = typename Config::ClusterShape; |
| 103 | using PerSmTileShape_MNK = typename Config::PerSmTileShape_MNK; |
| 104 | |
| 105 | using CollectiveEpilogue = |
| 106 | typename cutlass::epilogue::collective::CollectiveBuilder< |
| 107 | ArchTag, OperatorClass, PerSmTileShape_MNK, ClusterShape, |
| 108 | cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator, |
| 109 | ElementAccumulator, ElementC, LayoutCTag, AlignmentC, ElementD, |
| 110 | LayoutDTag, AlignmentD, |
| 111 | typename Config::EpilogueSchedule>::CollectiveOp; |
| 112 | |
| 113 | using CollectiveMainloop = |
| 114 | typename cutlass::gemm::collective::CollectiveBuilder< |
| 115 | ArchTag, OperatorClass, ElementA, LayoutATag, AlignmentA, ElementB, |
| 116 | LayoutBTag, AlignmentB, ElementAccumulator, MmaTileShape, |
| 117 | ClusterShape, |
| 118 | cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>( |
| 119 | sizeof(typename CollectiveEpilogue::SharedStorage))>, |
| 120 | typename Config::KernelSchedule>::CollectiveOp; |
| 121 | |
| 122 | using TileScheduler = typename Config::TileScheduler; |
| 123 | using GemmKernel = |
| 124 | cutlass::gemm::kernel::GemmUniversal<Shape<int, int, int, int>, |
| 125 | CollectiveMainloop, |
| 126 | CollectiveEpilogue, TileScheduler>; |
| 127 | using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>; |
| 128 | using StrideA = typename Gemm::GemmKernel::StrideA; |
| 129 | using LayoutA = decltype(cute::make_layout(make_shape(0, 0, 0), StrideA{})); |
| 130 | using LayoutSFA = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFA; |
| 131 | using StrideB = typename Gemm::GemmKernel::StrideB; |
| 132 | using LayoutB = decltype(cute::make_layout(make_shape(0, 0, 0), StrideB{})); |
| 133 | using LayoutSFB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutSFB; |
| 134 | using StrideC = typename Gemm::GemmKernel::StrideC; |
| 135 | using LayoutC = decltype(cute::make_layout(make_shape(0, 0, 0), StrideC{})); |
| 136 | using StrideD = typename Gemm::GemmKernel::StrideD; |
| 137 | using LayoutD = decltype(cute::make_layout(make_shape(0, 0, 0), StrideD{})); |
| 138 | }; |
| 139 | |
| 140 | template <typename Config> |
| 141 | typename Config::Gemm::Arguments args_from_options( |
| 142 | torch::stable::Tensor& D, torch::stable::Tensor const& A, |
| 143 | torch::stable::Tensor const& B, torch::stable::Tensor const& A_sf, |
| 144 | torch::stable::Tensor const& B_sf, torch::stable::Tensor const& alpha, |
| 145 | int64_t M, int64_t N, int64_t K) { |
| 146 | using ElementA = typename Config::Gemm::ElementA; |
| 147 | using ElementB = typename Config::Gemm::ElementB; |
| 148 | using ElementSFA = cutlass::float_ue4m3_t; |
| 149 | using ElementSFB = cutlass::float_ue4m3_t; |
| 150 | using ElementD = typename Config::Gemm::ElementD; |
| 151 | using ElementCompute = float; |
| 152 | using StrideA = typename Config::StrideA; |
| 153 | using StrideB = typename Config::StrideB; |
| 154 | using StrideD = typename Config::StrideD; |
| 155 | using Sm100BlkScaledConfig = typename Config::Gemm::GemmKernel:: |
| 156 | CollectiveMainloop::Sm1xxBlkScaledConfig; |
| 157 | |
| 158 | int m = static_cast<int>(M); |
| 159 | int n = static_cast<int>(N); |
| 160 | int k = static_cast<int>(K); |
| 161 | auto stride_A = cutlass::make_cute_packed_stride(StrideA{}, {m, k, 1}); |
| 162 | auto stride_B = cutlass::make_cute_packed_stride(StrideB{}, {n, k, 1}); |
| 163 | auto stride_D = cutlass::make_cute_packed_stride(StrideD{}, {m, n, 1}); |
| 164 | |
| 165 | auto layout_SFA = Sm100BlkScaledConfig::tile_atom_to_shape_SFA( |
| 166 | cute::make_shape(m, n, k, 1)); |
| 167 | auto layout_SFB = Sm100BlkScaledConfig::tile_atom_to_shape_SFB( |
| 168 | cute::make_shape(m, n, k, 1)); |
| 169 | |
| 170 | typename Config::Gemm::Arguments arguments{ |
| 171 | cutlass::gemm::GemmUniversalMode::kGemm, |
| 172 | {m, n, k, 1}, |
| 173 | {// Mainloop arguments |
| 174 | static_cast<ElementA const*>(A.data_ptr()), stride_A, |
| 175 | static_cast<ElementB const*>(B.data_ptr()), stride_B, |
| 176 | static_cast<ElementSFA const*>(A_sf.data_ptr()), layout_SFA, |
| 177 | static_cast<ElementSFB const*>(B_sf.data_ptr()), layout_SFB}, |
| 178 | { // Epilogue arguments |
| 179 | {}, // epilogue.thread |
| 180 | static_cast<ElementD const*>(D.data_ptr()), |
| 181 | stride_D, |
| 182 | static_cast<ElementD*>(D.data_ptr()), |
| 183 | stride_D}}; |
| 184 | auto& fusion_args = arguments.epilogue.thread; |
| 185 | fusion_args.alpha_ptr = static_cast<ElementCompute const*>(alpha.data_ptr()); |
| 186 | return arguments; |
| 187 | } |
| 188 | |
| 189 | template <typename Config> |
| 190 | void runGemm(torch::stable::Tensor& D, torch::stable::Tensor const& A, |
| 191 | torch::stable::Tensor const& B, torch::stable::Tensor const& A_sf, |
| 192 | torch::stable::Tensor const& B_sf, |
| 193 | torch::stable::Tensor const& alpha, int64_t m, int64_t n, |
| 194 | int64_t k, cudaStream_t stream) { |
| 195 | typename Config::Gemm gemm; |
| 196 | |
| 197 | auto arguments = |
| 198 | args_from_options<Config>(D, A, B, A_sf, B_sf, alpha, m, n, k); |
| 199 | |
| 200 | size_t workspace_size = Config::Gemm::get_workspace_size(arguments); |
| 201 | auto workspace = |
| 202 | torch::stable::empty(workspace_size, torch::headeronly::ScalarType::Byte, |
| 203 | std::nullopt, A.device()); |
| 204 | |
| 205 | CUTLASS_CHECK(gemm.can_implement(arguments)); |
| 206 | |
| 207 | CUTLASS_CHECK(gemm.initialize(arguments, workspace.data_ptr(), stream)); |
| 208 | |
| 209 | CUTLASS_CHECK(gemm.run(arguments, workspace.data_ptr(), stream)); |
| 210 | } |
| 211 | |
| 212 | // Dispatch function to select appropriate config based on M |
| 213 | template <typename OutType> |
| 214 | void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D, |
| 215 | torch::stable::Tensor const& A, |
| 216 | torch::stable::Tensor const& B, |
| 217 | torch::stable::Tensor const& A_sf, |
| 218 | torch::stable::Tensor const& B_sf, |
| 219 | torch::stable::Tensor const& alpha, int64_t m, |
| 220 | int64_t n, int64_t k, cudaStream_t stream) { |
| 221 | if (vllm::vllm_is_batch_invariant()) { |
| 222 | using BiGemm = Fp4GemmSm100<sm100_fp4_config_default, OutType>; |
| 223 | static_assert( |
| 224 | cute::is_same_v<typename BiGemm::TileScheduler, |
| 225 | cutlass::gemm::PersistentScheduler>, |
| 226 | "batch_invariant requires a persistent tile scheduler; stream-K or " |
| 227 | "split-K would break numerical invariance"); |
| 228 | runGemm<BiGemm>(D, A, B, A_sf, B_sf, alpha, m, n, k, stream); |
| 229 | return; |
| 230 | } |
| 231 | |
| 232 | uint32_t const mp2 = std::max(static_cast<uint32_t>(16), next_pow_2(m)); |
| 233 | |
| 234 | if (mp2 <= 16) { |
| 235 | // m in [1, 16] |
| 236 | runGemm<Fp4GemmSm100<sm100_fp4_config_M16, OutType>>( |
| 237 | D, A, B, A_sf, B_sf, alpha, m, n, k, stream); |
| 238 | } else if (mp2 <= 256) { |
| 239 | // m in (16, 256] |
| 240 | runGemm<Fp4GemmSm100<sm100_fp4_config_M256, OutType>>( |
| 241 | D, A, B, A_sf, B_sf, alpha, m, n, k, stream); |
| 242 | } else { |
| 243 | // m in (256, inf) |
| 244 | runGemm<Fp4GemmSm100<sm100_fp4_config_default, OutType>>( |
| 245 | D, A, B, A_sf, B_sf, alpha, m, n, k, stream); |
| 246 | } |
| 247 | } |
| 248 | |
| 249 | #else |
| 250 | template <typename OutType> |
| 251 | void cutlass_fp4_gemm_dispatch(torch::stable::Tensor& D, |
| 252 | torch::stable::Tensor const& A, |
| 253 | torch::stable::Tensor const& B, |
| 254 | torch::stable::Tensor const& A_sf, |
| 255 | torch::stable::Tensor const& B_sf, |
| 256 | torch::stable::Tensor const& alpha, int64_t m, |
| 257 | int64_t n, int64_t k, cudaStream_t stream) { |
| 258 | STD_TORCH_CHECK(false, |
| 259 | "Unsupported CUTLASS version. Set VLLM_CUTLASS_SRC_DIR to " |
| 260 | "a CUTLASS 3.8 source directory to enable support."); |
| 261 | } |
| 262 | #endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) |
| 263 | |
| 264 | #define CHECK_TYPE(x, st, m) \ |
| 265 | STD_TORCH_CHECK(x.scalar_type() == st, \ |
| 266 | ": Inconsistency of torch::stable::Tensor type:", m) |
| 267 | #define CHECK_TH_CUDA(x, m) \ |
| 268 | STD_TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor") |
| 269 | #define CHECK_CONTIGUOUS(x, m) \ |
| 270 | STD_TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous") |
| 271 | #define CHECK_INPUT(x, st, m) \ |
| 272 | CHECK_TH_CUDA(x, m); \ |
| 273 | CHECK_CONTIGUOUS(x, m); \ |
| 274 | CHECK_TYPE(x, st, m) |
| 275 | |
| 276 | constexpr auto FLOAT4_E2M1X2 = torch::headeronly::ScalarType::Byte; |
| 277 | constexpr auto SF_DTYPE = torch::headeronly::ScalarType::Float8_e4m3fn; |
| 278 | |
| 279 | void cutlass_scaled_fp4_mm_sm100a(torch::stable::Tensor& D, |
| 280 | torch::stable::Tensor const& A, |
| 281 | torch::stable::Tensor const& B, |
| 282 | torch::stable::Tensor const& A_sf, |
| 283 | torch::stable::Tensor const& B_sf, |
| 284 | torch::stable::Tensor const& alpha) { |
| 285 | CHECK_INPUT(A, FLOAT4_E2M1X2, "a"); |
| 286 | CHECK_INPUT(B, FLOAT4_E2M1X2, "b"); |
| 287 | |
| 288 | CHECK_INPUT(A_sf, SF_DTYPE, "scale_a"); |
| 289 | CHECK_INPUT(B_sf, SF_DTYPE, "scale_b"); |
| 290 | |
| 291 | CHECK_INPUT(alpha, torch::headeronly::ScalarType::Float, "alpha"); |
| 292 | |
| 293 | STD_TORCH_CHECK(A.dim() == 2, "a must be a matrix"); |
| 294 | STD_TORCH_CHECK(B.dim() == 2, "b must be a matrix"); |
| 295 | STD_TORCH_CHECK(A.size(1) == B.size(1), |
| 296 | "a and b shapes cannot be multiplied (", A.size(0), "x", |
| 297 | A.size(1), " and ", B.size(0), "x", B.size(1), ")"); |
| 298 | |
| 299 | auto const m = A.size(0); |
| 300 | auto const n = B.size(0); |
| 301 | auto const k = A.size(1) * 2; |
| 302 | |
| 303 | constexpr int alignment = 32; |
| 304 | STD_TORCH_CHECK(k % alignment == 0, "Expected k to be divisible by ", |
| 305 | alignment, ", but got a shape: (", A.size(0), "x", A.size(1), |
| 306 | "), k: ", k, "."); |
| 307 | STD_TORCH_CHECK(n % alignment == 0, "Expected n to be divisible by ", |
| 308 | alignment, ", but got b shape: (", B.size(0), "x", B.size(1), |
| 309 | ")."); |
| 310 | |
| 311 | auto round_up = [](int x, int y) { return (x + y - 1) / y * y; }; |
| 312 | int rounded_m = round_up(m, 128); |
| 313 | int rounded_n = round_up(n, 128); |
| 314 | // Since k is divisible by 32 (alignment), k / 16 is guaranteed to be an |
| 315 | // integer. |
| 316 | int rounded_k = round_up(k / 16, 4); |
| 317 | |
| 318 | STD_TORCH_CHECK(A_sf.dim() == 2, "scale_a must be a matrix"); |
| 319 | STD_TORCH_CHECK(B_sf.dim() == 2, "scale_b must be a matrix"); |
| 320 | STD_TORCH_CHECK(A_sf.size(1) == B_sf.size(1), |
| 321 | "scale_a and scale_b shapes cannot be multiplied (", |
| 322 | A_sf.size(0), "x", A_sf.size(1), " and ", B_sf.size(0), "x", |
| 323 | B_sf.size(1), ")"); |
| 324 | STD_TORCH_CHECK(A_sf.size(0) == rounded_m && A_sf.size(1) == rounded_k, |
| 325 | "scale_a must be padded and swizzled to a shape (", rounded_m, |
| 326 | "x", rounded_k, "), but got a shape (", A_sf.size(0), "x", |
| 327 | A_sf.size(1), ")"); |
| 328 | STD_TORCH_CHECK(B_sf.size(0) == rounded_n && B_sf.size(1) == rounded_k, |
| 329 | "scale_b must be padded and swizzled to a shape (", rounded_n, |
| 330 | "x", rounded_k, "), but got a shape (", B_sf.size(0), "x", |
| 331 | B_sf.size(1), ")"); |
| 332 | |
| 333 | auto out_dtype = D.scalar_type(); |
| 334 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 335 | A.get_device_index()); |
| 336 | const cudaStream_t stream = get_current_cuda_stream(A.get_device_index()); |
| 337 | |
| 338 | if (out_dtype == torch::headeronly::ScalarType::Half) { |
| 339 | cutlass_fp4_gemm_dispatch<cutlass::half_t>(D, A, B, A_sf, B_sf, alpha, m, n, |
| 340 | k, stream); |
| 341 | } else if (out_dtype == torch::headeronly::ScalarType::BFloat16) { |
| 342 | cutlass_fp4_gemm_dispatch<cutlass::bfloat16_t>(D, A, B, A_sf, B_sf, alpha, |
| 343 | m, n, k, stream); |
| 344 | } else { |
| 345 | STD_TORCH_CHECK(false, "Unsupported output data type of nvfp4 mm (", |
| 346 | out_dtype, ")"); |
| 347 | } |
| 348 | } |