Match NVFP4 block scale padding

This commit is contained in:
Daniel Maddern 2026-08-15 02:24:09 +07:00
parent 0e855b8e55
commit 57b84b8d6b

View file

@ -122,7 +122,7 @@ __global__ void quantize_nvfp4_bf16_kernel(
if (row < scale_rows) {
block_scale[row * scale_cols + block_col] = scale_byte;
}
if (row < q_rows) {
if (row < q_rows && block_col * 8 < q_cols) {
#pragma unroll
for (int pair = 0; pair < 8; ++pair) {
const float denom = tensor_scale * block_scale_value;
@ -271,7 +271,7 @@ std::vector<torch::Tensor> quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::
const int64_t q_rows = pad_16x ? roundup(rows, 16) : rows;
const int64_t q_cols = pad_16x ? roundup(cols, 16) / 2 : cols / 2;
const int64_t scale_rows = roundup(q_rows, 128);
const int64_t scale_cols = q_cols / 8;
const int64_t scale_cols = roundup(q_cols / 8, 4);
auto qdata = torch::empty({q_rows, q_cols}, input.options().dtype(torch::kUInt8));
auto block_scale = torch::empty({scale_rows, scale_cols}, input.options().dtype(torch::kFloat8_e4m3fn));
auto stream = at::cuda::getCurrentCUDAStream();