frost(gemm): Add Rubin kernel pipelines

benchmark/gemm/frost/benchmark_block_scale_matmul.py

@@ -25,6 +25,7 @@
 
 from cudnn.gemm.frost.compiler import jit_from_cudnn_graph
 from cudnn.gemm.frost.tile_config import CATALOG as _CATALOG
+from cudnn.gemm.frost.tile_config import by_name as _by_name
 
 
 def _build_spec_map():
@@ -36,8 +37,8 @@ def _build_spec_map():
         kb_want = 384 if cfg.pipeline == "sm103" else 128
         if cfg.cta_tile_m % 128 or cfg.cta_tile_n % 128 or cfg.cta_tile_k_bytes != kb_want:
             continue
-        # sm103 has 1ctamma + 2ctamma CLC templates (no static variants).
-        scheds = (("clc", ""),) if cfg.pipeline == "sm103" else (("clc", ""), ("static", "_static"))
+        # Only sm100 has static-scheduler variants; sm103 / sm107 are CLC-only.
+        scheds = (("clc", ""), ("static", "_static")) if cfg.pipeline == "sm100" else (("clc", ""),)
         for cg in (1, 2):
             if cg == 2 and (cfg.cgrp_size_m % 2 or cfg.cta_tile_m == 64):
                 continue
@@ -48,6 +49,27 @@ def _build_spec_map():
 
 _SPEC_MAP = _build_spec_map()
 
+_LABEL_RE = re.compile(r"^(CONFIG_sm\d+_\d+x\d+x\d+_\d+x\d+x\d+_cluster\d+x\d+)_([12])ctamma(_static)?$")
+
+
+def _spec_for(name):
+    """(geometry cfg, cta_group, scheduler) for a --configs label, or None.
+
+    The sweep set comes from the registry funnel over CATALOG; a label naming a
+    geometry outside it (e.g. a num_mma_m > 1 tile, which `by_name` synthesizes) is
+    still runnable, so parse it rather than reporting UNKNOWN_CONFIG."""
+    spec = _SPEC_MAP.get(name)
+    if spec is not None:
+        return spec
+    m = _LABEL_RE.match(name)
+    if m is None:
+        return None
+    try:
+        cfg = _by_name(m.group(1))
+    except (KeyError, NotImplementedError):
+        return None
+    return cfg, int(m.group(2)), "static" if m.group(3) else "clc"
+
 
 def _vp_bs(handles, a, b, c, sfa, sfb):
     """Block-scale single-GEMM variant-pack dict keyed by the graph's tensors."""
@@ -57,7 +79,7 @@ def _vp_bs(handles, a, b, c, sfa, sfb):
 
 def _build_plan(g, cfg, name):
     """JIT-compile the recorded graph with a forced tile config."""
-    return jit_from_cudnn_graph(g, config=cfg, cta_group=_SPEC_MAP[name][1], scheduler=_SPEC_MAP[name][2])
+    return jit_from_cudnn_graph(g, config=cfg, cta_group=_spec_for(name)[1], scheduler=_spec_for(name)[2])
 
 
 # Combo table (input dtype family + scale dtype + block size)
@@ -508,10 +530,10 @@ def _nsys_worker(shape, combo, configs, warmup, iters, ref_mode, nbuf) -> None:
     torch.cuda.synchronize()
 
     # 2. each block-scale config.
-    name_to_cfg = {lbl: sp[0] for lbl, sp in _SPEC_MAP.items()}
     config_names = configs or list(_SPEC_MAP)
     for name in config_names:
-        cfg = name_to_cfg.get(name)
+        spec = _spec_for(name)
+        cfg = spec[0] if spec else None
         if cfg is None:
             continue
         try:
@@ -603,7 +625,6 @@ def main() -> int:
 
     flops = 2 * B * M * N * K
     config_names = [c.strip() for c in args.configs.split(",")] if args.configs else list(_SPEC_MAP)
-    name_to_cfg = {lbl: sp[0] for lbl, sp in _SPEC_MAP.items()}
 
     print(f"\n=== block-scale matmul B={B} {M}x{N}x{K}  (~{flops / 1e9:.1f} GFLOP) — " f"{combo} in / BF16 out ===")
 
@@ -644,7 +665,8 @@ def _fmt_row(name: str, tflops: float, ms: float, note: str, ref_tflops: float)
             print("  reference kernel: not detected in nsys output")
 
         for name in config_names:
-            cfg = name_to_cfg.get(name)
+            spec = _spec_for(name)
+            cfg = spec[0] if spec else None
             if cfg is None:
                 rows.append((name, 0.0, float("inf"), "UNKNOWN_CONFIG"))
                 continue
