From 19bb6d3fa0c53be5099c32eeb713fbf1e691e993 Mon Sep 17 00:00:00 2001 From: Brandon Zhang <31413216+brandonfzhang@users.noreply.github.com> Date: Tue, 16 Jun 2026 12:40:18 -0700 Subject: [PATCH] fix(sdpa_benchmark): use sampled SM clock + per-arch MMA throughput for SOL% The MMA SOL% reported by benchmark_single_sdpa.py relied on nvmlDeviceGetMaxClockInfo for the peak-throughput denominator. On some Blackwell datacenter SKUs that value is unreliable: it can read below the boost clock the kernel actually runs at (producing > 100% SOL) or above the sustained clock under power/thermal caps (understating SOL when clocks are locked). Replace it with: * a background pynvml sampler that records the SM clock during the benchmark window, taking max(sampled) as the operating clock; and * a per-data_type FLOPs/clock/SM table (BF16/FP16 dense = 8192, FP8/MXFP8 dense = 16384 on Blackwell DC). Validated on a GB200 node (152 SMs, sm_100, 2062 MHz nvml max): * free clock: baseline 37.5%, patched 37.3% (agree when nvml is correct) * locked 1200: baseline 28.0%, patched 48.3% * locked 900: baseline 21.5%, patched 49.1% Patched SOL is clock-invariant by construction. Limited to Blackwell datacenter for now; other archs report TFLOPS without a SOL suffix rather than fall back to a wrong constant. --- .../benchmark_single_sdpa.py | 103 ++++++++++++++++-- 1 file changed, 94 insertions(+), 9 deletions(-) diff --git a/benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py b/benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py index c3dcbd163..ef45cb8fb 100755 --- a/benchmark/sdpa_benchmark_training/benchmark_single_sdpa.py +++ b/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