feat: Add full-dimension Hadamard rotation kernel

Add kHadamardRotateFull kernel that rotates across the entire last
dimension (512-8192), matching the approach used by QuIP#, QuaRot, and
SpinQuant for maximal outlier suppression. The existing block-diagonal
kernel (block_size 32-256) remains for use cases where block rotation
suffices.

Kernel design: one thread block per row with 3-4 butterfly levels:
1. In-thread butterfly (strides 1, 2, 4)
2. Warp shuffle butterfly (strides 8-128)
3. Cross-warp butterfly via shared memory (strides 256+)
4. Cross-chunk butterfly in registers (dims > 2048)

API: hadamard_rotate(data, block_size=0) for full-dimension mode.
Signs vector has dim//32 words (one per full row, not per block).

144 tests passing (79 existing + 65 new full-dimension tests).

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: TimDettmers

bitsandbytes/_ops.py

@@ -614,6 +614,43 @@ def _(data: torch.Tensor, block_size: int, signs: Optional[torch.Tensor]) -> tor
     return data
 
 
+# Full-dimension Hadamard rotation (in-place, for kbit quantization outlier spreading)
+# Unlike hadamard_rotate_ which uses block-diagonal Hadamard, this rotates across
+# the entire last dimension of the input tensor.
+
+torch.library.define(
+    "bitsandbytes::hadamard_rotate_full_",
+    "(Tensor(a!) data, int dim, Tensor? signs) -> Tensor(a!)",
+)
+
+
+@register_fake("bitsandbytes::hadamard_rotate_full_")
+def _(data: torch.Tensor, dim: int, signs: Optional[torch.Tensor]) -> torch.Tensor:
+    supported_dims = (512, 1024, 2048, 4096, 8192)
+    torch._check(
+        dim in supported_dims,
+        lambda: f"dim must be one of {supported_dims}, got {dim}",
+    )
+    torch._check(
+        data.numel() % dim == 0,
+        lambda: f"data.numel() ({data.numel()}) must be divisible by dim ({dim})",
+    )
+    torch._check(
+        data.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"hadamard_rotate_full only supports float16/bfloat16, got {data.dtype}",
+    )
+    if signs is not None:
+        torch._check(
+            signs.dtype == torch.int32,
+            lambda: f"signs must be int32, got {signs.dtype}",
+        )
+        torch._check(
+            signs.numel() == dim // 32,
+            lambda: f"signs must have {dim // 32} elements for dim={dim}, got {signs.numel()}",
+        )
+    return data
+
+
 # K-bit fused dequant + GEMM (production: fp16 + bf16)
 
 torch.library.define(

bitsandbytes/backends/cuda/ops.py

@@ -1026,6 +1026,34 @@ def _(data: torch.Tensor, block_size: int, signs: Optional[torch.Tensor]) -> tor
     return data
 
 
+@register_kernel("bitsandbytes::hadamard_rotate_full_", "cuda")
+def _(data: torch.Tensor, dim: int, signs: Optional[torch.Tensor]) -> torch.Tensor:
+    supported_dims = (512, 1024, 2048, 4096, 8192)
+    torch._check(
+        dim in supported_dims,
+        lambda: f"dim must be one of {supported_dims}, got {dim}",
+    )
+    torch._check(
+        data.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"hadamard_rotate_full only supports float16/bfloat16, got {data.dtype}",
+    )
+
+    num_rows = data.numel() // dim
+    tname = _KBIT_DTYPE_SUFFIX[data.dtype]
+    signs_ptr = get_ptr(signs) if signs is not None else None
+    with _cuda_device_of(data):
+        fn = getattr(lib, f"chadamard_rotate_full_{tname}")
+        fn(
+            get_ptr(data),
+            ct.c_int(num_rows),
+            ct.c_int(dim),
+            signs_ptr,
+            _get_tensor_stream(data),
+        )
+
+    return data
+
+
 def _kbit_gemm_prod_check(A, B_packed, B_absmax, codebook, N, k, k_chunks):
     torch._check(k >= 2 and k <= 5, lambda: f"k must be 2-5, got {k}")
     torch._check(

bitsandbytes/functional.py

@@ -1140,25 +1140,47 @@ def hadamard_rotate(
     block_size: int = 32,
     signs: Optional[Tensor] = None,
 ) -> Tensor:
-    """Apply in-place randomized Walsh-Hadamard rotation (H*D) to contiguous blocks.
+    """Apply in-place randomized Walsh-Hadamard rotation (H*D).
 
     Spreads outliers across quantization blocks, improving kbit accuracy.
     Since H*D is orthogonal, rotating both weights and activations with the
     same signs preserves the GEMM result: (H*D)(A) @ (H*D)(B)^T = A @ B^T.
 
+    Two modes:
+
+    **Block-diagonal** (block_size in {32, 64, 128, 256}): Applies independent
+    Hadamard rotations to contiguous blocks of ``block_size`` elements across
+    the flattened tensor. Fast and parallel, but only spreads outliers within
+    each block.
+
+    **Full-dimension** (block_size=0): Applies the Hadamard rotation across
+    the entire last dimension of the tensor. Matches the approach used by
+    QuIP#, QuaRot, and SpinQuant for maximal outlier suppression. The last
+    dimension must be a power of 2 in {512, 1024, 2048, 4096, 8192}.
+
     Args:
         data: Input tensor (float16 or bfloat16). Modified in-place.
-        block_size: Rotation block size (32, 64, 128, or 256).
-        signs: Optional int32 tensor of block_size//32 words. Each bit controls
-            the sign flip for one element within the block. If None, no sign
-            flips are applied (plain Hadamard). Generate once per model with
-            ``torch.randint(0, 2**32, (block_size // 32,), dtype=torch.int32)``.
+        block_size: Rotation block size (32, 64, 128, 256) for block-diagonal
+            mode, or 0 for full-dimension mode.
+        signs: Optional int32 tensor of sign-flip bits. For block-diagonal
+            mode: ``block_size // 32`` words (repeated per block). For
+            full-dimension mode: ``dim // 32`` words where ``dim`` is the
+            last dimension. Each bit controls the sign flip for one element.
+            If None, no sign flips (plain Hadamard). Generate once per model
+            with ``torch.randint(0, 2**32, (n_words,), dtype=torch.int32)``.
 
     Returns:
         The input tensor, rotated in-place.
     """
-    data_flat = data.contiguous().view(-1)
-    torch.ops.bitsandbytes.hadamard_rotate_(data_flat, block_size, signs)
+    if block_size == 0:
+        # Full-dimension mode: rotate across the entire last dimension.
+        dim = data.shape[-1]
+        data_flat = data.contiguous().view(-1)
+        torch.ops.bitsandbytes.hadamard_rotate_full_(data_flat, dim, signs)
+    else:
+        # Block-diagonal mode: independent rotations per block.
+        data_flat = data.contiguous().view(-1)
+        torch.ops.bitsandbytes.hadamard_rotate_(data_flat, block_size, signs)
     return data
 
 

csrc/ops.cu

@@ -1127,6 +1127,200 @@ INSTANTIATE_HADAMARD(256)
 
 #undef INSTANTIATE_HADAMARD
 
+// ===========================================================================
+// Full-dimension Hadamard rotation kernel.
+// One thread block processes one row of DIM elements using 3-4 butterfly levels:
+//   1. In-thread butterfly (strides 1..kNElts/2)
+//   2. Warp shuffle butterfly (strides kNElts..kNElts*16)
+//   3. Cross-warp butterfly via shared memory (strides across warps)
+//   4. Cross-chunk butterfly in registers (when kNChunks > 1)
+//
+// Grid: (num_rows,). Signs: DIM/32 uint32 words (one per full row, not per block).
+// ===========================================================================
+
+template <int kLogDim, int kNThreads, typename T>
+__global__ void kHadamardRotateFull(T* __restrict__ data, const int num_rows, const unsigned int* __restrict__ signs) {
+    constexpr int DIM = 1 << kLogDim;
+    constexpr int kNElts = 8; // elements per thread per chunk
+    constexpr int kNChunks = DIM / (kNThreads * kNElts);
+    constexpr int kNWarps = kNThreads / 32;
+
+    static_assert(DIM == kNThreads * kNElts * kNChunks, "dimension decomposition mismatch");
+    static_assert(kNElts == 8, "kNElts must be 8");
+    static_assert((kNThreads & (kNThreads - 1)) == 0, "kNThreads must be power of 2");
+
+    const int row = blockIdx.x;
+    if (row >= num_rows)
+        return;
+
+    T* row_data = data + (long long)row * DIM;
+
+    // Shared memory for cross-warp butterfly (only needed when kNWarps > 1).
+    // Use char[] to match other kernels in this TU, then cast to float*.
+    extern __shared__ char smem_raw[];
+    float* smem = reinterpret_cast<float*>(smem_raw);
+
+    const int tid = threadIdx.x;
+    const int warp_id = tid / 32;
+    const int lane_id = tid % 32;
+
+    // ---- Load elements (contiguous per thread) ----
+    float vals[kNChunks][kNElts];
+#pragma unroll
+    for (int c = 0; c < kNChunks; c++) {
+        const int base = c * kNThreads * kNElts + tid * kNElts;
+#pragma unroll
+        for (int i = 0; i < kNElts; i++) {
+            vals[c][i] = (float)row_data[base + i];
+        }
+    }
+
+    // ---- Apply sign flips (D matrix) before butterfly ----
+    // 8 contiguous elements at position 'base' always fit within one uint32 word
+    // since base is always a multiple of 8.
+    if (signs != nullptr) {
+#pragma unroll
+        for (int c = 0; c < kNChunks; c++) {
+            const int linear = c * kNThreads + tid; // which group of 8
+            const int word_idx = linear / 4;
+            const int byte_pos = (linear % 4) * 8;
+            const unsigned int byte_bits = (signs[word_idx] >> byte_pos) & 0xFFu;
+#pragma unroll
+            for (int i = 0; i < kNElts; i++) {
+                if (byte_bits & (1u << i))
+                    vals[c][i] = -vals[c][i];
+            }
+        }
+    }
+
+    // ---- Level 1: In-thread butterfly (strides 1, 2, 4) ----
+#pragma unroll
+    for (int c = 0; c < kNChunks; c++) {
+#pragma unroll
+        for (int s = 1; s < kNElts; s <<= 1) {
+#pragma unroll
+            for (int i = 0; i < kNElts; i++) {
+                int partner = i ^ s;
+                if (partner > i) {
+                    float a = vals[c][i], b = vals[c][partner];
+                    vals[c][i] = a + b;
+                    vals[c][partner] = a - b;
+                }
+            }
+        }
+    }
+
+    // ---- Level 2: Warp shuffle butterfly (shfl_xor s=1..16) ----
+#pragma unroll
+    for (int s = 1; s <= 16; s <<= 1) {
+#pragma unroll
+        for (int c = 0; c < kNChunks; c++) {
+#pragma unroll
+            for (int i = 0; i < kNElts; i++) {
+                float other = __shfl_xor_sync(0xFFFFFFFF, vals[c][i], s);
+                vals[c][i] = (lane_id & s) ? (other - vals[c][i]) : (vals[c][i] + other);
+            }
+        }
+    }
+
+    // ---- Level 3: Cross-warp butterfly via shared memory ----
+    if constexpr (kNWarps > 1) {
+        constexpr int VALS_PER_THREAD = kNChunks * kNElts;
+        // smem layout: smem[tid * VALS_PER_THREAD + c * kNElts + i]
+#pragma unroll
+        for (int ws = 1; ws < kNWarps; ws <<= 1) {
+            // Write my values to shared memory
+#pragma unroll
+            for (int c = 0; c < kNChunks; c++) {
+#pragma unroll
+                for (int i = 0; i < kNElts; i++) {
+                    smem[tid * VALS_PER_THREAD + c * kNElts + i] = vals[c][i];
+                }
+            }
+            __syncthreads();
+
+            // Read partner warp's values
+            const int partner_tid = (warp_id ^ ws) * 32 + lane_id;
+            const bool negate = (warp_id & ws) != 0;
+#pragma unroll
+            for (int c = 0; c < kNChunks; c++) {
+#pragma unroll
+                for (int i = 0; i < kNElts; i++) {
+                    float pval = smem[partner_tid * VALS_PER_THREAD + c * kNElts + i];
+                    vals[c][i] = negate ? (pval - vals[c][i]) : (vals[c][i] + pval);
+                }
+            }
+            __syncthreads();
+        }
+    }
+
+    // ---- Level 4: Cross-chunk butterfly (in-register, no communication) ----
+    if constexpr (kNChunks > 1) {
+#pragma unroll
+        for (int cs = 1; cs < kNChunks; cs <<= 1) {
+#pragma unroll
+            for (int c = 0; c < kNChunks; c++) {
+                int pc = c ^ cs;
+                if (pc > c) {
+#pragma unroll
+                    for (int i = 0; i < kNElts; i++) {
+                        float a = vals[c][i], b = vals[pc][i];
+                        vals[c][i] = a + b;
+                        vals[pc][i] = a - b;
+                    }
+                }
+            }
+        }
+    }
+
+    // ---- Normalize by 1/sqrt(DIM) ----
+    const float norm = rsqrtf((float)DIM);
+#pragma unroll
+    for (int c = 0; c < kNChunks; c++) {
+#pragma unroll
+        for (int i = 0; i < kNElts; i++)
+            vals[c][i] *= norm;
+    }
+
+    // ---- Store back ----
+#pragma unroll
+    for (int c = 0; c < kNChunks; c++) {
+        const int base = c * kNThreads * kNElts + tid * kNElts;
+#pragma unroll
+        for (int i = 0; i < kNElts; i++) {
+            row_data[base + i] = (T)vals[c][i];
+        }
+    }
+}
+
+// ---- Full-dimension Hadamard launch wrapper ----
+// kLogDim must match the dimension. kNThreads is the thread block size.
+
+template <int kLogDim, int kNThreads, typename T>
+void hadamardRotateFull(T* data, int num_rows, const unsigned int* signs, cudaStream_t stream) {
+    constexpr int DIM = 1 << kLogDim;
+    constexpr int kNElts = 8;
+    constexpr int kNChunks = DIM / (kNThreads * kNElts);
+    constexpr int smem_bytes = kNThreads * kNChunks * kNElts * sizeof(float);
+    kHadamardRotateFull<kLogDim, kNThreads, T><<<num_rows, kNThreads, smem_bytes, stream>>>(data, num_rows, signs);
+    CUDA_CHECK_RETURN(cudaPeekAtLastError());
+}
+
+// Explicit instantiations: dim 512..8192, 2 dtypes
+#define INSTANTIATE_HADAMARD_FULL(LOG_DIM, NTHREADS)                                                                   \
+    template void hadamardRotateFull<LOG_DIM, NTHREADS, half>(half*, int, const unsigned int*, cudaStream_t);          \
+    template void hadamardRotateFull<LOG_DIM, NTHREADS, __nv_bfloat16>(                                                \
+        __nv_bfloat16*, int, const unsigned int*, cudaStream_t                                                         \
+    );
+
+INSTANTIATE_HADAMARD_FULL(9, 64)   // dim=512
+INSTANTIATE_HADAMARD_FULL(10, 128) // dim=1024
+INSTANTIATE_HADAMARD_FULL(11, 256) // dim=2048
+INSTANTIATE_HADAMARD_FULL(12, 256) // dim=4096
+INSTANTIATE_HADAMARD_FULL(13, 256) // dim=8192
+
+#undef INSTANTIATE_HADAMARD_FULL
+
 // Datacenter GPU detection: Hopper (sm_90) and Blackwell datacenter (sm_100).
 // NOTE: sm_120 (RTX 5090, Blackwell consumer) lacks TMA/wgmma — must NOT match.
 #if defined(__CUDA_ARCH__)

csrc/pythonInterface.cpp

@@ -827,6 +827,39 @@ MAKE_HADAMARD_ROTATE(bf16, __nv_bfloat16)
 
 #undef MAKE_HADAMARD_ROTATE
 
+// Forward declarations of full-dimension hadamard rotation template
+template <int kLogDim, int kNThreads, typename T>
+void hadamardRotateFull(T* data, int num_rows, const unsigned int* signs, cudaStream_t stream);
+
+// Unmangled full-dimension hadamard rotation wrappers (dispatch dim at runtime)
+#define MAKE_HADAMARD_ROTATE_FULL(tname, T)                                                                            \
+    void hadamard_rotate_full_##tname(                                                                                 \
+        T* data, int num_rows, int dim, const unsigned int* signs, cudaStream_t stream                                 \
+    ) {                                                                                                                \
+        switch (dim) {                                                                                                 \
+        case 512:                                                                                                      \
+            hadamardRotateFull<9, 64, T>(data, num_rows, signs, stream);                                               \
+            break;                                                                                                     \
+        case 1024:                                                                                                     \
+            hadamardRotateFull<10, 128, T>(data, num_rows, signs, stream);                                             \
+            break;                                                                                                     \
+        case 2048:                                                                                                     \
+            hadamardRotateFull<11, 256, T>(data, num_rows, signs, stream);                                             \
+            break;                                                                                                     \
+        case 4096:                                                                                                     \
+            hadamardRotateFull<12, 256, T>(data, num_rows, signs, stream);                                             \
+            break;                                                                                                     \
+        case 8192:                                                                                                     \
+            hadamardRotateFull<13, 256, T>(data, num_rows, signs, stream);                                             \
+            break;                                                                                                     \
+        }                                                                                                              \
+    }
+
+MAKE_HADAMARD_ROTATE_FULL(fp16, half)
+MAKE_HADAMARD_ROTATE_FULL(bf16, __nv_bfloat16)
+
+#undef MAKE_HADAMARD_ROTATE_FULL
+
 #endif // BUILD_CUDA || BUILD_HIP (kbit unmangled)
 
 extern "C" {
@@ -1707,5 +1740,16 @@ void chadamard_rotate_bf16(__nv_bfloat16* data, int n, int block_size, const uns
     hadamard_rotate_bf16(data, n, block_size, signs, stream);
 }
 
+// Full-dimension Hadamard rotation extern C wrappers
+void chadamard_rotate_full_fp16(half* data, int num_rows, int dim, const unsigned int* signs, cudaStream_t stream) {
+    hadamard_rotate_full_fp16(data, num_rows, dim, signs, stream);
+}
+
+void chadamard_rotate_full_bf16(
+    __nv_bfloat16* data, int num_rows, int dim, const unsigned int* signs, cudaStream_t stream
+) {
+    hadamard_rotate_full_bf16(data, num_rows, dim, signs, stream);
+}
+
 #endif
 }