| 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 <cuda_runtime_api.h> |
| 20 | #include <cuda_runtime.h> |
| 21 | |
| 22 | #include <cuda_fp8.h> |
| 23 | |
| 24 | #include "libtorch_stable/torch_utils.h" |
| 25 | #include "libtorch_stable/dispatch_utils.h" |
| 26 | #include "../../cuda_vec_utils.cuh" |
| 27 | |
| 28 | #include "cuda_utils.h" |
| 29 | #include "libtorch_stable/launch_bounds_utils.h" |
| 30 | |
| 31 | // Define before including nvfp4_utils.cuh so the header |
| 32 | // can use this macro during compilation. |
| 33 | #define NVFP4_ENABLE_ELTS16 1 |
| 34 | #include "nvfp4_utils.cuh" |
| 35 | |
| 36 | namespace vllm { |
| 37 | |
| 38 | // Use UE4M3 by default. |
| 39 | template <class Type, bool UE8M0_SF = false, bool PADDED_OUTPUT = false> |
| 40 | __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) |
| 41 | cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, int32_t outputCols, |
| 42 | int32_t num_padded_cols, Type const* __restrict__ in, |
| 43 | float const* __restrict__ SFScale, |
| 44 | uint32_t* __restrict__ out, uint32_t* __restrict__ SFout) { |
| 45 | using PackedVec = vllm::PackedVec<Type, CVT_FP4_PACK16>; |
| 46 | |
| 47 | static constexpr int CVT_FP4_NUM_THREADS_PER_SF = |
| 48 | (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); |
| 49 | static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, |
| 50 | "Vec size is not matched."); |
| 51 | |
| 52 | #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) |
| 53 | cudaGridDependencySynchronize(); |
| 54 | cudaTriggerProgrammaticLaunchCompletion(); |
| 55 | #endif |
| 56 | |
| 57 | // Precompute SF layout parameter (constant for entire kernel). |
| 58 | int32_t const numKTiles = ((PADDED_OUTPUT ? outputCols : numCols) + 63) / 64; |
| 59 | |
| 60 | int sf_m = round_up<int>(numRows, 128); |
| 61 | int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x; |
| 62 | int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD; |
| 63 | |
| 64 | // Get the global scaling factor, which will be applied to the SF. |
| 65 | // Note SFScale is the same as next GEMM's alpha, which is |
| 66 | // (448.f / (Alpha_A / 6.f)). |
| 67 | float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0]; |
| 68 | |
| 69 | // Iterate over all rows and cols including padded ones - |
| 70 | // ensures we visit every single scale factor address to initialize it. |
| 71 | for (int rowIdx = blockIdx.x; rowIdx < sf_m; rowIdx += gridDim.x) { |
| 72 | if (colIdx < num_padded_cols) { |
| 73 | PackedVec in_vec; |
| 74 | int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; |
| 75 | |
| 76 | // If we are outside valid columns, feed zeros |
| 77 | bool valid = (rowIdx < numRows) && (elem_idx < numCols); |
| 78 | if constexpr (CVT_FP4_PACK16) { |
| 79 | ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec), |
| 80 | &reinterpret_cast<const uint32_t*>(in)[inOffset * 8], |
| 81 | valid); |
| 82 | } else { |
| 83 | ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec), |
| 84 | &reinterpret_cast<const uint32_t*>(in)[inOffset * 4], |
| 85 | valid); |
| 86 | } |
| 87 | |
| 88 | auto sf_out = |
| 89 | cvt_quant_to_fp4_get_sf_out_offset<uint32_t, |
| 90 | CVT_FP4_NUM_THREADS_PER_SF>( |
| 91 | rowIdx, colIdx, numKTiles, SFout); |
| 92 | |
| 93 | fp4_packed_t out_val{}; |
| 94 | if (valid) { |
| 95 | out_val = |
| 96 | cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>( |
| 97 | in_vec, global_scale, sf_out); |
| 98 | } else if (sf_out != nullptr) { |
| 99 | // The swizzled scale tensor is allocated with torch.empty. Explicitly |
| 100 | // initialize padded rows and columns in this kernel so downstream |
| 101 | // NVFP4 GEMMs never observe stale scale factors. |
| 102 | *sf_out = 0; |
| 103 | } |
| 104 | |
| 105 | if constexpr (PADDED_OUTPUT) { |
| 106 | bool valid_output = (rowIdx < numRows) && (elem_idx < outputCols); |
| 107 | if (valid_output) { |
| 108 | if constexpr (CVT_FP4_PACK16) { |
| 109 | int64_t outOffset = rowIdx * (outputCols / 8) + colIdx * 2; |
| 110 | uint64_t packed64 = |
| 111 | (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo); |
| 112 | reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64; |
| 113 | } else { |
| 114 | int64_t outOffset = |
| 115 | rowIdx * (outputCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; |
| 116 | reinterpret_cast<fp4_packed_t*>(out)[outOffset] = out_val; |
| 117 | } |
| 118 | } |
| 119 | } else if (valid) { |
| 120 | // #32520 fast path: input and output widths are identical, so avoid |
| 121 | // the padded-output predicate and use the original address arithmetic. |
| 122 | if constexpr (CVT_FP4_PACK16) { |
| 123 | int64_t outOffset = rowIdx * (numCols / 8) + colIdx * 2; |
| 124 | uint64_t packed64 = |
| 125 | (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo); |
| 126 | reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64; |
| 127 | } else { |
| 128 | reinterpret_cast<fp4_packed_t*>(out)[inOffset] = out_val; |
| 129 | } |
| 130 | } |
| 131 | } |
| 132 | } |
| 133 | } |
| 134 | |
| 135 | // Use UE4M3 by default. |
| 136 | template <class Type, bool UE8M0_SF = false> |
| 137 | __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) |
| 138 | cvt_fp16_to_fp4_sf_major(int32_t numRows, int32_t numCols, |
| 139 | int32_t outputCols, int32_t sf_n_unpadded, |
| 140 | int32_t num_packed_cols, |
| 141 | Type const* __restrict__ in, |
| 142 | float const* __restrict__ SFScale, |
| 143 | uint32_t* __restrict__ out, |
| 144 | uint32_t* __restrict__ SFout) { |
| 145 | using PackedVec = PackedVec<Type, CVT_FP4_PACK16>; |
| 146 | |
| 147 | static constexpr int CVT_FP4_NUM_THREADS_PER_SF = |
| 148 | (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD); |
| 149 | static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD, |
| 150 | "Vec size is not matched."); |
| 151 | |
| 152 | #if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) |
| 153 | cudaGridDependencySynchronize(); |
| 154 | cudaTriggerProgrammaticLaunchCompletion(); |
| 155 | #endif |
| 156 | |
| 157 | int32_t const colIdx = blockDim.x * blockIdx.y + threadIdx.x; |
| 158 | int elem_idx = colIdx * CVT_FP4_ELTS_PER_THREAD; |
| 159 | |
| 160 | // Get the global scaling factor, which will be applied to the SF. |
| 161 | // Note SFScale is the same as next GEMM's alpha, which is |
| 162 | // (448.f / (Alpha_A / 6.f)). |
| 163 | float const global_scale = (SFScale == nullptr) ? 1.0f : SFScale[0]; |
| 164 | |
| 165 | // Iterate over all rows and cols including padded ones - |
| 166 | // ensures we visit every single scale factor address to initialize it. |
| 167 | for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) { |
| 168 | if (colIdx < num_packed_cols) { |
| 169 | PackedVec in_vec; |
| 170 | int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; |
| 171 | |
| 172 | // If we are outside valid columns, feed zeros |
| 173 | bool valid = (rowIdx < numRows) && (elem_idx < numCols); |
| 174 | if constexpr (CVT_FP4_PACK16) { |
| 175 | ld256_cg_or_zero(reinterpret_cast<u32x8_t&>(in_vec), |
| 176 | &reinterpret_cast<const uint32_t*>(in)[inOffset * 8], |
| 177 | valid); |
| 178 | } else { |
| 179 | ld128_cg_or_zero(reinterpret_cast<uint4&>(in_vec), |
| 180 | &reinterpret_cast<const uint32_t*>(in)[inOffset * 4], |
| 181 | valid); |
| 182 | } |
| 183 | |
| 184 | auto sf_out = |
| 185 | sf_out_rowmajor_u8<uint32_t>(rowIdx, colIdx, sf_n_unpadded, SFout); |
| 186 | |
| 187 | auto out_val = |
| 188 | cvt_warp_fp16_to_fp4<Type, CVT_FP4_NUM_THREADS_PER_SF, UE8M0_SF>( |
| 189 | in_vec, global_scale, sf_out); |
| 190 | |
| 191 | if (rowIdx < numRows) { |
| 192 | if constexpr (CVT_FP4_PACK16) { |
| 193 | int64_t outOffset = rowIdx * (outputCols / 8) + colIdx * 2; |
| 194 | uint64_t packed64 = |
| 195 | (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo); |
| 196 | reinterpret_cast<uint64_t*>(out)[outOffset >> 1] = packed64; |
| 197 | } else { |
| 198 | int64_t outOffset = |
| 199 | rowIdx * (outputCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; |
| 200 | out[outOffset] = out_val; |
| 201 | } |
| 202 | } |
| 203 | } |
| 204 | } |
| 205 | } |
| 206 | |
| 207 | } // namespace vllm |
| 208 | |
| 209 | void scaled_fp4_quant_sm1xxa(torch::stable::Tensor const& output, |
| 210 | torch::stable::Tensor const& input, |
| 211 | torch::stable::Tensor const& output_sf, |
| 212 | torch::stable::Tensor const& input_sf, |
| 213 | bool is_sf_swizzled_layout) { |
| 214 | int32_t m = input.size(0); |
| 215 | int32_t n = input.size(1); |
| 216 | int32_t output_n = output.size(1) * 2; |
| 217 | |
| 218 | STD_TORCH_CHECK(n % 16 == 0, "The N dimension must be multiple of 16."); |
| 219 | STD_TORCH_CHECK(output_n % 16 == 0, |
| 220 | "The output tensor width must be a multiple of 16."); |
| 221 | STD_TORCH_CHECK( |
| 222 | input.scalar_type() == torch::headeronly::ScalarType::Half || |
| 223 | input.scalar_type() == torch::headeronly::ScalarType::BFloat16, |
| 224 | "Unsupported input data type for quantize_to_fp4."); |
| 225 | |
| 226 | int multiProcessorCount = |
| 227 | get_device_attribute(cudaDevAttrMultiProcessorCount, -1); |
| 228 | |
| 229 | auto input_sf_ptr = static_cast<float const*>(input_sf.data_ptr()); |
| 230 | auto sf_out = static_cast<int32_t*>(output_sf.data_ptr()); |
| 231 | auto output_ptr = static_cast<int64_t*>(output.data_ptr()); |
| 232 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 233 | input.get_device_index()); |
| 234 | auto stream = get_current_cuda_stream(input.get_device_index()); |
| 235 | |
| 236 | int output_sf_n_unpadded = int(output_n / CVT_FP4_SF_VEC_SIZE); |
| 237 | |
| 238 | // Grid, Block size. Each thread converts 8 values. |
| 239 | dim3 block(std::min(int(n / ELTS_PER_THREAD), 512)); |
| 240 | int const numBlocksPerSM = |
| 241 | vllm_runtime_blocks_per_sm(static_cast<int>(block.x)); |
| 242 | |
| 243 | if (is_sf_swizzled_layout) { |
| 244 | int sf_n_int = int(vllm::round_up(output_sf_n_unpadded, 4) / 4); |
| 245 | int32_t num_padded_cols = |
| 246 | sf_n_int * 4 * CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD; |
| 247 | |
| 248 | int grid_y = vllm::div_round_up(num_padded_cols, static_cast<int>(block.x)); |
| 249 | int grid_x = |
| 250 | std::min(vllm::computeEffectiveRows(m), |
| 251 | std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); |
| 252 | dim3 grid(grid_x, grid_y); |
| 253 | |
| 254 | VLLM_STABLE_DISPATCH_HALF_TYPES( |
| 255 | input.scalar_type(), "nvfp4_quant_kernel", [&] { |
| 256 | using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type; |
| 257 | auto input_ptr = static_cast<cuda_type const*>(input.data_ptr()); |
| 258 | cudaLaunchConfig_t config = {}; |
| 259 | config.gridDim = grid; |
| 260 | config.blockDim = block; |
| 261 | config.dynamicSmemBytes = 0; |
| 262 | config.stream = stream; |
| 263 | cudaLaunchAttribute attrs[1]; |
| 264 | attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 265 | attrs[0].val.programmaticStreamSerializationAllowed = 1; |
| 266 | config.numAttrs = 1; |
| 267 | config.attrs = attrs; |
| 268 | if (output_n == n) { |
| 269 | cudaLaunchKernelEx( |
| 270 | &config, vllm::cvt_fp16_to_fp4<cuda_type, false, false>, m, n, |
| 271 | output_n, num_padded_cols, input_ptr, input_sf_ptr, |
| 272 | reinterpret_cast<uint32_t*>(output_ptr), |
| 273 | reinterpret_cast<uint32_t*>(sf_out)); |
| 274 | } else { |
| 275 | cudaLaunchKernelEx( |
| 276 | &config, vllm::cvt_fp16_to_fp4<cuda_type, false, true>, m, n, |
| 277 | output_n, num_padded_cols, input_ptr, input_sf_ptr, |
| 278 | reinterpret_cast<uint32_t*>(output_ptr), |
| 279 | reinterpret_cast<uint32_t*>(sf_out)); |
| 280 | } |
| 281 | }); |
| 282 | } else { |
| 283 | int num_packed_cols = output_n / CVT_FP4_ELTS_PER_THREAD; |
| 284 | int grid_y = vllm::div_round_up(num_packed_cols, static_cast<int>(block.x)); |
| 285 | int grid_x = std::min( |
| 286 | m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); |
| 287 | dim3 grid(grid_x, grid_y); |
| 288 | |
| 289 | VLLM_STABLE_DISPATCH_HALF_TYPES( |
| 290 | input.scalar_type(), "nvfp4_quant_kernel", [&] { |
| 291 | using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type; |
| 292 | auto input_ptr = static_cast<cuda_type const*>(input.data_ptr()); |
| 293 | cudaLaunchConfig_t config = {}; |
| 294 | config.gridDim = grid; |
| 295 | config.blockDim = block; |
| 296 | config.dynamicSmemBytes = 0; |
| 297 | config.stream = stream; |
| 298 | cudaLaunchAttribute attrs[1]; |
| 299 | attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; |
| 300 | attrs[0].val.programmaticStreamSerializationAllowed = 1; |
| 301 | config.numAttrs = 1; |
| 302 | config.attrs = attrs; |
| 303 | cudaLaunchKernelEx( |
| 304 | &config, vllm::cvt_fp16_to_fp4_sf_major<cuda_type, false>, m, n, |
| 305 | output_n, output_sf_n_unpadded, num_packed_cols, input_ptr, |
| 306 | input_sf_ptr, reinterpret_cast<uint32_t*>(output_ptr), |
| 307 | reinterpret_cast<uint32_t*>(sf_out)); |
| 308 | }); |
| 309 | } |
| 310 | } |