V8 grouped scalar GEMV + inline work distribution for grouped MMA

Grouped scalar GEMV: ported V8 optimizations (64 threads, 2 warps,
vectorized int4 A loads, __launch_bounds__, M_VAL dispatch 1-4) and
switched from tiled/E4M4 to flat layout with float32 absmax.

Grouped MMA: replaced cudaMemcpy/cudaMalloc/cudaFree work_offsets
computation with inline linear scan over expert_offsets. Caller now
passes max_M directly, eliminating device-to-host sync per call.

Benchmark suite: added grouped_mma kernel type to ncu_driver and
bench_ncu.sh. model_summary.py now shows 5 kernel columns (MMA,
Scalar, Grouped, Grp MMA, fp16) with per-k TOTAL rows.

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

Author: TimDettmers

benchmarks/bench_ncu.sh

@@ -87,6 +87,13 @@ else
     echo "(no M<=4 values requested)" | tee "$RESULTS_DIR/grouped.txt"
 fi
 
+# ---- Grouped MMA kernel (all M values) ----
+echo ""
+echo "=== Grouped MMA (${NUM_EXPERTS} experts, all M) ==="
+printf "%-8s %2s %2s %10s\n" "shape" "k" "M" "avg_us"
+echo "---"
+run_ncu_bench grouped_mma "kbit_grouped_gemm_prod" "['moe_gu','moe_dn']" "$ALL_M" | tee "$RESULTS_DIR/grouped_mma.txt"
+
 # ---- cuBLAS fp16 baselines (CUDA events, all M values) ----
 echo ""
 echo "=== cuBLAS fp16 (dense mm + MoE bmm) ==="

benchmarks/model_summary.py

@@ -5,7 +5,7 @@
 side by side, the best kernel, and speedup vs cuBLAS fp16.
 
 Dense shapes have MMA, Scalar, fp16 columns.