@@ -687,7 +709,8 @@ def _fmt_row(name: str, tflops: float, ms: float, note: str, ref_tflops: float)
 
         ctx_dead = False
         for name in config_names:
-            cfg = name_to_cfg.get(name)
+            spec = _spec_for(name)
+            cfg = spec[0] if spec else None
             if cfg is None:
                 row = (name, 0.0, float("inf"), "UNKNOWN_CONFIG")
             elif ctx_dead:

benchmark/gemm/frost/benchmark_block_scale_matmul_swiglu.py

@@ -11,6 +11,7 @@
 from __future__ import annotations
 
 import argparse
+import re
 import sys
 from typing import Callable
 
@@ -21,6 +22,7 @@
 from types import SimpleNamespace
 
 from cudnn.gemm.frost.compiler import jit_from_cudnn_graph
+from cudnn.gemm.frost.tile_config import by_name as _by_name
 from cudnn.gemm.frost.graph_analyzer import analyze
 from cudnn.gemm.frost.kernel_registry import candidates as _registry_candidates
 
@@ -204,6 +206,27 @@ def _build_spec_map():
 
 _SPEC_MAP = _build_spec_map()
 
+_LABEL_RE = re.compile(r"^(CONFIG_sm\d+_\d+x\d+x\d+_\d+x\d+x\d+_cluster\d+x\d+)_([12])ctamma(_static)?$")
+
+
+def _spec_for(name):
+    """(geometry cfg, cta_group, scheduler) for a --configs label, or None.
+
+    The sweep set comes from the registry funnel over CATALOG; a label naming a
+    geometry outside it (e.g. a num_mma_m > 1 tile, which `by_name` synthesizes) is
+    still runnable, so parse it rather than reporting it unsweepable."""
+    spec = _SPEC_MAP.get(name)
+    if spec is not None:
+        return spec
+    m = _LABEL_RE.match(name)
+    if m is None:
+        return None
+    try:
+        cfg = _by_name(m.group(1))
+    except (KeyError, NotImplementedError):
+        return None
+    return cfg, int(m.group(2)), "static" if m.group(3) else "clc"
+
 
 def main() -> int:
     p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
@@ -244,10 +267,11 @@ def main() -> int:
 
     best = None
     for name in config_names:
-        if name not in _SPEC_MAP:
+        spec = _spec_for(name)
+        if spec is None:
             print(f"  {name:62s} UNKNOWN (not a sweepable block-scale strategy)")
             continue
-        cfg, cta_group, sched = _SPEC_MAP[name]
+        cfg, cta_group, sched = spec
         try:
             g, h = _graph(B, M, N, K)
             plan = _build_plan(g, cfg, cta_group, sched)

benchmark/gemm/frost/benchmark_matmul_mixed_input.py

@@ -27,6 +27,7 @@
 import torch
 
 from cudnn.gemm.frost.compiler import jit_from_cudnn_graph
+from cudnn.gemm.frost.tile_config import by_name as _by_name
 from cudnn.gemm.frost.graph_analyzer import analyze
 from cudnn.gemm.frost.kernel_registry import candidates as _candidates
 
@@ -81,6 +82,27 @@ def _build_spec_map():
 
 _SPEC_MAP = _build_spec_map()
 
+_LABEL_RE = re.compile(r"^(CONFIG_sm\d+_\d+x\d+x\d+_\d+x\d+x\d+_cluster\d+x\d+)_([12])ctamma(_static)?$")
+
+
+def _spec_for(name):
+    """(geometry cfg, cta_group, scheduler) for a --configs label, or None.
+
+    The sweep set comes from the registry funnel over CATALOG; a label naming a
+    geometry outside it (e.g. a num_mma_m > 1 tile, which `by_name` synthesizes) is
+    still runnable, so parse it rather than reporting UNKNOWN_CONFIG."""
+    spec = _SPEC_MAP.get(name)
+    if spec is not None:
+        return spec
+    m = _LABEL_RE.match(name)
+    if m is None:
+        return None
+    try:
+        cfg = _by_name(m.group(1))
+    except (KeyError, NotImplementedError):
+        return None
+    return cfg, int(m.group(2)), "static" if m.group(3) else "clc"
+
 
 def _vp(handles, a, b, c):
     """Variant-pack dict {tensor: buffer}; `a` is the narrow (load-dtype) A root operand."""
@@ -90,7 +112,7 @@ def _vp(handles, a, b, c):
 
 def _build_plan(g, cfg, name):
     """JIT-compile the graph with a forced tile config -> callable kernel."""
