diff --git a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu index 985cb72ac..b20d17001 100644 --- a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu +++ b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu @@ -3,7 +3,9 @@ #include #include +#include #include +#include #include "device_guard.h" #include "gpu_stream.h" @@ -29,7 +31,7 @@ void launch_qdq_mxfp4( const torch::stable::Tensor& a, int64_t group_size, void* out_ptr ) { quark::DeviceGuard guard(a); - int numel = a.numel(); + int64_t numel = a.numel(); int block_size; if (numel % kBlockSizeLarge == 0) { @@ -51,7 +53,15 @@ void launch_qdq_mxfp4( a.is_contiguous(), "Expected qdq_mxfp4 input to be contiguous!" ); - dim3 dimGrid(numel / block_size, 1, 1); + // `numel` is a multiple of `block_size`, so this division is exact. Clamp to + // the maximum grid dimension; `qdq_mxfp4_kernel` is a grid-stride loop, so a + // clamped grid just means each thread handles more than one element. For any + // tensor that fits in memory today the clamp is inactive and every thread + // still handles exactly one element. + const int64_t num_blocks = numel / block_size; + const int64_t max_blocks = static_cast(std::numeric_limits::max()); + + dim3 dimGrid(static_cast(std::min(num_blocks, max_blocks)), 1, 1); dim3 dimBlock(block_size, 1, 1); // < 1024: we are good! const cudaStream_t stream = getCurrentStream(); @@ -61,14 +71,14 @@ void launch_qdq_mxfp4( __half, FLOAT16_EXP_BITS, FLOAT16_MANTISSA_BITS, FLOAT16_EXP_BIAS, FLOAT16_VAL_TO_ADD, FLOAT16_SIGN_EXPONENT_MASK> <<>>( - (__half*)a.data_ptr(), (__half*)out_ptr + (__half*)a.data_ptr(), (__half*)out_ptr, numel ); } else if (a.scalar_type() == torch::headeronly::ScalarType::BFloat16) { qdq_mxfp4_kernel< __nv_bfloat16, BFLOAT16_EXP_BITS, BFLOAT16_MANTISSA_BITS, BFLOAT16_EXP_BIAS, BFLOAT16_VAL_TO_ADD, BFLOAT16_SIGN_EXPONENT_MASK> <<>>( - (__nv_bfloat16*)a.data_ptr(), (__nv_bfloat16*)out_ptr + (__nv_bfloat16*)a.data_ptr(), (__nv_bfloat16*)out_ptr, numel ); } else { STD_TORCH_CHECK(false, "Wrong input dtype in qdq_mxfp4!"); diff --git a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake_kernels.cuh b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake_kernels.cuh index df8e92cfa..22790b335 100644 --- a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake_kernels.cuh +++ b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake_kernels.cuh @@ -149,53 +149,64 @@ __device__ float_type fp16_to_fp4_simulate(float_type* val) { template < typename float_type, uint32_t half_exp_bits, uint32_t half_mantissa_bits, uint32_t half_exp_bias, uint16_t val_to_add, uint16_t sign_exponent_mask> -__global__ void qdq_mxfp4_kernel(float_type* inp, float_type* out) { - // Each thread handles one element. - - int idx = blockIdx.x * blockDim.x + threadIdx.x; - - float_type elem = inp[idx]; - float_type block_max = habs_impl(elem); - - // Compute the max across one warp via butterfly shuffles. Each thread - // handles a single value, so we need `log2(WARP_SIZE)` rounds of shuffle - // (5 rounds for a 32-lane warp). - for (int i = 1; i < static_cast(WARP_SIZE); i *= 2) { - block_max = - hmax_impl(block_max, habs_impl(shfl_xor_bf16_or_half(block_max, i))); +__global__ void qdq_mxfp4_kernel( + float_type* inp, float_type* out, int64_t numel +) { + // Each thread handles one element per grid-stride iteration. + // + // `numel` is a multiple of `blockDim.x` (the caller enforces that it is a + // multiple of 64 or 128), and the stride is a multiple of `blockDim.x` too, + // so the loop bound always falls on a block boundary. Every block that runs + // an iteration therefore has all of its threads active, which keeps the + // warps below fully populated -- `shfl_xor_bf16_or_half` requires all 32 + // lanes of a warp to participate. + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + + for (int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + idx < numel; idx += stride) { + float_type elem = inp[idx]; + float_type block_max = habs_impl(elem); + + // Compute the max across one warp via butterfly shuffles. Each thread + // handles a single value, so we need `log2(WARP_SIZE)` rounds of shuffle + // (5 rounds for a 32-lane warp). + for (int i = 1; i < static_cast(WARP_SIZE); i *= 2) { + block_max = + hmax_impl(block_max, habs_impl(shfl_xor_bf16_or_half(block_max, i))); + } + + // TODO: fix as well in quantize kernel. + // Apply rounding strategy to block_max. + // cannot take the address of an rvalue so need this intermediate + // `block_max_uint` variable? + uint16_t block_max_uint = + (*(uint16_t*)(&block_max) + val_to_add) & sign_exponent_mask; + + block_max = *(float_type*)(&block_max_uint); + + // Pick the largest power-of-two scale s.t. block_max / scale fits in fp4. + // Max fp4 magnitude is 6.0 (unbiased exp = FP4_MAX_NORMAL_EXP_UNBIASED), + // so we want the scale's unbiased exp to be `floor(log2(block_max)) - + // FP4_MAX_NORMAL_EXP_UNBIASED`. + uint8_t scale_exp = max( + 0, FLOAT8_E8M0_MAX_EXP + + min( + bf16_or_half2int_rn(hfloor_impl(hlog2_impl(block_max)) + ) - FP4_MAX_NORMAL_EXP_UNBIASED, + FLOAT8_E8M0_MAX_EXP + ) + ); + float_type scale = float_to_bf16_or_half( + powf(2.0, scale_exp - FLOAT8_E8M0_MAX_EXP) + ); + + elem = hdiv_impl(elem, scale); + + float_type elem_fp4 = fp16_to_fp4_simulate< + float_type, half_exp_bits, half_mantissa_bits, half_exp_bias>(&elem); + + out[idx] = hmul_impl(elem_fp4, scale); } - - // TODO: fix as well in quantize kernel. - // Apply rounding strategy to block_max. - // cannot take the address of an rvalue so need this intermediate - // `block_max_uint` variable? - uint16_t block_max_uint = - (*(uint16_t*)(&block_max) + val_to_add) & sign_exponent_mask; - - block_max = *(float_type*)(&block_max_uint); - - // Pick the largest power-of-two scale s.t. block_max / scale fits in fp4. - // Max fp4 magnitude is 6.0 (unbiased exp = FP4_MAX_NORMAL_EXP_UNBIASED), so - // we want the scale's unbiased exp to be `floor(log2(block_max)) - - // FP4_MAX_NORMAL_EXP_UNBIASED`. - uint8_t scale_exp = max( - 0, FLOAT8_E8M0_MAX_EXP + - min( - bf16_or_half2int_rn(hfloor_impl(hlog2_impl(block_max))) - - FP4_MAX_NORMAL_EXP_UNBIASED, - FLOAT8_E8M0_MAX_EXP - ) - ); - float_type scale = float_to_bf16_or_half( - powf(2.0, scale_exp - FLOAT8_E8M0_MAX_EXP) - ); - - elem = hdiv_impl(elem, scale); - - float_type elem_fp4 = fp16_to_fp4_simulate< - float_type, half_exp_bits, half_mantissa_bits, half_exp_bias>(&elem); - - out[idx] = hmul_impl(elem_fp4, scale); } #endif // USE_CUDA