-MoE shapes have Grouped, fp16 (bmm) columns.
+MoE shapes have Grouped (scalar), Grp MMA, fp16 (bmm) columns.
 """
 import os, sys
 
@@ -64,9 +64,10 @@ def main():
     mma = parse_results(os.path.join(results_dir, "mma.txt"))
     scalar = parse_results(os.path.join(results_dir, "scalar.txt"))
     grouped = parse_results(os.path.join(results_dir, "grouped.txt"))
+    grouped_mma = parse_results(os.path.join(results_dir, "grouped_mma.txt"))
     cublas_dense, cublas_moe = parse_cublas(os.path.join(results_dir, "cublas.txt"))
 
-    if not mma and not scalar and not grouped:
+    if not mma and not scalar and not grouped and not grouped_mma:
         print("No benchmark results found. Run bench_ncu.sh first.")
         return
 
@@ -78,25 +79,22 @@ def main():
 
     # Collect all M values
     all_M = set()
-    for key in list(mma.keys()) + list(scalar.keys()) + list(grouped.keys()):
-        all_M.add(key[2])
+    for d in [mma, scalar, grouped, grouped_mma]:
+        for key in d:
+            all_M.add(key[2])
     all_M = sorted(all_M)
 
-    # Column widths
+    # Column widths — 6 kernel columns
     SEP = "+"
-    HDR = (f"{SEP}--------+-----+-------+--------+---------+-------+--------+---------{SEP}")
-    TOP = (f"{SEP}========+=====+=======+========+=========+=======+========+========={SEP}")
+    HDR = f"{SEP}--------+-----+-------+--------+---------+---------+-------+--------+---------{SEP}"
+    TOP = f"{SEP}========+=====+=======+========+=========+=========+=======+========+========={SEP}"
 
     for M in all_M:
         print(f"\n  M={M}:")
         print(f"  {TOP}")
-        print(f"  | {'shape':<6} | {'k':>3} | {'MMA':>5} | {'Scalar':>6} | {'Grouped':>7} | {'fp16':>5} | {'Best':>6} | {'vs fp16':>7} |")
+        print(f"  | {'shape':<6} | {'k':>3} | {'MMA':>5} | {'Scalar':>6} | {'Grouped':>7} | {'Grp MMA':>7} | {'fp16':>5} | {'Best':>6} | {'vs fp16':>7} |")
         print(f"  {HDR}")
 
-        total_best = 0.0
-        total_fp16 = 0.0
-        all_complete = True
-
         for shape in all_shapes:
             is_moe = shape in moe_shapes
 
@@ -105,6 +103,7 @@ def main():
                 m_us = mma.get((shape, k, M)) if not is_moe else None
                 s_us = scalar.get((shape, k, M)) if not is_moe else None
                 g_us = grouped.get((shape, k, M)) if is_moe else None
+                gm_us = grouped_mma.get((shape, k, M)) if is_moe else None
                 fp16 = cublas_moe.get((shape, M)) if is_moe else cublas_dense.get((shape, M))
 
                 # Find best kbit kernel
@@ -115,40 +114,32 @@ def main():
                     candidates["Scalar"] = s_us
                 if g_us is not None:
                     candidates["Grouped"] = g_us
+                if gm_us is not None:
+                    candidates["Grp MMA"] = gm_us
 
                 if candidates:
                     best_name = min(candidates, key=candidates.get)
                     best_us = candidates[best_name]
                 else:
                     best_name, best_us = None, None
 
-                # Compare against fp16.
-                # "Best" = fastest kbit kernel. "vs fp16" = fp16 / Best.
-                # >1.00x means kbit wins, <1.00x means fp16 wins (slowdown).
-                # When no kbit kernel exists, Best falls back to fp16 and shows "-".
                 if best_us is not None and fp16 is not None:
                     speedup = f"{fp16 / best_us:5.2f}x"
-                    total_best += best_us
-                    total_fp16 += fp16
                 elif fp16 is not None and best_us is None:
-                    # No kbit kernel for this config — fp16 only
                     best_name = "-"
                     best_us = fp16
                     speedup = "     -"
-                    total_best += fp16
-                    total_fp16 += fp16
                 else:
                     speedup = "   N/A"
-                    all_complete = False
 
                 best_str = best_name if best_name else "N/A"
 
-                print(f"  | {shape:<6} | {k:>3} | {fmt(m_us)} | {fmt(s_us):>6} | {fmt(g_us):>7} | {fmt(fp16)} | {best_str:>6} | {speedup:>7} |")
+                print(f"  | {shape:<6} | {k:>3} | {fmt(m_us)} | {fmt(s_us):>6} | {fmt(g_us):>7} | {fmt(gm_us):>7} | {fmt(fp16)} | {best_str:>7} | {speedup:>7} |")
 
             print(f"  {HDR}")
 
         # Per-k total rows: sum best_us and fp16 across all shapes for each k
-        print(f"  | {'TOTAL':<6} |     |       |        |         |       |        |         |")
+        print(f"  | {'TOTAL':<6} |     |       |        |         |         |       |        |         |")
         for k in k_bits:
             k_best = 0.0
             k_fp16 = 0.0
@@ -158,6 +149,7 @@ def main():
                 m_us = mma.get((shape, k, M)) if not is_moe else None
                 s_us = scalar.get((shape, k, M)) if not is_moe else None
                 g_us = grouped.get((shape, k, M)) if is_moe else None
+                gm_us = grouped_mma.get((shape, k, M)) if is_moe else None
                 fp16 = cublas_moe.get((shape, M)) if is_moe else cublas_dense.get((shape, M))
 
                 candidates = {}
@@ -167,6 +159,8 @@ def main():
                     candidates["Scalar"] = s_us
                 if g_us is not None:
                     candidates["Grouped"] = g_us
+                if gm_us is not None:
+                    candidates["Grp MMA"] = gm_us
 
                 if candidates:
                     best_us = min(candidates.values())
@@ -182,9 +176,9 @@ def main():
 
             if k_complete and k_best > 0 and k_fp16 > 0:
                 overall = k_fp16 / k_best
-                print(f"  |  k={k:<3} | {k:>3} |       |        |         |       | {k_best:6.1f} | {overall:5.2f}x  |")
+                print(f"  |  k={k:<3} | {k:>3} |       |        |         |         |       | {k_best:6.1f} | {overall:5.2f}x  |")
             else:
-                print(f"  |  k={k:<3} | {k:>3} |       |        |         |       |   N/A  |    N/A  |")
+                print(f"  |  k={k:<3} | {k:>3} |       |        |         |         |       |   N/A  |    N/A  |")
         print(f"  {TOP}")
 
 

benchmarks/ncu_driver.py

@@ -1,9 +1,9 @@
 """ncu kernel driver — runs all shape x k x M configs in a single process.
 
 Used by bench_ncu.sh. Env vars:
-  KERNEL: "mma", "scalar", or "grouped"
+  KERNEL: "mma", "scalar", "grouped", or "grouped_mma"
   M_VALS: comma-separated M values (default "1,2,3,4,5,6,7,8")
-  NUM_EXPERTS: number of active experts for grouped kernel (default 8)
+  NUM_EXPERTS: number of active experts for grouped/grouped_mma kernel (default 8)
 
 Each config runs WARMUP + PROFILED kernel launches. ncu captures all
 matching launches; the sweep script skips warmup and averages profiled.
@@ -27,7 +27,7 @@
 m_vals = [int(x) for x in os.environ.get("M_VALS", "1,2,3,4,5,6,7,8").split(",")]
 NUM_EXPERTS = int(os.environ.get("NUM_EXPERTS", "8"))
 
-# Scalar and grouped kernels only support M<=4
+# Scalar and grouped scalar kernels only support M<=4
 if KERNEL in ("scalar", "grouped"):
     m_vals = [m for m in m_vals if m <= 4]
 
@@ -94,7 +94,52 @@
         torch.cuda.synchronize()
 
 elif KERNEL == "grouped":
-    # Pre-quantize MoE expert weights (NUM_EXPERTS copies per shape)
+    # Pre-quantize MoE expert weights (NUM_EXPERTS copies, flat layout)
+    moe_data = {}
+    for name, K_dim, N in moe_shapes:
+        for k in k_bits_list:
+            codebook = create_normal_float_codebook(k, device=dev)
+            packed_list = []
+            absmax_list = []
+            num_k_blocks = K_dim // 32
+            expected_packed = N * num_k_blocks * k
+            expected_absmax = N * num_k_blocks
+            for _ in range(NUM_EXPERTS):
+                W = torch.randn(K_dim * N, device=dev, dtype=torch.float32)
+                pf, af = torch.ops.bitsandbytes.quantize_kbit(W, codebook, k)
+                packed_list.append(pf[:expected_packed])
+                absmax_list.append(af.cuda()[:expected_absmax])
+            B_packed_all = torch.cat(packed_list, dim=0)
+            B_absmax_all = torch.cat(absmax_list, dim=0)
+            moe_data[(name, k)] = (K_dim, N, B_packed_all, B_absmax_all, codebook)
+
+    configs = []
+    for name, K_dim, N in moe_shapes:
+        for k in k_bits_list:
+            for M in m_vals:
+                configs.append((name, k, M))
+
+    for name, k, M in configs:
+        K_dim, N, B_packed_all, B_absmax_all, codebook = moe_data[(name, k)]
+        # M tokens per expert (all experts get same M for benchmarking)
+        total_tokens = M * NUM_EXPERTS
+        A_concat = torch.randn(total_tokens, K_dim, dtype=torch.float16, device=dev)
+        offsets = list(range(0, total_tokens + 1, M))
+        expert_offsets = torch.tensor(offsets, dtype=torch.int32, device=dev)
+
+        fn = lambda: torch.ops.bitsandbytes.kbit_grouped_scalar_gemv(
+            A_concat, B_packed_all, B_absmax_all, codebook,
+            expert_offsets, K_dim, N, k, NUM_EXPERTS, M)
+
+        for _ in range(WARMUP):
+            fn()
+        torch.cuda.synchronize()
+        for _ in range(PROFILED):
+            fn()
+        torch.cuda.synchronize()
+
+elif KERNEL == "grouped_mma":
+    # Pre-quantize MoE expert weights (NUM_EXPERTS copies, tiled layout for MMA)
     moe_data = {}
     for name, K_dim, N in moe_shapes:
         for k in k_bits_list:
@@ -119,15 +164,14 @@
 
     for name, k, M in configs:
         K_dim, N, B_packed_all, B_absmax_all, codebook = moe_data[(name, k)]
-        # M tokens per expert (all experts get same M for benchmarking)
         total_tokens = M * NUM_EXPERTS
         A_concat = torch.randn(total_tokens, K_dim, dtype=torch.float16, device=dev)
         offsets = list(range(0, total_tokens + 1, M))
         expert_offsets = torch.tensor(offsets, dtype=torch.int32, device=dev)
 