-    return jit_from_cudnn_graph(g, config=cfg, cta_group=_SPEC_MAP[name][1], scheduler=_SPEC_MAP[name][2])
+    return jit_from_cudnn_graph(g, config=cfg, cta_group=_spec_for(name)[1], scheduler=_spec_for(name)[2])
 
 
 # ---------------------------------------------------------------------------
@@ -378,10 +400,10 @@ def _nsys_worker(shape, configs, warmup, iters, nbuf, load_dt, tin_dt, tout_dt)
     torch.cuda.synchronize()
 
     # 2. each GEMM config.
-    name_to_cfg = {lbl: sp[0] for lbl, sp in _SPEC_MAP.items()}
     config_names = configs or list(_SPEC_MAP)
     for name in config_names:
-        cfg = name_to_cfg.get(name)
+        spec = _spec_for(name)
+        cfg = spec[0] if spec else None
         if cfg is None:
             continue
         try:
@@ -449,7 +471,6 @@ def main() -> int:
 
     flops = 2 * B * M * N * K
     config_names = [c.strip() for c in args.configs.split(",")] if args.configs else list(_SPEC_MAP)
-    name_to_cfg = {lbl: sp[0] for lbl, sp in _SPEC_MAP.items()}
 
     print(f"\n=== mixed-input matmul B={B} {M}x{N}x{K}  (~{flops / 1e9:.1f} GFLOP) " f"— A={load_dt} -> {tin_dt} @ {tin_dt} -> {tout_dt} ===")
 
@@ -493,7 +514,8 @@ def _fmt_row(name, tflops, ms, note, ref_tflops) -> str:
             print("  cuBLAS kernel: not detected in nsys output")
 
         for name in config_names:
-            cfg = name_to_cfg.get(name)
+            spec = _spec_for(name)
+            cfg = spec[0] if spec else None
             if cfg is None:
                 rows.append((name, 0.0, float("inf"), "UNKNOWN_CONFIG"))
                 continue
@@ -528,7 +550,8 @@ def _fmt_row(name, tflops, ms, note, ref_tflops) -> str:
 
         ctx_dead = False
         for name in config_names:
-            cfg = name_to_cfg.get(name)
+            spec = _spec_for(name)
+            cfg = spec[0] if spec else None
             if cfg is None:
                 row = (name, 0.0, float("inf"), "UNKNOWN_CONFIG")
             elif ctx_dead:

benchmark/gemm/frost/benchmark_matmul_swiglu.py

@@ -10,6 +10,7 @@
 from __future__ import annotations
 
 import argparse
+import re
 import sys
 import time
 from typing import Callable
@@ -21,6 +22,7 @@
 from types import SimpleNamespace
 
 from cudnn.gemm.frost.compiler import jit_from_cudnn_graph
+from cudnn.gemm.frost.tile_config import by_name as _by_name
 from cudnn.gemm.frost.graph_analyzer import analyze
 from cudnn.gemm.frost.kernel_registry import candidates as _registry_candidates
 
@@ -147,6 +149,27 @@ def _build_spec_map():
 
 _SPEC_MAP = _build_spec_map()
 
+_LABEL_RE = re.compile(r"^(CONFIG_sm\d+_\d+x\d+x\d+_\d+x\d+x\d+_cluster\d+x\d+)_([12])ctamma(_static)?$")
+
+
+def _spec_for(name):
+    """(geometry cfg, cta_group, scheduler) for a --configs label, or None.
+
+    The sweep set comes from the registry funnel over CATALOG; a label naming a
+    geometry outside it (e.g. a num_mma_m > 1 tile, which `by_name` synthesizes) is
+    still runnable, so parse it rather than reporting it unsweepable."""
+    spec = _SPEC_MAP.get(name)
+    if spec is not None:
+        return spec
+    m = _LABEL_RE.match(name)
+    if m is None:
+        return None
+    try:
+        cfg = _by_name(m.group(1))
+    except (KeyError, NotImplementedError):
+        return None
+    return cfg, int(m.group(2)), "static" if m.group(3) else "clc"
+
 
 # ---------------------------------------------------------------------------
 # Main
@@ -204,10 +227,11 @@ def main() -> int:
 
     best = None
     for label in config_names:
-        if label not in _SPEC_MAP:
+        spec = _spec_for(label)
+        if spec is None:
             print(f"  {label:62s} UNKNOWN (not a sweepable swiglu strategy)")
             continue
-        cfg, cta_group, sched = _SPEC_MAP[label]
+        cfg, cta_group, sched = spec
         try:
             g, h = _graph_swiglu(B, M, N, K, in_dt, out_dt)
             plan = _build_plan(g, cfg, cta_group, sched)

benchmark/gemm/frost/benchmark_moe_block_scale_matmul.py

