From 7f9b1560e410544855b8167d06821c321bbf53a6 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Sat, 15 Aug 2026 02:34:52 +0700 Subject: [PATCH] Add H3 NVFP4 scale column swizzles --- src/h3_blackwell_runtime/csrc/nvfp4_scale.cu | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu index d1a644f..5847167 100644 --- a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu @@ -122,7 +122,21 @@ __global__ void quantize_nvfp4_bf16_kernel( if (row < scale_rows) { int64_t scale_row = row; int64_t scale_col = block_col; - if (scale_cols >= 32) { + if (scale_cols == 336 || scale_cols == 896) { + const int64_t row_in_tile = row % 128; + const int64_t group = block_col / 4; + const int64_t pair = group / 2; + const int64_t phase = group % 2; + const int64_t col_stride = scale_cols == 336 ? 16 : 128; + const int64_t phase_offset = scale_cols == 336 ? 176 : 512; + int64_t swizzled_col_base = pair * col_stride + phase * phase_offset; + scale_row = (row / 128) * 128 + pair * (scale_cols == 336 ? 3 : 1) + phase + swizzled_col_base / scale_cols + ((row_in_tile % 32) / 2); + scale_col = (swizzled_col_base % scale_cols) + (block_col % 4) + (row_in_tile / 32) * 4 + (row_in_tile % 2) * 16; + if (scale_col >= scale_cols) { + scale_col -= scale_cols; + scale_row += 1; + } + } else if (scale_cols >= 32) { const int64_t row_in_tile = row % 128; scale_row = (row / 128) * 128 + (block_col / 4) * 16 + ((row_in_tile % 32) / 2); scale_col = (block_col % 4) + (row_in_tile / 32) * 4 + (row_in_tile % 2) * 16;