feat: Add VQ repack kernel (flat → tiled layout) and tiled GEMV support

- kRepackVQ<P_VAL> kernel: maps flat VQ byte layout to tile-interleaved layout
- Same tile geometry as kbit repack (TILE_K=64, TILE_N=128, BS=32)
- Full registration chain: ops.cu → pythonInterface → _ops.py → backends → functional.py
- repack_vq() Python wrapper in functional.py
- Verified: flat vs tiled GEMV produces bit-identical results for all (p,K,N,M)

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

Author: TimDettmers

bitsandbytes/_ops.py

@@ -744,6 +744,34 @@ def _(
     return packed_tiled, absmax_tiled
 
 
+# VQ repack: flat VQ byte layout -> tiled layout
+
+torch.library.define(
+    "bitsandbytes::repack_vq",
+    "(Tensor packed_flat, Tensor absmax_flat, int K_dim, int N, int p) -> (Tensor, Tensor)",
+)
+
+
+@register_fake("bitsandbytes::repack_vq")
+def _(
+    packed_flat: torch.Tensor, absmax_flat: torch.Tensor, K_dim: int, N: int, p: int
+) -> tuple[torch.Tensor, torch.Tensor]:
+    torch._check(p in (2, 4), lambda: f"p must be 2 or 4, got {p}")
+    TILE_K, TILE_N, BLOCKSIZE = 64, 128, 32
+    torch._check(N % TILE_N == 0, lambda: f"N ({N}) must be divisible by {TILE_N}")
+    torch._check(K_dim % BLOCKSIZE == 0, lambda: f"K_dim ({K_dim}) must be divisible by {BLOCKSIZE}")
+    K_dim_padded = ((K_dim + TILE_K - 1) // TILE_K) * TILE_K
+    k_tiles = K_dim_padded // TILE_K
+    n_tiles = N // TILE_N
+    k_blocks_per_tile = TILE_K // BLOCKSIZE
+    words_per_block = BLOCKSIZE // (p * 4)
+    total_words = k_tiles * n_tiles * TILE_N * k_blocks_per_tile * words_per_block
+    total_absmax = k_tiles * n_tiles * TILE_N * k_blocks_per_tile
+    packed_tiled = torch.empty(total_words, device=packed_flat.device, dtype=torch.int32)
+    absmax_tiled = torch.empty(total_absmax, device=packed_flat.device, dtype=torch.uint8)
+    return packed_tiled, absmax_tiled
+
+
 # Hadamard rotation (in-place, for kbit quantization outlier spreading)
 
 torch.library.define(

bitsandbytes/backends/cuda/ops.py

@@ -1210,6 +1210,51 @@ def _(
     return packed_tiled, absmax_tiled
 
 
+@register_kernel("bitsandbytes::repack_vq", "cuda")
+def _(
+    packed_flat: torch.Tensor,
+    absmax_flat: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    torch._check(p in (2, 4), lambda: f"p must be 2 or 4, got {p}")
+    torch._check(packed_flat.dtype == torch.int32, lambda: f"packed_flat must be int32, got {packed_flat.dtype}")
+    torch._check(
+        absmax_flat.dtype == torch.uint8, lambda: f"absmax_flat must be uint8 (E4M4), got {absmax_flat.dtype}"
+    )
+
+    TILE_K, TILE_N, BLOCKSIZE = 64, 128, 32
+    torch._check(N % TILE_N == 0, lambda: f"N ({N}) must be divisible by {TILE_N}")
+    torch._check(K_dim % BLOCKSIZE == 0, lambda: f"K_dim ({K_dim}) must be divisible by {BLOCKSIZE}")
+
+    K_dim_padded = ((K_dim + TILE_K - 1) // TILE_K) * TILE_K
+    k_tiles = K_dim_padded // TILE_K
+    n_tiles = N // TILE_N
+    k_blocks_per_tile = TILE_K // BLOCKSIZE
+    words_per_block = BLOCKSIZE // (p * 4)
+    total_words = k_tiles * n_tiles * TILE_N * k_blocks_per_tile * words_per_block
+    total_absmax = k_tiles * n_tiles * TILE_N * k_blocks_per_tile
+
+    # Zero-fill for padding regions (when K_dim is not multiple of TILE_K)
+    packed_tiled = torch.zeros(total_words, device=packed_flat.device, dtype=torch.int32)
+    absmax_tiled = torch.zeros(total_absmax, device=packed_flat.device, dtype=torch.uint8)
+
+    with _cuda_device_of(packed_flat):
+        fn = getattr(lib, f"crepack_vq_p{p}")
+        fn(
+            get_ptr(packed_flat),
+            get_ptr(absmax_flat),
+            get_ptr(packed_tiled),
+            get_ptr(absmax_tiled),
+            ct.c_int(K_dim),
+            ct.c_int(N),
+            _get_tensor_stream(packed_flat),
+        )
+
+    return packed_tiled, absmax_tiled
+
+
 @register_kernel("bitsandbytes::hadamard_rotate_", "cuda")
 def _(data: torch.Tensor, block_size: int, signs: Optional[torch.Tensor]) -> torch.Tensor:
     torch._check(

bitsandbytes/functional.py

@@ -1376,6 +1376,31 @@ def dequantize_vq(
     return result[:n]
 
 
+def repack_vq(
+    packed_flat: Tensor,
+    absmax_flat: Tensor,
+    K_dim: int,
+    N: int,
+    p: int = 2,
+) -> tuple[Tensor, Tensor]:
+    """Repack VQ quantized weights from flat to tiled layout.
+
+    Rearranges packed byte indices and absmax from flat column-major layout
+    to tile-interleaved layout used by vq_scalar_gemv_tiled and vq_gemm_prod.
+
+    Args:
+        packed_flat: int32 tensor of packed byte indices (from quantize_vq).
+        absmax_flat: uint8 E4M4 per-block absmax values.
+        K_dim: Reduction dimension.
+        N: Output dimension (must be multiple of 128).
+        p: VQ dimension (2 or 4).
+
+    Returns:
+        Tuple of (packed_tiled, absmax_tiled).
+    """
+    return torch.ops.bitsandbytes.repack_vq(packed_flat, absmax_flat, K_dim, N, p)
+
+
 def dequantize_kbit_tiled(
     packed: Tensor,
     absmax: Tensor,

csrc/ops.cu

@@ -1180,6 +1180,67 @@ void repackKbit(
     CUDA_CHECK_RETURN(cudaPeekAtLastError());
 }
 
+// ---- VQ Repack (flat VQ bytes -> tiled layout) ----
+// Same tile geometry as kbit repack but with VQ byte words instead of bit planes.
+// words_per_block = BS / (P_VAL * 4): p=2→4, p=4→2
+
+template <int P_VAL>
+__global__ void kRepackVQ(
+    const unsigned int* __restrict__ packed_flat, const unsigned char* __restrict__ absmax_flat,
+    unsigned int* __restrict__ packed_tiled, unsigned char* __restrict__ absmax_tiled, const int K_dim, const int N
+) {
+    constexpr int BS = 32;
+    constexpr int WORDS_PER_BLOCK = BS / (P_VAL * 4); // p=2: 4, p=4: 2
+    const int total_k_blocks = K_dim / BS;
+    const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+    if (idx >= N * total_k_blocks)
+        return;
+
+    const int n_idx = idx / total_k_blocks;
+    const int k_block_idx = idx % total_k_blocks;
+    const int k_start = k_block_idx * BS;
+
+    // Source: flat layout
+    const int flat_block_id = n_idx * total_k_blocks + k_block_idx;
+
+    // Destination: tiled layout
+    const int k_tile = k_start / KBIT_TILE_K;
+    const int n_tile = n_idx / KBIT_TILE_N;
+    const int col = n_idx % KBIT_TILE_N;
+    const int kb = (k_start % KBIT_TILE_K) / BS;
+
+    const int n_tiles = N / KBIT_TILE_N;
+    constexpr int KB_PER_TILE = KBIT_TILE_K / BS; // 2
+    constexpr int WORDS_PER_TILE = KBIT_TILE_N * KB_PER_TILE * WORDS_PER_BLOCK;
+    constexpr int ABS_PER_TILE = KBIT_TILE_N * KB_PER_TILE;
+
+    const int tile_base = k_tile * n_tiles + n_tile;
+    const int dst_word_base = tile_base * WORDS_PER_TILE + (col * KB_PER_TILE + kb) * WORDS_PER_BLOCK;
+    const int src_word_base = flat_block_id * WORDS_PER_BLOCK;
+
+#pragma unroll
+    for (int w = 0; w < WORDS_PER_BLOCK; w++)
+        packed_tiled[dst_word_base + w] = packed_flat[src_word_base + w];
+
+    const int dst_abs_idx = tile_base * ABS_PER_TILE + col * KB_PER_TILE + kb;
+    absmax_tiled[dst_abs_idx] = absmax_flat[flat_block_id];
+}
+
+// VQ Repack launcher
+template <int P_VAL>
+void repackVQ(
+    const unsigned int* packed_flat, const unsigned char* absmax_flat,
+    unsigned int* packed_tiled, unsigned char* absmax_tiled,
+    int K_dim, int N, cudaStream_t stream
+) {
+    int total_work = N * (K_dim / 32);
+    int block_size = 256;
+    int grid_size = (total_work + block_size - 1) / block_size;
+    kRepackVQ<P_VAL>
+        <<<grid_size, block_size, 0, stream>>>(packed_flat, absmax_flat, packed_tiled, absmax_tiled, K_dim, N);
+    CUDA_CHECK_RETURN(cudaPeekAtLastError());
+}
+
 // ===========================================================================
 // Hadamard rotation kernel (in-place, blocksize-templated)
 //
@@ -3500,6 +3561,14 @@ INSTANTIATE_KBIT_REPACK(3)
 INSTANTIATE_KBIT_REPACK(4)
 INSTANTIATE_KBIT_REPACK(5)
 
+// VQ repack: P_VAL
+#define INSTANTIATE_VQ_REPACK(P)                                                                                       \
+    template void repackVQ<P>(                                                                                         \
+        const unsigned int*, const unsigned char*, unsigned int*, unsigned char*, int, int, cudaStream_t                \
+    );
+INSTANTIATE_VQ_REPACK(2)
+INSTANTIATE_VQ_REPACK(4)
+
 // Production kernel instantiations — uint8 E4M4 absmax (default)
 #define INSTANTIATE_KBIT_GEMM_PROD_U8(K)                                                                               \
     template void kbitGemmProd<K, half, unsigned char>(                                                                \

csrc/pythonInterface.cpp

@@ -570,6 +570,21 @@ MAKE_KBIT_REPACK(3)
 MAKE_KBIT_REPACK(4)
 MAKE_KBIT_REPACK(5)
 
+// Forward declaration of VQ repack launcher
+template <int P>
+void repackVQ(const unsigned int*, const unsigned char*, unsigned int*, unsigned char*, int, int, cudaStream_t);
+
+#define MAKE_VQ_REPACK(P)                                                                                              \
+    void repack_vq_p##P(                                                                                               \
+        const unsigned int* packed_flat, const unsigned char* absmax_flat, unsigned int* packed_tiled,                  \
+        unsigned char* absmax_tiled, int K_dim, int N, cudaStream_t stream                                              \
+    ) {                                                                                                                \
+        repackVQ<P>(packed_flat, absmax_flat, packed_tiled, absmax_tiled, K_dim, N, stream);                            \
+    }
+
+MAKE_VQ_REPACK(2)
+MAKE_VQ_REPACK(4)
+
 // Forward declarations of GEMM launchers
 template <int K, typename scalar_t, typename ABSMAX_T>
 void kbitGemmProd(
@@ -1519,6 +1534,18 @@ MAKE_CKBIT_REPACK(3)
 MAKE_CKBIT_REPACK(4)
 MAKE_CKBIT_REPACK(5)
 
+// VQ repack extern C wrappers
+#define MAKE_CREPACK_VQ(P)                                                                                             \
+    void crepack_vq_p##P(                                                                                              \
+        const unsigned int* packed_flat, const unsigned char* absmax_flat, unsigned int* packed_tiled,                  \
+        unsigned char* absmax_tiled, int K_dim, int N, cudaStream_t stream                                              \
+    ) {                                                                                                                \
+        repack_vq_p##P(packed_flat, absmax_flat, packed_tiled, absmax_tiled, K_dim, N, stream);                         \
+    }
+
+MAKE_CREPACK_VQ(2)
+MAKE_CREPACK_VQ(4)
+
 // fp16 absmax - all output types
 MAKE_CKBIT_DEQUANT(fp16, half, fp16abs, half, 2)
 MAKE_CKBIT_DEQUANT(fp16, half, fp16abs, half, 3)