feat: Add VQ scalar GEMV kernel with flat layout support

- vq_scalar_gemv<P_VAL, M_VAL, TILED, scalar_t, ABSMAX_T> kernel in ops.cu
- Shared memory codebook: p=2 → half2[256] (1KB), p=4 → split half2[256]+half2[256] (2KB)
- Inner loop: per-byte iteration, codebook lookup, fp16→fp32 convert, FMA
- Supports M=1-4, fp16/bf16, E4M4 and float32 absmax
- 64 threads (2 warps), 40 registers M=1, 0 spills, full occupancy
- Correctness verified: all (p,dtype,M,K,N) combos pass rtol=1e-2
- Performance: 7.4 us for 2048x5120 M=1 p=2 (vs 2.1 us target - needs optimization)
- Full registration chain: ops.cu → pythonInterface.cpp → _ops.py → backends/cuda/ops.py

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

Author: TimDettmers

bitsandbytes/_ops.py

@@ -617,6 +617,106 @@ def _(
     return out
 
 
+# VQ scalar GEMV: byte-indexed codebook lookup GEMV for M=1-4
+
+torch.library.define(
+    "bitsandbytes::vq_scalar_gemv",
+    "(Tensor A, Tensor B_packed, Tensor B_absmax, Tensor codebook, int K_dim, int N, int p) -> Tensor",
+)
+
+
+@register_fake("bitsandbytes::vq_scalar_gemv")
+def _(
+    A: torch.Tensor,
+    B_packed: torch.Tensor,
+    B_absmax: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+) -> torch.Tensor:
+    torch._check(p in (2, 4), lambda: f"p must be 2 or 4, got {p}")
+    torch._check(A.dim() == 2 and A.shape[1] == K_dim, lambda: "A must be [M, K_dim]")
+    torch._check(A.shape[0] <= 4, lambda: f"vq_scalar_gemv supports M<=4, got {A.shape[0]}")
+    torch._check(A.dtype in (torch.float16, torch.bfloat16), lambda: f"A must be fp16 or bf16, got {A.dtype}")
+    M = A.shape[0]
+    return torch.empty(M, N, device=A.device, dtype=A.dtype)
+
+
+torch.library.define(
+    "bitsandbytes::vq_scalar_gemv.out",
+    "(Tensor A, Tensor B_packed, Tensor B_absmax, Tensor codebook, int K_dim, int N, int p, Tensor(a!) out) -> ()",
+)
+
+
+@register_fake("bitsandbytes::vq_scalar_gemv.out")
+def _(
+    A: torch.Tensor,
+    B_packed: torch.Tensor,
+    B_absmax: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+    out: torch.Tensor,
+) -> None:
+    pass
+
+
+# VQ scalar GEMV with tiled B layout
+
+torch.library.define(
+    "bitsandbytes::vq_scalar_gemv_tiled",
+    "(Tensor A, Tensor B_packed_tiled, Tensor B_absmax_tiled, Tensor codebook, int K_dim, int N, int p) -> Tensor",
+)
+
+
+@register_fake("bitsandbytes::vq_scalar_gemv_tiled")
+def _(
+    A: torch.Tensor,
+    B_packed_tiled: torch.Tensor,
+    B_absmax_tiled: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+) -> torch.Tensor:
+    torch._check(p in (2, 4), lambda: f"p must be 2 or 4, got {p}")
+    torch._check(A.dim() == 2 and A.shape[1] == K_dim, lambda: "A must be [M, K_dim]")
+    torch._check(A.shape[0] <= 4, lambda: f"vq_scalar_gemv_tiled supports M<=4, got {A.shape[0]}")
+    torch._check(A.dtype in (torch.float16, torch.bfloat16), lambda: f"A must be fp16 or bf16, got {A.dtype}")
+    M = A.shape[0]
+    return torch.empty(M, N, device=A.device, dtype=A.dtype)
+
+
+# VQ scalar GEMV tiled with pre-allocated output (CUDA graph compatible)
+
+torch.library.define(
+    "bitsandbytes::vq_scalar_gemv_tiled_",
+    "(Tensor A, Tensor B_packed_tiled, Tensor B_absmax_tiled, Tensor codebook, int K_dim, int N, int p, "
+    "Tensor(a!) out) -> Tensor(a!)",
+)
+
+
+@register_fake("bitsandbytes::vq_scalar_gemv_tiled_")
+def _(
+    A: torch.Tensor,
+    B_packed_tiled: torch.Tensor,
+    B_absmax_tiled: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    torch._check(p in (2, 4), lambda: f"p must be 2 or 4, got {p}")
+    torch._check(A.dim() == 2 and A.shape[1] == K_dim, lambda: "A must be [M, K_dim]")
+    torch._check(A.shape[0] <= 4, lambda: f"vq_scalar_gemv_tiled_ supports M<=4, got {A.shape[0]}")
+    torch._check(A.dtype in (torch.float16, torch.bfloat16), lambda: f"A must be fp16 or bf16, got {A.dtype}")
+    torch._check(out.dtype == A.dtype, lambda: f"out dtype {out.dtype} must match A dtype {A.dtype}")
+    return out
+
+
 # K-bit repack: flat bit-plane layout -> GEMM-tiled layout
 
 torch.library.define(

bitsandbytes/backends/cuda/ops.py

@@ -1059,6 +1059,113 @@ def _(
     return out
 
 
+def _vq_scalar_gemv_impl(
+    A: torch.Tensor,
+    B_packed: torch.Tensor,
+    B_absmax: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+    out: torch.Tensor,
+    tiled: bool = False,
+) -> None:
+    M = A.shape[0]
+    dtype_suffix = "fp16" if A.dtype == torch.float16 else "bf16"
+    tiled_str = "_tiled" if tiled else ""
+
+    with _cuda_device_of(A):
+        fn = getattr(lib, f"cvq_scalar_gemv{tiled_str}_{dtype_suffix}_p{p}")
+        fn(
+            get_ptr(A),
+            get_ptr(B_packed),
+            get_ptr(B_absmax),
+            get_ptr(codebook),
+            get_ptr(out),
+            ct.c_int(M),
+            ct.c_int(K_dim),
+            ct.c_int(N),
+            _get_tensor_stream(A),
+        )
+
+
+@register_kernel("bitsandbytes::vq_scalar_gemv", "cuda")
+def _(
+    A: torch.Tensor,
+    B_packed: torch.Tensor,
+    B_absmax: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+) -> torch.Tensor:
+    torch._check(p in (2, 4), lambda: f"p must be 2 or 4, got {p}")
+    torch._check(
+        A.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"vq_scalar_gemv supports float16 and bfloat16, got {A.dtype}",
+    )
+    M = A.shape[0]
+    out = torch.empty(M, N, device=A.device, dtype=A.dtype)
+    _vq_scalar_gemv_impl(A, B_packed, B_absmax, codebook, K_dim, N, p, out=out)
+    return out
+
+
+@register_kernel("bitsandbytes::vq_scalar_gemv.out", "cuda")
+def _(
+    A: torch.Tensor,
+    B_packed: torch.Tensor,
+    B_absmax: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+    out: torch.Tensor,
+) -> None:
+    _vq_scalar_gemv_impl(A, B_packed, B_absmax, codebook, K_dim, N, p, out=out)
+
+
+@register_kernel("bitsandbytes::vq_scalar_gemv_tiled", "cuda")
+def _(
+    A: torch.Tensor,
+    B_packed_tiled: torch.Tensor,
+    B_absmax_tiled: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+) -> torch.Tensor:
+    torch._check(p in (2, 4), lambda: f"p must be 2 or 4, got {p}")
+    torch._check(
+        A.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"vq_scalar_gemv_tiled supports float16 and bfloat16, got {A.dtype}",
+    )
+    M = A.shape[0]
+    out = torch.empty(M, N, device=A.device, dtype=A.dtype)
+    _vq_scalar_gemv_impl(A, B_packed_tiled, B_absmax_tiled, codebook, K_dim, N, p, out=out, tiled=True)
+    return out
+
+
+@register_kernel("bitsandbytes::vq_scalar_gemv_tiled_", "cuda")
+def _(
+    A: torch.Tensor,
+    B_packed_tiled: torch.Tensor,
+    B_absmax_tiled: torch.Tensor,
+    codebook: torch.Tensor,
+    K_dim: int,
+    N: int,
+    p: int,
+    out: torch.Tensor,
+) -> torch.Tensor:
+    torch._check(p in (2, 4), lambda: f"p must be 2 or 4, got {p}")
+    torch._check(
+        A.dtype in (torch.float16, torch.bfloat16),
+        lambda: f"vq_scalar_gemv_tiled_ supports float16 and bfloat16, got {A.dtype}",
+    )
+    M = A.shape[0]
+    _vq_scalar_gemv_impl(A, B_packed_tiled, B_absmax_tiled, codebook, K_dim, N, p, out=out, tiled=True)
+    return out
+
+
 @register_kernel("bitsandbytes::repack_kbit", "cuda")
 def _(
     packed_flat: torch.Tensor,