@@ -25,6 +25,7 @@
 import torch
 
 from cudnn.gemm.frost.compiler import jit_from_cudnn_graph
+from cudnn.gemm.frost.tile_config import by_name as _by_name
 from cudnn.gemm.frost.graph_analyzer import analyze
 from cudnn.gemm.frost.kernel_registry import candidates as _candidates
 from cudnn.gemm.frost.tile_config import TileConfig
@@ -38,7 +39,7 @@ def _vp_moe_bs(handles, token, weight, sfa, sfb, fto, output):
 
 def _build_plan(g, cfg, name):
     """JIT-compile the recorded graph with a forced tile config -> compiled kernel."""
-    return jit_from_cudnn_graph(g, config=cfg, cta_group=_SPEC_MAP[name][1], scheduler=_SPEC_MAP[name][2])
+    return jit_from_cudnn_graph(g, config=cfg, cta_group=_spec_for(name)[1], scheduler=_spec_for(name)[2])
 
 
 # combo : (is_fp4, block_size, a_dtype, sf_dtype)
@@ -109,6 +110,27 @@ def _build_spec_map():
 
 _SPEC_MAP = _build_spec_map()
 
+_LABEL_RE = re.compile(r"^(CONFIG_sm\d+_\d+x\d+x\d+_\d+x\d+x\d+_cluster\d+x\d+)_([12])ctamma(_static)?$")
+
+
+def _spec_for(name):
+    """(geometry cfg, cta_group, scheduler) for a --configs label, or None.
+
+    The sweep set comes from the registry funnel over CATALOG; a label naming a
+    geometry outside it (e.g. a num_mma_m > 1 tile, which `by_name` synthesizes) is
+    still runnable, so parse it rather than reporting UNKNOWN_CONFIG."""
+    spec = _SPEC_MAP.get(name)
+    if spec is not None:
+        return spec
+    m = _LABEL_RE.match(name)
+    if m is None:
+        return None
+    try:
+        cfg = _by_name(m.group(1))
+    except (KeyError, NotImplementedError):
+        return None
+    return cfg, int(m.group(2)), "static" if m.group(3) else "clc"
+
 
 def _offsets(S: int, E: int) -> torch.Tensor:
     """Even split: group g owns rows [g*group_m, (g+1)*group_m)."""
@@ -425,9 +447,9 @@ def _nsys_worker(shape, combo, configs, warmup, iters, nbuf, no_baseline=False)
     pool = _mkdata_pool(S, N, K, E, combo, nbuf)
 
     # 2. each MoE-block-scale config.
-    name_to_cfg = {lbl: sp[0] for lbl, sp in _SPEC_MAP.items()}
     for name in configs or list(_SPEC_MAP):
-        cfg = name_to_cfg.get(name)
+        spec = _spec_for(name)
+        cfg = spec[0] if spec else None
         if cfg is None:
             continue
         try:
@@ -501,7 +523,6 @@ def main() -> int:
 
     flops = 2 * S * N * K
     config_names = [c.strip() for c in args.configs.split(",")] if args.configs else list(_SPEC_MAP)
-    name_to_cfg = {lbl: sp[0] for lbl, sp in _SPEC_MAP.items()}
 
     print(f"\n=== moe_block_scale_matmul G={G} M={M} N={N} K={K}  " f"(S={S} tokens, ~{flops / 1e9:.1f} GFLOP) — {combo} ===")
 
@@ -536,11 +557,12 @@ def _fmt_row(name, tflops, ms, note, ref_tflops) -> str:
             cublas_tflops, cublas_ms = float("nan"), float("nan")
             print("  cuBLAS kernel: not detected in nsys output")
         for name in config_names:
-            cfg = name_to_cfg.get(name)
+            spec = _spec_for(name)
+            cfg = spec[0] if spec else None
             if cfg is None:
                 rows.append((name, 0.0, float("inf"), "UNKNOWN_CONFIG"))
                 continue
-            tok = _kernel_match_token(cfg, _SPEC_MAP[name][1])
+            tok = _kernel_match_token(cfg, _spec_for(name)[1])
             matches = [(k, v) for k, v in kern_times.items() if tok in k]
             if not matches:
                 rows.append((name, 0.0, float("inf"), "NO_KERNEL_IN_NSYS"))
@@ -587,7 +609,8 @@ def _fmt_row(name, tflops, ms, note, ref_tflops) -> str:
 
         ctx_dead = False
         for name in config_names:
-            cfg = name_to_cfg.get(name)
+            spec = _spec_for(name)
+            cfg = spec[0] if spec else None
             if cfg is None:
                 row = (name, 0.0, float("inf"), "UNKNOWN_CONFIG")
             elif ctx_dead: