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

Add H3 NVFP4 scale column swizzles

Browse source
This commit is contained in:
Daniel Maddern 2026-08-15 02:34:52 +07:00
parent 1200152d79
commit 7f9b1560e4
1 changed files with 15 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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;