diff --git a/mlx/backend/cuda/quantized/cuda_fp4.h b/mlx/backend/cuda/quantized/cuda_fp4.h deleted file mode 100644 index 4500be6b..00000000 --- a/mlx/backend/cuda/quantized/cuda_fp4.h +++ /dev/null @@ -1,100 +0,0 @@ -#pragma once - -struct __nv_fp8_e8m0 { - __device__ __nv_fp8_e8m0(float x) { - if (!std::isfinite(x)) { - __x = 0xFF; - return; - } - if (x < 0.0f) { - __x = 0x00; - return; - } - float le = std::log2f(x); - int n = static_cast(std::nearbyintf(le)); - - n = n < -127 ? -127 : n; - n = n > 127 ? 127 : n; - __x = static_cast(n + 127); - } - - __device__ operator float() { - if (__x == 0xFF) { - return std::numeric_limits::quiet_NaN(); - } - return std::ldexp(1.0f, static_cast(__x) - 127); - } - - uint8_t __x{0}; -}; - -struct __nv_fp4_e2m1 { - __device__ __nv_fp4_e2m1(float x) { - if (std::isnan(x)) { - __x = 0x7; - return; - } - - const uint8_t sign_bit = (std::signbit(x)) ? 0x8 : 0x0; - x = std::abs(x); - - if (x > 5.0f) { - __x = 0x7; - } else if (x >= 3.5f) { - __x = 0x6; - } else if (x > 2.5f) { - __x = 0x5; - } else if (x >= 1.75f) { - __x = 0x4; - } else if (x > 1.25f) { - __x = 0x3; - } else if (x >= 0.75f) { - __x = 0x2; - } else if (x > 0.25f) { - __x = 0x1; - } else { - __x = 0x0; - } - __x |= sign_bit; - } - - __device__ operator float() { - static const float LUT[16] = { - 0.0f, - 0.5f, - 1.0f, - 1.5f, - 2.0f, - 3.0f, - 4.0f, - 6.0f, - -0.0f, - -0.5f, - -1.0f, - -1.5f, - -2.0f, - -3.0f, - -4.0f, - -6.0f}; - - return LUT[__x]; - } - uint8_t __x{0}; -}; - -struct __nv_fp4x4_e2m1 { - __device__ operator float4() { - float4 out; - auto bits = __high & 0xf; - out.x = float(*(__nv_fp4_e2m1*)(&bits)); - bits = (__high >> 4) & 0xf; - out.y = float(*(__nv_fp4_e2m1*)(&bits)); - bits = (__low) & 0xf; - out.z = float(*(__nv_fp4_e2m1*)(&bits)); - bits = (__low >> 4) & 0xf; - out.w = float(*(__nv_fp4_e2m1*)(&bits)); - return out; - } - uint8_t __high{0}; - uint8_t __low{0}; -}; diff --git a/mlx/backend/cuda/quantized/fp_quantize.cu b/mlx/backend/cuda/quantized/fp_quantize.cu index a257ff65..5feb07db 100644 --- a/mlx/backend/cuda/quantized/fp_quantize.cu +++ b/mlx/backend/cuda/quantized/fp_quantize.cu @@ -11,8 +11,8 @@ #include #include -#include -#include +#include +#include constexpr float F8E4M3_MAX = 448.0f; constexpr float F4E2M1_MAX = 6.0f; @@ -24,9 +24,9 @@ template struct Dequantize { __device__ float operator()(uint8_t x) { if constexpr (bits == 8) { - return float(*(__nv_fp8_e4m3*)(&x)); + return float(*(cutlass::float_e4m3_t*)(&x)); } else { - return float(*(__nv_fp4_e2m1*)(&x)); + return float(*(cutlass::float_e2m1_t*)(&x)); } } }; @@ -45,7 +45,6 @@ __global__ void fp_quantize_dequantize( const float inv_scale_enc = use_global_scale ? 1.0f / scale_enc : 1.0f; using Tx2 = Vector2_t; - using Tx4 = Vector4_t; uint32_t rbits = 0; // reserved bits for future use auto block_size = cg::this_thread_block().dim_threads(); auto block_idx = cg::this_thread_block().group_index(); @@ -79,30 +78,35 @@ __global__ void fp_quantize_dequantize( scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; scale_dec_b *= scale_enc; // Convert to mx scale or nv scale - using ScaleType = - std::conditional_t; + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; auto s = ScaleType(scale_dec_b); float scale_enc_b = scale_enc / float(s); float scale_dec = float(s) * inv_scale_enc; AlignedVector w_hat; #pragma unroll - for (int i = 0; i < group_size / 4; i++) { - Tx4 w_Tx4 = *reinterpret_cast(&w_tile[i * 4]); - float4 dq; + for (int i = 0; i < group_size / 8; i++) { + auto& w = *reinterpret_cast*>(&w_tile[i * 8]); + cutlass::NumericArrayConverter fp32_t; + auto scaled = fp32_t(w) * scale_enc_b; + cutlass::Array dq; if constexpr (bits == 8) { - uint32_t quantized_val = - scale_cvt_Tx4_to_fp8x4(w_Tx4, scale_enc_b, rbits); - dq = dequant_fp8(quantized_val); + cutlass::NumericArrayConverter fp8_fp32; + auto quant = fp8_fp32(scaled); + cutlass::NumericArrayConverter fp32_fp8; + dq = fp32_fp8(quant); } else { - uint16_t quantized_val = - scale_cvt_Tx4_to_fp4x4(w_Tx4, scale_enc_b, rbits); - dq = dequant_fp4(quantized_val); + cutlass::NumericArrayConverter fp4_fp32; + auto quant = fp4_fp32(scaled); + cutlass::NumericArrayConverter fp32_fp4; + dq = fp32_fp4(quant); } - w_hat[i * 4] = static_cast(dq.x * scale_dec); - w_hat[i * 4 + 1] = static_cast(dq.y * scale_dec); - w_hat[i * 4 + 2] = static_cast(dq.z * scale_dec); - w_hat[i * 4 + 3] = static_cast(dq.w * scale_dec); + cutlass::NumericArrayConverter t_fp32; + *reinterpret_cast*>(&w_hat[i * 8]) = + t_fp32(dq * scale_dec); } store_vector(out, thread_idx, w_hat); } @@ -157,10 +161,12 @@ __global__ void fp_quantize_rowwise( scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; scale_dec_b *= scale_enc; // Convert to mx scale or nv scale - using ScaleType = - std::conditional_t; + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; auto s = ScaleType(scale_dec_b); - uint8_t q_scale = s.__x; + uint8_t q_scale = s.storage; float scale_enc_b = scale_enc / float(s); scales[thread_idx] = q_scale; @@ -256,11 +262,13 @@ __global__ void fp_quantize_columnwise( scale_dec_b /= bits == 4 ? F4E2M1_MAX : F8E4M3_MAX; scale_dec_b *= scale_enc; // Convert to mx scale or nv scale - using ScaleType = - std::conditional_t; + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; auto s = ScaleType(scale_dec_b); float scale_enc_b = scale_enc / float(s); - scales_smem[tidx][tidy] = s.__x; + scales_smem[tidx][tidy] = s.storage; int shared_idx = tidx * padded_local_cols + tidy * bytes_per_group; @@ -345,8 +353,10 @@ __global__ void fp_dequantize( } size_t gindex = oindex / group_size; - using ScaleType = - std::conditional_t; + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; auto scale = float(((ScaleType*)(scales))[gindex]) * inv_scale_enc; out += oindex; diff --git a/mlx/backend/cuda/quantized/mxfp8_quantize.cuh b/mlx/backend/cuda/quantized/mxfp8_quantize.cuh index 1edf3641..ee8494e2 100644 --- a/mlx/backend/cuda/quantized/mxfp8_quantize.cuh +++ b/mlx/backend/cuda/quantized/mxfp8_quantize.cuh @@ -1,32 +1,23 @@ #pragma once -#include -#include -#include #include "mlx/backend/cuda/vector_types.cuh" -namespace mlx::core::cu { +#include -// TODO implement fast path -template -__device__ __forceinline__ uint32_t -scale_cvt_Tx4_to_fp8x4_fallback(const Vector4_t input, const float scale) { - uint32_t out_fp8x4 = 0; - float4 scaled; - scaled.x = static_cast(input.x) * scale; - scaled.y = static_cast(input.y) * scale; - scaled.z = static_cast(input.z) * scale; - scaled.w = static_cast(input.w) * scale; - out_fp8x4 = __nv_fp8x4_e4m3(scaled).__x; - return out_fp8x4; -} +namespace mlx::core::cu { // Place holder for future fast path implementation template __device__ __forceinline__ uint32_t scale_cvt_Tx4_to_fp8x4( - const Vector4_t input, + const Vector4_t& input, const float scale, uint32_t rbits) { - return scale_cvt_Tx4_to_fp8x4_fallback(input, scale); + cutlass::NumericArrayConverter fp32_t; + auto scaled = + fp32_t(*reinterpret_cast*>(&input)) * scale; + cutlass::NumericArrayConverter fp8_fp32; + auto quant = fp8_fp32(scaled); + return *reinterpret_cast(&quant); } -} // namespace mlx::core::cu \ No newline at end of file + +} // namespace mlx::core::cu diff --git a/mlx/backend/cuda/quantized/nvfp4_quantize.cuh b/mlx/backend/cuda/quantized/nvfp4_quantize.cuh index da0df429..7a5616c7 100644 --- a/mlx/backend/cuda/quantized/nvfp4_quantize.cuh +++ b/mlx/backend/cuda/quantized/nvfp4_quantize.cuh @@ -1,10 +1,9 @@ #pragma once -#include -#include -#include #include "mlx/backend/cuda/vector_types.cuh" +#include + namespace mlx::core::cu { using bf16x4 = Vector4_t<__nv_bfloat16>; @@ -13,23 +12,15 @@ using f32x4 = Vector4_t; template __device__ __forceinline__ uint16_t -scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t input, const float scale) { +scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t& input, const float scale) { // Fallback implementation for architectures that do not support cvt // instructions or for cuda versions with no fp4 support (< 12.8) -> scalar - uint16_t out_fp4x4 = 0; - fp32x4 scaled; - scaled.x = static_cast(input.x) * scale; - scaled.y = static_cast(input.y) * scale; - scaled.z = static_cast(input.z) * scale; - scaled.w = static_cast(input.w) * scale; - uint8_t q0 = __nv_fp4_e2m1(scaled.x).__x; - uint8_t q1 = __nv_fp4_e2m1(scaled.y).__x; - uint8_t q2 = __nv_fp4_e2m1(scaled.z).__x; - uint8_t q3 = __nv_fp4_e2m1(scaled.w).__x; - out_fp4x4 = (static_cast(q3) << 12) | - (static_cast(q2) << 8) | (static_cast(q1) << 4) | - static_cast(q0); - return out_fp4x4; + cutlass::NumericArrayConverter fp32_t; + auto scaled = + fp32_t(*reinterpret_cast*>(&input)) * scale; + cutlass::NumericArrayConverter fp4_fp32; + auto quant = fp4_fp32(scaled); + return *reinterpret_cast(&quant); } #if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \ @@ -318,7 +309,7 @@ __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4_fast( template __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4( - const Vector4_t input, + const Vector4_t& input, const float scale, uint32_t rbits) { #if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \ @@ -331,4 +322,5 @@ __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4( return scale_cvt_Tx4_to_fp4x4_fallback(input, scale); #endif } -} // namespace mlx::core::cu \ No newline at end of file + +} // namespace mlx::core::cu diff --git a/mlx/backend/cuda/quantized/qmm/fp_qmv.cu b/mlx/backend/cuda/quantized/qmm/fp_qmv.cu index 33eb4adc..c6a1dbf1 100644 --- a/mlx/backend/cuda/quantized/qmm/fp_qmv.cu +++ b/mlx/backend/cuda/quantized/qmm/fp_qmv.cu @@ -9,6 +9,8 @@ #include #include +#include +#include namespace mlx::core { @@ -76,8 +78,10 @@ __device__ void fp_qmv_impl( vec += g_idx.x * cols; out += g_idx.x * rows; - using ScaleType = - std::conditional_t; + using ScaleType = std::conditional_t< + use_mx_scale, + cutlass::float_ue8m0_t, + cutlass::float_e4m3_t>; auto scales = (ScaleType*)(scales_); auto packed_cols = cols / vals_per_item; @@ -102,35 +106,41 @@ __device__ void fp_qmv_impl( for (int j = 0; j < n_per_step; ++j) { int k = n_per_step * i + j; if constexpr (bits == 8) { - auto v = dequant_fp8(local_mat[k]); + cutlass::NumericArrayConverter + converter; + auto v = converter( + *reinterpret_cast*>( + &local_mat[k])); local_sum.x += - v.x * static_cast(local_vec[vals_per_item * k]); + v[0] * static_cast(local_vec[vals_per_item * k]); local_sum.x += - v.y * static_cast(local_vec[vals_per_item * k + 1]); + v[1] * static_cast(local_vec[vals_per_item * k + 1]); local_sum.y += - v.z * static_cast(local_vec[vals_per_item * k + 2]); + v[2] * static_cast(local_vec[vals_per_item * k + 2]); local_sum.y += - v.w * static_cast(local_vec[vals_per_item * k + 3]); + v[3] * static_cast(local_vec[vals_per_item * k + 3]); } else { - auto v = dequant_fp4(local_mat[k]); + cutlass::NumericArrayConverter + converter; + auto v = converter( + *reinterpret_cast*>( + &local_mat[k])); local_sum.x += - v.x * static_cast(local_vec[vals_per_item * k]); + v[0] * static_cast(local_vec[vals_per_item * k]); local_sum.y += - v.y * static_cast(local_vec[vals_per_item * k + 1]); + v[1] * static_cast(local_vec[vals_per_item * k + 1]); local_sum.x += - v.z * static_cast(local_vec[vals_per_item * k + 2]); + v[2] * static_cast(local_vec[vals_per_item * k + 2]); local_sum.y += - v.w * static_cast(local_vec[vals_per_item * k + 3]); - - v = dequant_fp4(local_mat[k] >> 16); + v[3] * static_cast(local_vec[vals_per_item * k + 3]); local_sum.x += - v.x * static_cast(local_vec[vals_per_item * k + 4]); + v[4] * static_cast(local_vec[vals_per_item * k + 4]); local_sum.y += - v.y * static_cast(local_vec[vals_per_item * k + 5]); + v[5] * static_cast(local_vec[vals_per_item * k + 5]); local_sum.x += - v.z * static_cast(local_vec[vals_per_item * k + 6]); + v[6] * static_cast(local_vec[vals_per_item * k + 6]); local_sum.y += - v.w * static_cast(local_vec[vals_per_item * k + 7]); + v[7] * static_cast(local_vec[vals_per_item * k + 7]); } } sum += (local_sum.x + local_sum.y) * float(scales[i]); diff --git a/mlx/backend/cuda/quantized/quantized_utils.cuh b/mlx/backend/cuda/quantized/quantized_utils.cuh index 8cbe9b29..8a6c6f1d 100644 --- a/mlx/backend/cuda/quantized/quantized_utils.cuh +++ b/mlx/backend/cuda/quantized/quantized_utils.cuh @@ -1,22 +1,9 @@ // Copyright © 2025 Apple Inc. -#include -#include - namespace mlx::core { namespace cu { -inline __device__ float4 dequant_fp8(uint32_t bits) { - auto out = *(__nv_fp8x4_e4m3*)(&bits); - return out.operator float4(); -} - -inline __device__ float4 dequant_fp4(uint16_t bits) { - auto out = *(__nv_fp4x4_e2m1*)(&bits); - return out.operator float4(); -} - template inline constexpr __device__ short get_pack_factor() { return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);