From 7afb99e790a1ffff4be4a788c4c9186334f6cbfc Mon Sep 17 00:00:00 2001 From: zjin-lcf Date: Sat, 19 Sep 2026 15:47:27 -0500 Subject: [PATCH 1/2] Fix int32 overflow in qdq_mxfp4 for tensors above INT32_MAX elements `launch_qdq_mxfp4` truncated `a.numel()` into an `int`, and `qdq_mxfp4_kernel` indexed elements with an `int`. For a tensor with more than INT32_MAX elements this wraps negative, still passes the divisibility check, and launches with a garbage grid dimension, so the process dies with "Memory access fault by GPU node-N ... Reason: Unknown". Use 64-bit types for the element count and the element index, and check the grid dimension explicitly. This mirrors what the sibling mxfp4/dequantize.cu already does. Minimal reproducer (bf16, hidden 14336) on MI300A / ROCm 7.2: 2,132,795,392 elements (under INT32_MAX) -> OK 2,162,155,520 elements (over INT32_MAX) -> memory access fault With this change both cases pass. --- quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu | 12 ++++++++++-- .../kernel/hw_emulation/csrc/mxfp4/fake_kernels.cuh | 2 +- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu index 985cb72ac..98e1073dc 100644 --- a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu +++ b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu @@ -4,6 +4,7 @@ #include #include +#include #include "device_guard.h" #include "gpu_stream.h" @@ -29,7 +30,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 +52,14 @@ void launch_qdq_mxfp4( a.is_contiguous(), "Expected qdq_mxfp4 input to be contiguous!" ); - dim3 dimGrid(numel / block_size, 1, 1); + int64_t grid_size = numel / block_size; + + STD_TORCH_CHECK( + grid_size <= static_cast(std::numeric_limits::max()), + "Grid size exceeds CUDA maximum grid dimension" + ); + + dim3 dimGrid(grid_size, 1, 1); dim3 dimBlock(block_size, 1, 1); // < 1024: we are good! const cudaStream_t stream = getCurrentStream(); 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..128141bbd 100644 --- a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake_kernels.cuh +++ b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake_kernels.cuh @@ -152,7 +152,7 @@ template < __global__ void qdq_mxfp4_kernel(float_type* inp, float_type* out) { // Each thread handles one element. - int idx = blockIdx.x * blockDim.x + threadIdx.x; + int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; float_type elem = inp[idx]; float_type block_max = habs_impl(elem); From 467b738d31ff45a7eb015550d14d6f6912d6928e Mon Sep 17 00:00:00 2001 From: zjin-lcf Date: Sat, 19 Sep 2026 16:12:43 -0500 Subject: [PATCH 2/2] Use a grid-stride loop in qdq_mxfp4_kernel Addresses review feedback: rather than erroring out when the block count exceeds the maximum grid dimension, make the kernel a grid-stride loop and clamp the grid. This removes the ceiling entirely instead of moving it. The kernel derives each MXFP4 group's shared scale from a 32-lane butterfly shuffle, so every active warp must be fully populated. That holds here because the caller already requires `numel % block_size == 0`, and the stride is a multiple of block_size, so the loop bound always lands on a block boundary -- a block either runs an iteration with all of its threads or not at all. Verified on MI300A / ROCm 7.2 against the stock kernel: - bit-identical output on 8,388,608 elements, mixed magnitudes - bit-identical with the grid clamped to 1024 blocks, forcing 64 grid-stride iterations per thread (exercises the multi-iteration path) - 2,162,155,520 elements (over INT32_MAX): faulted before, passes now --- .../kernel/hw_emulation/csrc/mxfp4/fake.cu | 20 ++-- .../hw_emulation/csrc/mxfp4/fake_kernels.cuh | 103 ++++++++++-------- 2 files changed, 68 insertions(+), 55 deletions(-) diff --git a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu index 98e1073dc..b20d17001 100644 --- a/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu +++ b/quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu @@ -3,6 +3,7 @@ #include #include +#include #include #include @@ -52,14 +53,15 @@ void launch_qdq_mxfp4( a.is_contiguous(), "Expected qdq_mxfp4 input to be contiguous!" ); - int64_t grid_size = numel / block_size; + // `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()); - STD_TORCH_CHECK( - grid_size <= static_cast(std::numeric_limits::max()), - "Grid size exceeds CUDA maximum grid dimension" - ); - - dim3 dimGrid(grid_size, 1, 1); + 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(); @@ -69,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 128141bbd..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. - - int64_t idx = static_cast(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