fix(sdpa_benchmark): use sampled SM clock + per-arch MMA throughput for SOL%

benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py

@@ -21,10 +21,91 @@
 import numpy as np
 import functools
 import math
+import threading
+import time
 from typing import Optional, Dict, Any
 
 from torch.profiler import profile, record_function, ProfilerActivity
 
+
+# Dense MMA throughput (FLOPs / clock / SM) for Blackwell datacenter SKUs.
+# BF16/FP16 dense = 8192, FP8/MXFP8 dense = 16384 (MXFP8 uses the FP8
+# datapath with block scaling).
+# Keys match the strings accepted by the --data_type CLI flag.
+_BLACKWELL_DC_FLOPS_PER_CLOCK_PER_SM = {
+    "bfloat16":  8192,
+    "float16":   8192,
+    "fp8":      16384,
+    "mxfp8":    16384,
+}
+
+
+def _peak_flops_per_clock_per_sm(dtype_str):
+    """Return per-SM per-clock dense FLOPs for the current GPU + dtype.
+    Returns None on unsupported arch (anything other than Blackwell DC for now)."""
+    if not torch.cuda.is_available():
+        return None
+    props = torch.cuda.get_device_properties(torch.cuda.current_device())
+    if props.major != 10:  # only Blackwell DC is in scope
+        return None
+    return _BLACKWELL_DC_FLOPS_PER_CLOCK_PER_SM.get(dtype_str)
+
+
+class _SmClockSampler:
+    """Background thread that polls SM clock via NVML at ~1 kHz.
+
+    Used to capture the actual boost clock during the benchmark window.
+    `nvmlDeviceGetMaxClockInfo` is unreliable on some Blackwell datacenter
+    SKUs: it can report a value below the boost the kernel actually runs
+    at, producing nonsensical (>100%) SOL numbers downstream.
+    """
+
+    def __init__(self):
+        self._samples = []
+        self._stop = threading.Event()
+        self._thread = None
+        self._handle = None
+        self._pynvml = None
+
+    def start(self):
+        try:
+            import pynvml
+            pynvml.nvmlInit()
+            self._pynvml = pynvml
+            self._handle = pynvml.nvmlDeviceGetHandleByIndex(torch.cuda.current_device())
+        except Exception:
+            self._pynvml = None
+            return
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._thread.start()
+
+    def _run(self):
+        pynvml = self._pynvml
+        while not self._stop.is_set():
+            try:
+                self._samples.append(pynvml.nvmlDeviceGetClockInfo(self._handle, pynvml.NVML_CLOCK_SM))
+            except Exception:
+                pass
+            # Sample at ~1 kHz; kernels run much longer than this in aggregate
+            # across warmup + measurement iterations.
+            time.sleep(0.001)
+
+    def stop(self):
+        if self._thread is None:
+            return
+        self._stop.set()
+        self._thread.join()
+        try:
+            if self._pynvml is not None:
+                self._pynvml.nvmlShutdown()
+        except Exception:
+            pass
+
+    def peak_mhz(self):
+        """Return max sampled SM clock (MHz), or None if no samples."""
+        return max(self._samples) if self._samples else None
+
+
 try:
     import cutlass.cute as cute
     import cutlass
@@ -1282,6 +1363,11 @@ def tflops_per_sec(
 
     first_error = True  # For suppressing error message beyond first error
     sdpa_function = get_sdpa_function(args.sdpa_backend)
+
+    # Sample SM clock throughout the benchmark window so SOL% uses the actual
+    # boost clock the kernel ran at rather than nvml's (often-stale) max.
+    _clock_sampler = _SmClockSampler()
+    _clock_sampler.start()
     for i in range(total_iters):
         # FP8/MXFP8 needs randn in bfloat16 then convert (randn doesn't support fp8 well)
         randn_dtype = torch.bfloat16 if args.data_type in ("fp8", "mxfp8") else target_dtype
@@ -1611,6 +1697,8 @@ def tflops_per_sec(
         else:
             del query, key, value, output
 
+    _clock_sampler.stop()
+
     ## print results
     fwd_median_time = (
         np.median(np.array(forward_times[5:])) if len(forward_times) > 5 else (np.median(np.array(forward_times)) if len(forward_times) > 0 else 0.0)
@@ -1648,18 +1736,15 @@ def tflops_per_sec(
             args.sliding_window_size,
         )
 
-    # Compute MMA SOL%
+    # Compute MMA SOL% using the per-arch FLOPs/clk/SM table and the actual
+    # sampled boost clock observed during the benchmark window.
     _peak_mma_tflops = None
     try:
-        import pynvml
-
-        pynvml.nvmlInit()
-        _handle = pynvml.nvmlDeviceGetHandleByIndex(torch.cuda.current_device())
-        _max_clock_mhz = pynvml.nvmlDeviceGetMaxClockInfo(_handle, pynvml.NVML_CLOCK_SM)
-        pynvml.nvmlShutdown()
+        _flops_per_clk_per_sm = _peak_flops_per_clock_per_sm(args.data_type)
         _num_sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
-        _fma_per_clock = 8192 if args.data_type in ("fp8", "mxfp8") else 4096
-        _peak_mma_tflops = _fma_per_clock * 2 * _num_sms * _max_clock_mhz / 1e6
+        _sampled_mhz = _clock_sampler.peak_mhz()
+        if _flops_per_clk_per_sm is not None and _sampled_mhz is not None:
+            _peak_mma_tflops = _flops_per_clk_per_sm * _num_sms * _sampled_mhz / 1e6
     except Exception:
         pass