Big fix in cpu ternary multipliers and model inference output

Author: MohsenDehghankar

benchmarking/bit_1_58/bench_best_k.py

@@ -74,8 +74,8 @@ def bench_inference(fn, warmup=5, repeats=20):
 
 
 def bench_cpu(shapes, k_values, warmup, repeats):
-    from multiplier.bit_1_58.cpu.rsr_v3_3_nonsquare import (
-        RSRTernaryV3_3NonSquareMultiplier,
+    from multiplier.bit_1_58.cpu.rsr_nonsquare import (
+        RSRTernaryNonSquareMultiplier,
     )
 
     rows = []
@@ -105,7 +105,7 @@ def bench_cpu(shapes, k_values, warmup, repeats):
         best_k, best_t = None, float("inf")
         for k in k_values:
             try:
-                rsr = RSRTernaryV3_3NonSquareMultiplier(M, k)
+                rsr = RSRTernaryNonSquareMultiplier(M, k)
                 t_rsr = bench_inference(
                     lambda: rsr(v),
                     warmup=warmup,

benchmarking/bit_1_58/run_cpu_nonsquare.py

@@ -19,8 +19,8 @@
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", ".."))
 
-from multiplier.bit_1_58.pytorch import PytorchMultiplier
-from multiplier.bit_1_58.cpu.rsr_v3_3_nonsquare import RSRTernaryV3_3NonSquareMultiplier
+from multiplier.bit_1_58.pytorch import PytorchBF16Multiplier
+from multiplier.bit_1_58.cpu.rsr_nonsquare import RSRTernaryNonSquareMultiplier
 
 
 def bench_inference(multiplier, v, warmup=5, repeats=20):
@@ -49,15 +49,21 @@ def random_ternary(n_rows, n_cols):
 def main():
     # Non-square shapes: (n_rows, n_cols) — representative of neural network layers
     shapes = [
+        # BitNet-b1.58-2B-4T
+        (640, 2560),
+        (2560, 2560),
+        (2560, 6912),
+        (6912, 2560),
+        # Llama3-8B-1.58-100B-tokens
         (1024, 4096),
-        (4096, 1024),
-        (2048, 2048),
         (4096, 4096),
-        (4096, 11008),
-        (11008, 4096),
-        (8192, 8192),
         (4096, 14336),
         (14336, 4096),
+        # Falcon3-10B-Instruct-1.58bit
+        (1024, 3072),
+        (3072, 3072),
+        (3072, 23040),
+        (23040, 3072),
     ]
     k_values = [4, 8, 12]
     repeats = 20
@@ -83,7 +89,7 @@ def main():
 
         # Pytorch baseline (same for all k)
         try:
-            m_pt = PytorchMultiplier(M)
+            m_pt = PytorchBF16Multiplier(M)
             t_pytorch = bench_inference(m_pt, v, warmup=warmup, repeats=repeats)
         except Exception as e:
             print(f"  [error pytorch: {e}]")
@@ -96,7 +102,7 @@ def main():
 
         for k in k_values:
             try:
-                m_rsr = RSRTernaryV3_3NonSquareMultiplier(M, k)
+                m_rsr = RSRTernaryNonSquareMultiplier(M, k)
                 t_rsr = bench_inference(m_rsr, v, warmup=warmup, repeats=repeats)
             except Exception as e:
                 print(f"  [error rsr_ternary_nonsquare k={k}: {e}]")

benchmarking/llms/bench_inference.py

@@ -26,7 +26,9 @@ def _timed_generate(model, tokenizer, prompt, max_new_tokens, use_chat_template)
     torch.cuda.synchronize() if torch.cuda.is_available() else None
     t0 = time.perf_counter()
     text = generate_text(
-        model, tokenizer, prompt,
+        model,
+        tokenizer,
+        prompt,
         max_new_tokens=max_new_tokens,
         use_chat_template=use_chat_template,
     )
@@ -35,7 +37,9 @@ def _timed_generate(model, tokenizer, prompt, max_new_tokens, use_chat_template)
     return text, elapsed
 
 
-def bench_one(label, load_fn, prompt, max_new_tokens, use_chat_template, warmup, repeats):
+def bench_one(
+    label, load_fn, prompt, max_new_tokens, use_chat_template, warmup, repeats
+):
     print(f"\n{'='*60}")
     print(f"  {label}")
     print(f"{'='*60}")
@@ -48,13 +52,17 @@ def bench_one(label, load_fn, prompt, max_new_tokens, use_chat_template, warmup,
 
     # Warmup
     for i in range(warmup):
-        text, dt = _timed_generate(model, tokenizer, prompt, max_new_tokens, use_chat_template)
+        text, dt = _timed_generate(
+            model, tokenizer, prompt, max_new_tokens, use_chat_template
+        )
         print(f"  Warmup {i+1}: {dt:.3f}s")
 
     # Timed runs
     times = []
     for i in range(repeats):
-        text, dt = _timed_generate(model, tokenizer, prompt, max_new_tokens, use_chat_template)
+        text, dt = _timed_generate(
+            model, tokenizer, prompt, max_new_tokens, use_chat_template
+        )
         times.append(dt)
         n_tokens = len(tokenizer.encode(text))
         print(f"  Run {i+1}: {dt:.3f}s  ({n_tokens} tokens, {n_tokens/dt:.1f} tok/s)")
@@ -65,20 +73,37 @@ def bench_one(label, load_fn, prompt, max_new_tokens, use_chat_template, warmup,
     print(f"  Output: {text[:200]}...")
 
     import torch
+
     del model, tokenizer
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
-    import gc; gc.collect()
+    import gc
+
+    gc.collect()
+
+    return {
+        "label": label,
+        "avg_time": avg,
+        "tok_per_s": n_tokens / avg,
+        "n_tokens": n_tokens,
+    }
 
-    return {"label": label, "avg_time": avg, "tok_per_s": n_tokens / avg, "n_tokens": n_tokens}
 
+def _discover_model_dirs(model_dir, device):
+    """Find preprocessed model directories matching the given device suffix.
 
-def _discover_model_dirs(parent_dir, device):
-    """Find preprocessed model subdirectories matching the given device suffix."""
-    parent = Path(parent_dir)
+    Accepts either a single preprocessed model directory or a parent directory
+    containing multiple preprocessed model subdirectories.
+    """
+    p = Path(model_dir)
     suffix = f"_{device}"
+    # Single model directory: ends with the device suffix and has a config.
+    if p.name.endswith(suffix) and (p / "rsr_config.json").exists():
+        return [p]
+    # Otherwise treat as parent directory containing multiple models.
     dirs = sorted(
-        d for d in parent.iterdir()
+        d
+        for d in p.iterdir()
         if d.is_dir() and d.name.endswith(suffix) and (d / "rsr_config.json").exists()
     )
     return dirs
@@ -89,15 +114,26 @@ def main():
         description="Benchmark RSR vs HF ternary inference.",
         formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
-    parser.add_argument("--model-dir", required=True,
-                        help="Parent directory containing preprocessed model subdirectories")
-    parser.add_argument("--prompt", required=True)
+    parser.add_argument(
+        "--model-dir",
+        required=True,
+        help="Single preprocessed model directory or parent directory containing multiple",
+    )
+    parser.add_argument(
+        "--prompt",
+        required=False,
+        default="Write the numbers from one to two hundred in words separated by commas only:",
+    )
     parser.add_argument("--max-new-tokens", type=int, default=64)
     parser.add_argument("--warmup", type=int, default=1)
     parser.add_argument("--repeats", type=int, default=3)
     parser.add_argument("--no-chat-template", action="store_true")
-    parser.add_argument("--device", required=True, choices=["cpu", "cuda"],
-                        help="Device to benchmark on (also selects model variants)")
+    parser.add_argument(
+        "--device",
+        required=True,
+        choices=["cpu", "cuda"],
+        help="Device to benchmark on (also selects model variants)",
+    )
     _ALL_HF_DTYPES = {
         "cpu": ["float32", "bfloat16"],
         "cuda": ["float32", "bfloat16"],
@@ -110,10 +146,13 @@ def main():
         + [f"hf_{d}" for d in _ALL_HF_DTYPES["cuda"]]
         + [f"hf_{d}" for d in _EXTRA_DTYPES]
     )
-    parser.add_argument("--backends", nargs="+",
-                        default=None,
-                        choices=all_backend_names,
-                        help="Backends to benchmark (default: rsr + all HF dtypes for device)")
+    parser.add_argument(
+        "--backends",
+        nargs="+",
+        default=None,
+        choices=all_backend_names,
+        help="Backends to benchmark (default: rsr + all HF dtypes for device)",
+    )
     args = parser.parse_args()
 
     from integrations.hf.model_infer import load_preprocessed_model, load_hf_model
@@ -138,15 +177,23 @@ def main():
         if "rsr" in args.backends:
             md = str(model_dir)
             loaders["RSR"] = lambda md=md: load_preprocessed_model(
-                md, device=args.device,
+                md,
+                device=args.device,
+                dtype="bfloat16",
             )
         is_bf16_model = "bf16" in model_name or "bfloat16" in model_name
-        hf_dtypes = ["bfloat16"] if is_bf16_model else _ALL_HF_DTYPES[args.device] + _EXTRA_DTYPES
+        hf_dtypes = (
+            ["bfloat16"]
+            if is_bf16_model
+            else _ALL_HF_DTYPES[args.device] + _EXTRA_DTYPES
+        )
         for dtype in hf_dtypes:
             if f"hf_{dtype}" in args.backends:
                 md = str(model_dir)
                 loaders[f"HF {dtype}"] = lambda md=md, dt=dtype: load_hf_model(
-                    md, device=args.device, dtype=dt,
+                    md,
+                    device=args.device,
+                    dtype=dt,
                 )
         for label, load_fn in loaders.items():
             tasks.append((model_name, model_dir, label, load_fn))
@@ -155,24 +202,39 @@ def main():
     for i, (model_name, model_dir, label, load_fn) in enumerate(tasks):
         print(f"\n[{i+1}/{total}] {model_name} / {label}")
         try:
-            r = bench_one(label, load_fn, args.prompt, args.max_new_tokens,
-                          use_chat, args.warmup, args.repeats)
+            r = bench_one(
+                label,
+                load_fn,
+                args.prompt,
+                args.max_new_tokens,
+                use_chat,
+                args.warmup,
+                args.repeats,
+            )
         except Exception as exc:
             print(f"  FAILED: {exc}")
-            all_results.append({
-                "model": model_name, "label": label,
-                "avg_time": float("nan"), "tok_per_s": float("nan"),
-                "n_tokens": 0, "error": str(exc),
-            })
+            all_results.append(
+                {
+                    "model": model_name,
+                    "label": label,
+                    "avg_time": float("nan"),
+                    "tok_per_s": float("nan"),
+                    "n_tokens": 0,
+                    "error": str(exc),
+                }
+            )
             # Try to reset CUDA state; after a device-side assert the
             # context may be unrecoverable so we guard the cleanup too.
             import torch
+
             try:
                 if torch.cuda.is_available():
                     torch.cuda.empty_cache()
             except Exception:
                 pass
-            import gc; gc.collect()
+            import gc
+
+            gc.collect()
             continue
         r["model"] = model_name
         all_results.append(r)
@@ -187,7 +249,9 @@ def main():
         if "error" in r:
             print(f"  {r['model']:<30} {r['label']:<15} {'FAILED':>10} {'—':>10}")
         else:
-            print(f"  {r['model']:<30} {r['label']:<15} {r['avg_time']:>9.3f}s {r['tok_per_s']:>9.1f}")
+            print(
+                f"  {r['model']:<30} {r['label']:<15} {r['avg_time']:>9.3f}s {r['tok_per_s']:>9.1f}"
+            )
 
 
 if __name__ == "__main__":