Skip to content

Fix int32 overflow in qdq_mxfp4 causing GPU memory access fault above INT32_MAX elements #43

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zjin-lcf wants to merge 2 commits into
base: release/0.12
Choose a base branch
from
Open

Fix int32 overflow in qdq_mxfp4 causing GPU memory access fault above INT32_MAX elements #43

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7afb99e
Fix int32 overflow in qdq_mxfp4 for tensors above INT32_MAX elements
zjin-lcf Sep 19, 2026
467b738
Use a grid-stride loop in qdq_mxfp4_kernel
zjin-lcf Sep 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 14 additions & 4 deletions quark/torch/kernel/hw_emulation/csrc/mxfp4/fake.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,9 @@
#include <torch/csrc/stable/tensor.h>
#include <torch/headeronly/core/ScalarType.h>

#include <algorithm>
#include <cstdint>
#include <limits>

#include "device_guard.h"
#include "gpu_stream.h"
Expand All @@ -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) {
Expand All @@ -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<int64_t>(std::numeric_limits<int>::max());

dim3 dimGrid(static_cast<unsigned int>(std::min(num_blocks, max_blocks)), 1, 1);
dim3 dimBlock(block_size, 1, 1); // < 1024: we are good!

const cudaStream_t stream = getCurrentStream();
Expand All @@ -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>
<<<dimGrid, dimBlock, 0, stream>>>(
(__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>
<<<dimGrid, dimBlock, 0, stream>>>(
(__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!");
Expand Down
103 changes: 57 additions & 46 deletions quark/torch/kernel/hw_emulation/csrc/mxfp4/fake_kernels.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<int>(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<int64_t>(gridDim.x) * blockDim.x;

for (int64_t idx = static_cast<int64_t>(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<int>(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<float_type>(hfloor_impl(hlog2_impl(block_max))
) - FP4_MAX_NORMAL_EXP_UNBIASED,
FLOAT8_E8M0_MAX_EXP
)
);
float_type scale = float_to_bf16_or_half<float_type>(
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<float_type>(hfloor_impl(hlog2_impl(block_max))) -
FP4_MAX_NORMAL_EXP_UNBIASED,
FLOAT8_E8M0_MAX_EXP
)
);
float_type scale = float_to_bf16_or_half<float_type>(
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