| 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 | #pragma once |
| 18 | |
| 19 | #include <cuda_runtime.h> |
| 20 | #include <cuda_fp8.h> |
| 21 | #include <utility> |
| 22 | |
| 23 | #include "../../cuda_vec_utils.cuh" |
| 24 | |
| 25 | #if defined(NVFP4_ENABLE_ELTS16) && defined(CUDART_VERSION) && \ |
| 26 | CUDART_VERSION >= 12090 |
| 27 | #define ELTS_PER_THREAD 16 |
| 28 | #define CVT_FP4_PACK16 1 |
| 29 | constexpr int CVT_FP4_ELTS_PER_THREAD = 16; |
| 30 | #else |
| 31 | #define ELTS_PER_THREAD 8 |
| 32 | #define CVT_FP4_PACK16 0 |
| 33 | constexpr int CVT_FP4_ELTS_PER_THREAD = 8; |
| 34 | #endif |
| 35 | |
| 36 | constexpr int CVT_FP4_SF_VEC_SIZE = 16; |
| 37 | |
| 38 | namespace vllm { |
| 39 | |
| 40 | template <typename Int> |
| 41 | __host__ __device__ inline Int round_up(Int x, Int y) { |
| 42 | static_assert(std::is_integral_v<Int>, |
| 43 | "round_up argument must be integral type"); |
| 44 | return ((x + y - 1) / y) * y; |
| 45 | } |
| 46 | |
| 47 | template <typename Int> |
| 48 | __host__ __device__ __forceinline__ Int div_round_up(Int x, Int y) { |
| 49 | return (x + y - 1) / y; |
| 50 | } |
| 51 | |
| 52 | // Compute effective rows for grid configuration with swizzled SF layouts. |
| 53 | inline int computeEffectiveRows(int m) { |
| 54 | constexpr int ROW_TILE = 128; |
| 55 | return round_up(m, ROW_TILE); |
| 56 | } |
| 57 | |
| 58 | // Compute the shape of the swizzled SF output tensor. |
| 59 | // Returns (rounded_m, rounded_n / 4) where: |
| 60 | // rounded_m = round_up(m, 128) |
| 61 | // rounded_n = round_up(n / CVT_FP4_SF_VEC_SIZE, 4) |
| 62 | inline std::pair<int64_t, int64_t> computeSwizzledSFShape(int64_t m, |
| 63 | int64_t n) { |
| 64 | int64_t rounded_m = round_up(m, static_cast<int64_t>(128)); |
| 65 | int64_t scale_n = n / CVT_FP4_SF_VEC_SIZE; |
| 66 | int64_t rounded_n = round_up(scale_n, static_cast<int64_t>(4)); |
| 67 | return {rounded_m, rounded_n / 4}; |
| 68 | } |
| 69 | |
| 70 | // Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t). |
| 71 | inline __device__ uint32_t fp32_vec8_to_e2m1(float (&array)[8]) { |
| 72 | uint32_t val; |
| 73 | asm volatile( |
| 74 | "{\n" |
| 75 | ".reg .b8 byte0;\n" |
| 76 | ".reg .b8 byte1;\n" |
| 77 | ".reg .b8 byte2;\n" |
| 78 | ".reg .b8 byte3;\n" |
| 79 | "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" |
| 80 | "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" |
| 81 | "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" |
| 82 | "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" |
| 83 | "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" |
| 84 | "}" |
| 85 | : "=r"(val) |
| 86 | : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]), |
| 87 | "f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7])); |
| 88 | return val; |
| 89 | } |
| 90 | |
| 91 | // Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t). |
| 92 | __device__ __forceinline__ uint32_t fp32_vec8_to_e2m1(float2 (&array)[4]) { |
| 93 | uint32_t val; |
| 94 | asm volatile( |
| 95 | "{\n" |
| 96 | ".reg .b8 byte0;\n" |
| 97 | ".reg .b8 byte1;\n" |
| 98 | ".reg .b8 byte2;\n" |
| 99 | ".reg .b8 byte3;\n" |
| 100 | "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" |
| 101 | "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" |
| 102 | "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n" |
| 103 | "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n" |
| 104 | "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" |
| 105 | "}\n" |
| 106 | : "=r"(val) |
| 107 | : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y), |
| 108 | "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y)); |
| 109 | return val; |
| 110 | } |
| 111 | |
| 112 | struct u32x2 { |
| 113 | uint32_t lo, hi; |
| 114 | }; |
| 115 | |
| 116 | using fp4_packed_t = std::conditional_t<CVT_FP4_PACK16, u32x2, uint32_t>; |
| 117 | |
| 118 | __device__ __forceinline__ u32x2 fp32_vec16_to_e2m1(float2 (&array)[8]) { |
| 119 | u32x2 out; |
| 120 | asm volatile( |
| 121 | "{\n" |
| 122 | ".reg .b8 b0;\n" |
| 123 | ".reg .b8 b1;\n" |
| 124 | ".reg .b8 b2;\n" |
| 125 | ".reg .b8 b3;\n" |
| 126 | ".reg .b8 b4;\n" |
| 127 | ".reg .b8 b5;\n" |
| 128 | ".reg .b8 b6;\n" |
| 129 | ".reg .b8 b7;\n" |
| 130 | "cvt.rn.satfinite.e2m1x2.f32 b0, %3, %2;\n" |
| 131 | "cvt.rn.satfinite.e2m1x2.f32 b1, %5, %4;\n" |
| 132 | "cvt.rn.satfinite.e2m1x2.f32 b2, %7, %6;\n" |
| 133 | "cvt.rn.satfinite.e2m1x2.f32 b3, %9, %8;\n" |
| 134 | "cvt.rn.satfinite.e2m1x2.f32 b4, %11, %10;\n" |
| 135 | "cvt.rn.satfinite.e2m1x2.f32 b5, %13, %12;\n" |
| 136 | "cvt.rn.satfinite.e2m1x2.f32 b6, %15, %14;\n" |
| 137 | "cvt.rn.satfinite.e2m1x2.f32 b7, %17, %16;\n" |
| 138 | "mov.b32 %0, {b0, b1, b2, b3};\n" |
| 139 | "mov.b32 %1, {b4, b5, b6, b7};\n" |
| 140 | "}\n" |
| 141 | : "=r"(out.lo), "=r"(out.hi) |
| 142 | : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y), |
| 143 | "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y), |
| 144 | "f"(array[4].x), "f"(array[4].y), "f"(array[5].x), "f"(array[5].y), |
| 145 | "f"(array[6].x), "f"(array[6].y), "f"(array[7].x), "f"(array[7].y)); |
| 146 | return out; |
| 147 | } |
| 148 | |
| 149 | __device__ __forceinline__ uint32_t pack_fp4(float2 (&v)[4]) { |
| 150 | return fp32_vec8_to_e2m1(v); |
| 151 | } |
| 152 | |
| 153 | __device__ __forceinline__ u32x2 pack_fp4(float2 (&v)[8]) { |
| 154 | return fp32_vec16_to_e2m1(v); |
| 155 | } |
| 156 | |
| 157 | // Fast reciprocal. |
| 158 | __device__ __forceinline__ float reciprocal_approximate_ftz(float a) { |
| 159 | float b; |
| 160 | asm volatile("rcp.approx.ftz.f32 %0, %1;" : "=f"(b) : "f"(a)); |
| 161 | return b; |
| 162 | } |
| 163 | |
| 164 | // Compute SF output offset for swizzled tensor core layout. |
| 165 | // SF layout: [numMTiles, numKTiles, 32, 4, 4] |
| 166 | // Caller must precompute: numKTiles = (numCols + 63) / 64 |
| 167 | template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF> |
| 168 | __device__ __forceinline__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset( |
| 169 | int rowIdx, int colIdx, int32_t numKTiles, SFType* SFout) { |
| 170 | static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 || |
| 171 | CVT_FP4_NUM_THREADS_PER_SF == 2); |
| 172 | |
| 173 | // One pair of threads write one SF to global memory. |
| 174 | // TODO: stage through smem for packed STG.32 |
| 175 | // is it better than STG.8 from 4 threads ? |
| 176 | if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF != 0) { |
| 177 | return nullptr; |
| 178 | } |
| 179 | |
| 180 | // SF vector index (16 elements share one SF in the K dimension). |
| 181 | int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF; |
| 182 | int32_t mIdx = rowIdx; |
| 183 | |
| 184 | // Decompose indices using bitwise ops (all divisors are powers of 2). |
| 185 | // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)] |
| 186 | int32_t mTileIdx = mIdx >> 7; // mIdx / 128 |
| 187 | int32_t outerMIdx = mIdx & 31; // mIdx % 32 |
| 188 | int32_t innerMIdx = (mIdx >> 5) & 3; // (mIdx / 32) % 4 |
| 189 | int32_t kTileIdx = kIdx >> 2; // kIdx / 4 |
| 190 | int32_t innerKIdx = kIdx & 3; // kIdx % 4 |
| 191 | |
| 192 | // Compute global SF offset: mTileIdx * (numKTiles * 512) + kTileIdx * 512 + |
| 193 | // outerMIdx * 16 + innerMIdx * 4 + innerKIdx |
| 194 | // Use bitwise OR for non-overlapping lower bits. |
| 195 | int64_t SFOffset = (static_cast<int64_t>(mTileIdx) * numKTiles + kTileIdx) |
| 196 | << 9 | |
| 197 | (outerMIdx << 4) | (innerMIdx << 2) | innerKIdx; |
| 198 | |
| 199 | return reinterpret_cast<uint8_t*>(SFout) + SFOffset; |
| 200 | } |
| 201 | |
| 202 | template <class SFType> |
| 203 | __device__ __forceinline__ uint8_t* sf_out_rowmajor_u8(int row, int pack, |
| 204 | int packs_per_row_sf, |
| 205 | SFType* SFout) { |
| 206 | constexpr int PACK = CVT_FP4_ELTS_PER_THREAD; |
| 207 | constexpr int THREADS_PER_SF = |
| 208 | CVT_FP4_SF_VEC_SIZE / PACK; // 1 if PACK=16, 2 else PACK=8 |
| 209 | |
| 210 | if (threadIdx.x % THREADS_PER_SF != 0) return nullptr; |
| 211 | |
| 212 | int sf_col = |
| 213 | pack / THREADS_PER_SF; // PACK=16 => sf_col=pack; PACK=8 => sf_col=pack/2 |
| 214 | int64_t off = (int64_t)row * packs_per_row_sf + sf_col; |
| 215 | |
| 216 | return (uint8_t*)SFout + off; |
| 217 | } |
| 218 | |
| 219 | // Quantizes the provided PackedVec into the uint32_t output |
| 220 | template <class Type, int CVT_FP4_NUM_THREADS_PER_SF, bool UE8M0_SF = false> |
| 221 | __device__ __forceinline__ fp4_packed_t cvt_warp_fp16_to_fp4( |
| 222 | PackedVec<Type, CVT_FP4_PACK16>& vec, float SFScaleVal, uint8_t* SFout) { |
| 223 | // Get absolute maximum values among the local 8 values. |
| 224 | auto localMax = __habs2(vec.elts[0]); |
| 225 | |
| 226 | // Local maximum value. |
| 227 | #pragma unroll |
| 228 | for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { |
| 229 | localMax = __hmax2(localMax, __habs2(vec.elts[i])); |
| 230 | } |
| 231 | |
| 232 | // Get the absolute maximum among all 16 values (two threads). |
| 233 | |
| 234 | if constexpr (CVT_FP4_NUM_THREADS_PER_SF == 2) { |
| 235 | localMax = __hmax2(__shfl_xor_sync(0xffffffffu, localMax, 1), localMax); |
| 236 | } |
| 237 | // Get the final absolute maximum values. |
| 238 | float vecMax = float(__hmax(localMax.x, localMax.y)); |
| 239 | |
| 240 | // 8 bits representation of the SF. |
| 241 | float SFValue; |
| 242 | uint8_t fp8SFVal; |
| 243 | |
| 244 | if constexpr (UE8M0_SF) { |
| 245 | // OCP MX spec E8M0 scale computation (MXFP4 path): |
| 246 | // scale_exp = biased_exponent(round_up(vecMax)) - 2 |
| 247 | // -2 because max E2M1 value is 6.0 ≈ 2^2.58; we use 2^2=4 as the |
| 248 | // safe divisor so that max_val / scale <= 6.0 for values near 2^n. |
| 249 | uint32_t max_bits = __float_as_uint(vecMax); |
| 250 | // Add rounding bias at mantissa bit 21 (equivalent to bf16 val_to_add=32 |
| 251 | // at bit 5). Threshold: values with mantissa >= 0.75 (i.e. >= 1.75*2^n) |
| 252 | // round up to the next power of 2. |
| 253 | uint32_t rounded_bits = (max_bits + (1u << 21)) & 0xFF800000u; |
| 254 | uint32_t biased_exp = (rounded_bits >> 23) & 0xFFu; |
| 255 | uint32_t scale_exp = (biased_exp > 2u) ? (biased_exp - 2u) : 0u; |
| 256 | scale_exp = min(scale_exp, 254u); |
| 257 | fp8SFVal = static_cast<uint8_t>(scale_exp); |
| 258 | // Reconstruct scale as float32: scale = 2^(scale_exp - 127) |
| 259 | uint32_t sf_bits = scale_exp << 23; |
| 260 | SFValue = __uint_as_float(sf_bits); |
| 261 | } else { |
| 262 | // NVFP4 path: scale = max / 6.0, stored as E4M3. |
| 263 | SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f)); |
| 264 | // Here SFValue is always positive, so E4M3 is the same as UE4M3. |
| 265 | __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue); |
| 266 | reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp; |
| 267 | // Convert back to fp32. |
| 268 | SFValue = float(tmp); |
| 269 | } |
| 270 | |
| 271 | // Write the SF to global memory (STG.8). |
| 272 | if (SFout) *SFout = fp8SFVal; |
| 273 | |
| 274 | // Get the output scale (= 1 / SFValue for the MXFP4/UE8M0 path where |
| 275 | // SFScaleVal=1). Use exact division for UE8M0 to ensure bit-exact scaling |
| 276 | // that matches the reference QDQ implementation (dividing by a power-of-2 |
| 277 | // scale is exact in IEEE 754). |
| 278 | float outputScale; |
| 279 | if constexpr (UE8M0_SF) { |
| 280 | // SFValue is always a power of 2 for UE8M0, so 1/SFValue is exact. |
| 281 | outputScale = SFValue != 0.0f ? (1.0f / SFValue) : 0.0f; |
| 282 | } else { |
| 283 | // NVFP4 path: use fast approximate reciprocal (original behavior). |
| 284 | outputScale = SFValue != 0.0f |
| 285 | ? reciprocal_approximate_ftz( |
| 286 | SFValue * reciprocal_approximate_ftz(SFScaleVal)) |
| 287 | : 0.0f; |
| 288 | } |
| 289 | |
| 290 | // Convert the input to float. |
| 291 | float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2]; |
| 292 | |
| 293 | #pragma unroll |
| 294 | for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) { |
| 295 | fp2Vals[i] = cast_to_float2(vec.elts[i]); |
| 296 | fp2Vals[i].x *= outputScale; |
| 297 | fp2Vals[i].y *= outputScale; |
| 298 | } |
| 299 | |
| 300 | // Convert to e2m1 values. |
| 301 | return pack_fp4(fp2Vals); |
| 302 | } |
| 303 | |
| 304 | // silu in float32 |
| 305 | __device__ __forceinline__ float silu(float x) { |
| 306 | return __fdividef(x, (1.f + __expf(-x))); |
| 307 | } |
| 308 | |
| 309 | __device__ __forceinline__ float2 silu2(float2 x) { |
| 310 | return make_float2(silu(x.x), silu(x.y)); |
| 311 | } |
| 312 | |
| 313 | template <class Type> |
| 314 | __inline__ __device__ PackedVec<Type, CVT_FP4_PACK16> compute_silu_mul( |
| 315 | const PackedVec<Type, CVT_FP4_PACK16>& x_vec, |
| 316 | const PackedVec<Type, CVT_FP4_PACK16>& y_vec) { |
| 317 | PackedVec<Type, CVT_FP4_PACK16> result; |
| 318 | |
| 319 | #pragma unroll |
| 320 | for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) { |
| 321 | // silu_mul in float32 |
| 322 | using packed_t = typename PackedTypeConverter<Type>::Type; |
| 323 | float2 silu_vec = silu2(cast_to_float2(x_vec.elts[i])); |
| 324 | float2 y_f2 = cast_to_float2(y_vec.elts[i]); |
| 325 | result.elts[i] = cast_to_packed<packed_t>( |
| 326 | make_float2(silu_vec.x * y_f2.x, silu_vec.y * y_f2.y)); |
| 327 | } |
| 328 | return result; |
| 329 | } |
| 330 | |
| 331 | } // namespace vllm |