Add FlagGems single-op perf benchmark harness

pytest plugin set + driver script for compiler A/B perf comparison:
- reproducible runs: fixed seed, yaml-driven shapes, autotune record/replay
- per-shape ttgir dump of actually-used variants with readable naming
- cudagraph-based timing with documented fallback semantics

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

_bespoke_shape_plugin.py

@@ -0,0 +1,150 @@
+"""pytest plugin: make the shape yaml authoritative for ops whose inputs are
+*bespoke* (built inline / from TestParam objects / a sampled kit) and so are
+unreachable by `_shape_inject_plugin` and `_shape_iter_inject_plugin`.
+
+A per-op registry; each handler REUSES the class's own input construction and
+only swaps the shape source for our yaml (so we never measure a different
+kernel). Covered: upsample_bicubic2d_aa_backward, flash_mla_sparse_fwd,
+fused_deepseek_..._quant_insert, cutlass_scaled_mm. Ops with no yaml entry fall
+through unchanged. Couples to upstream internals; a mismatch surfaces as a loud
+benchmark FAIL, not silent corruption.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import yaml as _yaml
+
+_YAML_CACHE: dict | None = None
+
+
+def _load_yaml(path: str) -> dict:
+    if not path or not os.path.isfile(path):
+        return {}
+    try:
+        with open(path, "r") as f:
+            return _yaml.safe_load(f) or {}
+    except Exception as exc:
+        print(f"[bespoke-shape-plugin] failed to load {path}: {exc}", file=sys.stderr)
+        return {}
+
+
+def _yaml_shapes(op_name):
+    entry = (_YAML_CACHE or {}).get(op_name)
+    if isinstance(entry, dict):
+        shapes = entry.get("shapes")
+        if shapes:
+            return shapes
+    return None
+
+
+# --- per-op builders: reuse the class's own construction, swap the shape source ---
+
+def _upsample_builder(self, dtype, shapes, original):
+    # _cfgs entries are (N,C,Hi,Wi,Ho,Wo,ac,label); yaml carries the 7 numerics.
+    self._cfgs = [tuple(shape) + ("yaml",) for shape in shapes]
+    yield from original(self, dtype)
+
+
+def _fused_deepseek_builder(self, dtype, shapes, original):
+    TestParam = original.__globals__["TestParam"]
+    for shape in shapes:
+        num_tokens, num_heads, num_tokens_insert, block_size, max_pos, eps = shape
+        param = TestParam(
+            num_tokens=num_tokens,
+            num_heads=num_heads,
+            num_tokens_insert=num_tokens_insert,
+            block_size=block_size,
+            max_pos=max_pos,
+            eps=eps,
+        )
+        yield from type(self).make_input(param)
+
+
+def _flash_mla_sparse_builder(self, dtype, shapes, original):
+    TestParam = original.__globals__["TestParam"]
+    for shape in shapes:
+        s_q, s_kv, topk, h_q, d_qk = shape
+        param = TestParam(
+            s_q=s_q, s_kv=s_kv, topk=topk, h_q=h_q, d_qk=d_qk,
+            have_attn_sink=True,  # every hardcoded case uses True
+        )
+        yield from type(self).make_input_flashmla(param)
+
+
+_WRAP_BUILDERS = {
+    "UpsampleBicubic2dAaBackwardBenchmark": _upsample_builder,
+    "FusedDeepseekV4QnormRopeKVRopeQuantInsertBenchmark": _fused_deepseek_builder,
+    "FlashmlaSparseBenchmark": _flash_mla_sparse_builder,
+}
+
+
+def _wrap_get_input_iter(cls, builder) -> None:
+    original = cls.get_input_iter
+
+    def patched(self, dtype):
+        shapes = _yaml_shapes(getattr(self, "op_name", None))
+        if shapes:
+            yield from builder(self, dtype, shapes, original)
+        else:
+            yield from original(self, dtype)
+
+    cls.get_input_iter = patched
+
+
+def _patch_cutlass(cls) -> bool:
+    """Redirect CutlassScaledMMPerfKit's hardcoded mnk to the yaml M,N,K, reusing
+    the kit's own combination/sampling pipeline (the quant-mode sweep stays
+    intrinsic)."""
+    shapes = _yaml_shapes("cutlass_scaled_mm")
+    if not shapes:
+        return False
+    kit = getattr(sys.modules.get(cls.__module__), "CutlassScaledMMPerfKit", None)
+    if kit is None:
+        return False
+    import torch
+    from itertools import product
+
+    mnk = [tuple(shape) for shape in shapes]
+
+    def _get_all_combinations():
+        # Mirror upstream verbatim except mnk (from yaml); keep in sync with
+        # CutlassScaledMMPerfKit._get_all_combinations.
+        scale_shape_types = ["scalar", "vector", "matrix"]
+        if_use_bias = [True, False]
+        dtypes = [(torch.int8, torch.float16), (torch.float8_e4m3fn, torch.bfloat16)]
+        return product(mnk, scale_shape_types, scale_shape_types, if_use_bias, dtypes)
+
+    kit._get_all_combinations = staticmethod(_get_all_combinations)
+    return True
+
+
+def pytest_collection_finish(session):
+    global _YAML_CACHE
+    from benchmark import base as fg_base
+    from benchmark.conftest import Config as fg_Config
+
+    _YAML_CACHE = _load_yaml(getattr(fg_Config, "shape_file", "") or "")
+
+    def _subclasses(cls):
+        for sub in cls.__subclasses__():
+            yield sub
+            yield from _subclasses(sub)
+
+    covered = []
+    seen = set()
+    for cls in _subclasses(fg_base.Benchmark):
+        if cls in seen:
+            continue
+        seen.add(cls)
+        builder = _WRAP_BUILDERS.get(cls.__name__)
+        if builder is not None and "get_input_iter" in cls.__dict__:
+            _wrap_get_input_iter(cls, builder)
+            covered.append(cls.__name__)
+        if cls.__name__ == "CutlassScaledMMBenchmark" and _patch_cutlass(cls):
+            covered.append("CutlassScaledMMBenchmark(mnk)")
+
+    print(f"[bespoke-shape-plugin] yaml-driven inputs for: {', '.join(covered) or 'none'}",
+          file=sys.stderr, flush=True)