Add datacenter GPU macro, L2 prefetch hints, and generalized reduction

- Add BNB_DATACENTER_GPU compile-time macro for Hopper (sm_90) and
  Blackwell datacenter (sm_100). Consumer GPUs (sm_89, sm_120) are
  explicitly excluded.
- Add prefetch_l2() helper: issues prefetch.global.L2 on datacenter
  GPUs, compiles to no-op on consumer.
- Add L2 prefetch hints in MMA pipeline loop (prefetch tile kt+2)
- Add L2 prefetch hints in grouped MMA pipeline loop (same pattern)
- Add L2 prefetch hints in scalar GEMV inner loop (next iteration's B)
- Generalize scalar GEMV warp reduction to loop over NUM_WARPS instead
  of hardcoding 2-warp sum (cleaner, same behavior at NUM_WARPS=2)
- Update bench_tiled_vs_flat.py to benchmark v2 kernel alongside flat
  and tiled

Zero consumer regression: 174/174 tests pass, RTX 4090 benchmarks
unchanged. L2 prefetch effect on H100 is neutral for scalar GEMV
(expected — single-iteration-per-thread case).

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

Author: TimDettmers

benchmarks/bench_tiled_vs_flat.py

@@ -1,7 +1,7 @@
 """Benchmark tiled vs flat scalar GEMV with pre-allocated output buffers.
 
 Measures kernel-only time by pre-allocating all buffers before the timing loop.
-No allocations inside the measured region — fair comparison between flat and tiled.
+No allocations inside the measured region — fair comparison between flat, tiled, and tiled v2.
 
 Usage:
     python benchmarks/bench_tiled_vs_flat.py
@@ -40,11 +40,20 @@
 M_VALUES = [1, 2, 4]
 
 if args.graph:
-    print(f"{'shape':<8} {'K_dim':>5} {'N':>5} {'k':>2} {'M':>2}  {'flat_us':>8} {'±flat':>6} {'tiled_us':>8} {'±tiled':>6} {'diff%':>7}")
-    print("-" * 76)
+    print(
+        f"{'shape':<8} {'K_dim':>5} {'N':>5} {'k':>2} {'M':>2}"
+        f"  {'flat_us':>8} {'±flat':>6}"
+        f" {'tiled_us':>8} {'±tl':>4}"
+        f" {'v2_us':>8} {'±v2':>4}"
+        f" {'tl/fl%':>7} {'v2/fl%':>7}"
+    )
+    print("-" * 100)
 else:
-    print(f"{'shape':<8} {'K_dim':>5} {'N':>5} {'k':>2} {'M':>2}  {'flat_us':>8} {'tiled_us':>8} {'diff%':>7}")
-    print("-" * 60)
+    print(
+        f"{'shape':<8} {'K_dim':>5} {'N':>5} {'k':>2} {'M':>2}"
+        f"  {'flat_us':>8} {'tiled_us':>8} {'v2_us':>8} {'tl/fl%':>7} {'v2/fl%':>7}"
+    )
+    print("-" * 75)
 
 for name, K_dim, N in SHAPES:
     for k in K_VALUES:
@@ -63,16 +72,27 @@
             # Pre-allocate output buffers
             out_flat = torch.empty(M, N, dtype=torch.float16, device="cuda")
             out_tiled = torch.empty(M, N, dtype=torch.float16, device="cuda")
+            out_v2 = torch.empty(M, N, dtype=torch.float16, device="cuda")
+
+            # v2 workspace
+            n_tiles = N // 128
+            C_workspace = torch.zeros(M, N, dtype=torch.float32, device="cuda")
+            tile_counters = torch.zeros(n_tiles, dtype=torch.int32, device="cuda")
 
             if args.ncu:
-                # NCU mode: single call each, profiler captures kernel time
                 torch.ops.bitsandbytes.kbit_scalar_gemv.out(
                     A, packed_flat, absmax_flat, codebook, K_dim, N, k, out_flat
                 )
                 torch.ops.bitsandbytes.kbit_scalar_gemv_tiled_(
                     A, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out_tiled
                 )
-                print(f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}  {'ncu':>8} {'ncu':>8} {'ncu':>7}")
+                torch.ops.bitsandbytes.kbit_scalar_gemv_v2_(
+                    A, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out_v2, C_workspace, tile_counters
+                )
+                print(
+                    f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}"
+                    f"  {'ncu':>8} {'ncu':>8} {'ncu':>8} {'ncu':>7} {'ncu':>7}"
+                )
                 continue
 
             start = torch.cuda.Event(enable_timing=True)
@@ -88,11 +108,16 @@ def call_tiled():
                     A, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out_tiled
                 )
 
+            def call_v2():
+                torch.ops.bitsandbytes.kbit_scalar_gemv_v2_(
+                    A, packed_tiled, absmax_tiled, codebook, K_dim, N, k, out_v2, C_workspace, tile_counters
+                )
+
             if args.graph:
                 import statistics
 
                 # CUDA graph replay — measures kernel-only time
-                for fn in (call_flat, call_tiled):
+                for fn in (call_flat, call_tiled, call_v2):
                     for _ in range(3):
                         fn()
                 torch.cuda.synchronize()
@@ -117,38 +142,44 @@ def bench_graph(fn, trials, iters):
 
                 flat_us, flat_std = bench_graph(call_flat, args.trials, args.iters)
                 tiled_us, tiled_std = bench_graph(call_tiled, args.trials, args.iters)
+                v2_us, v2_std = bench_graph(call_v2, args.trials, args.iters)
             else:
-                # CUDA events timing (includes Python dispatch overhead)
-                for _ in range(args.warmup):
-                    call_flat()
-                torch.cuda.synchronize()
-                start.record()
-                for _ in range(args.iters):
-                    call_flat()
-                end.record()
-                torch.cuda.synchronize()
-                flat_us = start.elapsed_time(end) * 1000 / args.iters
+                def bench_events(fn):
+                    for _ in range(args.warmup):
+                        fn()
+                    torch.cuda.synchronize()
+                    start.record()
+                    for _ in range(args.iters):
+                        fn()
+                    end.record()
+                    torch.cuda.synchronize()
+                    return start.elapsed_time(end) * 1000 / args.iters
 
-                for _ in range(args.warmup):
-                    call_tiled()
-                torch.cuda.synchronize()
-                start.record()
-                for _ in range(args.iters):
-                    call_tiled()
-                end.record()
-                torch.cuda.synchronize()
-                tiled_us = start.elapsed_time(end) * 1000 / args.iters
+                flat_us = bench_events(call_flat)
+                tiled_us = bench_events(call_tiled)
+                v2_us = bench_events(call_v2)
 
-            diff_pct = (tiled_us - flat_us) / flat_us * 100
+            tl_pct = (tiled_us - flat_us) / flat_us * 100
+            v2_pct = (v2_us - flat_us) / flat_us * 100
             if args.graph:
                 print(
                     f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}"
-                    f"  {flat_us:>8.1f} {flat_std:>5.1f}σ {tiled_us:>8.1f} {tiled_std:>5.1f}σ {diff_pct:>+7.1f}%"
+                    f"  {flat_us:>8.1f} {flat_std:>5.1f}σ"
+                    f" {tiled_us:>8.1f} {tiled_std:>3.1f}σ"
+                    f" {v2_us:>8.1f} {v2_std:>3.1f}σ"
+                    f" {tl_pct:>+7.1f}% {v2_pct:>+7.1f}%"
                 )
             else:
-                print(f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}  {flat_us:>8.1f} {tiled_us:>8.1f} {diff_pct:>+7.1f}%")
+                print(
+                    f"{name:<8} {K_dim:>5} {N:>5} {k:>2} {M:>2}"
+                    f"  {flat_us:>8.1f} {tiled_us:>8.1f} {v2_us:>8.1f} {tl_pct:>+7.1f}% {v2_pct:>+7.1f}%"
+                )
 
         # Correctness check (once per shape/k)
         assert torch.equal(out_flat, out_tiled) or torch.allclose(out_flat, out_tiled, rtol=0.05, atol=0.1), (
-            f"MISMATCH {name} k={k}: max diff = {(out_flat - out_tiled).abs().max().item()}"
+            f"MISMATCH flat vs tiled {name} k={k}: max diff = {(out_flat - out_tiled).abs().max().item()}"
+        )
+        # v2 uses split-K so small FP diffs are expected
+        assert torch.allclose(out_flat.float(), out_v2.float(), rtol=0.1, atol=1.0), (
+            f"MISMATCH flat vs v2 {name} k={k}: max diff = {(out_flat.float() - out_v2.float()).abs().max().item()}"
         )

csrc/ops.cu

@@ -1011,6 +1011,23 @@ void repackKbit(
     CUDA_CHECK_RETURN(cudaPeekAtLastError());
 }
 
+// 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__)
+#define BNB_DATACENTER_GPU (__CUDA_ARCH__ == 900 || __CUDA_ARCH__ == 1000)
+#else
+#define BNB_DATACENTER_GPU 0
+#endif
+
+// L2 prefetch hint (datacenter GPUs only — consumer GPUs ignore it)
+__device__ __forceinline__ void prefetch_l2(const void* ptr) {
+#if BNB_DATACENTER_GPU
+    asm volatile("prefetch.global.L2 [%0];" ::"l"(ptr));
+#else
+    (void)ptr;
+#endif
+}
+
 // cp.async helpers (sm_80+) — used by production MMA and grouped MMA kernels
 __device__ __forceinline__ void cp_async_cg_16(void* __restrict__ smem, const void* __restrict__ gmem) {
     uint32_t smem_addr = static_cast<uint32_t>(__cvta_generic_to_shared(smem));
@@ -1299,6 +1316,12 @@ __global__ void __launch_bounds__(TILE_N_VAL <= 64 ? 128 : 256, TILE_N_VAL <= 64
             if (kt + 1 < kt_end) {
                 fetch_tile((kt + 1 - kt_start) % 2, kt + 1);
                 cp_async_fence();
+                // L2 prefetch for tile kt+2 (warms L2 before next fetch_tile issues cp.async)
+                if (kt + 2 < kt_end) {
+                    const int pf_tile = (kt + 2) * n_tiles + n_tile;
+                    prefetch_l2(B_packed + pf_tile * B_STAGE_WORDS);
+                    prefetch_l2(B_absmax + pf_tile * ABS_STAGE_ELEMS);
+                }
                 cp_async_wait<1>();
             } else {
                 cp_async_wait<0>();
@@ -1736,6 +1759,12 @@ __global__ void kbit_grouped_gemm_prod(
             if (kt + 1 < kt_end) {
                 fetch_tile((kt - kt_start + 1) % 2, kt + 1);
                 cp_async_fence();
+                // L2 prefetch for tile kt+2
+                if (kt + 2 < kt_end) {
+                    const int pf_tile = (kt + 2) * n_tiles + n_tile;
+                    prefetch_l2(B_packed + pf_tile * B_STAGE_WORDS);
+                    prefetch_l2(B_absmax + pf_tile * ABS_STAGE_ELEMS);
+                }
                 cp_async_wait<1>();
             } else {
                 cp_async_wait<0>();
@@ -2009,6 +2038,21 @@ __global__ void __launch_bounds__(64, M_VAL <= 2 ? 24 : 16) kbit_scalar_gemv(
             abs_idx = tile_base * ABS_PER_TILE + col_in_tile * KB_PER_TILE + kb;
         }
 
+        // L2 prefetch for next iteration's B data
+        {
+            const int next_block_idx = block_idx + BLOCK_SIZE;
+            if (next_block_idx < num_k_blocks) {
+                if constexpr (!TILED) {
+                    prefetch_l2(&B_col[next_block_idx * K_BITS]);
+                } else {
+                    const int nk_tile = next_block_idx / KB_PER_TILE;
+                    const int nkb = next_block_idx % KB_PER_TILE;
+                    const int ntb = nk_tile * n_tiles + n_tile;
+                    prefetch_l2(&B_packed[ntb * WORDS_PER_TILE + (col_in_tile * KB_PER_TILE + nkb) * K_BITS]);
+                }
+            }
+        }
+
         // Load k bit-plane words (guarded; invalid threads get 0)
         // Vector loads for power-of-2 K_BITS, scalar for others.
         const unsigned int* B_src = TILED ? B_packed : B_col;
@@ -2092,12 +2136,15 @@ __global__ void __launch_bounds__(64, M_VAL <= 2 ? 24 : 16) kbit_scalar_gemv(
     }
     __syncthreads();
 
-    // Thread 0 sums both warps and writes output
+    // Thread 0 sums all warps and writes output
     if (threadIdx.x == 0) {
 #pragma unroll
         for (int m = 0; m < M_VAL; m++) {
             if (m < M) {
-                float sum = s_partial[0 * M_MAX + m] + s_partial[1 * M_MAX + m];
+                float sum = 0.0f;
+#pragma unroll
+                for (int w = 0; w < NUM_WARPS; w++)
+                    sum += s_partial[w * M_MAX + m];
                 C[m * N + col] = ScalarOps<scalar_t>::from_float(sum);
             }
         }