From 1200152d798405d49f682e8ccc30906a84679425 Mon Sep 17 00:00:00 2001
From: Daniel Maddern <daniel.maddern@illuvium.io>
Date: Sat, 15 Aug 2026 02:30:41 +0700
Subject: [PATCH] Match wide NVFP4 block scale swizzle

---
 src/h3_blackwell_runtime/csrc/nvfp4_scale.cu | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu
index b06d508..d1a644f 100644
--- a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu
+++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu
@@ -120,8 +120,14 @@ __global__ void quantize_nvfp4_bf16_kernel(
       block_scale_value = decode_fp8_e4m3(scale_byte);
     }
     if (row < scale_rows) {
-      const int64_t tile_row = (row / 128) * 128 + (row % 32) * 4 + ((row % 128) / 32);
-      block_scale[tile_row * scale_cols + block_col] = scale_byte;
+      int64_t scale_row = row;
+      int64_t scale_col = block_col;
+      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;
+      }
+      block_scale[scale_row * scale_cols + scale_col] = scale_byte;
     }
     if (row < q_rows && block_col * 8 < q_cols) {
       #pragma unroll