-        fn = lambda: torch.ops.bitsandbytes.kbit_grouped_scalar_gemv(
+        fn = lambda: torch.ops.bitsandbytes.kbit_grouped_gemm(
             A_concat, B_packed_all, B_absmax_all, codebook,
-            expert_offsets, K_dim, N, k, NUM_EXPERTS)
+            expert_offsets, K_dim, N, k, NUM_EXPERTS, M)
 
         for _ in range(WARMUP):
             fn()

bitsandbytes/_ops.py

@@ -606,7 +606,7 @@ def _(
 torch.library.define(
     "bitsandbytes::kbit_grouped_gemm",
     "(Tensor A_concat, Tensor B_packed_all, Tensor B_absmax_all, Tensor codebook, "
-    "Tensor expert_offsets, int K_dim, int N, int k, int num_experts) -> Tensor",
+    "Tensor expert_offsets, int K_dim, int N, int k, int num_experts, int max_M) -> Tensor",
 )
 
 
@@ -621,6 +621,7 @@ def _(
     N: int,
     k: int,
     num_experts: int,
+    max_M: int,
 ) -> torch.Tensor:
     torch._check(k >= 2 and k <= 5, lambda: f"k must be 2-5, got {k}")
     torch._check(A_concat.dim() == 2 and A_concat.shape[1] == K_dim, lambda: "A_concat must be [total_M, K_dim]")
@@ -682,7 +683,7 @@ def _(
 torch.library.define(
     "bitsandbytes::kbit_grouped_scalar_gemv",
     "(Tensor A_concat, Tensor B_packed_all, Tensor B_absmax_all, Tensor codebook, "
-    "Tensor expert_offsets, int K_dim, int N, int k, int num_experts) -> Tensor",
+    "Tensor expert_offsets, int K_dim, int N, int k, int num_experts, int max_M) -> Tensor",
 )
 
 
@@ -697,6 +698,7 @@ def _(
     N: int,
     k: int,
     num_experts: int,
+    max_M: int,
 ) -> torch.Tensor:
     torch._check(k >= 2 and k <= 5, lambda: f"k must be 2-5, got {k}")
     torch._check(A_concat.dim() == 2 and A_concat.shape[1] == K_dim, lambda: "A_concat must be [total_M, K_dim]")

bitsandbytes/backends/cuda/ops.py

@@ -1085,6 +1085,7 @@ def _(
     N: int,
     k: int,
     num_experts: int,
+    max_M: int,
 ) -> torch.Tensor:
     torch._check(k >= 2 and k <= 5, lambda: f"k must be 2-5, got {k}")
     torch._check(
@@ -1114,6 +1115,7 @@ def _(
             ct.c_int(K_dim),
             ct.c_int(N),
             ct.c_int(num_experts),
+            ct.c_int(max_M),
         )
 
     return C_concat
@@ -1193,17 +1195,17 @@ def _(
     N: int,
     k: int,
     num_experts: int,
+    max_M: int,
 ) -> torch.Tensor:
     torch._check(k >= 2 and k <= 5, lambda: f"k must be 2-5, got {k}")
     torch._check(
         A_concat.dtype in (torch.float16, torch.bfloat16),
         lambda: f"kbit_grouped_scalar_gemv supports float16 and bfloat16, got {A_concat.dtype}",
     )
     torch._check(B_packed_all.dtype == torch.int32, lambda: f"B_packed must be int32, got {B_packed_all.dtype}")
-    torch._check(B_absmax_all.dtype == torch.uint8, lambda: f"B_absmax must be uint8 (E4M4), got {B_absmax_all.dtype}")
+    torch._check(B_absmax_all.dtype == torch.float32, lambda: f"B_absmax must be float32, got {B_absmax_all.dtype}")
     torch._check(codebook.dtype == torch.float32, lambda: f"codebook must be float32, got {codebook.dtype}")
     torch._check(expert_offsets.dtype == torch.int32, lambda: f"expert_offsets must be int32, got {expert_offsets.dtype}")
-    torch._check(N % 128 == 0, lambda: f"N ({N}) must be divisible by 128")
 
     total_M = A_concat.shape[0]
     C_concat = torch.empty(total_M, N, device=A_concat.device, dtype=A_concat.dtype)
@@ -1222,6 +1224,7 @@ def _(
             ct.c_int(K_dim),
             ct.c_int(N),
             ct.c_int(num_experts),
+            ct.c_int(max_M),
         )
 
     return C_concat