L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Remove custom fp4/fp8 classes (#3212)

Browse Source
This commit is contained in:
Cheng
2026-03-10 16:08:01 +09:00
committed by GitHub
parent d2702a4fc1
commit 8d022bcb86
6 changed files with 89 additions and 199 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/cuda/quantized/cuda_fp4.h
-100
View File
@@ -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<int>(std::nearbyintf(le));
n = n < -127 ? -127 : n;
n = n > 127 ? 127 : n;
__x = static_cast<uint8_t>(n + 127);
}
__device__ operator float() {
if (__x == 0xFF) {
return std::numeric_limits<float>::quiet_NaN();
}
return std::ldexp(1.0f, static_cast<int>(__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};
};
mlx/backend/cuda/quantized/fp_quantize.cu
+38 -28
View File
@@ -11,8 +11,8 @@
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <cuda_fp4.h>
#include <cuda_fp8.h>
#include <cutlass/float8.h>
#include <cutlass/numeric_conversion.h>
constexpr float F8E4M3_MAX = 448.0f;
constexpr float F4E2M1_MAX = 6.0f;
@@ -24,9 +24,9 @@ template <int bits>
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<T>;
using Tx4 = Vector4_t<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<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
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<T, group_size> w_hat;
#pragma unroll
for (int i = 0; i < group_size / 4; i++) {
Tx4 w_Tx4 = *reinterpret_cast<Tx4*>(&w_tile[i * 4]);
float4 dq;
for (int i = 0; i < group_size / 8; i++) {
auto& w = *reinterpret_cast<cutlass::Array<T, 8>*>(&w_tile[i * 8]);
cutlass::NumericArrayConverter<float, T, 8> fp32_t;
auto scaled = fp32_t(w) * scale_enc_b;
cutlass::Array<float, 8> dq;
if constexpr (bits == 8) {
uint32_t quantized_val =
scale_cvt_Tx4_to_fp8x4<T, USE_SR>(w_Tx4, scale_enc_b, rbits);
dq = dequant_fp8(quantized_val);
cutlass::NumericArrayConverter<cutlass::float_e4m3_t, float, 8> fp8_fp32;
auto quant = fp8_fp32(scaled);
cutlass::NumericArrayConverter<float, cutlass::float_e4m3_t, 8> fp32_fp8;
dq = fp32_fp8(quant);
} else {
uint16_t quantized_val =
scale_cvt_Tx4_to_fp4x4<T, USE_SR>(w_Tx4, scale_enc_b, rbits);
dq = dequant_fp4(quantized_val);
cutlass::NumericArrayConverter<cutlass::float_e2m1_t, float, 8> fp4_fp32;
auto quant = fp4_fp32(scaled);
cutlass::NumericArrayConverter<float, cutlass::float_e2m1_t, 8> fp32_fp4;
dq = fp32_fp4(quant);
}
w_hat[i * 4] = static_cast<T>(dq.x * scale_dec);
w_hat[i * 4 + 1] = static_cast<T>(dq.y * scale_dec);
w_hat[i * 4 + 2] = static_cast<T>(dq.z * scale_dec);
w_hat[i * 4 + 3] = static_cast<T>(dq.w * scale_dec);
cutlass::NumericArrayConverter<T, float, 8> t_fp32;
*reinterpret_cast<cutlass::Array<T, 8>*>(&w_hat[i * 8]) =
t_fp32(dq * scale_dec);
}
store_vector<group_size>(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<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
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<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
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<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
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;
mlx/backend/cuda/quantized/mxfp8_quantize.cuh
+11 -20
View File
@@ -1,32 +1,23 @@
#pragma once
#include <cuda.h>
#include <cuda_fp8.h>
#include <cuda_runtime.h>
#include "mlx/backend/cuda/vector_types.cuh"
namespace mlx::core::cu {
#include <cutlass/numeric_conversion.h>
// TODO implement fast path
template <typename T>
__device__ __forceinline__ uint32_t
scale_cvt_Tx4_to_fp8x4_fallback(const Vector4_t<T> input, const float scale) {
uint32_t out_fp8x4 = 0;
float4 scaled;
scaled.x = static_cast<float>(input.x) * scale;
scaled.y = static_cast<float>(input.y) * scale;
scaled.z = static_cast<float>(input.z) * scale;
scaled.w = static_cast<float>(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 <typename T, bool USE_SR>
__device__ __forceinline__ uint32_t scale_cvt_Tx4_to_fp8x4(
const Vector4_t<T> input,
const Vector4_t<T>& input,
const float scale,
uint32_t rbits) {
return scale_cvt_Tx4_to_fp8x4_fallback(input, scale);
cutlass::NumericArrayConverter<float, T, 4> fp32_t;
auto scaled =
fp32_t(*reinterpret_cast<const cutlass::Array<T, 4>*>(&input)) * scale;
cutlass::NumericArrayConverter<cutlass::float_e4m3_t, float, 4> fp8_fp32;
auto quant = fp8_fp32(scaled);
return *reinterpret_cast<uint32_t*>(&quant);
}
} // namespace mlx::core::cu
} // namespace mlx::core::cu
mlx/backend/cuda/quantized/nvfp4_quantize.cuh
+12 -20
View File
@@ -1,10 +1,9 @@
#pragma once
#include <cuda.h>
#include <cuda_fp4.h>
#include <cuda_runtime.h>
#include "mlx/backend/cuda/vector_types.cuh"
#include <cutlass/numeric_conversion.h>
namespace mlx::core::cu {
using bf16x4 = Vector4_t<__nv_bfloat16>;
@@ -13,23 +12,15 @@ using f32x4 = Vector4_t<float>;
template <typename T>
__device__ __forceinline__ uint16_t
scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t<T> input, const float scale) {
scale_cvt_Tx4_to_fp4x4_fallback(const Vector4_t<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<float>(input.x) * scale;
scaled.y = static_cast<float>(input.y) * scale;
scaled.z = static_cast<float>(input.z) * scale;
scaled.w = static_cast<float>(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<uint16_t>(q3) << 12) |
(static_cast<uint16_t>(q2) << 8) | (static_cast<uint16_t>(q1) << 4) |
static_cast<uint16_t>(q0);
return out_fp4x4;
cutlass::NumericArrayConverter<float, T, 4> fp32_t;
auto scaled =
fp32_t(*reinterpret_cast<const cutlass::Array<T, 4>*>(&input)) * scale;
cutlass::NumericArrayConverter<cutlass::float_e2m1_t, float, 4> fp4_fp32;
auto quant = fp4_fp32(scaled);
return *reinterpret_cast<uint16_t*>(&quant);
}
#if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \
@@ -318,7 +309,7 @@ __device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4_fast(
template <typename T, bool USE_SR>
__device__ __forceinline__ uint16_t scale_cvt_Tx4_to_fp4x4(
const Vector4_t<T> input,
const Vector4_t<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
} // namespace mlx::core::cu
mlx/backend/cuda/quantized/qmm/fp_qmv.cu
+28 -18
View File
@@ -9,6 +9,8 @@
#include <cooperative_groups.h>
#include <cooperative_groups/reduce.h>
#include <cutlass/float8.h>
#include <cutlass/numeric_conversion.h>
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<use_mx_scale, __nv_fp8_e8m0, __nv_fp8_e4m3>;
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<float, cutlass::float_e4m3_t, 4>
converter;
auto v = converter(
*reinterpret_cast<cutlass::Array<cutlass::float_e4m3_t, 4>*>(
&local_mat[k]));
local_sum.x +=
v.x * static_cast<float>(local_vec[vals_per_item * k]);
v[0] * static_cast<float>(local_vec[vals_per_item * k]);
local_sum.x +=
v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
v[1] * static_cast<float>(local_vec[vals_per_item * k + 1]);
local_sum.y +=
v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
v[2] * static_cast<float>(local_vec[vals_per_item * k + 2]);
local_sum.y +=
v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
v[3] * static_cast<float>(local_vec[vals_per_item * k + 3]);
} else {
auto v = dequant_fp4(local_mat[k]);
cutlass::NumericArrayConverter<float, cutlass::float_e2m1_t, 8>
converter;
auto v = converter(
*reinterpret_cast<cutlass::Array<cutlass::float_e2m1_t, 8>*>(
&local_mat[k]));
local_sum.x +=
v.x * static_cast<float>(local_vec[vals_per_item * k]);
v[0] * static_cast<float>(local_vec[vals_per_item * k]);
local_sum.y +=
v.y * static_cast<float>(local_vec[vals_per_item * k + 1]);
v[1] * static_cast<float>(local_vec[vals_per_item * k + 1]);
local_sum.x +=
v.z * static_cast<float>(local_vec[vals_per_item * k + 2]);
v[2] * static_cast<float>(local_vec[vals_per_item * k + 2]);
local_sum.y +=
v.w * static_cast<float>(local_vec[vals_per_item * k + 3]);
v = dequant_fp4(local_mat[k] >> 16);
v[3] * static_cast<float>(local_vec[vals_per_item * k + 3]);
local_sum.x +=
v.x * static_cast<float>(local_vec[vals_per_item * k + 4]);
v[4] * static_cast<float>(local_vec[vals_per_item * k + 4]);
local_sum.y +=
v.y * static_cast<float>(local_vec[vals_per_item * k + 5]);
v[5] * static_cast<float>(local_vec[vals_per_item * k + 5]);
local_sum.x +=
v.z * static_cast<float>(local_vec[vals_per_item * k + 6]);
v[6] * static_cast<float>(local_vec[vals_per_item * k + 6]);
local_sum.y +=
v.w * static_cast<float>(local_vec[vals_per_item * k + 7]);
v[7] * static_cast<float>(local_vec[vals_per_item * k + 7]);
}
}
sum += (local_sum.x + local_sum.y) * float(scales[i]);
mlx/backend/cuda/quantized/quantized_utils.cuh
-13
View File
@@ -1,22 +1,9 @@
// Copyright © 2025 Apple Inc.
#include <cuda_fp4.h>
#include <cuda_fp8.h>
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 <int bits, int wsize = 8>
inline constexpr __device__ short get_pack_factor() {
return (bits == 3 || bits == 5) ? 8 : (bits == 6 ? 4 : wsize / bits);