Align softmax_online scaffold with issue #38

bench/benchmark_online_softmax.py

@@ -1,11 +1,16 @@
 """Benchmark softmax_online op against torch.softmax and torch.compile(torch.softmax)."""
 
 import argparse
-import os
 
 import torch
 
-from forge_cute_py.ops.softmax_online import softmax_fwd, softmax_bwd
+from forge_cute_py.ops.softmax_online import (
+    get_softmax_online_backend,
+    list_softmax_online_backends,
+    set_softmax_online_backend,
+    softmax_bwd,
+    softmax_fwd,
+)
 from forge_cute_py.util.bench import do_bench, estimate_bandwidth, summarize_times
 
 SHORT_M = [128, 512, 2048, 8192]
@@ -26,6 +31,7 @@ def parse_str_list(s: str) -> list[str]:
 
 
 def main():
+    backend_choices = list_softmax_online_backends()
     parser = argparse.ArgumentParser(description="Benchmark softmax_online op")
     parser.add_argument(
         "--long", action="store_true", help="Use long-N benchmark suite (small M, large N)"
@@ -35,14 +41,9 @@ def main():
     parser.add_argument("--dtypes", type=parse_str_list, default=DEFAULT_DTYPES)
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iterations", type=int, default=100)
-    parser.add_argument(
-        "--impl",
-        choices=["auto", "ref", "kernel"],
-        default="auto",
-        help="softmax_online backend mode (FORGE_SOFTMAX_IMPL)",
-    )
+    parser.add_argument("--backend", choices=backend_choices, default="ref")
     args = parser.parse_args()
-    os.environ["FORGE_SOFTMAX_IMPL"] = args.impl
+    set_softmax_online_backend(args.backend)
 
     if args.m_sizes is None:
         args.m_sizes = LONG_M if args.long else SHORT_M
@@ -54,7 +55,9 @@ def main():
 
     gpu_name = torch.cuda.get_device_name(0)
     suite = "long" if args.long else "short"
-    print(f"softmax_online benchmarks [{suite}] ({gpu_name}) [impl={args.impl}]")
+    print(
+        f"softmax_online benchmarks [{suite}] ({gpu_name}) [backend={get_softmax_online_backend()}]"
+    )
     print()
 
     header = (

bench/run.py

@@ -1,6 +1,5 @@
 import argparse
 import json
-import os
 import sys
 from pathlib import Path
 
@@ -153,14 +152,18 @@ def fn():
                 "reason": str(exc),
             }
         except Exception as exc:
-            impl = os.getenv("FORGE_SOFTMAX_IMPL", "auto")
+            backend = (
+                ops.get_softmax_online_backend()
+                if hasattr(ops, "get_softmax_online_backend")
+                else "unknown"
+            )
             return {
                 "status": "skipped",
                 "op": op_name,
                 "shape": shape,
                 "dtype": str(dtype).replace("torch.", ""),
                 "dim": dim,
-                "reason": f"softmax_online failed (impl={impl}): {exc}",
+                "reason": f"softmax_online failed (backend={backend}): {exc}",
             }
         return {
             "status": "ok",