daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Match NVFP4 block scale padding

Browse source
This commit is contained in:
Daniel Maddern 2026-08-15 02:24:09 +07:00
parent 0e855b8e55
commit 57b84b8d6b
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/h3_blackwell_runtime/csrc/nvfp4_scale.cu
View file

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